complete state handling

2 files changed, 27 insertions, 6 deletions

diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index b79fc7b3..a33962a6 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -617,7 +617,7 @@ class HTTPResponse(HTTPMessage):
         timestamp_end: Timestamp indicating when request transmission ended
     """
     def __init__(self, httpversion, code, msg, headers, content, timestamp_start, timestamp_end):
-        assert isinstance(headers, ODictCaseless)
+        assert isinstance(headers, ODictCaseless) or headers is None
         HTTPMessage.__init__(self)
 
         self.httpversion = httpversion
@@ -643,7 +643,7 @@ class HTTPResponse(HTTPMessage):
 
     @classmethod
     def _from_state(cls, state):
-        f = cls(None, None, None, None, None, None, None, None)
+        f = cls(None, None, None, None, None, None, None)
         f._load_state(state)
         return f
 
@@ -948,6 +948,19 @@ class HTTPHandler(ProtocolHandler):
 
             if flow.request.form_in == "authority":
                 self.ssl_upgrade(flow.request)
+
+            flow.server_conn = self.c.server_conn
+
+            """
+            FIXME: Remove state test
+            d = flow._get_state()
+            print d
+            flow._load_state(d)
+            print flow._get_state()
+            copy = HTTPFlow._from_state(d)
+            print copy._get_state()
+            """
+
             return True
         except (HttpAuthenticationError, http.HttpError, ProxyError, tcp.NetLibError), e:
             self.handle_error(e, flow)
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index afd1d3fc..6d5dd236 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -36,7 +36,11 @@ class ProxyConfig:
 
 class ClientConnection(tcp.BaseHandler, stateobject.SimpleStateObject):
     def __init__(self, client_connection, address, server):
-        tcp.BaseHandler.__init__(self, client_connection, address, server)
+        if client_connection:  # Eventually, this object is restored from state
+            tcp.BaseHandler.__init__(self, client_connection, address, server)
+        else:
+            self.address = None
+            self.clientcert = None
 
         self.timestamp_start = utils.timestamp()
         self.timestamp_end = None
@@ -52,7 +56,9 @@ class ClientConnection(tcp.BaseHandler, stateobject.SimpleStateObject):
 
     @classmethod
     def _from_state(cls, state):
-        raise NotImplementedError # FIXME
+        f = cls(None, None, None)
+        f._load_state(state)
+        return f
 
     def convert_to_ssl(self, *args, **kwargs):
         tcp.BaseHandler.convert_to_ssl(self, *args, **kwargs)
@@ -86,7 +92,9 @@ class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
 
     @classmethod
     def _from_state(cls, state):
-        raise NotImplementedError  # FIXME
+        f = cls(None)
+        f._load_state(state)
+        return f
 
     def connect(self):
         self.timestamp_start = utils.timestamp()
@@ -163,7 +171,7 @@ class ConnectionHandler:
     def del_server_connection(self):
         if self.server_conn and self.server_conn.connection:
             self.server_conn.finish()
-            self.log("serverdisconnect", ["%s:%s" % self.server_conn.address])
+            self.log("serverdisconnect", ["%s:%s" % (self.server_conn.address.host, self.server_conn.address.port)])
             self.channel.tell("serverdisconnect", self)
         self.server_conn = None
         self.sni = None