Use configured size limit to keep previews in check.

3 files changed, 14 insertions, 15 deletions

diff --git a/libpathod/app.py b/libpathod/app.py
index 21ae9e0d..4b5758a0 100644
--- a/libpathod/app.py
+++ b/libpathod/app.py
@@ -62,7 +62,6 @@ def onelog(lid):
     return render_template("onelog.html", section="log", alog=l, lid=lid)
 
 
-SANITY = 1024*1024
 @app.route('/preview')
 def preview():
     spec = request.args["spec"]
@@ -78,11 +77,8 @@ def preview():
         args["syntaxerror"] = str(v)
         args["marked"] = v.marked()
         return render_template("preview.html", **args)
-    if r.length() > SANITY:
-        error = "Refusing to preview a response of %s bytes. This is for your own good."%r.length()
-        args["error"] = error
-    else:
-        s = cStringIO.StringIO()
-        r.serve(s)
-        args["output"] = s.getvalue()
+
+    s = cStringIO.StringIO()
+    r.serve(s, check=app.config["pathod"].check_size)
+    args["output"] = s.getvalue()
     return render_template("preview.html", **args)
diff --git a/libpathod/pathod.py b/libpathod/pathod.py
index 28484c5d..f6b5e0f9 100644
--- a/libpathod/pathod.py
+++ b/libpathod/pathod.py
@@ -83,7 +83,7 @@ class PathodHandler(tcp.BaseHandler):
             httpversion = httpversion,
         )
         if crafted:
-            response_log = crafted.serve(self.wfile, self.check_size)
+            response_log = crafted.serve(self.wfile, self.server.check_size)
             self.server.add_log(
                 dict(
                     type = "crafted",
@@ -107,11 +107,6 @@ class PathodHandler(tcp.BaseHandler):
             self.debug("%s %s"%(method, path))
         return True
 
-    def check_size(self, req, actions):
-        if self.server.sizelimit and req.effective_length(actions) > self.server.sizelimit:
-            return "Response too large."
-        return False
-
     def handle(self):
         if self.server.ssloptions:
             try:
@@ -179,6 +174,14 @@ class Pathod(tcp.TCPServer):
                     raise PathodError("Invalid page spec in anchor: '%s', %s"%(i[1], str(v)))
                 self.anchors.append((arex, aresp))
 
+    def check_size(self, req, actions):
+        """
+            A policy check that verifies the request size is withing limits.
+        """
+        if self.sizelimit and req.effective_length(actions) > self.sizelimit:
+            return "Response too large."
+        return False
+
     @property
     def request_settings(self):
         return dict(
diff --git a/libpathod/rparse.py b/libpathod/rparse.py
index b09010b5..f81bb1ed 100644
--- a/libpathod/rparse.py
+++ b/libpathod/rparse.py
@@ -708,7 +708,7 @@ class PathodErrorResponse(Response):
         Response.__init__(self)
         self.code = 800
         self.msg = LiteralGenerator(msg)
-        self.body = LiteralGenerator(body or msg)
+        self.body = LiteralGenerator("pathod error: " + (body or msg))
         self.headers = [
             (
                 LiteralGenerator("Content-Type"),