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

@@ -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

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)

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:

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

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",

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()

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):