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

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):

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()

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()

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.

pathod

@@ -24,6 +24,11 @@ if __name__ == "__main__":
         action="store_true",
         help='Serve with SSL.'
     )
+    parser.add_argument(
+        "--limit-size", dest='sizelimit', default=None,
+        type=str,
+        help='Size limit of served responses. Understands size suffixes, i.e. 100k.'
+    )
     parser.add_argument(
         "--keyfile", dest='ssl_keyfile', default=None,
         type=str,
@@ -67,12 +72,20 @@ if __name__ == "__main__":
     if not args.debug:
         logging.disable(logging.DEBUG)
 
+    sizelimit = None
+    if args.sizelimit:
+        try:
+            sizelimit = utils.parse_size(args.sizelimit)
+        except ValueError, v:
+            parser.error(v)
+
     try:
         pd = pathod.Pathod(
             (args.address, args.port),
             ssloptions = ssl,
             staticdir = args.staticdir,
-            anchors = alst
+            anchors = alst,
+            sizelimit = sizelimit,
         )
     except pathod.PathodError, v:
         parser.error(str(v))

test/test_pathod.py

@@ -46,7 +46,8 @@ class _DaemonTests:
         self.d = test.Daemon(
             staticdir=tutils.test_data.path("data"),
             anchors=[("/anchor/.*", "202")],
-            ssl = self.SSL
+            ssl = self.SSL,
+            sizelimit=1*1024*1024
         )
 
     @classmethod
@@ -73,6 +74,12 @@ class _DaemonTests:
             c.settimeout(timeout)
         return c.request(spec)
 
+    def test_sizelimit(self):
+        r = self.get("200:b@1g")
+        assert r.status_code == 800
+        l = self.d.log()[0]
+        assert "too large" in l["response"]["error"]
+
     def test_preline(self):
         v = self.pathoc(r"get:'/p/200':i0,'\r\n'")
         assert v[1] == 200

test/test_utils.py

@@ -1,6 +1,12 @@
 from libpathod import utils
 import tutils
 
+def test_parse_size():
+    assert utils.parse_size("100") == 100
+    assert utils.parse_size("100k") == 100 * 1024
+    tutils.raises("invalid size spec",  utils.parse_size, "foo")
+    tutils.raises("invalid size spec",  utils.parse_size, "100kk")
+
 
 def test_parse_anchor_spec():
     assert utils.parse_anchor_spec("foo=200") == ("foo", "200")