diff options
| -rw-r--r-- | libpathod/pathod.py | 6 | ||||
| -rw-r--r-- | libpathod/rparse.py | 11 | ||||
| -rw-r--r-- | libpathod/test.py | 11 | ||||
| -rw-r--r-- | libpathod/utils.py | 22 | ||||
| -rwxr-xr-x | pathod | 15 | ||||
| -rw-r--r-- | test/test_pathod.py | 9 | ||||
| -rw-r--r-- | test/test_utils.py | 6 |
7 files changed, 62 insertions, 18 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. @@ -25,6 +25,11 @@ if __name__ == "__main__": 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, help='SSL key file. If not specified, a default key is used.' @@ -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)) diff --git a/test/test_pathod.py b/test/test_pathod.py index d917e25c..8e1e7490 100644 --- a/test/test_pathod.py +++ b/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 diff --git a/test/test_utils.py b/test/test_utils.py index 72c892f0..8ca2da49 100644 --- a/test/test_utils.py +++ b/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") |