diff --git a/mitmproxy/addons/clientplayback_sansio.py b/mitmproxy/addons/clientplayback_sansio.py index b362aa7a4..a9b5c10f2 100644 --- a/mitmproxy/addons/clientplayback_sansio.py +++ b/mitmproxy/addons/clientplayback_sansio.py @@ -51,7 +51,7 @@ class MockServer(layers.http.HttpConnection): layers.http.ResponseProtocolError, )): pass - else: + else: # pragma: no cover ctx.log(f"Unexpected event during replay: {events}") @@ -130,6 +130,7 @@ class ClientPlayback: await h.replay() except Exception: ctx.log(f"Client replay has crashed!\n{traceback.format_exc()}", "error") + self.queue.task_done() self.inflight = None def check(self, f: flow.Flow) -> typing.Optional[str]: @@ -179,6 +180,7 @@ class ClientPlayback: except asyncio.QueueEmpty: break else: + self.queue.task_done() f.revert() updated.append(f) diff --git a/setup.cfg b/setup.cfg index 5006359f1..554d29b19 100644 --- a/setup.cfg +++ b/setup.cfg @@ -56,9 +56,11 @@ exclude = [tool:individual_coverage] exclude = + mitmproxy/addons/clientplayback.py mitmproxy/addons/onboardingapp/app.py mitmproxy/addons/session.py mitmproxy/addons/termlog.py + mitmproxy/connections.py mitmproxy/contentviews/base.py mitmproxy/controller.py mitmproxy/ctx.py diff --git a/test/mitmproxy/addons/test_clientplayback_sansio.py b/test/mitmproxy/addons/test_clientplayback_sansio.py new file mode 100644 index 000000000..4a35a8a3f --- /dev/null +++ b/test/mitmproxy/addons/test_clientplayback_sansio.py @@ -0,0 +1,139 @@ +import asyncio +from contextlib import asynccontextmanager + +import pytest + +from mitmproxy.addons.clientplayback_sansio import ClientPlayback, ReplayHandler +from mitmproxy.exceptions import CommandError, OptionsError +from mitmproxy.proxy2.context import Address +from mitmproxy.test import taddons, tflow + + +@asynccontextmanager +async def tcp_server(handle_conn) -> Address: + server = await asyncio.start_server(handle_conn, '127.0.0.1', 0) + await server.start_serving() + try: + yield server.sockets[0].getsockname() + finally: + server.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["regular", "upstream", "err"]) +async def test_playback(mode): + handler_ok = asyncio.Event() + + async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): + if mode == "err": + writer.close() + handler_ok.set() + return + if mode == "upstream": + conn_req = await reader.readuntil(b"\r\n\r\n") + assert conn_req == b'CONNECT address:22 HTTP/1.1\r\n\r\n' + writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n") + req = await reader.readuntil(b"data") + assert req == ( + b'GET /path HTTP/1.1\r\n' + b'header: qvalue\r\n' + b'content-length: 4\r\n' + b'\r\n' + b'data' + ) + writer.write(b"HTTP/1.1 204 No Content\r\n\r\n") + await writer.drain() + assert not await reader.read() + handler_ok.set() + + cp = ClientPlayback() + with taddons.context(cp) as tctx: + async with tcp_server(handler) as addr: + + cp.running() + flow = tflow.tflow() + flow.request.content = b"data" + if mode == "upstream": + tctx.options.mode = f"upstream:http://{addr[0]}:{addr[1]}" + else: + flow.request.host, flow.request.port = addr + cp.start_replay([flow]) + assert cp.count() == 1 + await cp.queue.join() + await handler_ok.wait() + cp.done() + if mode != "err": + assert flow.response.status_code == 204 + + +@pytest.mark.asyncio +async def test_playback_crash(monkeypatch): + async def raise_err(): + raise ValueError("oops") + + monkeypatch.setattr(ReplayHandler, "replay", raise_err) + cp = ClientPlayback() + with taddons.context(cp) as tctx: + cp.running() + cp.start_replay([tflow.tflow()]) + assert await tctx.master.await_log("Client replay has crashed!", level="error") + assert cp.count() == 0 + + +def test_check(): + cp = ClientPlayback() + f = tflow.tflow(resp=True) + f.live = True + assert "live flow" in cp.check(f) + + f = tflow.tflow(resp=True) + f.intercepted = True + assert "intercepted flow" in cp.check(f) + + f = tflow.tflow(resp=True) + f.request = None + assert "missing request" in cp.check(f) + + f = tflow.tflow(resp=True) + f.request.raw_content = None + assert "missing content" in cp.check(f) + + f = tflow.ttcpflow() + assert "Can only replay HTTP" in cp.check(f) + + +@pytest.mark.asyncio +async def test_start_stop(tdata): + cp = ClientPlayback() + with taddons.context(cp) as tctx: + cp.start_replay([tflow.tflow()]) + assert cp.count() == 1 + + cp.start_replay([tflow.twebsocketflow()]) + assert await tctx.master.await_log("Can only replay HTTP flows.", level="warn") + assert cp.count() == 1 + + cp.stop_replay() + assert cp.count() == 0 + + +def test_load(tdata): + cp = ClientPlayback() + with taddons.context(cp): + cp.load_file(tdata.path("mitmproxy/data/dumpfile-018.bin")) + assert cp.count() == 1 + + with pytest.raises(CommandError): + cp.load_file("/nonexistent") + assert cp.count() == 1 + + +def test_configure(tdata): + cp = ClientPlayback() + with taddons.context(cp) as tctx: + assert cp.count() == 0 + tctx.configure(cp, client_replay=[tdata.path("mitmproxy/data/dumpfile-018.bin")]) + assert cp.count() == 1 + tctx.configure(cp, client_replay=[]) + with pytest.raises(OptionsError): + tctx.configure(cp, client_replay=["nonexistent"])