From c7b5faf7dbaab518bbe9942f018861f738ebb2b0 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Mon, 23 Jul 2012 14:37:00 +1200
Subject: [PATCH] Add hooks for policy checks of served data.

---
 libpathod/pathod.py |  9 +++++++--
 libpathod/rparse.py | 49 +++++++++++++++++++++++++++++++++++++--------
 test/test_rparse.py |  7 +++++++
 3 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/libpathod/pathod.py b/libpathod/pathod.py
index 5180361be..2f9717df0 100644
--- a/libpathod/pathod.py
+++ b/libpathod/pathod.py
@@ -86,7 +86,7 @@ class PathodHandler(tcp.BaseHandler):
             httpversion = httpversion,
         )
         if crafted:
-            response_log = crafted.serve(self.wfile)
+            response_log = crafted.serve(self.wfile, self.check_size)
             if response_log["disconnect"]:
                 return
             self.server.add_log(
@@ -110,6 +110,9 @@ class PathodHandler(tcp.BaseHandler):
             self.debug("%s %s"%(method, path))
         return True
 
+    def check_size(self, req, actions):
+        return False
+
     def handle(self):
         if self.server.ssloptions:
             try:
@@ -145,7 +148,7 @@ class PathodHandler(tcp.BaseHandler):
 
 class Pathod(tcp.TCPServer):
     LOGBUF = 500
-    def __init__(self, addr, ssloptions=None, prefix="/p/", staticdir=None, anchors=None):
+    def __init__(self, addr, ssloptions=None, prefix="/p/", staticdir=None, anchors=None, sizelimit=None):
         """
             addr: (address, port) tuple. If port is 0, a free port will be
             automatically chosen.
@@ -153,11 +156,13 @@ class Pathod(tcp.TCPServer):
             prefix: string specifying the prefix at which to anchor response generation.
             staticdir: path to a directory of static resources, or None.
             anchors: A list of (regex, spec) tuples, or None.
+            sizelimit: Limit size of served data.
         """
         tcp.TCPServer.__init__(self, addr)
         self.ssloptions = ssloptions
         self.staticdir = staticdir
         self.prefix = prefix
+        self.sizelimit = sizelimit
         self.app = app.app
         self.app.config["pathod"] = self
         self.log = []
diff --git a/libpathod/rparse.py b/libpathod/rparse.py
index 98254abbc..8c70e1549 100644
--- a/libpathod/rparse.py
+++ b/libpathod/rparse.py
@@ -25,6 +25,18 @@ class ParseException(Exception):
 class ServerError(Exception): pass
 
 
+def actions_log(lst):
+    ret = []
+    for i in lst:
+        if i[1] == "inject":
+            ret.append(
+                [i[0], i[1], repr(i[2])]
+            )
+        else:
+            ret.append(i)
+    return ret
+
+
 def ready_actions(length, lst):
     ret = []
     for i in lst:
@@ -527,7 +539,15 @@ class Message:
                 l += len(i[2])
         return l
 
-    def serve(self, fp):
+    def serve(self, fp, check):
+        """
+            fp: The file pointer to write to.
+
+            check: A function called with the effective actions (after random
+            values have been calculated).  If it returns False service
+            proceeds, otherwise the return is treated as an error message to be
+            sent to the client, and service stops.
+        """
         started = time.time()
         if self.body and not utils.get_header("Content-Length", self.headers):
             self.headers.append(
@@ -554,13 +574,26 @@ class Message:
         vals.reverse()
         actions = ready_actions(self.length(), self.actions)
         actions.reverse()
+        if check:
+            ret = check(self, actions)
+            if ret:
+                err = InternalResponse(
+                    800,
+                    ret
+                )
+                err.serve(fp)
+                return dict(
+                    disconnect = True,
+                    actions = actions_log(actions),
+                    error = ret
+                )
         disconnect = write_values(fp, vals, actions[:])
         duration = time.time() - started
         ret = dict(
             disconnect = disconnect,
             started = started,
             duration = duration,
-            actions = actions,
+            actions = actions_log(actions),
         )
         for i in self.logattrs:
             v = getattr(self, i)
@@ -661,8 +694,8 @@ class CraftedRequest(Request):
         for i in tokens:
             i.accept(settings, self)
 
-    def serve(self, fp):
-        d = Request.serve(self, fp)
+    def serve(self, fp, check=None):
+        d = Request.serve(self, fp, check)
         d["spec"] = self.spec
         return d
 
@@ -674,8 +707,8 @@ class CraftedResponse(Response):
         for i in tokens:
             i.accept(settings, self)
 
-    def serve(self, fp):
-        d = Response.serve(self, fp)
+    def serve(self, fp, check=None):
+        d = Response.serve(self, fp, check)
         d["spec"] = self.spec
         return d
 
@@ -697,8 +730,8 @@ class InternalResponse(Response):
             )
         ]
 
-    def serve(self, fp):
-        d = Response.serve(self, fp)
+    def serve(self, fp, check=None):
+        d = Response.serve(self, fp, check)
         d["internal"] = True
         return d
 
diff --git a/test/test_rparse.py b/test/test_rparse.py
index 0cb3e3731..929d73428 100644
--- a/test/test_rparse.py
+++ b/test/test_rparse.py
@@ -419,6 +419,13 @@ class TestResponse:
         assert r.body[:]
         assert str(r)
 
+    def test_checkfunc(self):
+        s = cStringIO.StringIO()
+        r = rparse.parse_response({}, "400:b@100k")
+        def check(req, acts):
+            return "errmsg"
+        assert r.serve(s, check=check)["error"] == "errmsg"
+
     def test_render(self):
         s = cStringIO.StringIO()
         r = rparse.parse_response({}, "400'msg'")