From 9dcd15d350aa5e2ee2d6f7299d0183cabddeb992 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 4 Apr 2018 14:52:24 +1200 Subject: [PATCH] asyncio: ditch the handler context There are a few reasons for this. First, logs are now async, and can be called at any time. Second, the event loop is thread local, so there can only ever be one master per thread. These two things together completely obviate the need for a handler context. --- mitmproxy/addonmanager.py | 30 +++++++------------ mitmproxy/command.py | 6 +--- mitmproxy/master.py | 21 +++---------- mitmproxy/test/taddons.py | 11 ------- setup.py | 1 + test/mitmproxy/addons/test_script.py | 5 ++-- .../proxy/protocol/test_websocket.py | 1 + 7 files changed, 21 insertions(+), 54 deletions(-) diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 7c446c807..9d0e70699 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -8,7 +8,6 @@ from mitmproxy import exceptions from mitmproxy import eventsequence from mitmproxy import controller from mitmproxy import flow -from mitmproxy import log from . import ctx import pprint @@ -38,9 +37,6 @@ def cut_traceback(tb, func_name): class StreamLog: - """ - A class for redirecting output using contextlib. - """ def __init__(self, lg): self.log = lg @@ -183,9 +179,8 @@ class AddonManager: Add addons to the end of the chain, and run their load event. If any addon has sub-addons, they are registered. """ - with self.master.handlecontext(): - for i in addons: - self.chain.append(self.register(i)) + for i in addons: + self.chain.append(self.register(i)) def remove(self, addon): """ @@ -201,8 +196,7 @@ class AddonManager: raise exceptions.AddonManagerError("No such addon: %s" % n) self.chain = [i for i in self.chain if i is not a] del self.lookup[_get_name(a)] - with self.master.handlecontext(): - self.invoke_addon(a, "done") + self.invoke_addon(a, "done") def __len__(self): return len(self.chain) @@ -245,8 +239,7 @@ class AddonManager: def invoke_addon(self, addon, name, *args, **kwargs): """ - Invoke an event on an addon and all its children. This method must - run within an established handler context. + Invoke an event on an addon and all its children. """ if name not in eventsequence.Events: name = "event_" + name @@ -268,12 +261,11 @@ class AddonManager: def trigger(self, name, *args, **kwargs): """ - Establish a handler context and trigger an event across all addons + Trigger an event across all addons. """ - with self.master.handlecontext(): - for i in self.chain: - try: - with safecall(): - self.invoke_addon(i, name, *args, **kwargs) - except exceptions.AddonHalt: - return + for i in self.chain: + try: + with safecall(): + self.invoke_addon(i, name, *args, **kwargs) + except exceptions.AddonHalt: + return diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 114e882d8..35a7f42ae 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -95,11 +95,7 @@ class Command: Call the command with a list of arguments. At this point, all arguments are strings. """ - pargs = self.prepare_args(args) - - with self.manager.master.handlecontext(): - ret = self.func(*pargs) - + ret = self.func(*self.prepare_args(args)) if ret is None and self.returntype is None: return typ = mitmproxy.types.CommandTypes.get(self.returntype) diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 323f73dcc..4e17e5380 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -1,5 +1,4 @@ import threading -import contextlib import asyncio import logging @@ -58,6 +57,10 @@ class Master: self.waiting_flows = [] self.log = log.Log(self) + mitmproxy_ctx.master = self + mitmproxy_ctx.log = self.log + mitmproxy_ctx.options = self.options + @property def server(self): return self._server @@ -67,22 +70,6 @@ class Master: server.set_channel(self.channel) self._server = server - @contextlib.contextmanager - def handlecontext(self): - # Handlecontexts also have to nest - leave cleanup to the outermost - if mitmproxy_ctx.master: - yield - return - mitmproxy_ctx.master = self - mitmproxy_ctx.log = self.log - mitmproxy_ctx.options = self.options - try: - yield - finally: - mitmproxy_ctx.master = None - mitmproxy_ctx.log = None - mitmproxy_ctx.options = None - def start(self): self.should_exit.clear() if self.server: diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 6d439c4ab..73e574561 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -74,7 +74,6 @@ class context: options ) self.options = self.master.options - self.wrapped = None if loadcore: self.master.addons.add(core.Core()) @@ -82,20 +81,10 @@ class context: for a in addons: self.master.addons.add(a) - def ctx(self): - """ - Returns a new handler context. - """ - return self.master.handlecontext() - def __enter__(self): - self.wrapped = self.ctx() - self.wrapped.__enter__() return self def __exit__(self, exc_type, exc_value, traceback): - self.wrapped.__exit__(exc_type, exc_value, traceback) - self.wrapped = None return False @contextlib.contextmanager diff --git a/setup.py b/setup.py index 18fb3827c..3e9148471 100644 --- a/setup.py +++ b/setup.py @@ -88,6 +88,7 @@ setup( "flake8>=3.5, <3.6", "Flask>=0.10.1, <0.13", "mypy>=0.580,<0.581", + "pytest-asyncio>=0.8", "pytest-cov>=2.5.1,<3", "pytest-faulthandler>=1.3.1,<2", "pytest-timeout>=1.2.1,<2", diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 17b1e3723..3f0ce68c9 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -180,11 +180,12 @@ class TestScriptLoader: 'recorder responseheaders', 'recorder response' ] - def test_script_run_nonexistent(self): + @pytest.mark.asyncio + async def test_script_run_nonexistent(self): sc = script.ScriptLoader() with taddons.context(sc) as tctx: sc.script_run([tflow.tflow(resp=True)], "/") - tctx.master.has_log("/: No such script") + assert await tctx.master.await_log("/: No such script") def test_simple(self): sc = script.ScriptLoader() diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index 347c0dfca..3ce1436a0 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -47,6 +47,7 @@ class _WebSocketServerBase(net_tservers.ServerTestBase): class _WebSocketTestBase: + client = None @classmethod def setup_class(cls):