Use policy hook to apply a size limit in pathod, add corresponding cmdline arg.

4 files changed, 34 insertions, 16 deletions

diff --git a/libpathod/pathod.py b/libpathod/pathod.py
index 2f9717df..0247c204 100644
--- a/libpathod/pathod.py
+++ b/libpathod/pathod.py
@@ -87,8 +87,6 @@ class PathodHandler(tcp.BaseHandler):
         )
         if crafted:
             response_log = crafted.serve(self.wfile, self.check_size)
-            if response_log["disconnect"]:
-                return
             self.server.add_log(
                 dict(
                     type = "crafted",
@@ -96,6 +94,8 @@ class PathodHandler(tcp.BaseHandler):
                     response=response_log
                 )
             )
+            if response_log["disconnect"]:
+                return
         else:
             cc = wsgi.ClientConn(self.client_address)
             req = wsgi.Request(cc, "http", method, path, headers, content)
@@ -111,6 +111,8 @@ class PathodHandler(tcp.BaseHandler):
         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):
diff --git a/libpathod/rparse.py b/libpathod/rparse.py
index 8c70e154..7836ea51 100644
--- a/libpathod/rparse.py
+++ b/libpathod/rparse.py
@@ -213,20 +213,13 @@ class ValueNakedLiteral(_Value):
 
 
 class ValueGenerate:
-    UNITS = dict(
-        b = 1024**0,
-        k = 1024**1,
-        m = 1024**2,
-        g = 1024**3,
-        t = 1024**4,
-    )
     def __init__(self, usize, unit, datatype):
         if not unit:
             unit = "b"
         self.usize, self.unit, self.datatype = usize, unit, datatype
 
     def bytes(self):
-        return self.usize * self.UNITS[self.unit]
+        return self.usize * utils.SIZE_UNITS[self.unit]
 
     def get_generator(self, settings):
         return RandomGenerator(self.datatype, self.bytes())
@@ -235,7 +228,7 @@ class ValueGenerate:
     def expr(klass):
         e = pp.Literal("@").suppress() + v_integer
 
-        u = reduce(operator.or_, [pp.Literal(i) for i in klass.UNITS.keys()])
+        u = reduce(operator.or_, [pp.Literal(i) for i in utils.SIZE_UNITS.keys()])
         e = e + pp.Optional(u, default=None)
 
         s = pp.Literal(",").suppress()
diff --git a/libpathod/test.py b/libpathod/test.py
index 00e03823..b90c8de6 100644
--- a/libpathod/test.py
+++ b/libpathod/test.py
@@ -6,9 +6,9 @@ import tutils
 IFACE = "127.0.0.1"
 
 class Daemon:
-    def __init__(self, staticdir=None, anchors=(), ssl=None):
+    def __init__(self, staticdir=None, anchors=(), ssl=None, sizelimit=None):
         self.q = Queue.Queue()
-        self.thread = PaThread(self.q, staticdir, anchors, ssl)
+        self.thread = PaThread(self.q, staticdir, anchors, ssl, sizelimit)
         self.thread.start()
         self.port = self.q.get(True, 5)
         self.urlbase = "%s://%s:%s"%("https" if ssl else "http", IFACE, self.port)
@@ -43,9 +43,9 @@ class Daemon:
 
 
 class PaThread(threading.Thread):
-    def __init__(self, q, staticdir, anchors, ssl):
+    def __init__(self, q, staticdir, anchors, ssl, sizelimit):
         threading.Thread.__init__(self)
-        self.q, self.staticdir, self.anchors, self.ssl = q, staticdir, anchors, ssl
+        self.q, self.staticdir, self.anchors, self.ssl, self.sizelimit = q, staticdir, anchors, ssl, sizelimit
 
     def run(self):
         if self.ssl is True:
@@ -59,7 +59,8 @@ class PaThread(threading.Thread):
             (IFACE, 0),
             ssloptions = ssloptions,
             anchors = self.anchors,
-            staticdir = self.staticdir
+            staticdir = self.staticdir,
+            sizelimit = self.sizelimit
         )
         self.q.put(self.server.port)
         self.server.serve_forever()
diff --git a/libpathod/utils.py b/libpathod/utils.py
index c656a0d0..de83b19a 100644
--- a/libpathod/utils.py
+++ b/libpathod/utils.py
@@ -1,6 +1,28 @@
 import os, re
 import rparse
 
+SIZE_UNITS = dict(
+    b = 1024**0,
+    k = 1024**1,
+    m = 1024**2,
+    g = 1024**3,
+    t = 1024**4,
+)
+
+def parse_size(s):
+    try:
+        return int(s)
+    except ValueError:
+        pass
+    for i in SIZE_UNITS.keys():
+        if s.endswith(i):
+            try:
+                return int(s[:-1]) * SIZE_UNITS[i]
+            except ValueError:
+                break
+    raise ValueError("Invalid size specification.")
+
+
 def get_header(val, headers):
     """
         Header keys may be Values, so we have to "generate" them as we try the match.