diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py index 366929a5d..7573f2a5c 100644 --- a/mitmproxy/script/concurrent.py +++ b/mitmproxy/script/concurrent.py @@ -29,4 +29,8 @@ def concurrent(fn): "script.concurrent (%s)" % fn.__name__, target=run ).start() - return _concurrent + # Support @concurrent for class-based addons + if "." in fn.__qualname__: + return staticmethod(_concurrent) + else: + return _concurrent diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py new file mode 100644 index 000000000..bd047c99a --- /dev/null +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py @@ -0,0 +1,13 @@ +import time +from mitmproxy.script import concurrent + + +class ConcurrentClass: + + @concurrent + def request(flow): + time.sleep(0.1) + + +def start(): + return ConcurrentClass() diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index fb932d9a2..e81c023dd 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -44,3 +44,21 @@ class TestConcurrent(tservers.MasterTest): ) sc.start() assert "decorator not supported" in tctx.master.event_log[0][1] + + def test_concurrent_class(self): + with taddons.context() as tctx: + sc = script.Script( + tutils.test_data.path( + "mitmproxy/data/addonscripts/concurrent_decorator_class.py" + ) + ) + sc.start() + + f1, f2 = tflow.tflow(), tflow.tflow() + tctx.cycle(sc, f1) + tctx.cycle(sc, f2) + start = time.time() + while time.time() - start < 5: + if f1.reply.state == f2.reply.state == "committed": + return + raise ValueError("Script never acked")