Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-26 10:16:27 +00:00
Code Issues Projects Releases Wiki Activity

simplify responseheader scripthook

Browse Source
This commit is contained in:
Maximilian Hils 2014-07-20 11:17:53 +02:00
parent 0ba5a2cf19
commit 7398db80db
4 changed files with 15 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
libmproxy/flow.py
View File

@ -706,11 +706,9 @@ class FlowMaster(controller.Master):
self.process_new_request(f) self.process_new_request(f)
return f return f
def handle_responseheaders(self, r): def handle_responseheaders(self, f):
f = self.state.add_response(r) self.run_script_hook("responseheaders", f)
if f: f.reply()
self.run_script_hook("responseheaders", f)
r.reply()
return f return f
def handle_response(self, r): def handle_response(self, r):

20
libmproxy/protocol/http.py
View File

@ -870,7 +870,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_content=(not stream)) body_size_limit=self.c.config.body_size_limit, include_body=(not stream))
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")
@ -890,8 +890,6 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
def handle_flow(self): def handle_flow(self):
flow = HTTPFlow(self.c.client_conn, self.c.server_conn, self.change_server) flow = HTTPFlow(self.c.client_conn, self.c.server_conn, self.change_server)
flow.stream_expecting_body = False
flow.stream = False
try: try:
req = HTTPRequest.from_stream(self.c.client_conn.rfile, req = HTTPRequest.from_stream(self.c.client_conn.rfile,
body_size_limit=self.c.config.body_size_limit) body_size_limit=self.c.config.body_size_limit)
@ -917,20 +915,14 @@ 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, stream=True)
flow.response.stream = False
if flow.response.content == None:
flow.stream_expecting_body = True
flow.response.content = "" # set this to empty string or other things get really confused,
# flow.stream_expecting_body now contains the state info of whether or not
# body still remains to be read
# 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
responseheaders_reply = self.c.channel.ask("responseheaders", flow.response) self.c.channel.ask("responseheaders", flow)
# hm - do we need to do something with responseheaders_reply??
# 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 flow.stream_expecting_body and not flow.stream: 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, 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
@ -943,7 +935,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
disconnected_while_streaming = False disconnected_while_streaming = False
if not flow.stream or not flow.stream_expecting_body: 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())
else: else:

6
test/scripts/all.py
View File

@ -7,15 +7,15 @@ def serverconnect(ctx, cc):
ctx.log("XSERVERCONNECT") ctx.log("XSERVERCONNECT")
log.append("serverconnect") log.append("serverconnect")
def request(ctx, r): def request(ctx, f):
ctx.log("XREQUEST") ctx.log("XREQUEST")
log.append("request") log.append("request")
def response(ctx, r): def response(ctx, f):
ctx.log("XRESPONSE") ctx.log("XRESPONSE")
log.append("response") log.append("response")
def responseheaders(ctx, r): def responseheaders(ctx, f):
ctx.log("XRESPONSEHEADERS") ctx.log("XRESPONSEHEADERS")
log.append("responseheaders") log.append("responseheaders")

7
test/test_server.py
View File

@ -387,10 +387,9 @@ class MasterStreamRequest(tservers.TestMaster):
""" """
Enables the stream flag on the flow for all requests Enables the stream flag on the flow for all requests
""" """
def handle_responseheaders(self, r): def handle_responseheaders(self, f):
f = self.state.add_response(r) f.response.stream = True
f.stream = True f.reply()
r.reply()
return f return f
class TestStreamRequest(tservers.HTTPProxTest): class TestStreamRequest(tservers.HTTPProxTest):