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.