push failing tests down to 43

8 files changed, 228 insertions, 248 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index b4b939c7..b1971469 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -5,137 +5,20 @@
 import hashlib, Cookie, cookielib, copy, re, urlparse, threading
 import time, urllib
 import types
-import tnetstring, filt, script, utils, encoding, proxy
+import tnetstring, filt, script, utils, encoding
 from email.utils import parsedate_tz, formatdate, mktime_tz
 from netlib import odict, http, certutils, wsgi
 from .proxy import ClientConnection, ServerConnection
 import controller, version, protocol, stateobject
 import app
-
-
-HDR_FORM_URLENCODED = "application/x-www-form-urlencoded"
-CONTENT_MISSING = 0
+from .protocol import KILL
+from .protocol.http import HTTPResponse, CONTENT_MISSING
+from .proxy import RequestReplayThread
 
 ODict = odict.ODict
 ODictCaseless = odict.ODictCaseless
 
 
-class BackreferenceMixin(object):
-    """
-    If an attribute from the _backrefattr tuple is set,
-    this mixin sets a reference back on the attribute object.
-    Example:
-        e = Error()
-        f = Flow()
-        f.error = e
-        assert f is e.flow
-    """
-    _backrefattr = tuple()
-
-    def __setattr__(self, key, value):
-        super(BackreferenceMixin, self).__setattr__(key, value)
-        if key in self._backrefattr and value is not None:
-            setattr(value, self._backrefname, self)
-
-
-class Error(stateobject.SimpleStateObject):
-    """
-        An Error.
-
-        This is distinct from an HTTP error response (say, a code 500), which
-        is represented by a normal Response object. This class is responsible
-        for indicating errors that fall outside of normal HTTP communications,
-        like interrupted connections, timeouts, protocol errors.
-
-        Exposes the following attributes:
-
-            flow: Flow object
-            msg: Message describing the error
-            timestamp: Seconds since the epoch
-    """
-    def __init__(self, msg, timestamp=None):
-        """
-        @type msg: str
-        @type timestamp: float
-        """
-        self.msg = msg
-        self.timestamp = timestamp or utils.timestamp()
-
-    _stateobject_attributes = dict(
-        msg=str,
-        timestamp=float
-    )
-
-    @classmethod
-    def _from_state(cls, state):
-        f = cls(None)  # the default implementation assumes an empty constructor. Override accordingly.
-        f._load_state(state)
-        return f
-
-    def copy(self):
-        c = copy.copy(self)
-        return c
-
-
-class Flow(stateobject.SimpleStateObject, BackreferenceMixin):
-    def __init__(self, conntype, client_conn, server_conn):
-        self.conntype = conntype
-        self.client_conn = client_conn
-        self.server_conn = server_conn
-        self.error = None
-
-    _backrefattr = ("error",)
-    _backrefname = "flow"
-
-    _stateobject_attributes = dict(
-        error=Error,
-        client_conn=ClientConnection,
-        server_conn=ServerConnection,
-        conntype=str
-    )
-
-    def _get_state(self):
-        d = super(Flow, self)._get_state()
-        d.update(version=version.IVERSION)
-        return d
-
-    @classmethod
-    def _from_state(cls, state):
-        f = cls(None, None, None)
-        f._load_state(state)
-        return f
-
-    def copy(self):
-        f = copy.copy(self)
-        if self.error:
-            f.error = self.error.copy()
-        return f
-
-    def modified(self):
-        """
-            Has this Flow been modified?
-        """
-        if self._backup:
-            return self._backup != self._get_state()
-        else:
-            return False
-
-    def backup(self, force=False):
-        """
-            Save a backup of this Flow, which can be reverted to using a
-            call to .revert().
-        """
-        if not self._backup:
-            self._backup = self._get_state()
-
-    def revert(self):
-        """
-            Revert to the last backed up state.
-        """
-        if self._backup:
-            self._load_state(self._backup)
-            self._backup = None
-
 
 class AppRegistry:
     def __init__(self):
@@ -660,10 +543,8 @@ class FlowMaster(controller.Master):
             rflow = self.server_playback.next_flow(flow)
             if not rflow:
                 return None
-            # FIXME
-            response = Response._from_state(flow.request, rflow.response._get_state())
-            response._set_replay()
-            flow.response = response
+            response = HTTPResponse._from_state(rflow.response._get_state())
+            response.is_replay = True
             if self.refresh_server_playback:
                 response.refresh()
             flow.request.reply(response)
@@ -742,13 +623,13 @@ class FlowMaster(controller.Master):
         if f.request.content == CONTENT_MISSING:
             return "Can't replay request with missing content..."
         if f.request:
-            f.request._set_replay()
+            f.request.is_replay = True
             if f.request.content:
                 f.request.headers["Content-Length"] = [str(len(f.request.content))]
             f.response = None
             f.error = None
             self.process_new_request(f)
-            rt = proxy.RequestReplayThread(
+            rt = RequestReplayThread(
                     self.server.config,
                     f,
                     self.masterq,
@@ -791,7 +672,7 @@ class FlowMaster(controller.Master):
                 err = app.serve(r, r.wfile, **{"mitmproxy.master": self})
                 if err:
                     self.add_event("Error in wsgi app. %s"%err, "error")
-                r.reply(proxy.KILL)
+                r.reply(KILL)
                 return
         f = self.state.add_request(r)
         self.replacehooks.run(f)
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 8c44461e..be60f374 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -6,7 +6,7 @@ from netlib.odict import ODict, ODictCaseless
 from . import ProtocolHandler, ConnectionTypeChange, KILL
 from .. import encoding, utils, version, filt, controller, stateobject
 from ..proxy import ProxyError, AddressPriority
-from ..flow import Flow, Error
+from .primitives import Flow, Error
 
 
 HDR_FORM_URLENCODED = "application/x-www-form-urlencoded"
@@ -340,7 +340,7 @@ class HTTPRequest(HTTPMessage):
             Raises an Exception if the request cannot be assembled.
         """
         if self.content == CONTENT_MISSING:
-            raise RuntimeError("Cannot assemble flow with CONTENT_MISSING")
+            raise ProxyError(502, "Cannot assemble flow with CONTENT_MISSING")
         head = self._assemble_head(form)
         if self.content:
             return head + self.content
@@ -444,6 +444,8 @@ class HTTPRequest(HTTPMessage):
             If hostheader is True, we use the value specified in the request
             Host header to construct the URL.
         """
+        raise NotImplementedError
+        # FIXME: Take server_conn into account.
         host = None
         if hostheader:
             host = self.headers.get_first("host")
@@ -462,6 +464,8 @@ class HTTPRequest(HTTPMessage):
 
             Returns False if the URL was invalid, True if the request succeeded.
         """
+        raise NotImplementedError
+        # FIXME: Needs to update server_conn as well.
         parts = http.parse_url(url)
         if not parts:
             return False
@@ -595,7 +599,7 @@ class HTTPResponse(HTTPMessage):
             Raises an Exception if the request cannot be assembled.
         """
         if self.content == CONTENT_MISSING:
-            raise RuntimeError("Cannot assemble flow with CONTENT_MISSING")
+            raise ProxyError(502, "Cannot assemble flow with CONTENT_MISSING")
         head = self._assemble_head()
         if self.content:
             return head + self.content
@@ -711,7 +715,7 @@ class HTTPFlow(Flow):
         if self.request:
             f.request = self.request.copy()
         if self.response:
-            f.response = self.request.copy()
+            f.response = self.response.copy()
         return f
 
     def match(self, f):
@@ -795,8 +799,7 @@ class HTTPHandler(ProtocolHandler):
 
         for i in range(2):
             try:
-                self.c.server_conn.wfile.write(request_raw)
-                self.c.server_conn.wfile.flush()
+                self.c.server_conn.send(request_raw)
                 return HTTPResponse.from_stream(self.c.server_conn.rfile, request.method,
                                                 body_size_limit=self.c.config.body_size_limit)
             except (tcp.NetLibDisconnect, http.HttpErrorConnClosed), v:
@@ -821,6 +824,7 @@ class HTTPHandler(ProtocolHandler):
             flow.request = HTTPRequest.from_stream(self.c.client_conn.rfile,
                                                    body_size_limit=self.c.config.body_size_limit)
             self.c.log("request", [flow.request._assemble_first_line(flow.request.form_in)])
+            self.process_request(flow.request)
 
             request_reply = self.c.channel.ask("request" if LEGACY else "httprequest",
                                                flow.request if LEGACY else flow)
@@ -830,7 +834,6 @@ class HTTPHandler(ProtocolHandler):
             if isinstance(request_reply, HTTPResponse):
                 flow.response = request_reply
             else:
-                self.process_request(flow.request)
                 self.c.establish_server_connection()
                 flow.response = self.get_response_from_server(flow.request)
 
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
new file mode 100644
index 00000000..f77e097b
--- /dev/null
+++ b/libmproxy/protocol/primitives.py
@@ -0,0 +1,124 @@
+from .. import stateobject, utils, version
+from ..proxy import ServerConnection, ClientConnection
+import copy
+
+
+class _BackreferenceMixin(object):
+    """
+    If an attribute from the _backrefattr tuple is set,
+    this mixin sets a reference back on the attribute object.
+    Example:
+        e = Error()
+        f = Flow()
+        f.error = e
+        assert f is e.flow
+    """
+    _backrefattr = tuple()
+
+    def __setattr__(self, key, value):
+        super(_BackreferenceMixin, self).__setattr__(key, value)
+        if key in self._backrefattr and value is not None:
+            setattr(value, self._backrefname, self)
+
+
+class Error(stateobject.SimpleStateObject):
+    """
+        An Error.
+
+        This is distinct from an HTTP error response (say, a code 500), which
+        is represented by a normal Response object. This class is responsible
+        for indicating errors that fall outside of normal HTTP communications,
+        like interrupted connections, timeouts, protocol errors.
+
+        Exposes the following attributes:
+
+            flow: Flow object
+            msg: Message describing the error
+            timestamp: Seconds since the epoch
+    """
+    def __init__(self, msg, timestamp=None):
+        """
+        @type msg: str
+        @type timestamp: float
+        """
+        self.msg = msg
+        self.timestamp = timestamp or utils.timestamp()
+
+    _stateobject_attributes = dict(
+        msg=str,
+        timestamp=float
+    )
+
+    @classmethod
+    def _from_state(cls, state):
+        f = cls(None)  # the default implementation assumes an empty constructor. Override accordingly.
+        f._load_state(state)
+        return f
+
+    def copy(self):
+        c = copy.copy(self)
+        return c
+
+
+class Flow(stateobject.SimpleStateObject, _BackreferenceMixin):
+    def __init__(self, conntype, client_conn, server_conn):
+        self.conntype = conntype
+        self.client_conn = client_conn
+        self.server_conn = server_conn
+        self.error = None
+
+    _backrefattr = ("error",)
+    _backrefname = "flow"
+
+    _stateobject_attributes = dict(
+        error=Error,
+        client_conn=ClientConnection,
+        server_conn=ServerConnection,
+        conntype=str
+    )
+
+    def _get_state(self):
+        d = super(Flow, self)._get_state()
+        d.update(version=version.IVERSION)
+        return d
+
+    @classmethod
+    def _from_state(cls, state):
+        f = cls(None, None, None)
+        f._load_state(state)
+        return f
+
+    def copy(self):
+        f = copy.copy(self)
+
+        f.client_conn = self.client_conn.copy()
+        f.server_conn = self.server_conn.copy()
+
+        if self.error:
+            f.error = self.error.copy()
+        return f
+
+    def modified(self):
+        """
+            Has this Flow been modified?
+        """
+        if self._backup:
+            return self._backup != self._get_state()
+        else:
+            return False
+
+    def backup(self, force=False):
+        """
+            Save a backup of this Flow, which can be reverted to using a
+            call to .revert().
+        """
+        if not self._backup:
+            self._backup = self._get_state()
+
+    def revert(self):
+        """
+            Revert to the last backed up state.
+        """
+        if self._backup:
+            self._load_state(self._backup)
+            self._backup = None
\ No newline at end of file
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 4842a81f..feff2259 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -1,9 +1,8 @@
-import os, socket, time, threading
+import os, socket, time, threading, copy
 from OpenSSL import SSL
 from netlib import tcp, http, certutils, http_auth
 import utils, version, platform, controller, stateobject
 
-
 TRANSPARENT_SSL_PORTS = [443, 8443]
 
 
@@ -82,6 +81,9 @@ class ClientConnection(tcp.BaseHandler, stateobject.SimpleStateObject):
         self.address = tcp.Address(**state["address"]) if state["address"] else None
         self.clientcert = certutils.SSLCert.from_pem(state["clientcert"]) if state["clientcert"] else None
 
+    def copy(self):
+        return copy.copy(self)
+
     @classmethod
     def _from_state(cls, state):
         f = cls(None, None, None)
@@ -115,7 +117,9 @@ class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
         timestamp_ssl_setup=float,
         address=tcp.Address,
         source_address=tcp.Address,
-        cert=certutils.SSLCert
+        cert=certutils.SSLCert,
+        ssl_established=bool,
+        sni=str
     )
 
     def _get_state(self):
@@ -141,6 +145,9 @@ class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
         f._load_state(state)
         return f
 
+    def copy(self):
+        return copy.copy(self)
+
     def connect(self):
         self.timestamp_start = utils.timestamp()
         tcp.TCPClient.connect(self)
@@ -167,8 +174,10 @@ class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
         tcp.TCPClient.finish(self)
         self.timestamp_end = utils.timestamp()
 
+from . import protocol
+from .protocol.http import HTTPResponse
+
 
-"""
 class RequestReplayThread(threading.Thread):
     def __init__(self, config, flow, masterq):
         self.config, self.flow, self.channel = config, flow, controller.Channel(masterq)
@@ -177,24 +186,17 @@ class RequestReplayThread(threading.Thread):
     def run(self):
         try:
             r = self.flow.request
-            server = ServerConnection(self.config, r.scheme, r.host, r.port, r.host)
+            server = ServerConnection(self.flow.server_conn.address())
             server.connect()
-            server.send(r)
-            httpversion, code, msg, headers, content = http.read_response(
-                server.rfile, r.method, self.config.body_size_limit
-            )
-            response = flow.Response(
-                self.flow.request, httpversion, code, msg, headers, content, server.cert, 
-                server.rfile.first_byte_timestamp
-            )
-            self.channel.ask("response", response)
+            if self.flow.server_conn.ssl_established:
+                server.establish_ssl(self.config.clientcerts,
+                                     self.flow.server_conn.sni)
+            server.send(r._assemble())
+            self.flow.response = HTTPResponse.from_stream(server.rfile, r.method, body_size_limit=self.config.body_size_limit)
+            self.channel.ask("response", self.flow.response)
         except (ProxyError, http.HttpError, tcp.NetLibError), v:
-            err = flow.Error(str(v))
-            self.channel.ask("error", err)
-"""
-
-
-import protocol
+            self.flow.error = protocol.primitives.Error(str(v))
+            self.channel.ask("error", self.flow.error)
 
 class ConnectionHandler:
     def __init__(self, config, client_connection, client_address, server, channel, server_version):
diff --git a/test/test_filt.py b/test/test_filt.py
index 96fc58f9..452a4505 100644
--- a/test/test_filt.py
+++ b/test/test_filt.py
@@ -1,6 +1,7 @@
 import cStringIO
 from libmproxy import filt, flow
 from libmproxy.protocol import http
+from libmproxy.protocol.primitives import Error
 import tutils
 
 class TestParsing:
@@ -103,7 +104,7 @@ class TestMatching:
 
     def err(self):
         f = self.req()
-        f.error = flow.Error("msg")
+        f.error = Error("msg")
         return f
 
     def q(self, q, o):
diff --git a/test/test_flow.py b/test/test_flow.py
index 3bcb47d8..3e111fc1 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -1,7 +1,8 @@
 import Queue, time, os.path
 from cStringIO import StringIO
 import email.utils
-from libmproxy import filt, flow, controller, utils, tnetstring, proxy
+from libmproxy import filt, protocol, controller, utils, tnetstring, proxy, flow
+from libmproxy.protocol.primitives import Error
 import tutils
 
 
@@ -171,7 +172,10 @@ class TestServerPlaybackState:
 class TestFlow:
     def test_copy(self):
         f = tutils.tflow_full()
+        a0 = f._get_state()
         f2 = f.copy()
+        a = f._get_state()
+        b = f2._get_state()
         assert f == f2
         assert not f is f2
         assert f.request == f2.request
@@ -204,7 +208,6 @@ class TestFlow:
     def test_backup(self):
         f = tutils.tflow()
         f.response = tutils.tresp()
-        f.request = f.response.request
         f.request.content = "foo"
         assert not f.modified()
         f.backup()
@@ -221,18 +224,18 @@ class TestFlow:
         f.revert()
 
     def test_getset_state(self):
-        f = tutils.tflow()
-        f.response = tutils.tresp(f.request)
+        f = tutils.tflow_full()
         state = f._get_state()
-        assert f._get_state() == flow.Flow._from_state(state)._get_state()
+        assert f._get_state() == protocol.http.HTTPFlow._from_state(state)._get_state()
 
         f.response = None
-        f.error = flow.Error("error")
+        f.error = Error("error")
         state = f._get_state()
-        assert f._get_state() == flow.Flow._from_state(state)._get_state()
+        assert f._get_state() == protocol.http.HTTPFlow._from_state(state)._get_state()
 
-        f2 = tutils.tflow()
-        f2.error = flow.Error("e2")
+        f2 = f.copy()
+        assert f == f2
+        f2.error = Error("e2")
         assert not f == f2
         f._load_state(f2._get_state())
         assert f._get_state() == f2._get_state()
@@ -248,7 +251,6 @@ class TestFlow:
         assert f.request.reply.acked
         f.intercept()
         f.response = tutils.tresp()
-        f.request = f.response.request
         f.request.reply()
         assert not f.response.reply.acked
         f.kill(fm)
@@ -278,7 +280,6 @@ class TestFlow:
         f.accept_intercept()
         assert f.request.reply.acked
         f.response = tutils.tresp()
-        f.request = f.response.request
         f.intercept()
         f.request.reply()
         assert not f.response.reply.acked
@@ -369,16 +370,16 @@ class TestState:
         c = flow.State()
         req = tutils.treq()
         f = c.add_request(req)
-        e = flow.Error("message")
+        e = Error("message")
         assert c.add_error(e)
 
-        e = flow.Error("message")
+        e = Error("message")
         assert not c.add_error(e)
 
         c = flow.State()
         req = tutils.treq()
         f = c.add_request(req)
-        e = flow.Error("message")
+        e = Error("message")
         c.set_limit("~e")
         assert not c.view
         assert not c.view
@@ -435,7 +436,7 @@ class TestState:
     def _add_error(self, state):
         req = tutils.treq()
         f = state.add_request(req)
-        f.error = flow.Error("msg")
+        f.error = Error("msg")
 
     def test_clear(self):
         c = flow.State()
@@ -572,7 +573,7 @@ class TestFlowMaster:
         fm = flow.FlowMaster(None, s)
         assert not fm.load_script(tutils.test_data.path("scripts/reqerr.py"))
         req = tutils.treq()
-        fm.handle_clientconnect(req.client_conn)
+        fm.handle_clientconnect(req.flow.client_conn)
         assert fm.handle_request(req)
 
     def test_script(self):
@@ -580,7 +581,7 @@ class TestFlowMaster:
         fm = flow.FlowMaster(None, s)
         assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
         req = tutils.treq()
-        fm.handle_clientconnect(req.client_conn)
+        fm.handle_clientconnect(req.flow.client_conn)
         assert fm.scripts[0].ns["log"][-1] == "clientconnect"
         sc = proxy.ServerConnection((req.host, req.port))
         sc.reply = controller.DummyReply()
@@ -594,7 +595,7 @@ class TestFlowMaster:
         #load second script
         assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
         assert len(fm.scripts) == 2
-        dc = flow.ClientDisconnect(req.client_conn)
+        dc = flow.ClientDisconnect(req.flow.client_conn)
         dc.reply = controller.DummyReply()
         fm.handle_clientdisconnect(dc)
         assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
@@ -606,7 +607,7 @@ class TestFlowMaster:
         assert len(fm.scripts) == 0
 
         assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
-        err = flow.Error("msg")
+        err = Error("msg")
         err.reply = controller.DummyReply()
         fm.handle_error(err)
         assert fm.scripts[0].ns["log"][-1] == "error"
@@ -642,10 +643,9 @@ class TestFlowMaster:
 
         dc = flow.ClientDisconnect(req.flow.client_conn)
         dc.reply = controller.DummyReply()
-        req.client_conn.requestcount = 1
         fm.handle_clientdisconnect(dc)
 
-        err = flow.Error("msg")
+        err = Error("msg")
         err.reply = controller.DummyReply()
         fm.handle_error(err)
 
@@ -666,7 +666,7 @@ class TestFlowMaster:
         fm.tick(q)
         assert fm.state.flow_count()
 
-        err = flow.Error("error")
+        err = Error("error")
         err.reply = controller.DummyReply()
         fm.handle_error(err)
 
@@ -885,28 +885,6 @@ class TestRequest:
         assert not "if-modified-since" in r.headers
         assert not "if-none-match" in r.headers
 
-    def test_getset_state(self):
-        h = flow.ODictCaseless()
-        h["test"] = ["test"]
-        r = tutils.treq()
-        r.headers = h
-        state = r._get_state()
-        assert flow.Request._from_state(state) == r
-
-        r.client_conn = None
-        state = r._get_state()
-        assert flow.Request._from_state(state) == r
-
-        r2 = tutils.treq()
-        r2.headers = h
-        assert not r == r2
-        r._load_state(r2._get_state())
-        assert r == r2
-
-        r2.client_conn = None
-        r._load_state(r2._get_state())
-        assert not r.client_conn
-
     def test_replace(self):
         r = tutils.treq()
         r.path = "path/foo"
@@ -1072,21 +1050,6 @@ class TestResponse:
         c = "MOO=BAR; Expires=Tue, 08-Mar-2011 00:20:38 GMT; Path=foo.com; Secure"
         assert "00:21:38" in r._refresh_cookie(c, 60)
 
-    def test_getset_state(self):
-        h = flow.ODictCaseless()
-        h["test"] = ["test"]
-        c = flow.ClientConnect(("addr", 2222))
-        req = flow.Request(c, (1, 1), "host", 22, "https", "GET", "/", h, "content")
-        resp = flow.Response(req, (1, 1), 200, "msg", h.copy(), "content", None)
-
-        state = resp._get_state()
-        assert flow.Response._from_state(req, state) == resp
-
-        resp2 = flow.Response(req, (1, 1), 220, "foo", h.copy(), "test", None)
-        assert not resp == resp2
-        resp._load_state(resp2._get_state())
-        assert resp == resp2
-
     def test_replace(self):
         r = tutils.tresp()
         r.headers["Foo"] = ["fOo"]
@@ -1184,18 +1147,18 @@ class TestResponse:
         h["Content-Type"] = ["text/plain"]
         resp = tutils.tresp()
         resp.headers = h
-        assert resp.get_content_type()=="text/plain"
+        assert resp.headers.get_first("content-type")=="text/plain"
 
 
 class TestError:
     def test_getset_state(self):
-        e = flow.Error("Error")
+        e = Error("Error")
         state = e._get_state()
-        assert flow.Error._from_state(state) == e
+        assert Error._from_state(state) == e
 
         assert e.copy()
 
-        e2 = flow.Error("bar")
+        e2 = Error("bar")
         assert not e == e2
         e._load_state(e2._get_state())
         assert e == e2
@@ -1205,17 +1168,19 @@ class TestError:
         assert e3 == e
 
 
-class TestClientConnect:
+class TestClientConnection:
     def test_state(self):
-        c = flow.ClientConnect(("a", 22))
-        assert flow.ClientConnect._from_state(c._get_state()) == c
 
-        c2 = flow.ClientConnect(("a", 25))
+        c = tutils.tclient_conn()
+        assert proxy.ClientConnection._from_state(c._get_state()) == c
+
+        c2 = tutils.tclient_conn()
+        c2.address.address = (c2.address.host, 4242)
         assert not c == c2
 
-        c2.requestcount = 99
+        c2.timestamp_start = 42
         c._load_state(c2._get_state())
-        assert c.requestcount == 99
+        assert c.timestamp_start == 42
 
         c3 = c.copy()
         assert c3 == c
diff --git a/test/test_server.py b/test/test_server.py
index d3bf4676..f542062d 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -4,6 +4,7 @@ from netlib import tcp, http_auth, http
 from libpathod import pathoc, pathod
 import tutils, tservers
 from libmproxy import flow, proxy
+from libmproxy.protocol import KILL
 
 """
     Note that the choice of response code in these tests matters more than you
@@ -19,8 +20,8 @@ class CommonMixin:
 
     def test_replay(self):
         assert self.pathod("304").status_code == 304
-        assert len(self.master.state.view) == (2 if self.ssl else 1)
-        l = self.master.state.view[1 if self.ssl else 0]
+        assert len(self.master.state.view) == 1
+        l = self.master.state.view[0]
         assert l.response.code == 304
         l.request.path = "/p/305"
         rt = self.master.replay_request(l, block=True)
@@ -109,7 +110,7 @@ class TestHTTP(tservers.HTTPProxTest, CommonMixin, AppMixin):
     def test_proxy_ioerror(self):
         # Tests a difficult-to-trigger condition, where an IOError is raised
         # within our read loop.
-        with mock.patch("libmproxy.protocol.HTTPRequest.from_stream") as m:
+        with mock.patch("libmproxy.protocol.http.HTTPRequest.from_stream") as m:
             m.side_effect = IOError("error!")
             tutils.raises("server disconnect", self.pathod, "304")
 
@@ -240,10 +241,10 @@ class TestProxy(tservers.HTTPProxTest):
         f = self.pathod("304")
         assert f.status_code == 304
 
-        l = self.master.state.view[0]
-        assert l.request.client_conn.address
-        assert "host" in l.request.headers
-        assert l.response.code == 304
+        f = self.master.state.view[0]
+        assert f.client_conn.address
+        assert "host" in f.request.headers
+        assert f.response.code == 304
 
     def test_response_timestamps(self):
         # test that we notice at least 2 sec delay between timestamps
@@ -285,8 +286,7 @@ class TestProxy(tservers.HTTPProxTest):
         assert request.timestamp_end - request.timestamp_start <= 0.1
 
     def test_request_tcp_setup_timestamp_presence(self):
-        # tests that the first request in a tcp connection has a tcp_setup_timestamp
-        # while others do not
+        # tests that the client_conn a tcp connection has a tcp_setup_timestamp
         connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         connection.connect(("localhost", self.proxy.port))
         connection.send("GET http://localhost:%d/p/304:b@1k HTTP/1.1\r\n"%self.server.port)
@@ -297,18 +297,18 @@ class TestProxy(tservers.HTTPProxTest):
         connection.recv(5000)
         connection.close()
 
-        first_request = self.master.state.view[0].request
-        second_request = self.master.state.view[1].request
-        assert first_request.tcp_setup_timestamp
-        assert first_request.ssl_setup_timestamp == None
-        assert second_request.tcp_setup_timestamp == None
-        assert second_request.ssl_setup_timestamp == None
+        first_flow = self.master.state.view[0]
+        second_flow = self.master.state.view[1]
+        assert first_flow.server_conn.timestamp_tcp_setup
+        assert first_flow.server_conn.timestamp_ssl_setup is None
+        assert second_flow.server_conn.timestamp_tcp_setup
+        assert first_flow.server_conn.timestamp_tcp_setup == second_flow.server_conn.timestamp_tcp_setup
 
     def test_request_ip(self):
         f = self.pathod("200:b@100")
         assert f.status_code == 200
-        request = self.master.state.view[0].request
-        assert request.ip == "127.0.0.1"
+        f = self.master.state.view[0]
+        assert f.server_conn.peername == ("127.0.0.1", self.server.port)
 
 class TestProxySSL(tservers.HTTPProxTest):
     ssl=True
@@ -317,7 +317,7 @@ class TestProxySSL(tservers.HTTPProxTest):
         f = self.pathod("304:b@10k")
         assert f.status_code == 304
         first_request = self.master.state.view[0].request
-        assert first_request.ssl_setup_timestamp
+        assert first_request.flow.server_conn.timestamp_ssl_setup
 
 class MasterFakeResponse(tservers.TestMaster):
     def handle_request(self, m):
@@ -334,7 +334,7 @@ class TestFakeResponse(tservers.HTTPProxTest):
 
 class MasterKillRequest(tservers.TestMaster):
     def handle_request(self, m):
-        m.reply(proxy.KILL)
+        m.reply(KILL)
 
 
 class TestKillRequest(tservers.HTTPProxTest):
@@ -347,7 +347,7 @@ class TestKillRequest(tservers.HTTPProxTest):
 
 class MasterKillResponse(tservers.TestMaster):
     def handle_response(self, m):
-        m.reply(proxy.KILL)
+        m.reply(KILL)
 
 
 class TestKillResponse(tservers.HTTPProxTest):
diff --git a/test/tutils.py b/test/tutils.py
index 78bf5909..0d3b94f4 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -8,6 +8,7 @@ if os.name != "nt":
 from netlib import certutils
 from nose.plugins.skip import SkipTest
 from mock import Mock
+from time import time
 
 def _SkipWindows():
     raise SkipTest("Skipped on Windows.")
@@ -19,17 +20,20 @@ def SkipWindows(fn):
 
 
 def tclient_conn():
-    return proxy.ClientConnection._from_state(dict(
+    c = proxy.ClientConnection._from_state(dict(
         address=dict(address=("address", 22), use_ipv6=True),
         clientcert=None
     ))
+    c.reply = controller.DummyReply()
+    return c
 
 def tserver_conn():
-    return proxy.ServerConnection._from_state(dict(
+    c = proxy.ServerConnection._from_state(dict(
         address=dict(address=("address", 22), use_ipv6=True),
         source_address=dict(address=("address", 22), use_ipv6=True),
         cert=None
     ))
+    c.reply = controller.DummyReply()
 
 
 def treq(conn=None, content="content"):
@@ -58,7 +62,7 @@ def tresp(req=None, content="message"):
         address=dict(address=("address", 22), use_ipv6=True),
         source_address=None,
         cert=cert.to_pem()))
-    f.response = http.HTTPResponse((1, 1), 200, "OK", headers, content, None, None)
+    f.response = http.HTTPResponse((1, 1), 200, "OK", headers, content, time(), time())
     f.response.reply = controller.DummyReply()
     return f.response