remove ServerConnectionMixin.reconnect

13 files changed, 48 insertions, 62 deletions

diff --git a/libmproxy/protocol/__init__.py b/libmproxy/protocol/__init__.py
index c582592b..b0e66dbd 100644
--- a/libmproxy/protocol/__init__.py
+++ b/libmproxy/protocol/__init__.py
@@ -1,11 +1,11 @@
 from __future__ import (absolute_import, print_function, division)
-from .base import Layer, ServerConnectionMixin, Log, Kill
+from .base import Layer, ServerConnectionMixin, Kill
 from .http import Http1Layer, Http2Layer
 from .tls import TlsLayer, is_tls_record_magic
 from .rawtcp import RawTCPLayer
 
 __all__ = [
-    "Layer", "ServerConnectionMixin", "Log", "Kill",
+    "Layer", "ServerConnectionMixin", "Kill",
     "Http1Layer", "Http2Layer",
     "TlsLayer", "is_tls_record_magic",
     "RawTCPLayer"
diff --git a/libmproxy/protocol/base.py b/libmproxy/protocol/base.py
index 40ec0536..f1718065 100644
--- a/libmproxy/protocol/base.py
+++ b/libmproxy/protocol/base.py
@@ -81,15 +81,6 @@ class Layer(_LayerCodeCompletion):
         """
         return getattr(self.ctx, name)
 
-    def log(self, msg, level, subs=()):
-        full_msg = [
-            "{}: {}".format(repr(self.client_conn.address), msg)
-        ]
-        for i in subs:
-            full_msg.append("  -> " + i)
-        full_msg = "\n".join(full_msg)
-        self.channel.tell("log", Log(full_msg, level))
-
     @property
     def layers(self):
         return [self] + self.ctx.layers
@@ -106,15 +97,9 @@ class ServerConnectionMixin(object):
     def __init__(self, server_address=None):
         super(ServerConnectionMixin, self).__init__()
         self.server_conn = ServerConnection(server_address)
-        self._check_self_connect()
-
-    def reconnect(self):
-        address = self.server_conn.address
-        self._disconnect()
-        self.server_conn.address = address
-        self.connect()
+        self.__check_self_connect()
 
-    def _check_self_connect(self):
+    def __check_self_connect(self):
         """
         We try to protect the proxy from _accidentally_ connecting to itself,
         e.g. because of a failed transparent lookup or an invalid configuration.
@@ -134,10 +119,10 @@ class ServerConnectionMixin(object):
     def set_server(self, address, server_tls=None, sni=None, depth=1):
         if depth == 1:
             if self.server_conn:
-                self._disconnect()
+                self.disconnect()
             self.log("Set new server address: " + repr(address), "debug")
             self.server_conn.address = address
-            self._check_self_connect()
+            self.__check_self_connect()
             if server_tls:
                 raise ProtocolException(
                     "Cannot upgrade to TLS, no TLS layer on the protocol stack."
@@ -145,15 +130,16 @@ class ServerConnectionMixin(object):
         else:
             self.ctx.set_server(address, server_tls, sni, depth - 1)
 
-    def _disconnect(self):
+    def disconnect(self):
         """
         Deletes (and closes) an existing server connection.
         """
         self.log("serverdisconnect", "debug", [repr(self.server_conn.address)])
+        address = self.server_conn.address
         self.server_conn.finish()
         self.server_conn.close()
         self.channel.tell("serverdisconnect", self.server_conn)
-        self.server_conn = ServerConnection(None)
+        self.server_conn = ServerConnection(address)
 
     def connect(self):
         if not self.server_conn.address:
@@ -167,12 +153,6 @@ class ServerConnectionMixin(object):
                 "Server connection to %s failed: %s" % (repr(self.server_conn.address), e), e)
 
 
-class Log(object):
-    def __init__(self, msg, level="info"):
-        self.msg = msg
-        self.level = level
-
-
 class Kill(Exception):
     """
     Kill a connection.
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index dbf46973..7768b9ad 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -143,10 +143,6 @@ class Http1Layer(_StreamingHttpLayer):
         self.ctx.connect()
         self.server_protocol = HTTP1Protocol(self.server_conn)
 
-    def reconnect(self):
-        self.ctx.reconnect()
-        self.server_protocol = HTTP1Protocol(self.server_conn)
-
     def set_server(self, *args, **kwargs):
         self.ctx.set_server(*args, **kwargs)
         self.server_protocol = HTTP1Protocol(self.server_conn)
@@ -202,12 +198,6 @@ class Http2Layer(_HttpLayer):
                                              unhandled_frame_cb=self.handle_unexpected_frame_from_server)
         self.server_protocol.perform_connection_preface()
 
-    def reconnect(self):
-        self.ctx.reconnect()
-        self.server_protocol = HTTP2Protocol(self.server_conn, is_server=False,
-                                             unhandled_frame_cb=self.handle_unexpected_frame_from_server)
-        self.server_protocol.perform_connection_preface()
-
     def set_server(self, *args, **kwargs):
         self.ctx.set_server(*args, **kwargs)
         self.server_protocol = HTTP2Protocol(self.server_conn, is_server=False,
@@ -314,14 +304,10 @@ class UpstreamConnectLayer(Layer):
         else:
             pass  # swallow the message
 
-    def reconnect(self):
-        self.ctx.reconnect()
-        self._send_connect_request()
-
     def set_server(self, address, server_tls=None, sni=None, depth=1):
         if depth == 1:
             if self.ctx.server_conn:
-                self.ctx.reconnect()
+                self.ctx.disconnect()
             address = Address.wrap(address)
             self.connect_request.host = address.host
             self.connect_request.port = address.port
@@ -459,7 +445,8 @@ class HttpLayer(Layer):
             # > server detects timeout, disconnects
             # > read (100-n)% of large request
             # > send large request upstream
-            self.reconnect()
+            self.disconnect()
+            self.connect()
             get_response()
 
         # call the appropriate script hook - this is an opportunity for an
@@ -534,10 +521,12 @@ class HttpLayer(Layer):
             """
             # This is a very ugly (untested) workaround to solve a very ugly problem.
             if self.server_conn and self.server_conn.tls_established and not ssl:
-                self.reconnect()
+                self.disconnect()
+                self.connect()
             elif ssl and not hasattr(self, "connected_to") or self.connected_to != address:
                 if self.server_conn.tls_established:
-                    self.reconnect()
+                    self.disconnect()
+                    self.connect()
 
                 self.send_request(make_connect_request(address))
                 tls_layer = TlsLayer(self, False, True)
diff --git a/libmproxy/protocol/http_replay.py b/libmproxy/protocol/http_replay.py
index 2759a019..a9ee5506 100644
--- a/libmproxy/protocol/http_replay.py
+++ b/libmproxy/protocol/http_replay.py
@@ -6,7 +6,7 @@ from netlib.http.http1 import HTTP1Protocol
 from netlib.tcp import NetLibError
 from ..controller import Channel
 from ..models import Error, HTTPResponse, ServerConnection, make_connect_request
-from .base import Log, Kill
+from .base import Kill
 
 
 # TODO: Doesn't really belong into libmproxy.protocol...
@@ -89,8 +89,9 @@ class RequestReplayThread(threading.Thread):
             if self.channel:
                 self.channel.ask("error", self.flow)
         except Kill:
-            # KillSignal should only be raised if there's a channel in the
+            # Kill should only be raised if there's a channel in the
             # first place.
+            from ..proxy.root_context import Log
             self.channel.tell("log", Log("Connection killed", "info"))
         finally:
             r.form_out = form_out_backup
diff --git a/libmproxy/protocol/tls.py b/libmproxy/protocol/tls.py
index 2b37c5f4..00e016ea 100644
--- a/libmproxy/protocol/tls.py
+++ b/libmproxy/protocol/tls.py
@@ -338,11 +338,6 @@ class TlsLayer(Layer):
         if self._server_tls and not self.server_conn.tls_established:
             self._establish_tls_with_server()
 
-    def reconnect(self):
-        self.ctx.reconnect()
-        if self._server_tls and not self.server_conn.tls_established:
-            self._establish_tls_with_server()
-
     def set_server(self, address, server_tls=None, sni=None, depth=1):
         if depth == 1 and server_tls is not None:
             self.ctx.set_server(address, None, None, 1)
diff --git a/libmproxy/proxy/__init__.py b/libmproxy/proxy/__init__.py
index d5297cb1..be7f5207 100644
--- a/libmproxy/proxy/__init__.py
+++ b/libmproxy/proxy/__init__.py
@@ -2,8 +2,10 @@ from __future__ import (absolute_import, print_function, division)
 
 from .server import ProxyServer, DummyServer
 from .config import ProxyConfig
+from .root_context import RootContext, Log
 
 __all__ = [
     "ProxyServer", "DummyServer",
     "ProxyConfig",
+    "RootContext", "Log",
 ]
diff --git a/libmproxy/proxy/modes/http_proxy.py b/libmproxy/proxy/modes/http_proxy.py
index 90c54cc6..c7502c24 100644
--- a/libmproxy/proxy/modes/http_proxy.py
+++ b/libmproxy/proxy/modes/http_proxy.py
@@ -10,7 +10,7 @@ class HttpProxy(Layer, ServerConnectionMixin):
             layer()
         finally:
             if self.server_conn:
-                self._disconnect()
+                self.disconnect()
 
 
 class HttpUpstreamProxy(Layer, ServerConnectionMixin):
@@ -23,4 +23,4 @@ class HttpUpstreamProxy(Layer, ServerConnectionMixin):
             layer()
         finally:
             if self.server_conn:
-                self._disconnect()
+                self.disconnect()
diff --git a/libmproxy/proxy/modes/reverse_proxy.py b/libmproxy/proxy/modes/reverse_proxy.py
index b57ac5eb..28f4e6f8 100644
--- a/libmproxy/proxy/modes/reverse_proxy.py
+++ b/libmproxy/proxy/modes/reverse_proxy.py
@@ -14,4 +14,4 @@ class ReverseProxy(Layer, ServerConnectionMixin):
             layer()
         finally:
             if self.server_conn:
-                self._disconnect()
+                self.disconnect()
diff --git a/libmproxy/proxy/modes/socks_proxy.py b/libmproxy/proxy/modes/socks_proxy.py
index ebaf939e..0efeab67 100644
--- a/libmproxy/proxy/modes/socks_proxy.py
+++ b/libmproxy/proxy/modes/socks_proxy.py
@@ -57,4 +57,4 @@ class Socks5Proxy(Layer, ServerConnectionMixin):
             layer()
         finally:
             if self.server_conn:
-                self._disconnect()
+                self.disconnect()
diff --git a/libmproxy/proxy/modes/transparent_proxy.py b/libmproxy/proxy/modes/transparent_proxy.py
index 96ad86c4..d99485c9 100644
--- a/libmproxy/proxy/modes/transparent_proxy.py
+++ b/libmproxy/proxy/modes/transparent_proxy.py
@@ -21,4 +21,4 @@ class TransparentProxy(Layer, ServerConnectionMixin):
             layer()
         finally:
             if self.server_conn:
-                self._disconnect()
+                self.disconnect()
diff --git a/libmproxy/proxy/root_context.py b/libmproxy/proxy/root_context.py
index 35909612..88df8e47 100644
--- a/libmproxy/proxy/root_context.py
+++ b/libmproxy/proxy/root_context.py
@@ -85,9 +85,28 @@ class RootContext(object):
         # d = top_layer.client_conn.rfile.peek(len(HTTP2Protocol.CLIENT_CONNECTION_PREFACE))
         # is_http2_magic = (d == HTTP2Protocol.CLIENT_CONNECTION_PREFACE)
 
+    def log(self, msg, level, subs=()):
+        """
+        Send a log message to the master.
+        """
+
+        full_msg = [
+            "{}: {}".format(repr(self.client_conn.address), msg)
+        ]
+        for i in subs:
+            full_msg.append("  -> " + i)
+        full_msg = "\n".join(full_msg)
+        self.channel.tell("log", Log(full_msg, level))
+
     @property
     def layers(self):
         return []
 
     def __repr__(self):
         return "RootContext"
+
+
+class Log(object):
+    def __init__(self, msg, level="info"):
+        self.msg = msg
+        self.level = level
\ No newline at end of file
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index e9e8df09..5d067b45 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -8,10 +8,10 @@ from netlib import tcp
 from netlib.http.http1 import HTTP1Protocol
 from netlib.tcp import NetLibError
 from ..exceptions import ProtocolException, ServerException
-from ..protocol import Log, Kill
+from ..protocol import Kill
 from ..models import ClientConnection, make_error_response
 from .modes import HttpUpstreamProxy, HttpProxy, ReverseProxy, TransparentProxy, Socks5Proxy
-from .root_context import RootContext
+from .root_context import RootContext, Log
 
 
 class DummyServer:
diff --git a/test/test_dump.py b/test/test_dump.py
index a0ad6cb4..9f055bac 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -6,7 +6,7 @@ import netlib.tutils
 from netlib.http.semantics import CONTENT_MISSING
 
 from libmproxy import dump, flow
-from libmproxy.protocol import Log
+from libmproxy.proxy import Log
 import tutils
 import mock