From 5da27a9905302a5e43fdf4db8a7b7b784544bed2 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 19 Feb 2011 17:00:24 +1300 Subject: Refactor Flow primitives to remove HTTP1.0 assumption. This is a big patch removing the assumption that there's one connection per Request/Response pair. It touches pretty much every part of mitmproxy, so expect glitches until everything is ironed out. --- test/test_console.py | 32 ++++--------- test/test_dump.py | 2 +- test/test_filt.py | 2 +- test/test_flow.py | 129 ++++++++++++++++++++------------------------------- test/test_proxy.py | 10 ++-- test/utils.py | 6 +-- 6 files changed, 68 insertions(+), 113 deletions(-) (limited to 'test') diff --git a/test/test_console.py b/test/test_console.py index 93312824..6baab4ba 100644 --- a/test/test_console.py +++ b/test/test_console.py @@ -10,11 +10,9 @@ class uState(libpry.AutoTree): connect -> request -> response """ - bc = proxy.ClientConnection(("address", 22)) c = console.ConsoleState() - f = flow.Flow(bc) - c.add_browserconnect(f) - assert c.lookup(bc) + f = self._add_request(c) + assert f.request in c.flow_map assert c.get_focus() == (f, 0) def test_focus(self): @@ -24,18 +22,14 @@ class uState(libpry.AutoTree): connect -> request -> response """ c = console.ConsoleState() + f = self._add_request(c) - bc = proxy.ClientConnection(("address", 22)) - f = flow.Flow(bc) - c.add_browserconnect(f) assert c.get_focus() == (f, 0) assert c.get_from_pos(0) == (f, 0) assert c.get_from_pos(1) == (None, None) assert c.get_next(0) == (None, None) - bc2 = proxy.ClientConnection(("address", 22)) - f2 = flow.Flow(bc2) - c.add_browserconnect(f2) + f2 = self._add_request(c) assert c.get_focus() == (f, 1) assert c.get_next(0) == (f, 1) assert c.get_prev(1) == (f2, 0) @@ -52,25 +46,14 @@ class uState(libpry.AutoTree): assert c.get_focus() == (None, None) def _add_request(self, state): - f = utils.tflow() - state.add_browserconnect(f) - q = utils.treq(f.client_conn) - state.add_request(q) - return f + r = utils.treq() + return state.add_request(r) def _add_response(self, state): f = self._add_request(state) r = utils.tresp(f.request) state.add_response(r) - def test_add_request(self): - c = console.ConsoleState() - f = utils.tflow() - c.add_browserconnect(f) - q = utils.treq(f.client_conn) - c.focus = None - assert c.add_request(q) - def test_add_response(self): c = console.ConsoleState() f = self._add_request(c) @@ -118,11 +101,12 @@ class uformat_flow(libpry.AutoTree): assert ('method', '[edited] ') in console.format_flow(f, True) assert ('method', '[edited] ') in console.format_flow(f, True, True) - f.client_conn = proxy.ClientConnection(None) + f.request.set_replay() assert ('method', '[replay] ') in console.format_flow(f, True) assert ('method', '[replay] ') in console.format_flow(f, True, True) + class uPathCompleter(libpry.AutoTree): def test_lookup_construction(self): c = console._PathCompleter() diff --git a/test/test_dump.py b/test/test_dump.py index 978bf138..7b223645 100644 --- a/test/test_dump.py +++ b/test/test_dump.py @@ -10,7 +10,7 @@ class uDumpMaster(libpry.AutoTree): req = utils.treq() cc = req.client_conn resp = utils.tresp(req) - m.handle_clientconnection(cc) + m.handle_clientconnect(cc) m.handle_request(req) m.handle_response(resp) diff --git a/test/test_filt.py b/test/test_filt.py index 6f8579d3..791b9b39 100644 --- a/test/test_filt.py +++ b/test/test_filt.py @@ -72,7 +72,7 @@ class uParsing(libpry.AutoTree): class uMatching(libpry.AutoTree): def req(self): - conn = proxy.ClientConnection(("one", 2222)) + conn = proxy.ClientConnect(("one", 2222)) headers = utils.Headers() headers["header"] = ["qvalue"] return proxy.Request( diff --git a/test/test_flow.py b/test/test_flow.py index 35d336e8..3998943c 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -46,38 +46,6 @@ class uFlow(libpry.AutoTree): state = f.get_state() assert f == flow.Flow.from_state(state) - def test_simple(self): - f = utils.tflow() - assert console.format_flow(f, True) - assert console.format_flow(f, False) - - f.request = utils.treq() - assert console.format_flow(f, True) - assert console.format_flow(f, False) - - f.response = utils.tresp() - f.response.headers["content-type"] = ["text/html"] - assert console.format_flow(f, True) - assert console.format_flow(f, False) - f.response.code = 404 - assert console.format_flow(f, True) - assert console.format_flow(f, False) - - assert console.format_flow(f, True) - assert console.format_flow(f, False) - - f.client_conn.set_replay() - assert console.format_flow(f, True) - assert console.format_flow(f, False) - - f.response = None - assert console.format_flow(f, True) - assert console.format_flow(f, False) - - f.error = proxy.Error(200, "test") - assert console.format_flow(f, True) - assert console.format_flow(f, False) - def test_kill(self): f = utils.tflow() f.request = utils.treq() @@ -115,10 +83,10 @@ class uFlow(libpry.AutoTree): class uState(libpry.AutoTree): def test_backup(self): - bc = proxy.ClientConnection(("address", 22)) + bc = proxy.ClientConnect(("address", 22)) c = flow.State() - f = flow.Flow(bc) - c.add_browserconnect(f) + req = utils.treq() + f = c.add_request(req) f.backup() c.revert(f) @@ -129,92 +97,98 @@ class uState(libpry.AutoTree): connect -> request -> response """ - bc = proxy.ClientConnection(("address", 22)) + bc = proxy.ClientConnect(("address", 22)) c = flow.State() - f = flow.Flow(bc) - c.add_browserconnect(f) - assert c.lookup(bc) + c.clientconnect(bc) + assert len(c.client_connections) == 1 req = utils.treq(bc) - assert c.add_request(req) + f = c.add_request(req) + assert f assert len(c.flow_list) == 1 - assert c.lookup(req) + assert c.flow_map.get(req) newreq = utils.treq() assert c.add_request(newreq) - assert c.lookup(newreq) + assert c.flow_map.get(newreq) resp = utils.tresp(req) assert c.add_response(resp) assert len(c.flow_list) == 2 - assert c.lookup(resp) + assert c.flow_map.get(resp.request) newresp = utils.tresp() assert not c.add_response(newresp) - assert not c.lookup(newresp) + assert not c.flow_map.get(newresp.request) + + dc = proxy.ClientDisconnect(bc) + c.clientdisconnect(dc) + assert not c.client_connections def test_err(self): - bc = proxy.ClientConnection(("address", 22)) + bc = proxy.ClientConnect(("address", 22)) c = flow.State() - f = flow.Flow(bc) - c.add_browserconnect(f) - e = proxy.Error(bc, "message") + req = utils.treq() + f = c.add_request(req) + e = proxy.Error(f, "message") assert c.add_error(e) - e = proxy.Error(proxy.ClientConnection(("address", 22)), "message") + e = proxy.Error(utils.tflow(), "message") assert not c.add_error(e) def test_view(self): c = flow.State() - f = utils.tflow() - c.add_browserconnect(f) + req = utils.treq() + c.clientconnect(req.client_conn) + assert len(c.view) == 0 + + f = c.add_request(req) assert len(c.view) == 1 - c.set_limit(filt.parse("~q")) + + c.set_limit(filt.parse("~s")) assert len(c.view) == 0 + resp = utils.tresp(req) + c.add_response(resp) + assert len(c.view) == 1 c.set_limit(None) + assert len(c.view) == 1 - - f = utils.tflow() - req = utils.treq(f.client_conn) - c.add_browserconnect(f) + req = utils.treq() + c.clientconnect(req.client_conn) c.add_request(req) assert len(c.view) == 2 c.set_limit(filt.parse("~q")) assert len(c.view) == 1 c.set_limit(filt.parse("~s")) - assert len(c.view) == 0 + assert len(c.view) == 1 def _add_request(self, state): - f = utils.tflow() - state.add_browserconnect(f) - q = utils.treq(f.client_conn) - state.add_request(q) + req = utils.treq() + f = state.add_request(req) return f def _add_response(self, state): - f = self._add_request(state) - r = utils.tresp(f.request) - state.add_response(r) + req = utils.treq() + f = state.add_request(req) + resp = utils.tresp(req) + state.add_response(resp) def _add_error(self, state): - f = utils.tflow() - f.error = proxy.Error(None, "msg") - state.add_browserconnect(f) - q = utils.treq(f.client_conn) - state.add_request(q) + req = utils.treq() + f = state.add_request(req) + f.error = proxy.Error(f, "msg") def test_kill_flow(self): c = flow.State() - f = utils.tflow() - c.add_browserconnect(f) + req = utils.treq() + f = c.add_request(req) c.kill_flow(f) assert not c.flow_list def test_clear(self): c = flow.State() - f = utils.tflow() - c.add_browserconnect(f) + f = self._add_request(c) f.intercepting = True c.clear() @@ -265,15 +239,12 @@ class uFlowMaster(libpry.AutoTree): def test_one(self): s = flow.State() f = flow.FlowMaster(None, s) - req = utils.treq() - f.handle_clientconnection(req.client_conn) - assert len(s.flow_list) == 1 + f.handle_request(req) assert len(s.flow_list) == 1 - f.handle_request(req) - resp = utils.tresp() - resp.request = req + + resp = utils.tresp(req) f.handle_response(resp) assert len(s.flow_list) == 1 diff --git a/test/test_proxy.py b/test/test_proxy.py index 7bb608ae..ba9d9bfa 100644 --- a/test/test_proxy.py +++ b/test/test_proxy.py @@ -213,7 +213,7 @@ class uRequest(libpry.AutoTree): def test_simple(self): h = utils.Headers() h["test"] = ["test"] - c = proxy.ClientConnection(("addr", 2222)) + c = proxy.ClientConnect(("addr", 2222)) r = proxy.Request(c, "host", 22, "https", "GET", "/", h, "content") u = r.url() assert r.set_url(u) @@ -225,17 +225,17 @@ class uRequest(libpry.AutoTree): def test_getset_state(self): h = utils.Headers() h["test"] = ["test"] - c = proxy.ClientConnection(("addr", 2222)) + c = proxy.ClientConnect(("addr", 2222)) r = proxy.Request(c, "host", 22, "https", "GET", "/", h, "content") state = r.get_state() - assert proxy.Request.from_state(c, state) == r + assert proxy.Request.from_state(state) == r class uResponse(libpry.AutoTree): def test_simple(self): h = utils.Headers() h["test"] = ["test"] - c = proxy.ClientConnection(("addr", 2222)) + c = proxy.ClientConnect(("addr", 2222)) req = proxy.Request(c, "host", 22, "https", "GET", "/", h, "content") resp = proxy.Response(req, 200, "msg", h.copy(), "content") assert resp.short() @@ -244,7 +244,7 @@ class uResponse(libpry.AutoTree): def test_getset_state(self): h = utils.Headers() h["test"] = ["test"] - c = proxy.ClientConnection(("addr", 2222)) + c = proxy.ClientConnect(("addr", 2222)) r = proxy.Request(c, "host", 22, "https", "GET", "/", h, "content") req = proxy.Request(c, "host", 22, "https", "GET", "/", h, "content") resp = proxy.Response(req, 200, "msg", h.copy(), "content") diff --git a/test/utils.py b/test/utils.py index b63c48dc..b1dc46d4 100644 --- a/test/utils.py +++ b/test/utils.py @@ -2,7 +2,7 @@ from libmproxy import proxy, utils, filt, flow def treq(conn=None): if not conn: - conn = proxy.ClientConnection(("address", 22)) + conn = proxy.ClientConnect(("address", 22)) headers = utils.Headers() headers["header"] = ["qvalue"] return proxy.Request(conn, "host", 80, "http", "GET", "/path", headers, "content") @@ -17,6 +17,6 @@ def tresp(req=None): def tflow(): - bc = proxy.ClientConnection(("address", 22)) - return flow.Flow(bc) + r = treq() + return flow.Flow(r) -- cgit v1.2.3