Merge pull request #1172 from cortesi/solidcore

First steps to solidifying the core

6 files changed, 104 insertions, 104 deletions

diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py
index f7bf615a..83ad428e 100644
--- a/test/mitmproxy/test_controller.py
+++ b/test/mitmproxy/test_controller.py
@@ -2,7 +2,7 @@ from threading import Thread, Event
 
 from mock import Mock
 
-from mitmproxy.controller import Reply, DummyReply, Channel, ServerThread, ServerMaster, Master
+from mitmproxy import controller
 from six.moves import queue
 
 from mitmproxy.exceptions import Kill
@@ -10,11 +10,15 @@ from mitmproxy.proxy import DummyServer
 from netlib.tutils import raises
 
 
+class TMsg:
+    pass
+
+
 class TestMaster(object):
     def test_simple(self):
-
-        class DummyMaster(Master):
-            def handle_panic(self, _):
+        class DummyMaster(controller.Master):
+            @controller.handler
+            def log(self, _):
                 m.should_exit.set()
 
             def tick(self, timeout):
@@ -23,14 +27,14 @@ class TestMaster(object):
 
         m = DummyMaster()
         assert not m.should_exit.is_set()
-        m.event_queue.put(("panic", 42))
+        msg = TMsg()
+        msg.reply = controller.DummyReply()
+        m.event_queue.put(("log", msg))
         m.run()
         assert m.should_exit.is_set()
 
-
-class TestServerMaster(object):
-    def test_simple(self):
-        m = ServerMaster()
+    def test_server_simple(self):
+        m = controller.Master()
         s = DummyServer(None)
         m.add_server(s)
         m.start()
@@ -42,7 +46,7 @@ class TestServerMaster(object):
 class TestServerThread(object):
     def test_simple(self):
         m = Mock()
-        t = ServerThread(m)
+        t = controller.ServerThread(m)
         t.run()
         assert m.serve_forever.called
 
@@ -50,7 +54,7 @@ class TestServerThread(object):
 class TestChannel(object):
     def test_tell(self):
         q = queue.Queue()
-        channel = Channel(q, Event())
+        channel = controller.Channel(q, Event())
         m = Mock()
         channel.tell("test", m)
         assert q.get() == ("test", m)
@@ -66,21 +70,21 @@ class TestChannel(object):
 
         Thread(target=reply).start()
 
-        channel = Channel(q, Event())
+        channel = controller.Channel(q, Event())
         assert channel.ask("test", Mock()) == 42
 
     def test_ask_shutdown(self):
         q = queue.Queue()
         done = Event()
         done.set()
-        channel = Channel(q, done)
+        channel = controller.Channel(q, done)
         with raises(Kill):
             channel.ask("test", Mock())
 
 
 class TestDummyReply(object):
     def test_simple(self):
-        reply = DummyReply()
+        reply = controller.DummyReply()
         assert not reply.acked
         reply()
         assert reply.acked
@@ -88,18 +92,18 @@ class TestDummyReply(object):
 
 class TestReply(object):
     def test_simple(self):
-        reply = Reply(42)
+        reply = controller.Reply(42)
         assert not reply.acked
         reply("foo")
         assert reply.acked
         assert reply.q.get() == "foo"
 
     def test_default(self):
-        reply = Reply(42)
+        reply = controller.Reply(42)
         reply()
         assert reply.q.get() == 42
 
     def test_reply_none(self):
-        reply = Reply(42)
+        reply = controller.Reply(42)
         reply(None)
         assert reply.q.get() is None
diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py
index ad4cee53..7d625c34 100644
--- a/test/mitmproxy/test_dump.py
+++ b/test/mitmproxy/test_dump.py
@@ -64,14 +64,14 @@ class TestDumpMaster:
         f = tutils.tflow(req=netlib.tutils.treq(content=content))
         l = Log("connect")
         l.reply = mock.MagicMock()
-        m.handle_log(l)
-        m.handle_clientconnect(f.client_conn)
-        m.handle_serverconnect(f.server_conn)
-        m.handle_request(f)
+        m.log(l)
+        m.clientconnect(f.client_conn)
+        m.serverconnect(f.server_conn)
+        m.request(f)
         if not f.error:
             f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=content))
-            f = m.handle_response(f)
-        m.handle_clientdisconnect(f.client_conn)
+            f = m.response(f)
+        m.clientdisconnect(f.client_conn)
         return f
 
     def _dummy_cycle(self, n, filt, content, **options):
@@ -95,8 +95,8 @@ class TestDumpMaster:
         o = dump.Options(flow_detail=1)
         m = dump.DumpMaster(None, o, outfile=cs)
         f = tutils.tflow(err=True)
-        m.handle_request(f)
-        assert m.handle_error(f)
+        m.request(f)
+        assert m.error(f)
         assert "error" in cs.getvalue()
 
     def test_missing_content(self):
@@ -105,10 +105,10 @@ class TestDumpMaster:
         m = dump.DumpMaster(None, o, outfile=cs)
         f = tutils.tflow()
         f.request.content = None
-        m.handle_request(f)
+        m.request(f)
         f.response = HTTPResponse.wrap(netlib.tutils.tresp())
         f.response.content = None
-        m.handle_response(f)
+        m.response(f)
         assert "content missing" in cs.getvalue()
 
     def test_replay(self):
@@ -160,7 +160,7 @@ class TestDumpMaster:
         assert o.verbosity == 2
 
     def test_filter(self):
-        assert not "GET" in self._dummy_cycle(1, "~u foo", "", verbosity=1)
+        assert "GET" not in self._dummy_cycle(1, "~u foo", "", verbosity=1)
 
     def test_app(self):
         o = dump.Options(app=True)
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 5441ea59..c5e39966 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -53,7 +53,7 @@ class TestStickyCookieState:
         assert s.domain_match("www.google.com", ".google.com")
         assert s.domain_match("google.com", ".google.com")
 
-    def test_handle_response(self):
+    def test_response(self):
         c = "SSID=mooo; domain=.google.com, FOO=bar; Domain=.google.com; Path=/; "\
             "Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; "
 
@@ -100,7 +100,7 @@ class TestStickyCookieState:
         assert len(s.jar[googlekey].keys()) == 1
         assert s.jar[googlekey]["somecookie"].items()[0][1] == "newvalue"
 
-    def test_handle_request(self):
+    def test_request(self):
         s, f = self._response("SSID=mooo", "www.google.com")
         assert "cookie" not in f.request.headers
         s.handle_request(f)
@@ -109,7 +109,7 @@ class TestStickyCookieState:
 
 class TestStickyAuthState:
 
-    def test_handle_response(self):
+    def test_response(self):
         s = flow.StickyAuthState(filt.parse(".*"))
         f = tutils.tflow(resp=True)
         f.request.headers["authorization"] = "foo"
@@ -460,25 +460,20 @@ class TestFlow(object):
         fm = flow.FlowMaster(None, s)
         f = tutils.tflow()
         f.intercept(mock.Mock())
-        assert not f.reply.acked
         f.kill(fm)
-        assert f.reply.acked
+        for i in s.view:
+            assert "killed" in str(i.error)
 
     def test_killall(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
 
         f = tutils.tflow()
-        fm.handle_request(f)
+        f.intercept(fm)
 
-        f = tutils.tflow()
-        fm.handle_request(f)
-
-        for i in s.view:
-            assert not i.reply.acked
         s.killall(fm)
         for i in s.view:
-            assert i.reply.acked
+            assert "killed" in str(i.error)
 
     def test_accept_intercept(self):
         f = tutils.tflow()
@@ -803,8 +798,8 @@ class TestFlowMaster:
         fm = flow.FlowMaster(None, s)
         fm.load_script(tutils.test_data.path("scripts/reqerr.py"))
         f = tutils.tflow()
-        fm.handle_clientconnect(f.client_conn)
-        assert fm.handle_request(f)
+        fm.clientconnect(f.client_conn)
+        assert fm.request(f)
 
     def test_script(self):
         s = flow.State()
@@ -812,18 +807,18 @@ class TestFlowMaster:
         fm.load_script(tutils.test_data.path("scripts/all.py"))
         f = tutils.tflow(resp=True)
 
-        fm.handle_clientconnect(f.client_conn)
+        fm.clientconnect(f.client_conn)
         assert fm.scripts[0].ns["log"][-1] == "clientconnect"
-        fm.handle_serverconnect(f.server_conn)
+        fm.serverconnect(f.server_conn)
         assert fm.scripts[0].ns["log"][-1] == "serverconnect"
-        fm.handle_request(f)
+        fm.request(f)
         assert fm.scripts[0].ns["log"][-1] == "request"
-        fm.handle_response(f)
+        fm.response(f)
         assert fm.scripts[0].ns["log"][-1] == "response"
         # load second script
         fm.load_script(tutils.test_data.path("scripts/all.py"))
         assert len(fm.scripts) == 2
-        fm.handle_clientdisconnect(f.server_conn)
+        fm.clientdisconnect(f.server_conn)
         assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
         assert fm.scripts[1].ns["log"][-1] == "clientdisconnect"
 
@@ -833,7 +828,7 @@ class TestFlowMaster:
         fm.load_script(tutils.test_data.path("scripts/all.py"))
 
         f.error = tutils.terr()
-        fm.handle_error(f)
+        fm.error(f)
         assert fm.scripts[0].ns["log"][-1] == "error"
 
     def test_duplicate_flow(self):
@@ -858,21 +853,20 @@ class TestFlowMaster:
         fm.anticache = True
         fm.anticomp = True
         f = tutils.tflow(req=None)
-        fm.handle_clientconnect(f.client_conn)
+        fm.clientconnect(f.client_conn)
         f.request = HTTPRequest.wrap(netlib.tutils.treq())
-        fm.handle_request(f)
+        fm.request(f)
         assert s.flow_count() == 1
 
         f.response = HTTPResponse.wrap(netlib.tutils.tresp())
-        fm.handle_response(f)
-        assert not fm.handle_response(None)
+        fm.response(f)
         assert s.flow_count() == 1
 
-        fm.handle_clientdisconnect(f.client_conn)
+        fm.clientdisconnect(f.client_conn)
 
         f.error = Error("msg")
         f.error.reply = controller.DummyReply()
-        fm.handle_error(f)
+        fm.error(f)
 
         fm.load_script(tutils.test_data.path("scripts/a.py"))
         fm.shutdown()
@@ -901,7 +895,7 @@ class TestFlowMaster:
         assert fm.state.flow_count()
 
         f.error = Error("error")
-        fm.handle_error(f)
+        fm.error(f)
 
     def test_server_playback(self):
         s = flow.State()
@@ -982,12 +976,12 @@ class TestFlowMaster:
         fm.set_stickycookie(".*")
         f = tutils.tflow(resp=True)
         f.response.headers["set-cookie"] = "foo=bar"
-        fm.handle_request(f)
-        fm.handle_response(f)
+        fm.request(f)
+        fm.response(f)
         assert fm.stickycookie_state.jar
         assert not "cookie" in f.request.headers
         f = f.copy()
-        fm.handle_request(f)
+        fm.request(f)
         assert f.request.headers["cookie"] == "foo=bar"
 
     def test_stickyauth(self):
@@ -1002,12 +996,12 @@ class TestFlowMaster:
         fm.set_stickyauth(".*")
         f = tutils.tflow(resp=True)
         f.request.headers["authorization"] = "foo"
-        fm.handle_request(f)
+        fm.request(f)
 
         f = tutils.tflow(resp=True)
         assert fm.stickyauth_state.hosts
         assert not "authorization" in f.request.headers
-        fm.handle_request(f)
+        fm.request(f)
         assert f.request.headers["authorization"] == "foo"
 
     def test_stream(self):
@@ -1023,15 +1017,15 @@ class TestFlowMaster:
             f = tutils.tflow(resp=True)
 
             fm.start_stream(file(p, "ab"), None)
-            fm.handle_request(f)
-            fm.handle_response(f)
+            fm.request(f)
+            fm.response(f)
             fm.stop_stream()
 
             assert r()[0].response
 
             f = tutils.tflow()
             fm.start_stream(file(p, "ab"), None)
-            fm.handle_request(f)
+            fm.request(f)
             fm.shutdown()
 
             assert not r()[1].response
diff --git a/test/mitmproxy/test_script.py b/test/mitmproxy/test_script.py
index f321d15c..dd6f51ae 100644
--- a/test/mitmproxy/test_script.py
+++ b/test/mitmproxy/test_script.py
@@ -7,7 +7,7 @@ def test_duplicate_flow():
     fm = flow.FlowMaster(None, s)
     fm.load_script(tutils.test_data.path("scripts/duplicate_flow.py"))
     f = tutils.tflow()
-    fm.handle_request(f)
+    fm.request(f)
     assert fm.state.flow_count() == 2
     assert not fm.state.view[0].request.is_replay
     assert fm.state.view[1].request.is_replay
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index 0701d52b..bb4949e1 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -12,6 +12,7 @@ from netlib.http import authentication, http1
 from netlib.tutils import raises
 from pathod import pathoc, pathod
 
+from mitmproxy import controller
 from mitmproxy.proxy.config import HostMatcher
 from mitmproxy.exceptions import Kill
 from mitmproxy.models import Error, HTTPResponse, HTTPFlow
@@ -190,8 +191,8 @@ class TcpMixin:
             assert i_cert == i2_cert == n_cert
 
         # Make sure that TCP messages are in the event log.
-        assert any("305" in m for m in self.master.log)
-        assert any("306" in m for m in self.master.log)
+        assert any("305" in m for m in self.master.tlog)
+        assert any("306" in m for m in self.master.tlog)
 
 
 class AppMixin:
@@ -260,7 +261,7 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin):
         p = self.pathoc()
         assert p.request(req % self.server.urlbase)
         assert p.request(req % self.server2.urlbase)
-        assert switched(self.proxy.log)
+        assert switched(self.proxy.tlog)
 
     def test_blank_leading_line(self):
         p = self.pathoc()
@@ -499,7 +500,7 @@ class TestHttps2Http(tservers.ReverseProxyTest):
     def test_sni(self):
         p = self.pathoc(ssl=True, sni="example.com")
         assert p.request("get:'/p/200'").status_code == 200
-        assert all("Error in handle_sni" not in msg for msg in self.proxy.log)
+        assert all("Error in handle_sni" not in msg for msg in self.proxy.tlog)
 
     def test_http(self):
         p = self.pathoc(ssl=False)
@@ -623,7 +624,8 @@ class TestProxySSL(tservers.HTTPProxyTest):
 class MasterRedirectRequest(tservers.TestMaster):
     redirect_port = None  # Set by TestRedirectRequest
 
-    def handle_request(self, f):
+    @controller.handler
+    def request(self, f):
         if f.request.path == "/p/201":
 
             # This part should have no impact, but it should also not cause any exceptions.
@@ -634,12 +636,13 @@ class MasterRedirectRequest(tservers.TestMaster):
 
             # This is the actual redirection.
             f.request.port = self.redirect_port
-        super(MasterRedirectRequest, self).handle_request(f)
+        super(MasterRedirectRequest, self).request(f)
 
-    def handle_response(self, f):
+    @controller.handler
+    def response(self, f):
         f.response.content = str(f.client_conn.address.port)
         f.response.headers["server-conn-id"] = str(f.server_conn.source_address.port)
-        super(MasterRedirectRequest, self).handle_response(f)
+        super(MasterRedirectRequest, self).response(f)
 
 
 class TestRedirectRequest(tservers.HTTPProxyTest):
@@ -689,10 +692,9 @@ class MasterStreamRequest(tservers.TestMaster):
     """
         Enables the stream flag on the flow for all requests
     """
-
-    def handle_responseheaders(self, f):
+    @controller.handler
+    def responseheaders(self, f):
         f.response.stream = True
-        f.reply()
 
 
 class TestStreamRequest(tservers.HTTPProxyTest):
@@ -739,8 +741,8 @@ class TestStreamRequest(tservers.HTTPProxyTest):
 
 
 class MasterFakeResponse(tservers.TestMaster):
-
-    def handle_request(self, f):
+    @controller.handler
+    def request(self, f):
         resp = HTTPResponse.wrap(netlib.tutils.tresp())
         f.reply(resp)
 
@@ -761,13 +763,14 @@ class TestServerConnect(tservers.HTTPProxyTest):
     def test_unnecessary_serverconnect(self):
         """A replayed/fake response with no_upstream_cert should not connect to an upstream server"""
         assert self.pathod("200").status_code == 200
-        for msg in self.proxy.tmaster.log:
+        for msg in self.proxy.tmaster.tlog:
             assert "serverconnect" not in msg
 
 
 class MasterKillRequest(tservers.TestMaster):
 
-    def handle_request(self, f):
+    @controller.handler
+    def request(self, f):
         f.reply(Kill)
 
 
@@ -783,7 +786,8 @@ class TestKillRequest(tservers.HTTPProxyTest):
 
 class MasterKillResponse(tservers.TestMaster):
 
-    def handle_response(self, f):
+    @controller.handler
+    def response(self, f):
         f.reply(Kill)
 
 
@@ -812,7 +816,8 @@ class TestTransparentResolveError(tservers.TransparentProxyTest):
 
 class MasterIncomplete(tservers.TestMaster):
 
-    def handle_request(self, f):
+    @controller.handler
+    def request(self, f):
         resp = HTTPResponse.wrap(netlib.tutils.tresp())
         resp.content = None
         f.reply(resp)
@@ -930,7 +935,9 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
             k = [0]  # variable scope workaround: put into array
             _func = getattr(master, attr)
 
-            def handler(f):
+            @controller.handler
+            def handler(*args):
+                f = args[-1]
                 k[0] += 1
                 if not (k[0] in exclude):
                     f.client_conn.finish()
@@ -940,13 +947,16 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
 
             setattr(master, attr, handler)
 
-        kill_requests(self.chain[1].tmaster, "handle_request",
-                      exclude=[
-                          # fail first request
-                          2,  # allow second request
-        ])
+        kill_requests(
+            self.chain[1].tmaster,
+            "request",
+            exclude = [
+                # fail first request
+                2,  # allow second request
+            ]
+        )
 
-        kill_requests(self.chain[0].tmaster, "handle_request",
+        kill_requests(self.chain[0].tmaster, "request",
                       exclude=[
                           1,  # CONNECT
                           # fail first request
@@ -1004,10 +1014,10 @@ class AddUpstreamCertsToClientChainMixin:
     ssl = True
     servercert = tutils.test_data.path("data/trusted-server.crt")
     ssloptions = pathod.SSLOptions(
-            cn="trusted-cert",
-            certs=[
-                ("trusted-cert", servercert)
-            ]
+        cn="trusted-cert",
+        certs=[
+            ("trusted-cert", servercert)
+        ]
     )
 
     def test_add_upstream_certs_to_client_chain(self):
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index c9d68cfd..24ebb476 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -39,19 +39,11 @@ class TestMaster(flow.FlowMaster):
         self.apps.add(errapp, "errapp", 80)
         self.clear_log()
 
-    def handle_request(self, f):
-        flow.FlowMaster.handle_request(self, f)
-        f.reply()
-
-    def handle_response(self, f):
-        flow.FlowMaster.handle_response(self, f)
-        f.reply()
-
     def clear_log(self):
-        self.log = []
+        self.tlog = []
 
     def add_event(self, message, level=None):
-        self.log.append(message)
+        self.tlog.append(message)
 
 
 class ProxyThread(threading.Thread):
@@ -68,8 +60,8 @@ class ProxyThread(threading.Thread):
         return self.tmaster.server.address.port
 
     @property
-    def log(self):
-        return self.tmaster.log
+    def tlog(self):
+        return self.tmaster.tlog
 
     def run(self):
         self.tmaster.run()