diff options
| -rw-r--r-- | mitmproxy/controller.py | 48 | ||||
| -rw-r--r-- | mitmproxy/flow/master.py | 4 | ||||
| -rw-r--r-- | mitmproxy/models/flow.py | 5 | ||||
| -rw-r--r-- | test/mitmproxy/test_server.py | 7 |
4 files changed, 37 insertions, 27 deletions
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 1498c3ad..1aac82db 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -145,23 +145,6 @@ class Channel(object): self.q.put((mtype, m)) -class DummyReply(object): - """ - A reply object that does nothing. 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 take(self): - self.taken = True - - def __call__(self, msg=False): - self.acked = True - - # Special value to distinguish the case where no reply was sent NO_REPLY = object() @@ -192,7 +175,7 @@ def handler(f): ret = f(*args, **kwargs) if handling and not message.reply.acked and not message.reply.taken: - message.reply() + message.reply.ack() return ret # Mark this function as a handler wrapper wrapper.func_dict["__handler"] = True @@ -215,6 +198,12 @@ class Reply(object): # Has a handler taken responsibility for ack-ing? self.handled = False + def ack(self): + self(NO_REPLY) + + def kill(self): + self(exceptions.Kill) + def take(self): self.taken = True @@ -231,3 +220,26 @@ class Reply(object): if not self.acked: # This will be ignored by the interpreter, but emit a warning raise exceptions.ControlException("Un-acked message") + + +class DummyReply(object): + """ + A reply object that does nothing. 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() + + def ack(self): + self() + + def take(self): + self.taken = True + + def __call__(self, msg=False): + self.acked = True diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index ec0bf36d..31475f5b 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -411,7 +411,7 @@ class FlowMaster(controller.Master): ) if err: self.add_event("Error in wsgi app. %s" % err, "error") - f.reply(exceptions.Kill) + f.reply.kill() return if f not in self.state.flows: # don't add again on replay self.state.add_flow(f) @@ -428,7 +428,7 @@ class FlowMaster(controller.Master): if self.stream_large_bodies: self.stream_large_bodies.run(f, False) except netlib.exceptions.HttpException: - f.reply(exceptions.Kill) + f.reply.kill() return self.run_script_hook("responseheaders", f) return f diff --git a/mitmproxy/models/flow.py b/mitmproxy/models/flow.py index e2dac221..de86e451 100644 --- a/mitmproxy/models/flow.py +++ b/mitmproxy/models/flow.py @@ -4,7 +4,6 @@ import time import copy import uuid -from mitmproxy import exceptions from mitmproxy import stateobject from mitmproxy import version from mitmproxy.models.connections import ClientConnection @@ -155,7 +154,7 @@ class Flow(stateobject.StateObject): """ self.error = Error("Connection killed") self.intercepted = False - self.reply(exceptions.Kill) + self.reply.kill() master.error(self) def intercept(self, master): @@ -175,5 +174,5 @@ class Flow(stateobject.StateObject): if not self.intercepted: return self.intercepted = False - self.reply() + self.reply.ack() master.handle_accept_intercept(self) diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index b58c4f44..1cd6cb0c 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -14,7 +14,6 @@ from pathod import pathoc, pathod from mitmproxy import controller from mitmproxy.proxy.config import HostMatcher -from mitmproxy.exceptions import Kill from mitmproxy.models import Error, HTTPResponse, HTTPFlow from . import tutils, tservers @@ -771,7 +770,7 @@ class MasterKillRequest(tservers.TestMaster): @controller.handler def request(self, f): - f.reply(Kill) + f.reply.kill() class TestKillRequest(tservers.HTTPProxyTest): @@ -788,7 +787,7 @@ class MasterKillResponse(tservers.TestMaster): @controller.handler def response(self, f): - f.reply(Kill) + f.reply.kill() class TestKillResponse(tservers.HTTPProxyTest): @@ -942,7 +941,7 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest): if not (k[0] in exclude): f.client_conn.finish() f.error = Error("terminated") - f.reply(Kill) + f.reply.kill() return _func(f) setattr(master, attr, handler) |