fix inline script redirects

4 files changed, 17 insertions, 11 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 3d9ef722..8605d7a1 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -8,6 +8,7 @@ import Cookie
 import cookielib
 import os
 import re
+from libmproxy.protocol2.http import RequestReplayThread
 
 from netlib import odict, wsgi, tcp
 from netlib.http.semantics import CONTENT_MISSING
@@ -934,7 +935,7 @@ class FlowMaster(controller.Master):
             f.response = None
             f.error = None
             self.process_new_request(f)
-            rt = http.RequestReplayThread(
+            rt = RequestReplayThread(
                 self.server.config,
                 f,
                 self.masterq if run_scripthooks else False,
diff --git a/libmproxy/protocol2/http.py b/libmproxy/protocol2/http.py
index 32c0116b..792cf266 100644
--- a/libmproxy/protocol2/http.py
+++ b/libmproxy/protocol2/http.py
@@ -212,10 +212,11 @@ class UpstreamConnectLayer(Layer):
         self.ctx.reconnect()
         self.send_to_server(self.connect_request)
 
-    def set_server(self, address, server_tls, sni, depth=1):
+    def set_server(self, address, server_tls=None, sni=None, depth=1):
         if depth == 1:
             if self.ctx.server_conn:
                 self.ctx.reconnect()
+            address = Address.wrap(address)
             self.connect_request.host = address.host
             self.connect_request.port = address.port
             self.server_conn.address = address
@@ -227,11 +228,16 @@ class HttpLayer(Layer):
     def __init__(self, ctx, mode):
         super(HttpLayer, self).__init__(ctx)
         self.mode = mode
+        self.__original_server_conn = None
+        "Contains the original destination in transparent mode, which needs to be restored"
+        "if an inline script modified the target server for a single http request"
 
     def __call__(self):
+        if self.mode == "transparent":
+            self.__original_server_conn = self.server_conn
         while True:
             try:
-                flow = HTTPFlow(self.client_conn, self.server_conn, live=True)
+                flow = HTTPFlow(self.client_conn, self.server_conn, live=self)
 
                 try:
                     request = self.read_from_client()
@@ -288,7 +294,7 @@ class HttpLayer(Layer):
                 flow.live = False
 
     def handle_regular_mode_connect(self, request):
-        self.set_server((request.host, request.port), False, None)
+        self.set_server((request.host, request.port))
         self.send_to_client(make_connect_response(request.httpversion))
         layer = self.ctx.next_layer(self)
         layer()
@@ -433,11 +439,10 @@ class HttpLayer(Layer):
             if flow.request.form_in == "authority":
                 flow.request.scheme = "http"  # pseudo value
         else:
-            flow.request.host = self.ctx.server_conn.address.host
-            flow.request.port = self.ctx.server_conn.address.port
-            flow.request.scheme = "https" if self.server_conn.tls_established else "http"
+            flow.request.host = self.__original_server_conn.address.host
+            flow.request.port = self.__original_server_conn.address.port
+            flow.request.scheme = "https" if self.__original_server_conn.tls_established else "http"
 
-        # TODO: Expose .set_server functionality to inline scripts
         request_reply = self.channel.ask("request", flow)
         if request_reply is None or request_reply == KILL:
             raise Kill()
diff --git a/libmproxy/protocol2/layer.py b/libmproxy/protocol2/layer.py
index 7cb76591..f72320ff 100644
--- a/libmproxy/protocol2/layer.py
+++ b/libmproxy/protocol2/layer.py
@@ -112,7 +112,7 @@ class ServerConnectionMixin(object):
         self.server_conn.address = address
         self.connect()
 
-    def set_server(self, address, server_tls, sni, depth=1):
+    def set_server(self, address, server_tls=None, sni=None, depth=1):
         if depth == 1:
             if self.server_conn:
                 self._disconnect()
diff --git a/libmproxy/protocol2/tls.py b/libmproxy/protocol2/tls.py
index 433dd65d..b1b80034 100644
--- a/libmproxy/protocol2/tls.py
+++ b/libmproxy/protocol2/tls.py
@@ -110,9 +110,9 @@ class TlsLayer(Layer):
         if self._server_tls and not self.server_conn.tls_established:
             self._establish_tls_with_server()
 
-    def set_server(self, address, server_tls, sni, depth=1):
+    def set_server(self, address, server_tls=None, sni=None, depth=1):
         self.ctx.set_server(address, server_tls, sni, depth)
-        if server_tls is not None:
+        if depth == 1 and server_tls is not None:
             self._sni_from_server_change = sni
             self._server_tls = server_tls