diff options
| -rw-r--r-- | mitmproxy/addons/session.py | 48 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_session.py | 68 |
2 files changed, 77 insertions, 39 deletions
diff --git a/mitmproxy/addons/session.py b/mitmproxy/addons/session.py index f9d3af3f..2e4d2147 100644 --- a/mitmproxy/addons/session.py +++ b/mitmproxy/addons/session.py @@ -241,8 +241,8 @@ class Session: self.started = True self.db_store = SessionDB(ctx.options.session_path) loop = asyncio.get_event_loop() - tasks = (self._writer, self._tweaker) - loop.create_task(asyncio.gather(*(t() for t in tasks))) + loop.create_task(self._writer()) + loop.create_task(self._tweaker()) def configure(self, updated): if "view_order" in updated: @@ -269,27 +269,30 @@ class Session: self._flush_period *= 1.1 self._flush_rate *= 0.9 - def load_view(self, ids=None): - flows = [] - ids_from_store = [] - if ids is None: - ids = [fid for _, fid in self._view] - for fid in ids: - # Flow could be at the same time in database and in hot storage. We want the most updated version. - if fid in self._hot_store: - flows.append(self._hot_store[fid]) - elif fid in self.db_store: - ids_from_store.append(fid) - else: - flows.append(None) - flows += self.db_store.retrieve_flows(ids_from_store) - return flows + def load_view(self): + ids = [fid for _, fid in self._view] + flows = self.load_storage(ids) + return sorted(flows, key=lambda f: self._generate_order(f)) - def load_storage(self): + def load_storage(self, ids=None): flows = [] - flows += self.db_store.retrieve_flows() - for flow in self._hot_store.values(): - flows.append(flow) + ids_from_store = [] + if ids is not None: + for fid in ids: + # A same flow could be at the same time in hot and db storage. We want the most updated version. + if fid in self._hot_store: + flows.append(self._hot_store[fid]) + elif fid in self.db_store: + ids_from_store.append(fid) + else: + flows.append(None) + flows += self.db_store.retrieve_flows(ids_from_store) + else: + for flow in self._hot_store.values(): + flows.append(flow) + for flow in self.db_store.retrieve_flows(): + if flow.id not in self._hot_store: + flows.append(flow) return flows def clear_storage(self): @@ -304,7 +307,7 @@ class Session: ln += 1 return ln + len(self.db_store) - def _generate_order(self, f: http.HTTPFlow) -> typing.Union[str, int, float]: + def _generate_order(self, f: http.HTTPFlow) -> typing.Optional[typing.Union[str, int, float]]: o = self.order if o == "time": return f.request.timestamp_start or 0 @@ -319,6 +322,7 @@ class Session: if f.response and f.response.raw_content: s += len(f.response.raw_content) return s + return None def set_order(self, order: str) -> None: if order not in orders: diff --git a/test/mitmproxy/addons/test_session.py b/test/mitmproxy/addons/test_session.py index 41e8a401..11a41a6a 100644 --- a/test/mitmproxy/addons/test_session.py +++ b/test/mitmproxy/addons/test_session.py @@ -24,9 +24,10 @@ class TestSession: @staticmethod def start_session(fp=None): s = session.Session() - tctx = taddons.context() - tctx.master.addons.add(s) - tctx.options.session_path = None + with taddons.context() as tctx: + tctx.master.addons.add(s) + tctx.options.session_path = None + tctx.options.view_filter = None if fp: s._flush_period = fp s.running() @@ -97,7 +98,10 @@ class TestSession: s.order = "size" assert s._generate_order(tf) == len(tf.request.raw_content) + len(tf.response.raw_content) - def test_simple(self): + s.order = "invalid" + assert not s._generate_order(tf) + + def test_storage_simple(self): s = session.Session() ctx.options = taddons.context() ctx.options.session_path = None @@ -106,8 +110,8 @@ class TestSession: assert s.store_count() == 0 s.request(f) assert s._view == [(1, f.id)] - assert s.load_view([f.id]) == [f] - assert s.load_view(['nonexistent']) == [None] + assert s.load_view() == [f] + assert s.load_storage(['nonexistent']) == [None] s.error(f) s.response(f) @@ -136,14 +140,17 @@ class TestSession: assert len(s._view) == 0 assert s.store_count() == 0 - def test_filter(self): + def test_storage_filter(self): s = self.start_session() s.request(self.tft(method="get")) s.request(self.tft(method="put")) s.request(self.tft(method="get")) s.request(self.tft(method="put")) assert len(s._view) == 4 - s.set_filter("~m get") + with taddons.context() as tctx: + tctx.master.addons.add(s) + tctx.options.view_filter = '~m get' + s.configure({"view_filter"}) assert [f.request.method for f in s.load_view()] == ["GET", "GET"] assert s.store_count() == 4 with pytest.raises(CommandError): @@ -152,19 +159,24 @@ class TestSession: assert len(s._view) == 4 @pytest.mark.asyncio - async def test_flush_withspecials(self): + async def test_storage_flush_with_specials(self): s = self.start_session(fp=0.5) f = self.tft() s.request(f) - await asyncio.sleep(2) + await asyncio.sleep(1) assert len(s._hot_store) == 0 - assert all([lflow.__dict__ == flow.__dict__ for lflow, flow in list(zip(s.load_view(), [f]))]) + f.response = http.HTTPResponse.wrap(tutils.tresp()) + s.response(f) + assert len(s._hot_store) == 1 + assert s.load_storage() == [f] + await asyncio.sleep(1) + assert all([lflow.__dict__ == flow.__dict__ for lflow, flow in list(zip(s.load_storage(), [f]))]) f.server_conn.via = tflow.tserver_conn() s.request(f) await asyncio.sleep(1) assert len(s._hot_store) == 0 - assert all([lflow.__dict__ == flow.__dict__ for lflow, flow in list(zip(s.load_view(), [f]))]) + assert all([lflow.__dict__ == flow.__dict__ for lflow, flow in list(zip(s.load_storage(), [f]))]) flows = [self.tft() for _ in range(500)] s.update(flows) @@ -174,23 +186,23 @@ class TestSession: assert s._flush_period < fp and s._flush_rate > fr @pytest.mark.asyncio - async def test_bodies(self): + async def test_storage_bodies(self): # Need to test for configure # Need to test for set_order s = self.start_session(fp=0.5) f = self.tft() f2 = self.tft(start=1) - f.request.content = b"A"*1001 + f.request.content = b"A" * 1001 s.request(f) s.request(f2) await asyncio.sleep(1.0) content = s.db_store.con.execute( "SELECT type_id, content FROM body WHERE body.flow_id == (?);", [f.id] ).fetchall()[0] - assert content == (1, b"A"*1001) + assert content == (1, b"A" * 1001) assert s.db_store.body_ledger == {f.id} - f.response = http.HTTPResponse.wrap(tutils.tresp(content=b"A"*1001)) - f2.response = http.HTTPResponse.wrap(tutils.tresp(content=b"A"*1001)) + f.response = http.HTTPResponse.wrap(tutils.tresp(content=b"A" * 1001)) + f2.response = http.HTTPResponse.wrap(tutils.tresp(content=b"A" * 1001)) # Content length is wrong for some reason -- quick fix f.response.headers['content-length'] = b"1001" f2.response.headers['content-length'] = b"1001" @@ -207,3 +219,25 @@ class TestSession: assert len(rows) == 1 assert s.db_store.body_ledger == {f.id} assert all([lf.__dict__ == rf.__dict__ for lf, rf in list(zip(s.load_view(), [f, f2]))]) + + @pytest.mark.asyncio + async def test_storage_order(self): + s = self.start_session(fp=0.5) + s.request(self.tft(method="GET", start=4)) + s.request(self.tft(method="PUT", start=2)) + s.request(self.tft(method="GET", start=3)) + s.request(self.tft(method="PUT", start=1)) + assert [i.request.timestamp_start for i in s.load_view()] == [1, 2, 3, 4] + await asyncio.sleep(1.0) + assert [i.request.timestamp_start for i in s.load_view()] == [1, 2, 3, 4] + with taddons.context() as tctx: + tctx.master.addons.add(s) + tctx.options.view_order = "method" + s.configure({"view_order"}) + assert [i.request.method for i in s.load_view()] == ["GET", "GET", "PUT", "PUT"] + + s.set_order("time") + assert [i.request.timestamp_start for i in s.load_view()] == [1, 2, 3, 4] + + with pytest.raises(CommandError): + s.set_order("not_an_order") |