diff options
| -rw-r--r-- | mitmproxy/controller.py | 28 | ||||
| -rw-r--r-- | mitmproxy/master.py | 85 | ||||
| -rw-r--r-- | mitmproxy/proxy/protocol/http_replay.py | 16 | ||||
| -rw-r--r-- | mitmproxy/tools/console/master.py | 11 | ||||
| -rw-r--r-- | mitmproxy/tools/main.py | 5 | ||||
| -rw-r--r-- | mitmproxy/tools/web/master.py | 4 | ||||
| -rw-r--r-- | release/README.md | 4 | ||||
| -rw-r--r-- | test/mitmproxy/data/addonscripts/shutdown.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/proxy/protocol/test_http2.py | 6 | ||||
| -rw-r--r-- | test/mitmproxy/proxy/protocol/test_websocket.py | 21 | ||||
| -rw-r--r-- | test/mitmproxy/proxy/test_server.py | 42 | ||||
| -rw-r--r-- | test/mitmproxy/test_controller.py | 79 | ||||
| -rw-r--r-- | test/mitmproxy/test_flow.py | 9 | ||||
| -rw-r--r-- | test/mitmproxy/test_fuzzing.py | 12 | ||||
| -rw-r--r-- | test/mitmproxy/tools/test_main.py | 10 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 60 |
16 files changed, 170 insertions, 224 deletions
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index beb210ca..79b049c9 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -1,4 +1,5 @@ import queue +import asyncio from mitmproxy import exceptions @@ -7,9 +8,10 @@ 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, q, should_exit): - self.q = q + def __init__(self, loop, q, should_exit): + self.loop = loop self.should_exit = should_exit + self._q = q def ask(self, mtype, m): """ @@ -20,18 +22,11 @@ class Channel: exceptions.Kill: All connections should be closed immediately. """ m.reply = Reply(m) - self.q.put((mtype, m)) - while not self.should_exit.is_set(): - try: - # The timeout is here so we can handle a should_exit event. - g = m.reply.q.get(timeout=0.5) - except queue.Empty: # pragma: no cover - continue - if g == exceptions.Kill: - raise exceptions.Kill() - return g - m.reply._state = "committed" # suppress error message in __del__ - raise exceptions.Kill() + asyncio.run_coroutine_threadsafe(self._q.put((mtype, m)), self.loop) + g = m.reply.q.get() + if g == exceptions.Kill: + raise exceptions.Kill() + return g def tell(self, mtype, m): """ @@ -39,7 +34,7 @@ class Channel: then return immediately. """ m.reply = DummyReply() - self.q.put((mtype, m)) + asyncio.run_coroutine_threadsafe(self._q.put((mtype, m)), self.loop) NO_REPLY = object() # special object we can distinguish from a valid "None" reply. @@ -52,7 +47,8 @@ class Reply: """ def __init__(self, obj): self.obj = obj - self.q = queue.Queue() # type: queue.Queue + # Spawn an event loop in the current thread + self.q = queue.Queue() self._state = "start" # "start" -> "taken" -> "committed" diff --git a/mitmproxy/master.py b/mitmproxy/master.py index a5e948f6..372bb289 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -1,6 +1,7 @@ import threading import contextlib -import queue +import asyncio +import logging from mitmproxy import addonmanager from mitmproxy import options @@ -18,6 +19,13 @@ from mitmproxy.coretypes import basethread from . import ctx as mitmproxy_ctx +# Conclusively preventing cross-thread races on proxy shutdown turns out to be +# very hard. We could build a thread sync infrastructure for this, or we could +# wait until we ditch threads and move all the protocols into the async loop. +# Until then, silence non-critical errors. +logging.getLogger('asyncio').setLevel(logging.CRITICAL) + + class ServerThread(basethread.BaseThread): def __init__(self, server): self.server = server @@ -35,11 +43,19 @@ class Master: The master handles mitmproxy's main event loop. """ 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, + ) + asyncio.ensure_future(self.main()) + asyncio.ensure_future(self.tick()) + self.options = opts or options.Options() # type: options.Options self.commands = command.CommandManager(self) self.addons = addonmanager.AddonManager(self) - self.event_queue = queue.Queue() - self.should_exit = threading.Event() self._server = None self.first_tick = True self.waiting_flows = [] @@ -50,9 +66,7 @@ class Master: @server.setter def server(self, server): - server.set_channel( - controller.Channel(self.event_queue, self.should_exit) - ) + server.set_channel(self.channel) self._server = server @contextlib.contextmanager @@ -71,7 +85,8 @@ class Master: mitmproxy_ctx.log = None mitmproxy_ctx.options = None - def tell(self, mtype, m): + # This is a vestigial function that will go away in a refactor very soon + def tell(self, mtype, m): # pragma: no cover m.reply = controller.DummyReply() self.event_queue.put((mtype, m)) @@ -86,38 +101,43 @@ class Master: if self.server: ServerThread(self.server).start() - def run(self): - self.start() - try: - while not self.should_exit.is_set(): - self.tick(0.1) - finally: - self.shutdown() + async def main(self): + while True: + try: + mtype, obj = await self.event_queue.get() + except RuntimeError: + return + if mtype not in eventsequence.Events: # pragma: no cover + raise exceptions.ControlException("Unknown event %s" % repr(mtype)) + self.addons.handle_lifecycle(mtype, obj) + self.event_queue.task_done() - def tick(self, timeout): + async def tick(self): if self.first_tick: self.first_tick = False self.addons.trigger("running") - self.addons.trigger("tick") - changed = False + while True: + if self.should_exit.is_set(): + asyncio.get_event_loop().stop() + return + self.addons.trigger("tick") + await asyncio.sleep(0.1) + + def run(self): + self.start() + asyncio.ensure_future(self.tick()) + loop = asyncio.get_event_loop() try: - mtype, obj = self.event_queue.get(timeout=timeout) - if mtype not in eventsequence.Events: - raise exceptions.ControlException( - "Unknown event %s" % repr(mtype) - ) - self.addons.handle_lifecycle(mtype, obj) - self.event_queue.task_done() - changed = True - except queue.Empty: - pass - return changed + loop.run_forever() + finally: + self.shutdown() + loop.close() + self.addons.trigger("done") def shutdown(self): if self.server: self.server.shutdown() self.should_exit.set() - self.addons.trigger("done") def _change_reverse_host(self, f): """ @@ -199,12 +219,7 @@ class Master: host = f.request.headers.pop(":authority") f.request.headers.insert(0, "host", host) - rt = http_replay.RequestReplayThread( - self.options, - f, - self.event_queue, - self.should_exit - ) + rt = http_replay.RequestReplayThread(self.options, f, self.channel) rt.start() # pragma: no cover if block: rt.join() diff --git a/mitmproxy/proxy/protocol/http_replay.py b/mitmproxy/proxy/protocol/http_replay.py index 0f3be1ea..b2cca2b1 100644 --- a/mitmproxy/proxy/protocol/http_replay.py +++ b/mitmproxy/proxy/protocol/http_replay.py @@ -1,7 +1,3 @@ -import queue -import threading -import typing - from mitmproxy import log from mitmproxy import controller from mitmproxy import exceptions @@ -25,20 +21,12 @@ class RequestReplayThread(basethread.BaseThread): self, opts: options.Options, f: http.HTTPFlow, - event_queue: typing.Optional[queue.Queue], - should_exit: threading.Event + channel: controller.Channel, ) -> None: - """ - event_queue can be a queue or None, if no scripthooks should be - processed. - """ self.options = opts self.f = f f.live = True - if event_queue: - self.channel = controller.Channel(event_queue, should_exit) - else: - self.channel = None + self.channel = channel super().__init__( "RequestReplay (%s)" % f.request.url ) diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index 5cc1cf43..de660b17 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -1,3 +1,4 @@ +import asyncio import mailcap import mimetypes import os @@ -182,12 +183,6 @@ class ConsoleMaster(master.Master): ) self.ui.clear() - def ticker(self, *userdata): - changed = self.tick(timeout=0) - if changed: - self.loop.draw_screen() - self.loop.set_alarm_in(0.01, self.ticker) - def inject_key(self, key): self.loop.process_input([key]) @@ -206,6 +201,7 @@ class ConsoleMaster(master.Master): ) self.loop = urwid.MainLoop( urwid.SolidFill("x"), + event_loop=urwid.AsyncioEventLoop(loop=asyncio.get_event_loop()), screen = self.ui, handle_mouse = self.options.console_mouse, ) @@ -214,8 +210,6 @@ class ConsoleMaster(master.Master): self.loop.widget = self.window self.window.refresh() - self.loop.set_alarm_in(0.01, self.ticker) - if self.start_err: def display_err(*_): self.sig_add_log(None, self.start_err) @@ -236,6 +230,7 @@ class ConsoleMaster(master.Master): finally: sys.stderr.flush() super().shutdown() + self.addons.trigger("done") def shutdown(self): raise urwid.ExitMainLoop diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 91488a1f..53c236bb 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -1,5 +1,6 @@ from __future__ import print_function # this is here for the version check to work on Python 2. +import asyncio import sys if sys.version_info < (3, 6): @@ -117,8 +118,10 @@ def run( def cleankill(*args, **kwargs): master.shutdown() - signal.signal(signal.SIGTERM, cleankill) + loop = asyncio.get_event_loop() + for signame in ('SIGINT', 'SIGTERM'): + loop.add_signal_handler(getattr(signal, signame), master.shutdown) master.run() except exceptions.OptionsError as e: print("%s: %s" % (sys.argv[0], e), file=sys.stderr) diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index 4c597f0e..b7eddcce 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -2,6 +2,8 @@ import webbrowser import tornado.httpserver import tornado.ioloop +from tornado.platform.asyncio import AsyncIOMainLoop + from mitmproxy import addons from mitmproxy import log from mitmproxy import master @@ -102,6 +104,7 @@ class WebMaster(master.Master): ) def run(self): # pragma: no cover + AsyncIOMainLoop().install() iol = tornado.ioloop.IOLoop.instance() @@ -109,7 +112,6 @@ class WebMaster(master.Master): http_server.listen(self.options.web_port, self.options.web_iface) iol.add_callback(self.start) - tornado.ioloop.PeriodicCallback(lambda: self.tick(timeout=0), 5).start() web_url = "http://{}:{}/".format(self.options.web_iface, self.options.web_port) self.add_log( diff --git a/release/README.md b/release/README.md index 0e9c373b..aaac8b0c 100644 --- a/release/README.md +++ b/release/README.md @@ -5,9 +5,9 @@ Make sure run all these steps on the correct branch you want to create a new rel - Update CHANGELOG - Verify that all CI tests pass - Tag the release and push to Github - - For alphas, betas, and release candidates, use lightweight tags. + - For alphas, betas, and release candidates, use lightweight tags. This is necessary so that the .devXXXX counter does not reset. - - For final releases, use annotated tags. + - For final releases, use annotated tags. This makes the .devXXXX counter reset. - Wait for tag CI to complete diff --git a/test/mitmproxy/data/addonscripts/shutdown.py b/test/mitmproxy/data/addonscripts/shutdown.py index 51a99b5c..3da4d03e 100644 --- a/test/mitmproxy/data/addonscripts/shutdown.py +++ b/test/mitmproxy/data/addonscripts/shutdown.py @@ -1,5 +1,5 @@ from mitmproxy import ctx -def running(): +def tick(): ctx.master.shutdown() diff --git a/test/mitmproxy/proxy/protocol/test_http2.py b/test/mitmproxy/proxy/protocol/test_http2.py index d9aa03b4..13f28728 100644 --- a/test/mitmproxy/proxy/protocol/test_http2.py +++ b/test/mitmproxy/proxy/protocol/test_http2.py @@ -10,7 +10,6 @@ import h2 from mitmproxy import options import mitmproxy.net -from mitmproxy.addons import core from ...net import tservers as net_tservers from mitmproxy import exceptions from mitmproxy.net.http import http1, http2 @@ -90,9 +89,7 @@ class _Http2TestBase: @classmethod def setup_class(cls): cls.options = cls.get_options() - tmaster = tservers.TestMaster(cls.options) - tmaster.addons.add(core.Core()) - cls.proxy = tservers.ProxyThread(tmaster) + cls.proxy = tservers.ProxyThread(tservers.TestMaster, cls.options) cls.proxy.start() @classmethod @@ -120,6 +117,7 @@ class _Http2TestBase: def teardown(self): if self.client: self.client.close() + self.server.server.wait_for_silence() def setup_connection(self): self.client = mitmproxy.net.tcp.TCPClient(("127.0.0.1", self.proxy.port)) diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index 661605b7..e5ed8e9d 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -3,10 +3,10 @@ import os import struct import tempfile import traceback +import time from mitmproxy import options from mitmproxy import exceptions -from mitmproxy.addons import core from mitmproxy.http import HTTPFlow from mitmproxy.websocket import WebSocketFlow @@ -52,9 +52,7 @@ class _WebSocketTestBase: @classmethod def setup_class(cls): cls.options = cls.get_options() - tmaster = tservers.TestMaster(cls.options) - tmaster.addons.add(core.Core()) - cls.proxy = tservers.ProxyThread(tmaster) + cls.proxy = tservers.ProxyThread(tservers.TestMaster, cls.options) cls.proxy.start() @classmethod @@ -163,7 +161,7 @@ class TestSimple(_WebSocketTest): def websocket_start(self, f): f.stream = streaming - self.master.addons.add(Stream()) + self.proxy.set_addons(Stream()) self.setup_connection() frame = websockets.Frame.from_file(self.client.rfile) @@ -204,7 +202,7 @@ class TestSimple(_WebSocketTest): def websocket_message(self, f): f.messages[-1].content = "foo" - self.master.addons.add(Addon()) + self.proxy.set_addons(Addon()) self.setup_connection() frame = websockets.Frame.from_file(self.client.rfile) @@ -235,7 +233,7 @@ class TestKillFlow(_WebSocketTest): def websocket_message(self, f): f.kill() - self.master.addons.add(KillFlow()) + self.proxy.set_addons(KillFlow()) self.setup_connection() with pytest.raises(exceptions.TcpDisconnect): @@ -329,7 +327,12 @@ class TestPong(_WebSocketTest): assert frame.header.opcode == websockets.OPCODE.PONG assert frame.payload == b'foobar' - assert self.master.has_log("Pong Received from server", "info") + for i in range(20): + if self.master.has_log("Pong Received from server", "info"): + break + time.sleep(0.01) + else: + raise AssertionError("No pong seen") class TestClose(_WebSocketTest): @@ -405,7 +408,7 @@ class TestStreaming(_WebSocketTest): def websocket_start(self, f): f.stream = streaming - self.master.addons.add(Stream()) + self.proxy.set_addons(Stream()) self.setup_connection() frame = None diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index 986dfb39..aed4a774 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -21,14 +21,6 @@ from pathod import pathod from .. import tservers from ...conftest import skip_appveyor -""" - Note that the choice of response code in these tests matters more than you - might think. libcurl treats a 304 response code differently from, say, a - 200 response code - it will correctly terminate a 304 response with no - content-length header, whereas it will block forever waiting for content - for a 200 response. -""" - class CommonMixin: @@ -284,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): """ @@ -591,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): @@ -739,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(): @@ -778,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 @@ -787,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 @@ -799,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") @@ -828,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 @@ -844,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") @@ -857,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 @@ -871,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 @@ -894,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 @@ -977,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) ) @@ -996,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.chain[1].set_addons(RewriteToHttps()) p = self.pathoc() with p.connect(): resp = p.request("get:'/p/418'") @@ -1071,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/test_controller.py b/test/mitmproxy/test_controller.py index e840380a..f7c64ed9 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -1,82 +1,31 @@ -from threading import Thread, Event -from unittest.mock import Mock +import asyncio import queue import pytest from mitmproxy.exceptions import Kill, ControlException from mitmproxy import controller -from mitmproxy import master -from mitmproxy import proxy from mitmproxy.test import taddons -class TMsg: - pass +@pytest.mark.asyncio +async def test_master(): + class TMsg: + pass + class tAddon: + def log(self, _): + ctx.master.should_exit.set() -class TestMaster: - def test_simple(self): - class tAddon: - def log(self, _): - ctx.master.should_exit.set() + with taddons.context(tAddon()) as ctx: + assert not ctx.master.should_exit.is_set() - with taddons.context() as ctx: - ctx.master.addons.add(tAddon()) - assert not ctx.master.should_exit.is_set() + async def test(): msg = TMsg() msg.reply = controller.DummyReply() - ctx.master.event_queue.put(("log", msg)) - ctx.master.run() - assert ctx.master.should_exit.is_set() - - def test_server_simple(self): - m = master.Master(None) - m.server = proxy.DummyServer() - m.start() - m.shutdown() - m.start() - m.shutdown() + await ctx.master.channel.tell("log", msg) - -class TestServerThread: - def test_simple(self): - m = Mock() - t = master.ServerThread(m) - t.run() - assert m.serve_forever.called - - -class TestChannel: - def test_tell(self): - q = queue.Queue() - channel = controller.Channel(q, Event()) - m = Mock(name="test_tell") - channel.tell("test", m) - assert q.get() == ("test", m) - assert m.reply - - def test_ask_simple(self): - q = queue.Queue() - - def reply(): - m, obj = q.get() - assert m == "test" - obj.reply.send(42) - obj.reply.take() - obj.reply.commit() - - Thread(target=reply).start() - - channel = controller.Channel(q, Event()) - assert channel.ask("test", Mock(name="test_ask_simple")) == 42 - - def test_ask_shutdown(self): - q = queue.Queue() - done = Event() - done.set() - channel = controller.Channel(q, done) - with pytest.raises(Kill): - channel.ask("test", Mock(name="test_ask_shutdown")) + asyncio.ensure_future(test()) + assert not ctx.master.should_exit.is_set() class TestReply: diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 8cc11a16..4042de5b 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -7,7 +7,7 @@ import mitmproxy.io from mitmproxy import flowfilter from mitmproxy import options from mitmproxy.io import tnetstring -from mitmproxy.exceptions import FlowReadException, ReplayException, ControlException +from mitmproxy.exceptions import FlowReadException, ReplayException from mitmproxy import flow from mitmproxy import http from mitmproxy.net import http as net_http @@ -169,9 +169,10 @@ class TestFlowMaster: f.error = flow.Error("msg") fm.addons.handle_lifecycle("error", f) - fm.tell("foo", f) - with pytest.raises(ControlException): - fm.tick(timeout=1) + # FIXME: This no longer works, because we consume on the main loop. + # fm.tell("foo", f) + # with pytest.raises(ControlException): + # fm.addons.trigger("unknown") fm.shutdown() diff --git a/test/mitmproxy/test_fuzzing.py b/test/mitmproxy/test_fuzzing.py index 905ba1cd..57d0ca55 100644 --- a/test/mitmproxy/test_fuzzing.py +++ b/test/mitmproxy/test_fuzzing.py @@ -25,14 +25,4 @@ class TestFuzzy(tservers.HTTPProxyTest): p = self.pathoc() with p.connect(): resp = p.request(req % self.server.port) - assert resp.status_code == 400 - - # def test_invalid_upstream(self): - # req = r"get:'http://localhost:%s/p/200:i10,\x27+\x27'" - # p = self.pathoc() - # assert p.request(req % self.server.port).status_code == 502 - - # def test_upstream_disconnect(self): - # req = r'200:d0' - # p = self.pathod(req) - # assert p.status_code == 502 + assert resp.status_code == 400
\ No newline at end of file diff --git a/test/mitmproxy/tools/test_main.py b/test/mitmproxy/tools/test_main.py index 88e2fe86..a514df74 100644 --- a/test/mitmproxy/tools/test_main.py +++ b/test/mitmproxy/tools/test_main.py @@ -1,19 +1,25 @@ +import pytest + from mitmproxy.test import tutils from mitmproxy.tools import main shutdown_script = tutils.test_data.path("mitmproxy/data/addonscripts/shutdown.py") -def test_mitmweb(): +@pytest.mark.asyncio +async def test_mitmweb(): main.mitmweb([ "--no-web-open-browser", "-q", + "-p", "0", "-s", shutdown_script ]) -def test_mitmdump(): +@pytest.mark.asyncio +async def test_mitmdump(): main.mitmdump([ "-q", + "-p", "0", "-s", shutdown_script ]) diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 0040b023..2d102a5d 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -2,7 +2,9 @@ import os.path import threading import tempfile import sys +import time from unittest import mock +import asyncio import mitmproxy.platform from mitmproxy.addons import core @@ -12,6 +14,7 @@ from mitmproxy import controller from mitmproxy import options from mitmproxy import exceptions from mitmproxy import io +from mitmproxy.utils import human import pathod.test import pathod.pathoc @@ -62,11 +65,6 @@ class TestState: if f not in self.flows: self.flows.append(f) - # TODO: add TCP support? - # def tcp_start(self, f): - # if f not in self.flows: - # self.flows.append(f) - class TestMaster(taddons.RecordingMaster): @@ -90,13 +88,12 @@ class TestMaster(taddons.RecordingMaster): class ProxyThread(threading.Thread): - def __init__(self, tmaster): + def __init__(self, masterclass, options): threading.Thread.__init__(self) - self.tmaster = tmaster - self.name = "ProxyThread (%s:%s)" % ( - tmaster.server.address[0], - tmaster.server.address[1], - ) + self.masterclass = masterclass + self.options = options + self.tmaster = None + self.event_loop = None controller.should_exit = False @property @@ -107,11 +104,27 @@ class ProxyThread(threading.Thread): def tlog(self): return self.tmaster.logs + def shutdown(self): + self.tmaster.shutdown() + def run(self): + 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)" % human.format_address(self.tmaster.server.address) self.tmaster.run() - def shutdown(self): - self.tmaster.shutdown() + def set_addons(self, *addons): + self.tmaster.reset(addons) + self.tmaster.addons.trigger("tick") + + def start(self): + super().start() + while True: + if self.tmaster: + break + time.sleep(0.01) class ProxyTestBase: @@ -132,9 +145,7 @@ class ProxyTestBase: ssloptions=cls.ssloptions) cls.options = cls.get_options() - tmaster = cls.masterclass(cls.options) - tmaster.addons.add(core.Core()) - cls.proxy = ProxyThread(tmaster) + cls.proxy = ProxyThread(cls.masterclass, cls.options) cls.proxy.start() @classmethod @@ -173,6 +184,9 @@ class ProxyTestBase: ssl_insecure=True, ) + def set_addons(self, *addons): + self.proxy.set_addons(*addons) + def addons(self): """ Can be over-ridden to add a standard set of addons to tests. @@ -327,8 +341,7 @@ class SocksModeTest(HTTPProxyTest): return opts -class ChainProxyTest(ProxyTestBase): - +class HTTPUpstreamProxyTest(HTTPProxyTest): """ Chain three instances of mitmproxy in a row to test upstream mode. Proxy order is cls.proxy -> cls.chain[0] -> cls.chain[1] @@ -344,11 +357,12 @@ class ChainProxyTest(ProxyTestBase): cls.chain = [] for _ in range(cls.n): opts = cls.get_options() - tmaster = cls.masterclass(opts) - tmaster.addons.add(core.Core()) - proxy = ProxyThread(tmaster) + proxy = ProxyThread(cls.masterclass, opts) proxy.start() cls.chain.insert(0, proxy) + while True: + if proxy.event_loop and proxy.event_loop.is_running(): + break super().setup_class() @@ -372,7 +386,3 @@ class ChainProxyTest(ProxyTestBase): mode="upstream:" + s, ) return opts - - -class HTTPUpstreamProxyTest(ChainProxyTest, HTTPProxyTest): - pass |