diff options
| -rw-r--r-- | libmproxy/flow.py | 8 | ||||
| -rw-r--r-- | libmproxy/protocol/http.py | 20 | ||||
| -rw-r--r-- | test/scripts/all.py | 6 | ||||
| -rw-r--r-- | test/test_server.py | 7 |
4 files changed, 15 insertions, 26 deletions
diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 5b99427a..55097756 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -706,11 +706,9 @@ class FlowMaster(controller.Master): self.process_new_request(f) return f - def handle_responseheaders(self, r): - f = self.state.add_response(r) - if f: - self.run_script_hook("responseheaders", f) - r.reply() + def handle_responseheaders(self, f): + self.run_script_hook("responseheaders", f) + f.reply() return f def handle_response(self, r): diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py index 242443ec..711cb06c 100644 --- a/libmproxy/protocol/http.py +++ b/libmproxy/protocol/http.py @@ -870,7 +870,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_content=(not stream)) + body_size_limit=self.c.config.body_size_limit, include_body=(not stream)) return res except (tcp.NetLibDisconnect, http.HttpErrorConnClosed), v: self.c.log("error in server communication: %s" % str(v), level="debug") @@ -890,8 +890,6 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin): def handle_flow(self): flow = HTTPFlow(self.c.client_conn, self.c.server_conn, self.change_server) - flow.stream_expecting_body = False - flow.stream = False try: req = HTTPRequest.from_stream(self.c.client_conn.rfile, body_size_limit=self.c.config.body_size_limit) @@ -917,20 +915,14 @@ 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) - - 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 + flow.response = self.get_response_from_server(flow.request, stream=True) + flow.response.stream = False # 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) - # hm - do we need to do something with responseheaders_reply?? + self.c.channel.ask("responseheaders", flow) # 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.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 - 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 self.c.client_conn.send(flow.response._assemble()) else: diff --git a/test/scripts/all.py b/test/scripts/all.py index 15a5fc02..3acaf694 100644 --- a/test/scripts/all.py +++ b/test/scripts/all.py @@ -7,15 +7,15 @@ def serverconnect(ctx, cc): ctx.log("XSERVERCONNECT") log.append("serverconnect") -def request(ctx, r): +def request(ctx, f): ctx.log("XREQUEST") log.append("request") -def response(ctx, r): +def response(ctx, f): ctx.log("XRESPONSE") log.append("response") -def responseheaders(ctx, r): +def responseheaders(ctx, f): ctx.log("XRESPONSEHEADERS") log.append("responseheaders") diff --git a/test/test_server.py b/test/test_server.py index 795a749f..71f00d96 100644 --- a/test/test_server.py +++ b/test/test_server.py @@ -387,10 +387,9 @@ class MasterStreamRequest(tservers.TestMaster): """ Enables the stream flag on the flow for all requests """ - def handle_responseheaders(self, r): - f = self.state.add_response(r) - f.stream = True - r.reply() + def handle_responseheaders(self, f): + f.response.stream = True + f.reply() return f class TestStreamRequest(tservers.HTTPProxTest): |