fix #328

3 files changed, 50 insertions, 41 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index dd897115..2540435e 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -704,7 +704,7 @@ class FlowMaster(controller.Master):
         return f
 
     def handle_request(self, r):
-        if r.flow.client_conn and r.flow.client_conn.wfile:
+        if r.flow.live:
             app = self.apps.get(r)
             if app:
                 err = app.serve(r, r.flow.client_conn.wfile, **{"mitmproxy.master": self})
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index fdc29df1..38a6cb49 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -5,7 +5,7 @@ import threading
 from netlib import http, tcp, http_status
 import netlib.utils
 from netlib.odict import ODict, ODictCaseless
-from .primitives import KILL, ProtocolHandler, TemporaryServerChangeMixin, Flow, Error
+from .primitives import KILL, ProtocolHandler, LiveConnection, Flow, Error
 from ..proxy.connection import ServerConnection
 from .. import encoding, utils, controller, stateobject, proxy
 
@@ -543,8 +543,8 @@ class HTTPRequest(HTTPMessage):
         self.path = path
 
         if host != self.get_host() or port != self.get_port():
-            if self.flow.change_server:
-                self.flow.change_server((host, port), ssl=is_ssl)
+            if self.flow.live:
+                self.flow.live.change_server((host, port), ssl=is_ssl)
             else:
                 # There's not live server connection, we're just changing the attributes here.
                 self.flow.server_conn = ServerConnection((host, port),
@@ -789,15 +789,15 @@ class HTTPFlow(Flow):
     The following additional attributes are exposed:
 
         intercepting: Is this flow currently being intercepted?
+        live: Does this flow have a live client connection?
     """
 
-    def __init__(self, client_conn, server_conn, change_server=None):
-        super(HTTPFlow, self).__init__("http", client_conn, server_conn)
+    def __init__(self, client_conn, server_conn, live=None):
+        super(HTTPFlow, self).__init__("http", client_conn, server_conn, live)
         self.request = None
         """@type: HTTPRequest"""
         self.response = None
         """@type: HTTPResponse"""
-        self.change_server = change_server  # Used by flow.request.set_url to change the server address
 
         self.intercepting = False  # FIXME: Should that rather be an attribute of Flow?
 
@@ -904,7 +904,7 @@ class HttpAuthenticationError(Exception):
         return "Proxy Authentication Required"
 
 
-class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
+class HTTPHandler(ProtocolHandler):
     def __init__(self, c):
         super(HTTPHandler, self).__init__(c)
         self.expected_form_in = c.config.http_form_in
@@ -943,7 +943,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
                     raise v
 
     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.live)
         try:
             req = HTTPRequest.from_stream(self.c.client_conn.rfile,
                                           body_size_limit=self.c.config.body_size_limit)
@@ -1038,7 +1038,8 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
 
             # If the user has changed the target server on this connection,
             # restore the original target server
-            self.restore_server()
+            flow.live.restore_server()
+            flow.live = None
 
             return True
         except (HttpAuthenticationError, http.HttpError, proxy.ProxyError, tcp.NetLibError), e:
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index 7b936f7f..a227d904 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -71,12 +71,14 @@ class Error(stateobject.SimpleStateObject):
 
 
 class Flow(stateobject.SimpleStateObject, BackreferenceMixin):
-    def __init__(self, conntype, client_conn, server_conn):
+    def __init__(self, conntype, client_conn, server_conn, live=None):
         self.conntype = conntype
         self.client_conn = client_conn
         """@type: ClientConnection"""
         self.server_conn = server_conn
         """@type: ServerConnection"""
+        self.live = live  # Used by flow.request.set_url to change the server address
+        """@type: LiveConnection"""
 
         self.error = None
         """@type: Error"""
@@ -140,6 +142,8 @@ class ProtocolHandler(object):
     def __init__(self, c):
         self.c = c
         """@type: libmproxy.proxy.server.ConnectionHandler"""
+        self.live = LiveConnection(c)
+        """@type: LiveConnection"""
 
     def handle_messages(self):
         """
@@ -164,46 +168,50 @@ class ProtocolHandler(object):
         raise error  # pragma: nocover
 
 
-class TemporaryServerChangeMixin(object):
+class LiveConnection(object):
     """
-    This mixin allows safe modification of the target server,
-    without any need to expose the ConnectionHandler to the Flow.
+    This facade allows protocol handlers to interface with a live connection,
+    without requiring the expose the ConnectionHandler.
     """
-    def change_server(self, address, ssl):
-        address = netlib.tcp.Address.wrap(address)
-        if address == self.c.server_conn.address():
-            return
-        priority = AddressPriority.MANUALLY_CHANGED
-
-        self.c.log("Temporarily change server connection: %s:%s -> %s:%s" % (
-            self.c.server_conn.address.host,
-            self.c.server_conn.address.port,
-            address.host,
-            address.port
-        ), "debug")
-
-        if not hasattr(self, "_backup_server_conn"):
-            self._backup_server_conn = self.c.server_conn
-            self.c.server_conn = None
-        else:  # This is at least the second temporary change. We can kill the current connection.
-            self.c.del_server_connection()
+    def __init__(self, c):
+        self._c = c
+        """@type: libmproxy.proxy.server.ConnectionHandler"""
 
-        self.c.set_server_address(address, priority)
-        self.c.establish_server_connection(ask=False)
-        if ssl:
-            self.c.establish_ssl(server=True)
+    def change_server(self, address, ssl, persistent_change=False):
+        address = netlib.tcp.Address.wrap(address)
+        if address != self._c.server_conn.address:
+
+            self._c.log("Change server connection: %s:%s -> %s:%s" % (
+                self._c.server_conn.address.host,
+                self._c.server_conn.address.port,
+                address.host,
+                address.port
+            ), "debug")
+
+            if not hasattr(self, "_backup_server_conn"):
+                self._backup_server_conn = self._c.server_conn
+                self._c.server_conn = None
+            else:  # This is at least the second temporary change. We can kill the current connection.
+                self._c.del_server_connection()
+
+            self._c.set_server_address(address, AddressPriority.MANUALLY_CHANGED)
+            self._c.establish_server_connection(ask=False)
+            if ssl:
+                self._c.establish_ssl(server=True)
+        if hasattr(self, "_backup_server_conn") and persistent_change:
+            del self._backup_server_conn
 
     def restore_server(self):
         if not hasattr(self, "_backup_server_conn"):
             return
 
-        self.c.log("Restore original server connection: %s:%s -> %s:%s" % (
-            self.c.server_conn.address.host,
-            self.c.server_conn.address.port,
+        self._c.log("Restore original server connection: %s:%s -> %s:%s" % (
+            self._c.server_conn.address.host,
+            self._c.server_conn.address.port,
             self._backup_server_conn.address.host,
             self._backup_server_conn.address.port
         ), "debug")
 
-        self.c.del_server_connection()
-        self.c.server_conn = self._backup_server_conn
+        self._c.del_server_connection()
+        self._c.server_conn = self._backup_server_conn
         del self._backup_server_conn
\ No newline at end of file