Merge pull request #1714 from cortesi/justno2

Test and interface improvements

examples/har_dump.py

@@ -145,7 +145,7 @@ def response(flow):
             "params": params
         }
 
-    if flow.server_conn:
+    if flow.server_conn.connected():
         entry["serverIPAddress"] = str(flow.server_conn.ip_address.address[0])
 
     HAR["log"]["entries"].append(entry)

mitmproxy/addons/dumper.py

@@ -98,7 +98,7 @@ class Dumper:
             self.echo("")
 
     def _echo_request_line(self, flow):
-        if flow.client_conn is not None:
+        if flow.client_conn:
             client = click.style(
                 strutils.escape_control_characters(
                     repr(flow.client_conn.address)

mitmproxy/connections.py

@@ -47,7 +47,7 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
         self.cipher_name = None
         self.tls_version = None
 
-    def __bool__(self):
+    def connected(self):
         return bool(self.connection) and not self.finished
 
     def __repr__(self):
@@ -146,7 +146,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
         self.timestamp_ssl_setup = None
         self.protocol = None
 
-    def __bool__(self):
+    def connected(self):
         return bool(self.connection) and not self.finished
 
     def __repr__(self):

mitmproxy/proxy/modes/http_proxy.py

@@ -8,7 +8,7 @@ class HttpProxy(protocol.Layer, protocol.ServerConnectionMixin):
         try:
             layer()
         finally:
-            if self.server_conn:
+            if self.server_conn.connected():
                 self.disconnect()
 
 
@@ -22,5 +22,5 @@ class HttpUpstreamProxy(protocol.Layer, protocol.ServerConnectionMixin):
         try:
             layer()
         finally:
-            if self.server_conn:
+            if self.server_conn.connected():
                 self.disconnect()

mitmproxy/proxy/modes/reverse_proxy.py

@@ -12,5 +12,5 @@ class ReverseProxy(protocol.Layer, protocol.ServerConnectionMixin):
         try:
             layer()
         finally:
-            if self.server_conn:
+            if self.server_conn.connected():
                 self.disconnect()

mitmproxy/proxy/modes/socks_proxy.py

@@ -56,5 +56,5 @@ class Socks5Proxy(protocol.Layer, protocol.ServerConnectionMixin):
         try:
             layer()
         finally:
-            if self.server_conn:
+            if self.server_conn.connected():
                 self.disconnect()

mitmproxy/proxy/modes/transparent_proxy.py

@@ -19,5 +19,5 @@ class TransparentProxy(protocol.Layer, protocol.ServerConnectionMixin):
         try:
             layer()
         finally:
-            if self.server_conn:
+            if self.server_conn.connected():
                 self.disconnect()

mitmproxy/proxy/protocol/base.py

@@ -101,7 +101,7 @@ class ServerConnectionMixin:
                 try:
                     # Do something.
                 finally:
-                    if self.server_conn:
+                    if self.server_conn.connected():
                         self.disconnect()
     """
 
@@ -139,7 +139,7 @@ class ServerConnectionMixin:
         """
         Sets a new server address. If there is an existing connection, it will be closed.
         """
-        if self.server_conn:
+        if self.server_conn.connected():
             self.disconnect()
         self.log("Set new server address: " + repr(address), "debug")
         self.server_conn.address = address

mitmproxy/proxy/protocol/http.py

@@ -76,8 +76,8 @@ class ConnectServerConnection:
     def __getattr__(self, item):
         return getattr(self.via, item)
 
-    def __bool__(self):
-        return bool(self.via)
+    def connected(self):
+        return self.via.connected()
 
 
 class UpstreamConnectLayer(base.Layer):
@@ -101,7 +101,7 @@ class UpstreamConnectLayer(base.Layer):
             raise exceptions.ProtocolException("Reconnect: Upstream server refuses CONNECT request")
 
     def connect(self):
-        if not self.server_conn:
+        if not self.server_conn.connected():
             self.ctx.connect()
             self._send_connect_request()
         else:
@@ -112,7 +112,7 @@ class UpstreamConnectLayer(base.Layer):
             self.ctx.set_server(address)
 
     def set_server(self, address):
-        if self.ctx.server_conn:
+        if self.ctx.server_conn.connected():
             self.ctx.disconnect()
         address = tcp.Address.wrap(address)
         self.connect_request.host = address.host
@@ -378,10 +378,10 @@ class HttpLayer(base.Layer):
                 self.set_server(address)
                 self.set_server_tls(tls, address.host)
             # Establish connection is neccessary.
-            if not self.server_conn:
+            if not self.server_conn.connected():
                 self.connect()
         else:
-            if not self.server_conn:
+            if not self.server_conn.connected():
                 self.connect()
             if tls:
                 raise exceptions.HttpProtocolException("Cannot change scheme in upstream proxy mode.")

mitmproxy/proxy/protocol/http2.py

@@ -96,7 +96,7 @@ class Http2Layer(base.Layer):
         self.client_conn.h2 = SafeH2Connection(self.client_conn, config=config)
 
     def _initiate_server_conn(self):
-        if self.server_conn:
+        if self.server_conn.connected():
             config = h2.config.H2Configuration(
                 client_side=True,
                 header_encoding=False,

mitmproxy/proxy/protocol/http_replay.py

@@ -115,5 +115,5 @@ class RequestReplayThread(basethread.BaseThread):
         finally:
             r.first_line_format = first_line_format_backup
             self.f.live = False
-            if server:
+            if server.connected():
                 server.finish()

mitmproxy/proxy/protocol/tls.py

@@ -364,7 +364,7 @@ class TlsLayer(base.Layer):
             )
         )
         establish_server_tls_now = (
-            (self.server_conn and self._server_tls) or
+            (self.server_conn.connected() and self._server_tls) or
             client_tls_requires_server_connection
         )
 
@@ -389,7 +389,7 @@ class TlsLayer(base.Layer):
             return "TlsLayer(inactive)"
 
     def connect(self):
-        if not self.server_conn:
+        if not self.server_conn.connected():
             self.ctx.connect()
         if self._server_tls and not self.server_conn.tls_established:
             self._establish_tls_with_server()

test/mitmproxy/data/addonscripts/concurrent_decorator.py

@@ -1,6 +1,7 @@
 import time
 from mitmproxy.script import concurrent
 
+
 @concurrent
 def request(flow):
     time.sleep(0.1)

test/mitmproxy/script/test_concurrent.py

@@ -1,10 +1,9 @@
 from mitmproxy.test import tflow
 from mitmproxy.test import tutils
+from mitmproxy.test import taddons
+
 from mitmproxy import controller
 from mitmproxy.addons import script
-from mitmproxy import options
-from mitmproxy import proxy
-from mitmproxy import master
 
 import time
 
@@ -21,29 +20,29 @@ class Thing:
 class TestConcurrent(mastertest.MasterTest):
     @ttutils.skip_appveyor
     def test_concurrent(self):
-        m = master.Master(options.Options(), proxy.DummyServer())
-        sc = script.Script(
-            tutils.test_data.path(
-                "mitmproxy/data/addonscripts/concurrent_decorator.py"
+        with taddons.context() as tctx:
+            sc = script.Script(
+                tutils.test_data.path(
+                    "mitmproxy/data/addonscripts/concurrent_decorator.py"
+                )
             )
-        )
-        m.addons.add(sc)
-        f1, f2 = tflow.tflow(), tflow.tflow()
-        m.request(f1)
-        m.request(f2)
-        start = time.time()
-        while time.time() - start < 5:
-            if f1.reply.state == f2.reply.state == "committed":
-                return
-        raise ValueError("Script never acked")
+            sc.start()
+
+            f1, f2 = tflow.tflow(), tflow.tflow()
+            tctx.cycle(sc, f1)
+            tctx.cycle(sc, f2)
+            start = time.time()
+            while time.time() - start < 5:
+                if f1.reply.state == f2.reply.state == "committed":
+                    return
+            raise ValueError("Script never acked")
 
     def test_concurrent_err(self):
-        m = mastertest.RecordingMaster(options.Options(), proxy.DummyServer())
-        sc = script.Script(
-            tutils.test_data.path(
-                "mitmproxy/data/addonscripts/concurrent_decorator_err.py"
+        with taddons.context() as tctx:
+            sc = script.Script(
+                tutils.test_data.path(
+                    "mitmproxy/data/addonscripts/concurrent_decorator_err.py"
+                )
             )
-        )
-        with m.handlecontext():
             sc.start()
-        assert "decorator not supported" in m.event_log[0][1]
+            assert "decorator not supported" in tctx.master.event_log[0][1]