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

unify stream handling

Browse Source
This commit is contained in:
Maximilian Hils 2014-07-21 14:09:24 +02:00
parent 7398db80db
commit 562ac9e721
2 changed files with 28 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

59
libmproxy/protocol/http.py
View File

@ -293,7 +293,8 @@ class HTTPRequest(HTTPMessage):
raise http.HttpError(400, "Invalid headers")
if include_body:
content = http.read_http_body(rfile, headers, body_size_limit, True)
content = http.read_http_body(rfile, headers, body_size_limit,
method, None, True)
timestamp_end = utils.timestamp()
return HTTPRequest(form_in, method, scheme, host, port, path, httpversion, headers,
@ -634,9 +635,9 @@ class HTTPResponse(HTTPMessage):
def _assemble_headers(self, preserve_transfer_encoding=False):
headers = self.headers.copy()
utils.del_all(headers,['Proxy-Connection'])
utils.del_all(headers, ['Proxy-Connection'])
if not preserve_transfer_encoding:
utils.del_all(headers,['Transfer-Encoding'])
utils.del_all(headers, ['Transfer-Encoding'])
if self.content:
headers["Content-Length"] = [str(len(self.content))]
@ -646,7 +647,8 @@ class HTTPResponse(HTTPMessage):
return str(headers)
def _assemble_head(self, preserve_transfer_encoding=False):
return '%s\r\n%s\r\n' % (self._assemble_first_line(), self._assemble_headers(preserve_transfer_encoding=preserve_transfer_encoding))
return '%s\r\n%s\r\n' % (
self._assemble_first_line(), self._assemble_headers(preserve_transfer_encoding=preserve_transfer_encoding))
def _assemble(self):
"""
@ -862,7 +864,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
pass
self.c.close = True
def get_response_from_server(self, request, stream=False):
def get_response_from_server(self, request, include_body=True):
self.c.establish_server_connection()
request_raw = request._assemble()
@ -870,7 +872,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
try:
self.c.server_conn.send(request_raw)
res = HTTPResponse.from_stream(self.c.server_conn.rfile, request.method,
body_size_limit=self.c.config.body_size_limit, include_body=(not stream))
body_size_limit=self.c.config.body_size_limit, include_body=include_body)
return res
except (tcp.NetLibDisconnect, http.HttpErrorConnClosed), v:
self.c.log("error in server communication: %s" % str(v), level="debug")
@ -915,7 +917,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
else:
# read initially in "stream" mode, so we can get the headers separately
flow.response = self.get_response_from_server(flow.request, stream=True)
flow.response = self.get_response_from_server(flow.request, include_body=False)
flow.response.stream = False
# call the appropriate script hook - this is an opportunity for an inline script to set flow.stream = True
@ -923,7 +925,9 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
# now get the rest of the request body, if body still needs to be read but not streaming this response
if not flow.response.stream and flow.response.content is None:
flow.response.content = http.read_http_body(self.c.server_conn.rfile, flow.response.headers, self.c.config.body_size_limit, False)
flow.response.content = http.read_http_body(self.c.server_conn.rfile, flow.response.headers,
self.c.config.body_size_limit,
flow.request.method, flow.response.code, False)
flow.server_conn = self.c.server_conn # no further manipulation of self.c.server_conn beyond this point
# we can safely set it as the final attribute value here.
@ -933,8 +937,6 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
if response_reply is None or response_reply == KILL:
return False
disconnected_while_streaming = False
if flow.response.content is not None:
# if not streaming or there is no body to be read, we'll already have the body, just send it
self.c.client_conn.send(flow.response._assemble())
@ -946,38 +948,19 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
h = flow.response._assemble_head(preserve_transfer_encoding=True)
self.c.client_conn.send(h)
# if chunked then we send back each chunk
if http.has_chunked_encoding(flow.response.headers):
while 1:
content = http.read_next_chunk(self.c.server_conn.rfile, flow.response.headers, False)
if not http.write_chunk(self.c.client_conn.wfile, content):
break
for chunk in http.read_http_body_chunked(self.c.server_conn.rfile,
flow.response.headers,
self.c.config.body_size_limit, "GET", 200, False, 4096):
for part in chunk:
self.c.client_conn.wfile.write(part)
self.c.client_conn.wfile.flush()
self.c.client_conn.wfile.flush()
else: # not chunked, we send back 4k at a time
clen = http.expected_http_body_size(flow.response.headers, False)
clen = clen if clen >= 0 else (64 * 1024 * 1024 * 1024) # arbitrary max of 64G if no length set
rcount = 0
blocksize = 4096
while 1:
bytes_to_read = min(blocksize, clen - rcount)
if bytes_to_read == 0:
break
content = self.c.server_conn.rfile.read(bytes_to_read)
if content == "": # check for EOF
disconnected_while_streaming = True
break
rcount += len(content)
self.c.client_conn.wfile.write(content)
self.c.client_conn.wfile.flush()
if rcount >= clen: # check for having read up to clen
break
flow.timestamp_end = utils.timestamp()
if (disconnected_while_streaming or http.connection_close(flow.request.httpversion, flow.request.headers) or
http.connection_close(flow.response.httpversion, flow.response.headers)):
if (http.connection_close(flow.request.httpversion, flow.request.headers) or
http.connection_close(flow.response.httpversion, flow.response.headers) or
http.expected_http_body_size(flow.response.headers, False, flow.request.method,
flow.response.code) == -1):
return False
if flow.request.form_in == "authority":

9
test/test_server.py
View File

@ -427,16 +427,15 @@ class TestStreamRequest(tservers.HTTPProxTest):
fconn = connection.makefile()
spec = '200:h"Transfer-Encoding"="chunked":r:b"4\\r\\nthis\\r\\n7\\r\\nisatest\\r\\n0\\r\\n\\r\\n"'
connection.send("GET %s/p/%s HTTP/1.1\r\n"%(self.server.urlbase, spec))
connection.send("\r\n");
connection.send("\r\n")
httpversion, code, msg, headers, content = http.read_response(fconn, "GET", 100000, include_body=False)
httpversion, code, msg, headers, content = http.read_response(fconn, "GET", None, include_body=False)
assert headers["Transfer-Encoding"][0] == 'chunked'
assert code == 200
assert http.read_next_chunk(fconn, headers, False) == "this"
assert http.read_next_chunk(fconn, headers, False) == "isatest"
assert http.read_next_chunk(fconn, headers, False) == None
chunks = list(content for _, content, _ in http.read_http_body_chunked(fconn, headers, None, "GET", 200, False))
assert chunks == ["this", "isatest", ""]
connection.close()