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

55
libmproxy/protocol/http.py
View File

@ -293,7 +293,8 @@ class HTTPRequest(HTTPMessage):
raise http.HttpError(400, "Invalid headers") raise http.HttpError(400, "Invalid headers")
if include_body: 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() timestamp_end = utils.timestamp()
return HTTPRequest(form_in, method, scheme, host, port, path, httpversion, headers, return HTTPRequest(form_in, method, scheme, host, port, path, httpversion, headers,
@ -646,7 +647,8 @@ class HTTPResponse(HTTPMessage):
return str(headers) return str(headers)
def _assemble_head(self, preserve_transfer_encoding=False): 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): def _assemble(self):
""" """
@ -862,7 +864,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
pass pass
self.c.close = True 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() self.c.establish_server_connection()
request_raw = request._assemble() request_raw = request._assemble()
@ -870,7 +872,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
try: try:
self.c.server_conn.send(request_raw) self.c.server_conn.send(request_raw)
res = HTTPResponse.from_stream(self.c.server_conn.rfile, request.method, 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 return res
except (tcp.NetLibDisconnect, http.HttpErrorConnClosed), v: except (tcp.NetLibDisconnect, http.HttpErrorConnClosed), v:
self.c.log("error in server communication: %s" % str(v), level="debug") self.c.log("error in server communication: %s" % str(v), level="debug")
@ -915,7 +917,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
else: else:
# read initially in "stream" mode, so we can get the headers separately # 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 flow.response.stream = False
# call the appropriate script hook - this is an opportunity for an inline script to set flow.stream = True # 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 # 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: 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 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. # 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: if response_reply is None or response_reply == KILL:
return False return False
disconnected_while_streaming = False
if flow.response.content is not None: 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 # 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()) 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) h = flow.response._assemble_head(preserve_transfer_encoding=True)
self.c.client_conn.send(h) self.c.client_conn.send(h)
# if chunked then we send back each chunk for chunk in http.read_http_body_chunked(self.c.server_conn.rfile,
if http.has_chunked_encoding(flow.response.headers): flow.response.headers,
while 1: self.c.config.body_size_limit, "GET", 200, False, 4096):
content = http.read_next_chunk(self.c.server_conn.rfile, flow.response.headers, False) for part in chunk:
if not http.write_chunk(self.c.client_conn.wfile, content): self.c.client_conn.wfile.write(part)
break
self.c.client_conn.wfile.flush() 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() flow.timestamp_end = utils.timestamp()
if (disconnected_while_streaming or http.connection_close(flow.request.httpversion, flow.request.headers) or if (http.connection_close(flow.request.httpversion, flow.request.headers) or
http.connection_close(flow.response.httpversion, flow.response.headers)): 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 return False
if flow.request.form_in == "authority": if flow.request.form_in == "authority":

9
test/test_server.py
View File

@ -427,16 +427,15 @@ class TestStreamRequest(tservers.HTTPProxTest):
fconn = connection.makefile() fconn = connection.makefile()
spec = '200:h"Transfer-Encoding"="chunked":r:b"4\\r\\nthis\\r\\n7\\r\\nisatest\\r\\n0\\r\\n\\r\\n"' 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("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 headers["Transfer-Encoding"][0] == 'chunked'
assert code == 200 assert code == 200
assert http.read_next_chunk(fconn, headers, False) == "this" chunks = list(content for _, content, _ in http.read_http_body_chunked(fconn, headers, None, "GET", 200, False))
assert http.read_next_chunk(fconn, headers, False) == "isatest" assert chunks == ["this", "isatest", ""]
assert http.read_next_chunk(fconn, headers, False) == None
connection.close() connection.close()