streamline HTTPHandler.handle_flow()

1 files changed, 96 insertions, 95 deletions

diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 9593c3cb..c67cb471 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -160,13 +160,6 @@ class HTTPMessage(stateobject.SimpleStateObject):
         c += self.headers.replace(pattern, repl, *args, **kwargs)
         return c
 
-    @classmethod
-    def from_stream(cls, rfile, include_body=True, body_size_limit=None):
-        """
-        Parse an HTTP message from a file stream
-        """
-        raise NotImplementedError()  # pragma: nocover
-
     def _assemble_first_line(self):
         """
         Returns the assembled request/response line
@@ -644,7 +637,8 @@ class HTTPResponse(HTTPMessage):
 
         if self.content:
             headers["Content-Length"] = [str(len(self.content))]
-        elif not preserve_transfer_encoding and 'Transfer-Encoding' in self.headers:  # add content-length for chuncked transfer-encoding with no content
+        # add content-length for chuncked transfer-encoding with no content
+        elif not preserve_transfer_encoding and 'Transfer-Encoding' in self.headers:
             headers["Content-Length"] = ["0"]
 
         return str(headers)
@@ -873,19 +867,21 @@ class HTTPHandler(ProtocolHandler):
         while self.handle_flow():
             pass
 
-    def get_response_from_server(self, request, include_body=True):
+    def get_response_from_server(self, flow):
         self.c.establish_server_connection()
-        request_raw = request.assemble()
+        request_raw = flow.request.assemble()
 
-        for i in range(2):
+        for attempt in (0, 1):
             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=include_body)
-                return res
+                # Only get the headers at first...
+                flow.response = HTTPResponse.from_stream(self.c.server_conn.rfile, flow.request.method,
+                                                         body_size_limit=self.c.config.body_size_limit,
+                                                         include_body=False)
+                break
             except (tcp.NetLibDisconnect, http.HttpErrorConnClosed), v:
                 self.c.log("error in server communication: %s" % repr(v), level="debug")
-                if i < 1:
+                if attempt == 0:
                     # In any case, we try to reconnect at least once.
                     # This is necessary because it might be possible that we already initiated an upstream connection
                     # after clientconnect that has already been expired, e.g consider the following event log:
@@ -899,13 +895,24 @@ class HTTPHandler(ProtocolHandler):
                 else:
                     raise
 
+        # call the appropriate script hook - this is an opportunity for an inline script to set flow.stream = True
+        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.response.stream:
+            flow.response.content = CONTENT_MISSING
+        else:
+            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)
+
     def handle_flow(self):
         flow = HTTPFlow(self.c.client_conn, self.c.server_conn, self.live)
         try:
             try:
                 req = HTTPRequest.from_stream(self.c.client_conn.rfile,
                                               body_size_limit=self.c.config.body_size_limit)
-            except tcp.NetLibDisconnect:  # specifically ignore disconnects that happen before/between requests.
+            except tcp.NetLibDisconnect:  # don't throw an error for disconnects that happen before/between requests.
                 return False
             self.c.log("request", "debug", [req._assemble_first_line(req.form_in)])
             ret = self.process_request(flow, req)
@@ -927,20 +934,7 @@ class HTTPHandler(ProtocolHandler):
             if isinstance(request_reply, HTTPResponse):
                 flow.response = request_reply
             else:
-
-                # read initially in "stream" mode, so we can get the headers separately
-                flow.response = self.get_response_from_server(flow.request, include_body=False)
-
-                # call the appropriate script hook - this is an opportunity for an inline script to set flow.stream = True
-                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.response.stream:
-                    flow.response.content = CONTENT_MISSING
-                else:
-                    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)
+                self.get_response_from_server(flow)
 
             # no further manipulation of self.c.server_conn beyond this point
             # we can safely set it as the final attribute value here.
@@ -951,67 +945,34 @@ class HTTPHandler(ProtocolHandler):
             if response_reply is None or response_reply == KILL:
                 return False
 
-            if not flow.response.stream:
-                # no streaming:
-                # we already received the full response from the server and can send it to the client straight away.
-                self.c.client_conn.send(flow.response.assemble())
-            else:
-                # streaming:
-                # First send the body and then transfer the response incrementally:
-                h = flow.response._assemble_head(preserve_transfer_encoding=True)
-                self.c.client_conn.send(h)
-                for chunk in http.read_http_body_chunked(self.c.server_conn.rfile,
-                                                         flow.response.headers,
-                                                         self.c.config.body_size_limit, flow.request.method,
-                                                         flow.response.code, False, 4096):
-                    for part in chunk:
-                        self.c.client_conn.wfile.write(part)
-                    self.c.client_conn.wfile.flush()
-                flow.response.timestamp_end = utils.timestamp()
-
-            flow.timestamp_end = utils.timestamp()
+            self.send_response_to_client(flow)
 
-            close_connection = (
-                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)
-            if close_connection:
-                if flow.request.form_in == "authority" and flow.response.code == 200:
-                    # Workaround for https://github.com/mitmproxy/mitmproxy/issues/313:
-                    # Some proxies (e.g. Charles) send a CONNECT response with HTTP/1.0 and no Content-Length header
-                    pass
-                else:
-                    return False
+            if self.check_close_connection(flow):
+                return False
 
             # We sent a CONNECT request to an upstream proxy.
             if flow.request.form_in == "authority" and flow.response.code == 200:
-                # TODO: Eventually add headers (space/usefulness tradeoff)
-                # Make sure to add state info before the actual upgrade happens.
-                # During the upgrade, we may receive an SNI indication from the client,
+                # TODO: Possibly add headers (memory consumption/usefulness tradeoff)
+                # Make sure to add state info before the actual processing of the CONNECT request happens.
+                # During an SSL upgrade, we may receive an SNI indication from the client,
                 # which resets the upstream connection. If this is the case, we must
                 # already re-issue the CONNECT request at this point.
                 self.c.server_conn.state.append(("http", {"state": "connect",
                                                           "host": flow.request.host,
                                                           "port": flow.request.port}))
-
-                if self.c.check_ignore_address((flow.request.host, flow.request.port)):
-                    self.c.log("Ignore host: %s:%s" % self.c.server_conn.address(), "info")
-                    TCPHandler(self.c).handle_messages()
+                if not self.process_connect_request((flow.request.host, flow.request.port)):
                     return False
-                else:
-                    if flow.request.port in self.c.config.ssl_ports:
-                        self.ssl_upgrade()
-                    self.skip_authentication = True
 
             # If the user has changed the target server on this connection,
             # restore the original target server
             flow.live.restore_server()
-            flow.live = None
 
-            return True
+            return True  # Next flow please.
         except (HttpAuthenticationError, http.HttpError, proxy.ProxyError, tcp.NetLibError), e:
             self.handle_error(e, flow)
+        finally:
+            flow.timestamp_end = utils.timestamp()
+            flow.live = None  # Connection is not live anymore.
         return False
 
     def handle_server_reconnect(self, state):
@@ -1060,16 +1021,6 @@ class HTTPHandler(ProtocolHandler):
         self.c.client_conn.wfile.write(html_content)
         self.c.client_conn.wfile.flush()
 
-    def ssl_upgrade(self):
-        """
-        Upgrade the connection to SSL after an authority (CONNECT) request has been made.
-        """
-        self.c.log("Received CONNECT request. Upgrading to SSL...", "debug")
-        self.expected_form_in = "relative"
-        self.expected_form_out = "relative"
-        self.c.establish_ssl(server=True, client=True)
-        self.c.log("Upgrade to SSL completed.", "debug")
-
     def process_request(self, flow, request):
         """
         @returns:
@@ -1114,16 +1065,7 @@ class HTTPHandler(ProtocolHandler):
                         ('Proxy-agent: %s\r\n' % self.c.server_version) +
                         '\r\n'
                     )
-
-                    if self.c.check_ignore_address(self.c.server_conn.address):
-                        self.c.log("Ignore host: %s:%s" % self.c.server_conn.address(), "info")
-                        TCPHandler(self.c).handle_messages()
-                        return False
-                    else:
-                        if self.c.server_conn.address.port in self.c.config.ssl_ports:
-                            self.ssl_upgrade()
-                        self.skip_authentication = True
-                        return True
+                    return self.process_connect_request(self.c.server_conn.address)
                 else:  # upstream proxy mode
                     return None
             else:
@@ -1140,7 +1082,6 @@ class HTTPHandler(ProtocolHandler):
                     self.c.set_server_address((request.host, request.port))
                     flow.server_conn = self.c.server_conn
 
-
             return None
 
         raise http.HttpError(400, "Invalid HTTP request form (expected: %s, got: %s)" %
@@ -1182,6 +1123,66 @@ class HTTPHandler(ProtocolHandler):
 
         flow.server_conn = self.c.server_conn
 
+    def send_response_to_client(self, flow):
+        if not flow.response.stream:
+            # no streaming:
+            # we already received the full response from the server and can send it to the client straight away.
+            self.c.client_conn.send(flow.response.assemble())
+        else:
+            # streaming:
+            # First send the body and then transfer the response incrementally:
+            h = flow.response._assemble_head(preserve_transfer_encoding=True)
+            self.c.client_conn.send(h)
+            for chunk in http.read_http_body_chunked(self.c.server_conn.rfile,
+                                                     flow.response.headers,
+                                                     self.c.config.body_size_limit, flow.request.method,
+                                                     flow.response.code, False, 4096):
+                for part in chunk:
+                    self.c.client_conn.wfile.write(part)
+                self.c.client_conn.wfile.flush()
+        flow.response.timestamp_end = utils.timestamp()
+
+    def check_close_connection(self, flow):
+        """
+        Checks if the connection should be closed depending on the HTTP semantics. Returns True, if so.
+        """
+        close_connection = (
+            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)
+        if close_connection:
+            if flow.request.form_in == "authority" and flow.response.code == 200:
+                # Workaround for https://github.com/mitmproxy/mitmproxy/issues/313:
+                # Some proxies (e.g. Charles) send a CONNECT response with HTTP/1.0 and no Content-Length header
+                pass
+            else:
+                return True
+        return False
+
+    def process_connect_request(self, address):
+        """
+        Process a CONNECT request.
+        Returns True if the CONNECT request has been processed successfully.
+        Returns False, if the connection should be closed immediately.
+        """
+        address = tcp.Address.wrap(address)
+        if self.c.check_ignore_address(address):
+            self.c.log("Ignore host: %s:%s" % address(), "info")
+            TCPHandler(self.c).handle_messages()
+            return False
+        else:
+            self.expected_form_in = "relative"
+            self.expected_form_out = "relative"
+            self.skip_authentication = True
+
+            if address.port in self.c.config.ssl_ports:
+                self.c.log("Received CONNECT request to SSL port. Upgrading to SSL...", "debug")
+                self.c.establish_ssl(server=True, client=True)
+                self.c.log("Upgrade to SSL completed.", "debug")
+
+            return True
+
     def authenticate(self, request):
         if self.c.config.authenticator:
             if self.c.config.authenticator.authenticate(request.headers):