From 0ee1b40c1701f3b60d799a93cb7f161da2816e5a Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 8 Aug 2016 19:20:06 -0700 Subject: improve reply semantics --- mitmproxy/controller.py | 114 ++++++++++++++++++++++++------------------------ 1 file changed, 57 insertions(+), 57 deletions(-) (limited to 'mitmproxy/controller.py') diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 6d7bd773..10044c0a 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -207,32 +207,29 @@ def handler(f): # acking is ours. If not, it's someone else's and we ignore it. handling = False # We're the first handler - ack responsibility is ours - if not message.reply.handled: + if message.reply.state == "unhandled": handling = True - message.reply.handled = True + message.reply.handle() with master.handlecontext(): ret = f(master, message) if handling: master.addons(f.__name__, message) - if handling and not message.reply.acked and not message.reply.taken: - message.reply.ack() - # Reset the handled flag - it's common for us to feed the same object # through handlers repeatedly, so we don't want this to persist across # calls. - if handling: - if not message.reply.taken: - message.reply.commit() - message.reply.handled = False + if handling and not message.reply.taken: + if message.reply.value == NO_REPLY: + message.reply.ack() + message.reply.commit() return ret # Mark this function as a handler wrapper wrapper.__dict__["__handler"] = True return wrapper -NO_REPLY = object() +NO_REPLY = object() # special object we can distinguish from a valid "None" reply. class Reply(object): @@ -244,67 +241,70 @@ class Reply(object): def __init__(self, obj): self.obj = obj self.q = queue.Queue() - # Has this message been acked? - self.acked = False - # What's the ack message? - self.ack_message = NO_REPLY - # Has the user taken responsibility for ack-ing? - self.taken = False - # Has a handler taken responsibility for ack-ing? - self.handled = False - def ack(self): - self.send(self.obj) + self.state = "unhandled" # "unhandled" -> "handled" -> "taken" -> "committed" + self.value = NO_REPLY # holds the reply value. May change before things are actually commited. - def kill(self): - self.send(exceptions.Kill) + def handle(self): + """ + Reply are handled by controller.handlers, which may be nested. The first handler takes + responsibility and handles the reply. + """ + if self.state != "unhandled": + raise exceptions.ControlException("Message is {}, but expected it to be unhandled.".format(self.state)) + self.state = "handled" def take(self): - self.taken = True - - def send(self, msg): - if self.acked: - raise exceptions.ControlException("Message already acked.") - if self.ack_message != NO_REPLY: - # We may have overrides for this later. - raise exceptions.ControlException("Message already has a response.") - self.acked = True - self.ack_message = msg - if self.taken: - self.commit() + """ + Scripts or other parties make "take" a reply out of a normal flow. + For example, intercepted flows are taken out so that the connection thread does not proceed. + """ + if self.state != "handled": + raise exceptions.ControlException("Message is {}, but expected it to be handled.".format(self.state)) + self.state = "taken" def commit(self): """ - This is called by the handler to actually send the ack message. + Ultimately, messages are commited. This is done either automatically by the handler + if the message is not taken or manually by the entity which called .take(). """ - if self.ack_message == NO_REPLY: - raise exceptions.ControlException("Message has no response.") - self.q.put(self.ack_message) + if self.state != "taken": + raise exceptions.ControlException("Message is {}, but expected it to be taken.".format(self.state)) + self.state = "committed" + self.q.put(self.value) + + def ack(self): + self.send(self.obj) + + def kill(self): + self.send(exceptions.Kill) + + def send(self, msg): + if self.state not in ("handled", "taken"): + raise exceptions.ControlException( + "Reply is currently {}, did not expect a call to .send().".format(self.state) + ) + if self.value is not NO_REPLY: + raise exceptions.ControlException("There is already a reply for this message.") + self.value = msg def __del__(self): - if not self.acked: + if self.state != "comitted": # This will be ignored by the interpreter, but emit a warning - raise exceptions.ControlException("Un-acked message: %s" % self.obj) + raise exceptions.ControlException("Uncomitted message: %s" % self.obj) -class DummyReply(object): +class DummyReply(Reply): """ - A reply object that does nothing. Useful when we need an object to seem - like it has a channel, and during testing. + A reply object that is not connected to anything. In contrast to regular Reply objects, + DummyReply objects are reset to "unhandled" after a commit so that they can be used + multiple times. Useful when we need an object to seem like it has a channel, + and during testing. """ def __init__(self): - self.acked = False - self.taken = False - self.handled = False - - def kill(self): - self.send(None) + super(DummyReply, self).__init__(None) - def ack(self): - self.send(None) - - def take(self): - self.taken = True - - def send(self, msg): - self.acked = True + def commit(self): + super(DummyReply, self).commit() + self.state = "unhandled" + self.value = NO_REPLY -- cgit v1.2.3