aboutsummaryrefslogtreecommitdiffstats
path: root/mitmproxy/controller.py
diff options
context:
space:
mode:
authorMaximilian Hils <git@maximilianhils.com>2016-08-08 19:20:06 -0700
committerMaximilian Hils <git@maximilianhils.com>2016-08-08 19:20:06 -0700
commit0ee1b40c1701f3b60d799a93cb7f161da2816e5a (patch)
tree13804a1d149e88b417be515cc524448173ca7f73 /mitmproxy/controller.py
parenta52a1df23c7593daf9979fab56c35e1ebccb6d1f (diff)
downloadmitmproxy-0ee1b40c1701f3b60d799a93cb7f161da2816e5a.tar.gz
mitmproxy-0ee1b40c1701f3b60d799a93cb7f161da2816e5a.tar.bz2
mitmproxy-0ee1b40c1701f3b60d799a93cb7f161da2816e5a.zip
improve reply semantics
Diffstat (limited to 'mitmproxy/controller.py')
-rw-r--r--mitmproxy/controller.py114
1 files changed, 57 insertions, 57 deletions
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