Merge pull request #3044 from mhils/issue-3038

fix http retry timeout

2 files changed, 27 insertions, 52 deletions

diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py
index 99286fa5..24b411e8 100644
--- a/mitmproxy/proxy/protocol/http.py
+++ b/mitmproxy/proxy/protocol/http.py
@@ -333,8 +333,20 @@ class HttpLayer(base.Layer):
                     f.request.scheme
                 )
 
-                try:
+                def get_response():
                     self.send_request_headers(f.request)
+                    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()
+
+                try:
+                    get_response()
                 except exceptions.NetlibException as e:
                     self.log(
                         "server communication error: %s" % repr(e),
@@ -357,22 +369,17 @@ class HttpLayer(base.Layer):
                         raise exceptions.ProtocolException(
                             "First and only attempt to get response via HTTP2 failed."
                         )
+                    elif f.request.stream:
+                        # We may have already consumed some request chunks already,
+                        # so all we can do is signal downstream that upstream closed the connection.
+                        self.send_error_response(408, "Request Timeout")
+                        f.error = flow.Error(repr(e))
+                        self.channel.ask("error", f)
+                        return False
 
                     self.disconnect()
                     self.connect()
-                    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()
+                    get_response()
 
                 # 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 aed4a774..f594fb40 100644
--- a/test/mitmproxy/proxy/test_server.py
+++ b/test/mitmproxy/proxy/test_server.py
@@ -229,28 +229,14 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin):
                 p.request("get:'%s'" % response)
 
     def test_reconnect(self):
-        req = "get:'%s/p/200:b@1'" % self.server.urlbase
+        req = "get:'%s/p/200:b@1:da'" % 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():
-            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
+            assert p.request(req)
+            # Server has disconnected. Mitmproxy should detect this, and reconnect.
+            assert p.request(req)
+            assert p.request(req)
 
     def test_get_connection_switching(self):
         req = "get:'%s/p/200:b@1'"
@@ -1045,22 +1031,6 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
         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].set_addons(RequestKiller([1, 2]))
         self.chain[1].set_addons(RequestKiller([1]))
 
@@ -1075,9 +1045,7 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
             assert len(self.chain[0].tmaster.state.flows) == 1
             assert len(self.chain[1].tmaster.state.flows) == 1
 
-            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\"'")
+            req = p.request("get:'/p/418:b\"content2\"'")
 
             assert req.status_code == 502