From cee4b724591dd614220c6df8fadd71cad00855b6 Mon Sep 17 00:00:00 2001
From: Robert Xiao <brx@cs.ubc.ca>
Date: Tue, 1 Feb 2022 23:40:39 -0800
Subject: [PATCH] Support async hooks. Fixes #4207.

---
 CHANGELOG.md                   |  1 +
 examples/addons/nonblocking.py | 22 ++++++++++++-----
 mitmproxy/addonmanager.py      | 43 ++++++++++++++++++++++++++++++----
 3 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7e9beb2c7..cdd53c9f7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,7 @@
 * Fix compatibility with BoringSSL (@pmoulton)
 * Change connection event hooks to be blocking.
   Processing will only resume once the event hook has finished. (@Prinzhorn)
+* Allow addon hooks to be async (@nneonneo, #4207)
 
 ## 28 September 2021: mitmproxy 7.0.4
 
diff --git a/examples/addons/nonblocking.py b/examples/addons/nonblocking.py
index e6b79f87d..2ee092980 100644
--- a/examples/addons/nonblocking.py
+++ b/examples/addons/nonblocking.py
@@ -1,16 +1,26 @@
 """
-Make events hooks non-blocking.
-
-When event hooks are decorated with @concurrent, they will be run in their own thread, freeing the main event loop.
-Please note that this generally opens the door to race conditions and decreases performance if not required.
+Make events hooks non-blocking using async or @concurrent
 """
+import asyncio
 import time
 
 from mitmproxy.script import concurrent
+from mitmproxy import ctx
 
 
-@concurrent  # Remove this and see what happens
-def request(flow):
+# Hooks can be async, which allows the hook to call async functions and perform async I/O
+# without blocking other requests. This is generally preferred for new addons.
+async def request(flow):
+    ctx.log.info(f"handle request: {flow.request.host}{flow.request.path}")
+    await asyncio.sleep(5)
+    ctx.log.info(f"start  request: {flow.request.host}{flow.request.path}")
+
+
+# Another option is to use @concurrent, which launches the hook in its own thread.
+# Please note that this generally opens the door to race conditions and decreases performance if not required.
+# Rename the function below to request(flow) to try it out.
+@concurrent  # Remove this to make it synchronous and see what happens
+def request_concurrent(flow):
     # This is ugly in mitmproxy's UI, but you don't want to use mitmproxy.ctx.log from a different thread.
     print(f"handle request: {flow.request.host}{flow.request.path}")
     time.sleep(5)
diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
index a0def682c..984e858d8 100644
--- a/mitmproxy/addonmanager.py
+++ b/mitmproxy/addonmanager.py
@@ -1,4 +1,5 @@
 import contextlib
+import inspect
 import pprint
 import sys
 import traceback
@@ -240,7 +241,7 @@ class AddonManager:
         if isinstance(message.reply, controller.DummyReply):
             message.reply.reset()
 
-        self.trigger(event)
+        await self.async_trigger(event)
 
         if message.reply.state == "start":
             message.reply.take()
@@ -250,18 +251,18 @@ class AddonManager:
                 message.reply.mark_reset()
 
         if isinstance(message, flow.Flow):
-            self.trigger(hooks.UpdateHook([message]))
+            await self.async_trigger(hooks.UpdateHook([message]))
 
-    def invoke_addon(self, addon, event: hooks.Hook):
+    def _iter_hooks(self, addon, event: hooks.Hook):
         """
-            Invoke an event on an addon and all its children.
+            Enumerate all hook callables belonging to the given addon
         """
         assert isinstance(event, hooks.Hook)
         for a in traverse([addon]):
             func = getattr(a, event.name, None)
             if func:
                 if callable(func):
-                    func(*event.args())
+                    yield a, func
                 elif isinstance(func, types.ModuleType):
                     # we gracefully exclude module imports with the same name as hooks.
                     # For example, a user may have "from mitmproxy import log" in an addon,
@@ -273,6 +274,38 @@ class AddonManager:
                         f"Addon handler {event.name} ({a}) not callable"
                     )
 
+    async def async_invoke_addon(self, addon, event: hooks.Hook):
+        """
+            Asynchronously invoke an event on an addon and all its children.
+        """
+        for addon, func in self._iter_hooks(addon, event):
+            res = func(*event.args())
+            # Support both async and sync hook functions
+            if inspect.isawaitable(res):
+                await res
+
+    def invoke_addon(self, addon, event: hooks.Hook):
+        """
+            Invoke an event on an addon and all its children.
+        """
+        for addon, func in self._iter_hooks(addon, event):
+            if inspect.iscoroutinefunction(func):
+                raise exceptions.AddonManagerError(
+                    f"Async handler {event.name} ({addon}) cannot be called from sync context"
+                )
+            func(*event.args())
+
+    async def async_trigger(self, event: hooks.Hook):
+        """
+            Asynchronously trigger an event across all addons.
+        """
+        for i in self.chain:
+            try:
+                with safecall():
+                    await self.async_invoke_addon(i, event)
+            except exceptions.AddonHalt:
+                return
+
     def trigger(self, event: hooks.Hook):
         """
             Trigger an event across all addons.