diff options
| -rw-r--r-- | mitmproxy/controller.py | 3 | ||||
| -rw-r--r-- | mitmproxy/master.py | 12 | ||||
| -rw-r--r-- | test/mitmproxy/proxy/test_server.py | 34 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 8 |
4 files changed, 33 insertions, 24 deletions
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 97fb0f51..79b049c9 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -8,8 +8,9 @@ class Channel: The only way for the proxy server to communicate with the master is to use the channel it has been given. """ - def __init__(self, loop, q): + def __init__(self, loop, q, should_exit): self.loop = loop + self.should_exit = should_exit self._q = q def ask(self, mtype, m): diff --git a/mitmproxy/master.py b/mitmproxy/master.py index c2d2b6a0..2f47ae09 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -1,7 +1,6 @@ import threading import contextlib import asyncio -import signal from mitmproxy import addonmanager from mitmproxy import options @@ -37,11 +36,16 @@ class Master: """ def __init__(self, opts): self.event_queue = asyncio.Queue() + self.should_exit = threading.Event() + self.channel = controller.Channel( + asyncio.get_event_loop(), + self.event_queue, + self.should_exit, + ) self.options = opts or options.Options() # type: options.Options self.commands = command.CommandManager(self) self.addons = addonmanager.AddonManager(self) - self.should_exit = threading.Event() self._server = None self.first_tick = True self.waiting_flows = [] @@ -52,7 +56,7 @@ class Master: @server.setter def server(self, server): - server.set_channel(controller.Channel(asyncio.get_event_loop(), self.event_queue)) + server.set_channel(self.channel) self._server = server @contextlib.contextmanager @@ -202,7 +206,7 @@ class Master: host = f.request.headers.pop(":authority") f.request.headers.insert(0, "host", host) - rt = http_replay.RequestReplayThread(self.options, f, self.server.channel) + rt = http_replay.RequestReplayThread(self.options, f, self.channel) rt.start() # pragma: no cover if block: rt.join() diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index 4cfaa523..9e276294 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -276,10 +276,9 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): s = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py") ) - self.master.addons.add(s) + self.set_addons(s) d = self.pathod('200:b"foo"') assert d.content == b"bar" - self.master.addons.remove(s) def test_first_line_rewrite(self): """ @@ -583,12 +582,11 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin): s = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py") ) - self.master.addons.add(s) + self.set_addons(s) self._tcpproxy_on() d = self.pathod('200:b"foo"') self._tcpproxy_off() assert d.content == b"bar" - self.master.addons.remove(s) class TestTransparentSSL(tservers.TransparentProxyTest, CommonMixin, TcpMixin): @@ -731,7 +729,7 @@ class TestRedirectRequest(tservers.HTTPProxyTest): This test verifies that the original destination is restored for the third request. """ - self.proxy.tmaster.addons.add(ARedirectRequest(self.server2.port)) + self.set_addons(ARedirectRequest(self.server2.port)) p = self.pathoc() with p.connect(): @@ -770,7 +768,7 @@ class AStreamRequest: class TestStreamRequest(tservers.HTTPProxyTest): def test_stream_simple(self): - self.proxy.tmaster.addons.add(AStreamRequest()) + self.set_addons(AStreamRequest()) p = self.pathoc() with p.connect(): # a request with 100k of data but without content-length @@ -779,7 +777,7 @@ class TestStreamRequest(tservers.HTTPProxyTest): assert len(r1.content) > 100000 def test_stream_multiple(self): - self.proxy.tmaster.addons.add(AStreamRequest()) + self.set_addons(AStreamRequest()) p = self.pathoc() with p.connect(): # simple request with streaming turned on @@ -791,7 +789,7 @@ class TestStreamRequest(tservers.HTTPProxyTest): assert r1.status_code == 201 def test_stream_chunked(self): - self.proxy.tmaster.addons.add(AStreamRequest()) + self.set_addons(AStreamRequest()) connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM) connection.connect(("127.0.0.1", self.proxy.port)) fconn = connection.makefile("rb") @@ -820,7 +818,7 @@ class AFakeResponse: class TestFakeResponse(tservers.HTTPProxyTest): def test_fake(self): - self.proxy.tmaster.addons.add(AFakeResponse()) + self.set_addons(AFakeResponse()) f = self.pathod("200") assert "header-response" in f.headers @@ -836,7 +834,7 @@ class TestServerConnect(tservers.HTTPProxyTest): def test_unnecessary_serverconnect(self): """A replayed/fake response with no upstream_cert should not connect to an upstream server""" - self.proxy.tmaster.addons.add(AFakeResponse()) + self.set_addons(AFakeResponse()) assert self.pathod("200").status_code == 200 assert not self.proxy.tmaster.has_log("serverconnect") @@ -849,7 +847,7 @@ class AKillRequest: class TestKillRequest(tservers.HTTPProxyTest): def test_kill(self): - self.proxy.tmaster.addons.add(AKillRequest()) + self.set_addons(AKillRequest()) with pytest.raises(exceptions.HttpReadDisconnect): self.pathod("200") # Nothing should have hit the server @@ -863,7 +861,7 @@ class AKillResponse: class TestKillResponse(tservers.HTTPProxyTest): def test_kill(self): - self.proxy.tmaster.addons.add(AKillResponse()) + self.set_addons(AKillResponse()) with pytest.raises(exceptions.HttpReadDisconnect): self.pathod("200") # The server should have seen a request @@ -886,7 +884,7 @@ class AIncomplete: class TestIncompleteResponse(tservers.HTTPProxyTest): def test_incomplete(self): - self.proxy.tmaster.addons.add(AIncomplete()) + self.set_addons(AIncomplete()) assert self.pathod("200").status_code == 502 @@ -969,7 +967,7 @@ class TestUpstreamProxySSL( def test_change_upstream_proxy_connect(self): # skip chain[0]. - self.proxy.tmaster.addons.add( + self.set_addons( UpstreamProxyChanger( ("127.0.0.1", self.chain[1].port) ) @@ -988,8 +986,8 @@ class TestUpstreamProxySSL( Client <- HTTPS -> Proxy <- HTTP -> Proxy <- HTTPS -> Server """ - self.proxy.tmaster.addons.add(RewriteToHttp()) - self.chain[1].tmaster.addons.add(RewriteToHttps()) + self.set_addons(RewriteToHttp()) + self.set_addons(RewriteToHttps()) p = self.pathoc() with p.connect(): resp = p.request("get:'/p/418'") @@ -1063,8 +1061,8 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest): http1obj.server_conn.wfile.write(headers) http1obj.server_conn.wfile.flush() - self.chain[0].tmaster.addons.add(RequestKiller([1, 2])) - self.chain[1].tmaster.addons.add(RequestKiller([1])) + self.chain[0].set_addons(RequestKiller([1, 2])) + self.chain[1].set_addons(RequestKiller([1])) p = self.pathoc() with p.connect(): diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 9e6cef97..f3270e5c 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -92,6 +92,7 @@ class ProxyThread(threading.Thread): self.masterclass = masterclass self.options = options self.tmaster = None + self.event_loop = None controller.should_exit = False @property @@ -106,7 +107,8 @@ class ProxyThread(threading.Thread): self.tmaster.shutdown() def run(self): - asyncio.set_event_loop(asyncio.new_event_loop()) + self.event_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self.event_loop) self.tmaster = self.masterclass(self.options) self.tmaster.addons.add(core.Core()) self.name = "ProxyThread (%s:%s)" % ( @@ -177,6 +179,10 @@ class ProxyTestBase: ssl_insecure=True, ) + def set_addons(self, *addons): + self.proxy.tmaster.reset(addons) + self.proxy.tmaster.addons.trigger("tick") + def addons(self): """ Can be over-ridden to add a standard set of addons to tests. |