diff options
| -rw-r--r-- | mitmproxy/proxy/protocol/http.py | 28 | ||||
| -rw-r--r-- | test/mitmproxy/proxy/test_server.py | 47 |
2 files changed, 55 insertions, 20 deletions
diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py index 93865c74..502280c1 100644 --- a/mitmproxy/proxy/protocol/http.py +++ b/mitmproxy/proxy/protocol/http.py @@ -329,20 +329,8 @@ class HttpLayer(base.Layer): f.request.scheme ) - def get_response(): - if f.request.stream: - self.send_request_headers(f.request) - chunks = self.read_request_body(f.request) - if callable(f.request.stream): - chunks = f.request.stream(chunks) - self.send_request_body(f.request, chunks) - else: - self.send_request(f.request) - - f.response = self.read_response_headers() - try: - get_response() + self.send_request_headers(f.request) except exceptions.NetlibException as e: self.log( "server communication error: %s" % repr(e), @@ -368,7 +356,19 @@ class HttpLayer(base.Layer): self.disconnect() self.connect() - get_response() + self.send_request_headers(f.request) + + # This is taken out of the try except block because when streaming + # we can't send the request body while retrying as the generator gets exhausted + if f.request.stream: + chunks = self.read_request_body(f.request) + if callable(f.request.stream): + chunks = f.request.stream(chunks) + self.send_request_body(f.request, chunks) + else: + self.send_request_body(f.request, [f.request.data.content]) + + f.response = self.read_response_headers() # call the appropriate script hook - this is an opportunity for # an inline script to set f.stream = True diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index bd61f600..4cae756a 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -239,13 +239,28 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): p.request("get:'%s'" % response) def test_reconnect(self): - req = "get:'%s/p/200:b@1:da'" % self.server.urlbase + req = "get:'%s/p/200:b@1'" % self.server.urlbase p = self.pathoc() + + class MockOnce: + call = 0 + + def mock_once(self, http1obj, req): + self.call += 1 + if self.call == 1: + raise exceptions.TcpDisconnect + else: + headers = http1.assemble_request_head(req) + http1obj.server_conn.wfile.write(headers) + http1obj.server_conn.wfile.flush() + with p.connect(): - assert p.request(req) - # Server has disconnected. Mitmproxy should detect this, and reconnect. - assert p.request(req) - assert p.request(req) + with mock.patch("mitmproxy.proxy.protocol.http1.Http1Layer.send_request_headers", + side_effect=MockOnce().mock_once, autospec=True): + # Server disconnects while sending headers but mitmproxy reconnects + resp = p.request(req) + assert resp + assert resp.status_code == 200 def test_get_connection_switching(self): req = "get:'%s/p/200:b@1'" @@ -1072,6 +1087,23 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest): proxified to an upstream http proxy, we need to send the CONNECT request again. """ + + class MockOnce: + call = 0 + + def mock_once(self, http1obj, req): + self.call += 1 + + if self.call == 2: + headers = http1.assemble_request_head(req) + http1obj.server_conn.wfile.write(headers) + http1obj.server_conn.wfile.flush() + raise exceptions.TcpDisconnect + else: + headers = http1.assemble_request_head(req) + http1obj.server_conn.wfile.write(headers) + http1obj.server_conn.wfile.flush() + self.chain[0].tmaster.addons.add(RequestKiller([1, 2])) self.chain[1].tmaster.addons.add(RequestKiller([1])) @@ -1086,7 +1118,10 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest): assert len(self.chain[0].tmaster.state.flows) == 1 assert len(self.chain[1].tmaster.state.flows) == 1 - req = p.request("get:'/p/418:b\"content2\"'") + with mock.patch("mitmproxy.proxy.protocol.http1.Http1Layer.send_request_headers", + side_effect=MockOnce().mock_once, autospec=True): + req = p.request("get:'/p/418:b\"content2\"'") + assert req.status_code == 502 assert len(self.proxy.tmaster.state.flows) == 2 |