Add client_playback_concurrency option (#4842)

* nowait

* docs, tests, flake8

* we ideally support other values in the future

Co-authored-by: Maximilian Hils <github@maximilianhils.com>
This commit is contained in:
Brad Dixon 2021-10-08 07:08:36 -04:00 committed by GitHub
parent 79896e23b7
commit 9346002e0f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 19 additions and 2 deletions

View File

@ -8,6 +8,7 @@
* Improve error message on TLS version mismatch (@mhils) * Improve error message on TLS version mismatch (@mhils)
* Windows: Switch to Python's default asyncio event loop, which increases the number of sockets * Windows: Switch to Python's default asyncio event loop, which increases the number of sockets
that can be processed simultaneously (@mhils) that can be processed simultaneously (@mhils)
* Add `client_replay_concurrency` option, which allows more than one client replay request to be in-flight at a time. (@rbdixon)
## 28 September 2021: mitmproxy 7.0.4 ## 28 September 2021: mitmproxy 7.0.4

View File

@ -133,6 +133,9 @@ class ClientPlayback:
self.inflight = await self.queue.get() self.inflight = await self.queue.get()
try: try:
h = ReplayHandler(self.inflight, self.options) h = ReplayHandler(self.inflight, self.options)
if ctx.options.client_replay_concurrency == -1:
asyncio_utils.create_task(h.replay(), name="client playback awaiting response")
else:
await h.replay() await h.replay()
except Exception: except Exception:
ctx.log(f"Client replay has crashed!\n{traceback.format_exc()}", "error") ctx.log(f"Client replay has crashed!\n{traceback.format_exc()}", "error")
@ -160,6 +163,10 @@ class ClientPlayback:
"client_replay", typing.Sequence[str], [], "client_replay", typing.Sequence[str], [],
"Replay client requests from a saved file." "Replay client requests from a saved file."
) )
loader.add_option(
"client_replay_concurrency", int, 1,
"Concurrency limit on in-flight client replay requests. Currently the only valid values are 1 and -1 (no limit)."
)
def configure(self, updated): def configure(self, updated):
if "client_replay" in updated and ctx.options.client_replay: if "client_replay" in updated and ctx.options.client_replay:
@ -169,6 +176,10 @@ class ClientPlayback:
raise exceptions.OptionsError(str(e)) raise exceptions.OptionsError(str(e))
self.start_replay(flows) self.start_replay(flows)
if "client_replay_concurrency" in updated:
if ctx.options.client_replay_concurrency not in [-1, 1]:
raise exceptions.OptionsError("Currently the only valid client_replay_concurrency values are -1 and 1.")
@command.command("replay.client.count") @command.command("replay.client.count")
def count(self) -> int: def count(self) -> int:
""" """

View File

@ -22,7 +22,8 @@ async def tcp_server(handle_conn) -> Address:
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["regular", "upstream", "err"]) @pytest.mark.parametrize("mode", ["regular", "upstream", "err"])
async def test_playback(mode): @pytest.mark.parametrize("concurrency", [-1, 1])
async def test_playback(mode, concurrency):
handler_ok = asyncio.Event() handler_ok = asyncio.Event()
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
@ -50,6 +51,7 @@ async def test_playback(mode):
cp = ClientPlayback() cp = ClientPlayback()
ps = Proxyserver() ps = Proxyserver()
with taddons.context(cp, ps) as tctx: with taddons.context(cp, ps) as tctx:
tctx.configure(cp, client_replay_concurrency=concurrency)
async with tcp_server(handler) as addr: async with tcp_server(handler) as addr:
cp.running() cp.running()
@ -140,3 +142,6 @@ def test_configure(tdata):
tctx.configure(cp, client_replay=[]) tctx.configure(cp, client_replay=[])
with pytest.raises(OptionsError): with pytest.raises(OptionsError):
tctx.configure(cp, client_replay=["nonexistent"]) tctx.configure(cp, client_replay=["nonexistent"])
tctx.configure(cp, client_replay_concurrency=-1)
with pytest.raises(OptionsError):
tctx.configure(cp, client_replay_concurrency=-2)