diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 35817a85d..6d7bd7734 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -222,7 +222,9 @@ def handler(f): # 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 message.reply.handled: + if handling: + if not message.reply.taken: + message.reply.commit() message.reply.handled = False return ret # Mark this function as a handler wrapper @@ -230,6 +232,9 @@ def handler(f): return wrapper +NO_REPLY = object() + + class Reply(object): """ Messages sent through a channel are decorated with a "reply" attribute. @@ -241,6 +246,8 @@ class Reply(object): 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? @@ -258,8 +265,21 @@ class Reply(object): 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.q.put(msg) + self.ack_message = msg + if self.taken: + self.commit() + + def commit(self): + """ + This is called by the handler to actually send the ack message. + """ + if self.ack_message == NO_REPLY: + raise exceptions.ControlException("Message has no response.") + self.q.put(self.ack_message) def __del__(self): if not self.acked: