diff options
| -rw-r--r-- | libpathod/pathod.py | 9 | ||||
| -rw-r--r-- | libpathod/rparse.py | 49 | ||||
| -rw-r--r-- | test/test_rparse.py | 7 |
3 files changed, 55 insertions, 10 deletions
diff --git a/libpathod/pathod.py b/libpathod/pathod.py index 5180361b..2f9717df 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 98254abb..8c70e154 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 0cb3e373..929d7342 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'") |