Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

fix HTTP retry if sending a request fails once

Browse Source
This commit is contained in:
Ujjwal Verma 2017-06-29 01:07:53 +05:30 committed by Thomas Kriechbaumer
parent 389d37ebff
commit 9e1902be62
2 changed files with 55 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
mitmproxy/proxy/protocol/http.py
View File

@ -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

47
test/mitmproxy/proxy/test_server.py
View File

@ -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