client replay: expad and consolidate tests

4 files changed, 128 insertions, 97 deletions

diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py
index 15fa7394..11d2453b 100644
--- a/mitmproxy/addons/clientplayback.py
+++ b/mitmproxy/addons/clientplayback.py
@@ -34,10 +34,10 @@ class RequestReplayThread(basethread.BaseThread):
 
     def run(self):
         while True:
-            f = self.queue.get(block=True, timeout=None)
+            f = self.queue.get()
             self.replay(f)
 
-    def replay(self, f):
+    def replay(self, f):  # pragma: no cover
         f.live = True
         r = f.request
         bsl = human.parse_size(self.options.body_size_limit)
@@ -118,12 +118,13 @@ class RequestReplayThread(basethread.BaseThread):
             f.live = False
             if server.connected():
                 server.finish()
+                server.close()
 
 
 class ClientPlayback:
     def __init__(self):
-        self.q: queue.Queue = queue.Queue()
-        self.thread: RequestReplayThread | None = None
+        self.q = queue.Queue()
+        self.thread: RequestReplayThread = None
 
     def check(self, f: http.HTTPFlow):
         if f.live:
@@ -184,23 +185,25 @@ class ClientPlayback:
         """
         lst = []
         for f in flows:
-            err = self.check(f)
+            hf = typing.cast(http.HTTPFlow, f)
+
+            err = self.check(hf)
             if err:
                 ctx.log.warn(err)
                 continue
 
-            lst.append(f)
+            lst.append(hf)
             # Prepare the flow for replay
-            f.backup()
-            f.request.is_replay = True
-            f.response = None
-            f.error = None
+            hf.backup()
+            hf.request.is_replay = True
+            hf.response = None
+            hf.error = None
             # https://github.com/mitmproxy/mitmproxy/issues/2197
-            if f.request.http_version == "HTTP/2.0":
-                f.request.http_version = "HTTP/1.1"
-                host = f.request.headers.pop(":authority")
-                f.request.headers.insert(0, "host", host)
-            self.q.put(f)
+            if hf.request.http_version == "HTTP/2.0":
+                hf.request.http_version = "HTTP/1.1"
+                host = hf.request.headers.pop(":authority")
+                hf.request.headers.insert(0, "host", host)
+            self.q.put(hf)
         ctx.master.addons.trigger("update", lst)
 
     @command.command("replay.client.file")
diff --git a/mitmproxy/net/tcp.py b/mitmproxy/net/tcp.py
index 22016291..18429daa 100644
--- a/mitmproxy/net/tcp.py
+++ b/mitmproxy/net/tcp.py
@@ -372,12 +372,11 @@ class TCPClient(_Connection):
         # Make sure to close the real socket, not the SSL proxy.
         # OpenSSL is really good at screwing up, i.e. when trying to recv from a failed connection,
         # it tries to renegotiate...
-        if not self.connection:
-            return
-        elif isinstance(self.connection, SSL.Connection):
-            close_socket(self.connection._socket)
-        else:
-            close_socket(self.connection)
+        if self.connection:
+            if isinstance(self.connection, SSL.Connection):
+                close_socket(self.connection._socket)
+            else:
+                close_socket(self.connection)
 
     def convert_to_tls(self, sni=None, alpn_protos=None, **sslctx_kwargs):
         context = tls.create_client_context(
diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py
index a63bec53..1b385e23 100644
--- a/test/mitmproxy/addons/test_clientplayback.py
+++ b/test/mitmproxy/addons/test_clientplayback.py
@@ -1,13 +1,16 @@
+import time
 import pytest
-from unittest import mock
 
-from mitmproxy.test import tflow
+from mitmproxy.test import tflow, tutils
 from mitmproxy import io
 from mitmproxy import exceptions
+from mitmproxy.net import http as net_http
 
 from mitmproxy.addons import clientplayback
 from mitmproxy.test import taddons
 
+from .. import tservers
+
 
 def tdump(path, flows):
     with open(path, "wb") as f:
@@ -21,77 +24,80 @@ class MockThread():
         return False
 
 
-class TestClientPlayback:
-    # @staticmethod
-    # def wait_until_not_live(flow):
-    #     """
-    #     Race condition: We don't want to replay the flow while it is still live.
-    #     """
-    #     s = time.time()
-    #     while flow.live:
-    #         time.sleep(0.001)
-    #         if time.time() - s > 5:
-    #             raise RuntimeError("Flow is live for too long.")
-
-    # def test_replay(self):
-    #     assert self.pathod("304").status_code == 304
-    #     assert len(self.master.state.flows) == 1
-    #     l = self.master.state.flows[-1]
-    #     assert l.response.status_code == 304
-    #     l.request.path = "/p/305"
-    #     self.wait_until_not_live(l)
-    #     rt = self.master.replay_request(l, block=True)
-    #     assert l.response.status_code == 305
-
-    #     # Disconnect error
-    #     l.request.path = "/p/305:d0"
-    #     rt = self.master.replay_request(l, block=True)
-    #     assert rt
-    #     if isinstance(self, tservers.HTTPUpstreamProxyTest):
-    #         assert l.response.status_code == 502
-    #     else:
-    #         assert l.error
-
-    #     # Port error
-    #     l.request.port = 1
-    #     # In upstream mode, we get a 502 response from the upstream proxy server.
-    #     # In upstream mode with ssl, the replay will fail as we cannot establish
-    #     # SSL with the upstream proxy.
-    #     rt = self.master.replay_request(l, block=True)
-    #     assert rt
-    #     if isinstance(self, tservers.HTTPUpstreamProxyTest):
-    #         assert l.response.status_code == 502
-    #     else:
-    #         assert l.error
-
-    # def test_replay(self):
-    #     opts = options.Options()
-    #     fm = master.Master(opts)
-    #     f = tflow.tflow(resp=True)
-    #     f.request.content = None
-    #     with pytest.raises(ReplayException, match="missing"):
-    #         fm.replay_request(f)
-
-    #     f.request = None
-    #     with pytest.raises(ReplayException, match="request"):
-    #         fm.replay_request(f)
-
-    #     f.intercepted = True
-    #     with pytest.raises(ReplayException, match="intercepted"):
-    #         fm.replay_request(f)
-
-    #     f.live = True
-    #     with pytest.raises(ReplayException, match="live"):
-    #         fm.replay_request(f)
-
-    #     req = tutils.treq(headers=net_http.Headers(((b":authority", b"foo"), (b"header", b"qvalue"), (b"content-length", b"7"))))
-    #     f = tflow.tflow(req=req)
-    #     f.request.http_version = "HTTP/2.0"
-    #     with mock.patch('mitmproxy.proxy.protocol.http_replay.RequestReplayThread.run'):
-    #         rt = fm.replay_request(f)
-    #         assert rt.f.request.http_version == "HTTP/1.1"
-    #         assert ":authority" not in rt.f.request.headers
+class TBase(tservers.HTTPProxyTest):
+    @staticmethod
+    def wait_response(flow):
+        """
+        Race condition: We don't want to replay the flow while it is still live.
+        """
+        s = time.time()
+        while True:
+            if flow.response or flow.error:
+                break
+            time.sleep(0.001)
+            if time.time() - s > 5:
+                raise RuntimeError("Flow is live for too long.")
+
+    @staticmethod
+    def reset(f):
+        f.live = False
+        f.repsonse = False
+        f.error = False
+
+    def addons(self):
+        return [clientplayback.ClientPlayback()]
+
+    def test_replay(self):
+        cr = self.master.addons.get("clientplayback")
+
+        assert self.pathod("304").status_code == 304
+        assert len(self.master.state.flows) == 1
+        l = self.master.state.flows[-1]
+        assert l.response.status_code == 304
+        l.request.path = "/p/305"
+        cr.start_replay([l])
+        self.wait_response(l)
+        assert l.response.status_code == 305
+
+        # Disconnect error
+        cr.stop_replay()
+        self.reset(l)
+        l.request.path = "/p/305:d0"
+        cr.start_replay([l])
+        self.wait_response(l)
+        if isinstance(self, tservers.HTTPUpstreamProxyTest):
+            assert l.response.status_code == 502
+        else:
+            assert l.error
+
+        # # Port error
+        cr.stop_replay()
+        self.reset(l)
+        l.request.port = 1
+        # In upstream mode, we get a 502 response from the upstream proxy server.
+        # In upstream mode with ssl, the replay will fail as we cannot establish
+        # SSL with the upstream proxy.
+        cr.start_replay([l])
+        self.wait_response(l)
+        if isinstance(self, tservers.HTTPUpstreamProxyTest):
+            assert l.response.status_code == 502
+        else:
+            assert l.error
+
+
+class TestHTTPProxy(TBase, tservers.HTTPProxyTest):
+    pass
+
+
+class TestHTTPSProxy(TBase, tservers.HTTPProxyTest):
+    ssl = True
+
+
+class TestUpstreamProxy(TBase, tservers.HTTPUpstreamProxyTest):
+    pass
 
+
+class TestClientPlayback:
     def test_load_file(self, tmpdir):
         cp = clientplayback.ClientPlayback()
         with taddons.context(cp):
@@ -133,13 +139,37 @@ class TestClientPlayback:
             f.request.raw_content = None
             assert "missing content" in cp.check(f)
 
-    def test_playback(self):
+    @pytest.mark.asyncio
+    async def test_playback(self):
         cp = clientplayback.ClientPlayback()
-        with taddons.context(cp):
+        with taddons.context(cp) as ctx:
             assert cp.count() == 0
             f = tflow.tflow(resp=True)
             cp.start_replay([f])
             assert cp.count() == 1
 
             cp.stop_replay()
-            assert cp.count() == 0
\ No newline at end of file
+            assert cp.count() == 0
+
+            f.live = True
+            cp.start_replay([f])
+            assert cp.count() == 0
+            await ctx.master.await_log("live")
+
+    def test_http2(self):
+        cp = clientplayback.ClientPlayback()
+        with taddons.context(cp):
+            req = tutils.treq(
+                headers = net_http.Headers(
+                    (
+                        (b":authority", b"foo"),
+                        (b"header", b"qvalue"),
+                        (b"content-length", b"7")
+                    )
+                )
+            )
+            f = tflow.tflow(req=req)
+            f.request.http_version = "HTTP/2.0"
+            cp.start_replay([f])
+            assert f.request.http_version == "HTTP/1.1"
+            assert ":authority" not in f.request.headers
diff --git a/test/mitmproxy/tools/web/test_app.py b/test/mitmproxy/tools/web/test_app.py
index b272861c..3d18987d 100644
--- a/test/mitmproxy/tools/web/test_app.py
+++ b/test/mitmproxy/tools/web/test_app.py
@@ -9,7 +9,6 @@ import tornado.testing
 from tornado import httpclient
 from tornado import websocket
 
-from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy.test import tflow
 from mitmproxy.tools.web import app