From 9e1902be62ba5a6c8d6d1496b3cb18f7deb0e583 Mon Sep 17 00:00:00 2001
From: Ujjwal Verma <ujjwalverma1111@gmail.com>
Date: Thu, 29 Jun 2017 01:07:53 +0530
Subject: [PATCH] fix HTTP retry if sending a request fails once

---
 mitmproxy/proxy/protocol/http.py    | 28 ++++++++---------
 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 93865c74d..502280c14 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 bd61f6004..4cae756a7 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