From 8f2cee722592ac100a4b159d2945044331afcfbc Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Fri, 15 Nov 2019 21:00:47 +0100 Subject: fix #3469 --- mitmproxy/addons/clientplayback.py | 20 +++++++------- mitmproxy/addons/serverplayback.py | 55 +++++++++++++++++++++++++------------- 2 files changed, 46 insertions(+), 29 deletions(-) diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py index 7bdaeb33..7adefd7a 100644 --- a/mitmproxy/addons/clientplayback.py +++ b/mitmproxy/addons/clientplayback.py @@ -1,23 +1,23 @@ import queue import threading -import typing import time +import typing -from mitmproxy import log +import mitmproxy.types +from mitmproxy import command +from mitmproxy import connections from mitmproxy import controller +from mitmproxy import ctx from mitmproxy import exceptions -from mitmproxy import http from mitmproxy import flow +from mitmproxy import http +from mitmproxy import io +from mitmproxy import log from mitmproxy import options -from mitmproxy import connections +from mitmproxy.coretypes import basethread from mitmproxy.net import server_spec, tls from mitmproxy.net.http import http1 -from mitmproxy.coretypes import basethread from mitmproxy.utils import human -from mitmproxy import ctx -from mitmproxy import io -from mitmproxy import command -import mitmproxy.types class RequestReplayThread(basethread.BaseThread): @@ -117,7 +117,7 @@ class RequestReplayThread(basethread.BaseThread): finally: r.first_line_format = first_line_format_backup f.live = False - if server.connected(): + if server and server.connected(): server.finish() server.close() diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index 0818696f..f2e381b2 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -1,16 +1,19 @@ import hashlib -import urllib import typing +import urllib -from mitmproxy import ctx -from mitmproxy import flow +import mitmproxy.types +from mitmproxy import command +from mitmproxy import ctx, http from mitmproxy import exceptions +from mitmproxy import flow from mitmproxy import io -from mitmproxy import command -import mitmproxy.types class ServerPlayback: + flowmap: typing.Dict[typing.Hashable, typing.List[http.HTTPFlow]] + configured: bool + def __init__(self): self.flowmap = {} self.configured = False @@ -82,10 +85,10 @@ class ServerPlayback: Replay server responses from flows. """ self.flowmap = {} - for i in flows: - if i.response: # type: ignore - l = self.flowmap.setdefault(self._hash(i), []) - l.append(i) + for f in flows: + if isinstance(f, http.HTTPFlow): + lst = self.flowmap.setdefault(self._hash(f), []) + lst.append(f) ctx.master.addons.trigger("update", []) @command.command("replay.server.file") @@ -108,12 +111,11 @@ class ServerPlayback: def count(self) -> int: return sum([len(i) for i in self.flowmap.values()]) - def _hash(self, flow): + def _hash(self, flow: http.HTTPFlow) -> typing.Hashable: """ Calculates a loose hash of the flow request. """ r = flow.request - _, _, path, _, query, _ = urllib.parse.urlparse(r.url) queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True) @@ -158,20 +160,33 @@ class ServerPlayback: repr(key).encode("utf8", "surrogateescape") ).digest() - def next_flow(self, request): + def next_flow(self, flow: http.HTTPFlow) -> typing.Optional[http.HTTPFlow]: """ Returns the next flow object, or None if no matching flow was found. """ - hsh = self._hash(request) - if hsh in self.flowmap: + request = flow.request + hash = self._hash(flow) + if hash in self.flowmap: if ctx.options.server_replay_nopop: - return self.flowmap[hsh][0] + return next(( + flow + for flow in self.flowmap[hash] + if flow.response + ), None) else: - ret = self.flowmap[hsh].pop(0) - if not self.flowmap[hsh]: - del self.flowmap[hsh] + ret = self.flowmap[hash].pop(0) + while not ret.response: + if self.flowmap[hash]: + ret = self.flowmap[hash].pop(0) + else: + del self.flowmap[hash] + return None + if not self.flowmap[hash]: + del self.flowmap[hash] return ret + else: + return None def configure(self, updated): if not self.configured and ctx.options.server_replay: @@ -182,10 +197,11 @@ class ServerPlayback: raise exceptions.OptionsError(str(e)) self.load_flows(flows) - def request(self, f): + def request(self, f: http.HTTPFlow) -> None: if self.flowmap: rflow = self.next_flow(f) if rflow: + assert rflow.response response = rflow.response.copy() response.is_replay = True if ctx.options.server_replay_refresh: @@ -197,4 +213,5 @@ class ServerPlayback: f.request.url ) ) + assert f.reply f.reply.kill() -- cgit v1.2.3 From 248034c528c158c95e87f5a888176ab14e11a6dc Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Fri, 15 Nov 2019 21:17:29 +0100 Subject: tests++ --- mitmproxy/addons/serverplayback.py | 1 - test/mitmproxy/addons/test_serverplayback.py | 34 +++++++++++++++++++++++----- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index f2e381b2..0e1c2309 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -165,7 +165,6 @@ class ServerPlayback: Returns the next flow object, or None if no matching flow was found. """ - request = flow.request hash = self._hash(flow) if hash in self.flowmap: if ctx.options.server_replay_nopop: diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py index c6a0c1f4..2e42fa03 100644 --- a/test/mitmproxy/addons/test_serverplayback.py +++ b/test/mitmproxy/addons/test_serverplayback.py @@ -1,13 +1,13 @@ import urllib -import pytest -from mitmproxy.test import taddons -from mitmproxy.test import tflow +import pytest import mitmproxy.test.tutils -from mitmproxy.addons import serverplayback from mitmproxy import exceptions from mitmproxy import io +from mitmproxy.addons import serverplayback +from mitmproxy.test import taddons +from mitmproxy.test import tflow def tdump(path, flows): @@ -321,7 +321,7 @@ def test_server_playback_full(): with taddons.context(s) as tctx: tctx.configure( s, - server_replay_refresh = True, + server_replay_refresh=True, ) f = tflow.tflow() @@ -345,7 +345,7 @@ def test_server_playback_kill(): with taddons.context(s) as tctx: tctx.configure( s, - server_replay_refresh = True, + server_replay_refresh=True, server_replay_kill_extra=True ) @@ -357,3 +357,25 @@ def test_server_playback_kill(): f.request.host = "nonexistent" tctx.cycle(s, f) assert f.reply.value == exceptions.Kill + + +def test_server_playback_response_deleted(): + """ + The server playback addon holds references to flows that can be modified by the user in the meantime. + One thing that can happen is that users remove the response object. This happens for example when doing a client + replay at the same time. + """ + sp = serverplayback.ServerPlayback() + with taddons.context(sp) as tctx: + tctx.configure(sp) + f1 = tflow.tflow(resp=True) + f2 = tflow.tflow(resp=True) + + assert not sp.flowmap + + sp.load_flows([f1, f2]) + assert sp.flowmap + + f1.response = f2.response = None + assert not sp.next_flow(f1) + assert not sp.flowmap -- cgit v1.2.3