diff --git a/examples/har_dump.py b/examples/har_dump.py index ab7bf1e14..aeb154d2f 100644 --- a/examples/har_dump.py +++ b/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) diff --git a/mitmproxy/addons/dumper.py b/mitmproxy/addons/dumper.py index 7e5f91560..89a9eab8b 100644 --- a/mitmproxy/addons/dumper.py +++ b/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) diff --git a/mitmproxy/connections.py b/mitmproxy/connections.py index 143f576b9..f3a752221 100644 --- a/mitmproxy/connections.py +++ b/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): diff --git a/mitmproxy/proxy/modes/http_proxy.py b/mitmproxy/proxy/modes/http_proxy.py index d4fb9d687..1acc42f90 100644 --- a/mitmproxy/proxy/modes/http_proxy.py +++ b/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() diff --git a/mitmproxy/proxy/modes/reverse_proxy.py b/mitmproxy/proxy/modes/reverse_proxy.py index 87d5d8c8e..8942477d7 100644 --- a/mitmproxy/proxy/modes/reverse_proxy.py +++ b/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() diff --git a/mitmproxy/proxy/modes/socks_proxy.py b/mitmproxy/proxy/modes/socks_proxy.py index adcd8fc1d..042580373 100644 --- a/mitmproxy/proxy/modes/socks_proxy.py +++ b/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() diff --git a/mitmproxy/proxy/modes/transparent_proxy.py b/mitmproxy/proxy/modes/transparent_proxy.py index ae532235a..95bfecedd 100644 --- a/mitmproxy/proxy/modes/transparent_proxy.py +++ b/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() diff --git a/mitmproxy/proxy/protocol/base.py b/mitmproxy/proxy/protocol/base.py index 97e90051b..4d67bb71a 100644 --- a/mitmproxy/proxy/protocol/base.py +++ b/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 diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py index 542f6a94a..1e2c627ec 100644 --- a/mitmproxy/proxy/protocol/http.py +++ b/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.") diff --git a/mitmproxy/proxy/protocol/http2.py b/mitmproxy/proxy/protocol/http2.py index f440d100e..7a28ac2ef 100644 --- a/mitmproxy/proxy/protocol/http2.py +++ b/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, diff --git a/mitmproxy/proxy/protocol/http_replay.py b/mitmproxy/proxy/protocol/http_replay.py index c37badd3f..2e4c91b0d 100644 --- a/mitmproxy/proxy/protocol/http_replay.py +++ b/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() diff --git a/mitmproxy/proxy/protocol/tls.py b/mitmproxy/proxy/protocol/tls.py index 8a344faf5..796477b21 100644 --- a/mitmproxy/proxy/protocol/tls.py +++ b/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() diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator.py b/test/mitmproxy/data/addonscripts/concurrent_decorator.py index a56c2af1c..162c00f4b 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator.py +++ b/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) diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index 71b87d288..bb760f926 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/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]