Merge pull request #2051 from MatthewShao/fix-#1928-@concurrent-in-OOP

Fix #1928, @concurrent decorator in class

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

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

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