Get rid of ReplayConnection - we now have only one ClientConnection class.

3 files changed, 29 insertions, 30 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 9442b0dd..04e46064 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -47,7 +47,7 @@ def format_keyvals(lst, key="key", val="text", space=5, indent=0):
 def format_flow(f, focus, extended=False, padding=3):
     if not f.request and not f.response:
         txt = [
-            ("title", " Connection from %s..."%(f.client_conn.address)),
+            ("title", " Connection from %s..."%(f.client_conn.address[0])),
         ]
     else:
         if extended:
@@ -1263,7 +1263,7 @@ class ConsoleMaster(controller.Master):
         self.refresh_connection(f)
 
     # Handlers
-    def handle_browserconnection(self, r):
+    def handle_clientconnection(self, r):
         f = flow.Flow(r)
         self.state.add_browserconnect(f)
         r.ack()
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 22ed79e4..e88b8f16 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -13,10 +13,6 @@ class RunException(Exception):
         self.errout = errout
 
 
-class ReplayConnection:
-    pass
-
-
 # begin nocover
 class ReplayThread(threading.Thread):
     def __init__(self, flow, masterq):
@@ -110,6 +106,7 @@ class Flow:
             self.response = proxy.Response.from_state(self.request, state["response"])
         if state["error"]:
             self.error = proxy.Error.from_state(state["error"])
+        self.client_conn = self.request.client_conn
 
     @classmethod
     def from_state(klass, state):
@@ -131,7 +128,7 @@ class Flow:
     def backup(self):
         if not self._backup:
             self._backup = [
-                self.client_conn.copy() if self.client_conn else None,
+                self.client_conn.copy(),
                 self.request.copy() if self.request else None,
                 self.response.copy() if self.response else None,
                 self.error.copy() if self.error else None,
@@ -152,7 +149,7 @@ class Flow:
         return False
 
     def is_replay(self):
-        return isinstance(self.client_conn, ReplayConnection)
+        return self.client_conn.is_replay()
 
     def kill(self):
         if self.request and not self.request.acked:
@@ -231,6 +228,8 @@ class State:
         data = bson.loads(js)
         data = [Flow.from_state(i) for i in data["flows"]]
         self.flow_list.extend(data)
+        for i in data:
+            self.flow_map[i.client_conn] = i
 
     def set_limit(self, limit):
         """
@@ -246,7 +245,7 @@ class State:
             return tuple(self.flow_list[:])
 
     def get_client_conn(self, itm):
-        if isinstance(itm, (proxy.ClientConnection, ReplayConnection)):
+        if isinstance(itm, proxy.ClientConnection):
             return itm
         elif hasattr(itm, "client_conn"):
             return itm.client_conn
@@ -284,19 +283,11 @@ class State:
         self.delete_flow(f)
 
     def revert(self, f):
-        """
-            Replaces the matching client_conn object with a ReplayConnection object.
-        """
         conn = self.get_client_conn(f)
-        if conn in self.flow_map:
-            del self.flow_map[conn]
         f.revert()
-        self.flow_map[f.client_conn] = f
 
     def replay(self, f, masterq):
         """
-            Replaces the matching client_conn object with a ReplayConnection object.
-
             Returns None if successful, or error message if not.
         """
         #begin nocover
@@ -305,16 +296,11 @@ class State:
         if f.request:
             f.backup()
             conn = self.get_client_conn(f)
-            if conn in self.flow_map:
-                del self.flow_map[conn]
-            rp = ReplayConnection()
-            f.client_conn = rp
-            f.request.client_conn = rp
+            f.client_conn.set_replay()
             if f.request.content:
                 f.request.headers["content-length"] = [str(len(f.request.content))]
             f.response = None
             f.error = None
-            self.flow_map[rp] = f
             rt = ReplayThread(f, masterq)
             rt.start()
         #end nocover
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 66be9525..938c4d21 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -106,7 +106,7 @@ class Request(controller.Msg):
     @classmethod
     def from_state(klass, state):
         return klass(
-            None,
+            ClientConnection(None),
             state["host"],
             state["port"],
             state["scheme"],
@@ -222,10 +222,23 @@ class Response(controller.Msg):
 
 
 class ClientConnection(controller.Msg):
-    def __init__(self, address, port):
-        self.address, self.port = address, port
+    def __init__(self, address):
+        """
+            address is an (address, port) tuple, or None if this connection has
+            been replayed from within mitmproxy.
+        """
+        self.address = address
         controller.Msg.__init__(self)
 
+    def set_replay(self):
+        self.address = None
+
+    def is_replay(self):
+        if self.address:
+            return False
+        else:
+            return True
+
     def copy(self):
         return copy.copy(self)
 
@@ -350,10 +363,10 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
 
     def handle(self):
         server = None
-        bc = ClientConnection(*self.client_address)
-        bc.send(self.mqueue)
+        cc = ClientConnection(self.client_address)
+        cc.send(self.mqueue)
         try:
-            request = self.read_request(bc)
+            request = self.read_request(cc)
             request = request.send(self.mqueue)
             if request is None:
                 self.finish()
@@ -369,7 +382,7 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
         except IOError:
             pass
         except ProxyError, e:
-            err = Error(bc, e.msg)
+            err = Error(cc, e.msg)
             err.send(self.mqueue)
             self.send_error(e.code, e.msg)
         if server: