Flow.intercept: use an Event instead of the reply system

This is patch 3/4 of the reply-ectomy.

mitmproxy/addons/clientplayback.py

@@ -98,6 +98,8 @@ class ReplayHandler(server.ConnectionHandler):
         data.reply = AsyncReply(data)
         await ctx.master.addons.handle_lifecycle(hook)
         await data.reply.done.wait()
+        if isinstance(data, flow.Flow):
+            await data.wait_for_resume()
         if isinstance(hook, (layers.http.HttpResponseHook, layers.http.HttpErrorHook)):
             if self.transports:
                 # close server connections

mitmproxy/addons/proxyserver.py

@@ -51,6 +51,8 @@ class ProxyConnectionHandler(server.StreamConnectionHandler):
             await self.master.addons.handle_lifecycle(hook)
             await data.reply.done.wait()
             data.reply = None
+            if isinstance(data, flow.Flow):
+                await data.wait_for_resume()
 
     def log(self, message: str, level: str = "info") -> None:
         x = log.LogEntry(self.log_prefix + message, level)

mitmproxy/flow.py

@@ -1,3 +1,4 @@
+import asyncio
 import time
 import typing  # noqa
 import uuid
@@ -118,6 +119,7 @@ class Flow(stateobject.StateObject):
         self.live = live
 
         self.intercepted: bool = False
+        self._resume_event: typing.Optional[asyncio.Event] = None
         self._backup: typing.Optional[Flow] = None
         self.reply: typing.Optional[controller.Reply] = None
         self.marked: str = ""
@@ -213,7 +215,18 @@ class Flow(stateobject.StateObject):
         if self.intercepted:
             return
         self.intercepted = True
-        self.reply.take()
+        if self._resume_event is not None:
+            self._resume_event.clear()
+
+    async def wait_for_resume(self):
+        """
+        Wait until this Flow is resumed.
+        """
+        if not self.intercepted:
+            return
+        if self._resume_event is None:
+            self._resume_event = asyncio.Event()
+        await self._resume_event.wait()
 
     def resume(self):
         """
@@ -222,9 +235,8 @@ class Flow(stateobject.StateObject):
         if not self.intercepted:
             return
         self.intercepted = False
-        # If a flow is intercepted and then duplicated, the duplicated one is not taken.
+        if self._resume_event is not None:
-        if self.reply.state == "taken":
+            self._resume_event.set()
-            self.reply.commit()
 
     @property
     def timestamp_start(self) -> float:

test/mitmproxy/test_http.py

@@ -1,3 +1,4 @@
+import asyncio
 import email
 import time
 import json
@@ -719,16 +720,46 @@ class TestHTTPFlow:
     def test_intercept(self):
         f = tflow()
         f.intercept()
-        assert f.reply.state == "taken"
+        assert f.intercepted
         f.intercept()
-        assert f.reply.state == "taken"
+        assert f.intercepted
 
     def test_resume(self):
         f = tflow()
-        f.intercept()
-        assert f.reply.state == "taken"
         f.resume()
-        assert f.reply.state == "committed"
+        assert not f.intercepted
+        f.intercept()
+        assert f.intercepted
+        f.resume()
+        assert not f.intercepted
+
+    @pytest.mark.asyncio
+    async def test_wait_for_resume(self):
+        f = tflow()
+        await f.wait_for_resume()
+
+        f = tflow()
+        f.intercept()
+        f.resume()
+        await f.wait_for_resume()
+
+        f = tflow()
+        f.intercept()
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(f.wait_for_resume(), 0.2)
+        f.resume()
+        await f.wait_for_resume()
+
+        f = tflow()
+        f.intercept()
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(f.wait_for_resume(), 0.2)
+        f.resume()
+        f.intercept()
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(f.wait_for_resume(), 0.2)
+        f.resume()
+        await f.wait_for_resume()
 
     def test_resume_duplicated(self):
         f = tflow()