From e9c834a30ddab5a8b6ef7b31c8ebebcf7e955371 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Sat, 26 Jul 2014 12:02:18 +0200
Subject: [PATCH] fix #259

---
 libmproxy/flow.py                |  9 ++-------
 libmproxy/protocol/http.py       |  2 ++
 libmproxy/protocol/primitives.py |  1 +
 libmproxy/proxy/server.py        | 12 +++++++++---
 test/test_dump.py                |  2 +-
 test/test_flow.py                |  2 +-
 6 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index b6b490222..5062c2027 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -672,12 +672,7 @@ class FlowMaster(controller.Master):
         self.run_script_hook("clientdisconnect", r)
         r.reply()
 
-    def handle_serverconnection(self, sc):
-        # To unify the mitmproxy script API, we call the script hook
-        # "serverconnect" rather than "serverconnection".  As things are handled
-        # differently in libmproxy (ClientConnect + ClientDisconnect vs
-        # ServerConnection class), there is no "serverdisonnect" event at the
-        # moment.
+    def handle_serverconnect(self, sc):
         self.run_script_hook("serverconnect", sc)
         sc.reply()
 
@@ -791,7 +786,7 @@ class FilteredFlowWriter:
 
 
 class RequestReplayThread(threading.Thread):
-    name="RequestReplayThread"
+    name = "RequestReplayThread"
 
     def __init__(self, config, flow, masterq, should_exit):
         self.config, self.flow, self.channel = config, flow, controller.Channel(masterq, should_exit)
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 8a3210785..8f9d5f5d4 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -1035,12 +1035,14 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
                 if not self.c.config.get_upstream_server:
                     self.c.set_server_address((request.host, request.port),
                                               proxy.AddressPriority.FROM_PROTOCOL)
+                    self.c.establish_server_connection()
                     flow.server_conn = self.c.server_conn  # Update server_conn attribute on the flow
                     self.c.client_conn.send(
                         'HTTP/1.1 200 Connection established\r\n' +
                         ('Proxy-agent: %s\r\n' % self.c.server_version) +
                         '\r\n'
                     )
+
                     self.ssl_upgrade()
                     self.skip_authentication = True
                     return False
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index 8c0ea5db7..5743bb6aa 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -181,6 +181,7 @@ class TemporaryServerChangeMixin(object):
             self.c.del_server_connection()
 
         self.c.set_server_address(address, priority)
+        self.c.establish_server_connection(ask=False)
         if ssl:
             self.c.establish_ssl(server=True)
 
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index 72228f16e..613697756 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -146,15 +146,20 @@ class ConnectionHandler:
         self.log("Set new server address: %s:%s" % (address.host, address.port), "debug")
         self.server_conn = ServerConnection(address, priority)
 
-    def establish_server_connection(self):
+    def establish_server_connection(self, ask=True):
         """
         Establishes a new server connection.
         If there is already an existing server connection, the function returns immediately.
+
+        By default, this function ".ask"s the proxy master. This is deadly if this function is already called from the
+        master (e.g. via change_server), because this navigates us in a simple deadlock (the master is single-threaded).
+        In these scenarios, ask=False can be passed to suppress the call to the master.
         """
         if self.server_conn.connection:
             return
         self.log("serverconnect", "debug", ["%s:%s" % self.server_conn.address()[:2]])
-        self.channel.tell("serverconnect", self)
+        if ask:
+            self.channel.ask("serverconnect", self)
         try:
             self.server_conn.connect()
         except tcp.NetLibError, v:
@@ -185,9 +190,10 @@ class ConnectionHandler:
             self.log("Establish SSL", "debug", subs)
 
         if server:
+            if not self.server_conn or not self.server_conn.connection:
+                raise ProxyError(502, "No server connection.")
             if self.server_conn.ssl_established:
                 raise ProxyError(502, "SSL to Server already established.")
-            self.establish_server_connection()  # make sure there is a server connection.
             self.server_conn.establish_ssl(self.config.clientcerts, self.sni)
         if client:
             if self.client_conn.ssl_established:
diff --git a/test/test_dump.py b/test/test_dump.py
index 604bc5421..6f70450fb 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -30,7 +30,7 @@ class TestDumpMaster:
         m.handle_clientconnect(cc)
         sc = proxy.connection.ServerConnection((req.get_host(), req.get_port()), None)
         sc.reply = mock.MagicMock()
-        m.handle_serverconnection(sc)
+        m.handle_serverconnect(sc)
         m.handle_request(req)
         resp = tutils.tresp(req, content=content)
         f = m.handle_response(resp)
diff --git a/test/test_flow.py b/test/test_flow.py
index 890130035..88e7b9d70 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -589,7 +589,7 @@ class TestFlowMaster:
         assert fm.scripts[0].ns["log"][-1] == "clientconnect"
         sc = ServerConnection((req.get_host(), req.get_port()), None)
         sc.reply = controller.DummyReply()
-        fm.handle_serverconnection(sc)
+        fm.handle_serverconnect(sc)
         assert fm.scripts[0].ns["log"][-1] == "serverconnect"
         f = fm.handle_request(req)
         assert fm.scripts[0].ns["log"][-1] == "request"