From b97b1f17cf1f61b75682ae902f8bdea5fb85763c Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Sun, 24 Aug 2014 14:22:11 +0200 Subject: [PATCH] fix #328 --- libmproxy/flow.py | 2 +- libmproxy/protocol/http.py | 19 +++++----- libmproxy/protocol/primitives.py | 64 ++++++++++++++++++-------------- test/test_server.py | 4 +- 4 files changed, 49 insertions(+), 40 deletions(-) diff --git a/libmproxy/flow.py b/libmproxy/flow.py index dd8971159..2540435ee 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -704,7 +704,7 @@ class FlowMaster(controller.Master): return f def handle_request(self, r): - if r.flow.client_conn and r.flow.client_conn.wfile: + if r.flow.live: app = self.apps.get(r) if app: err = app.serve(r, r.flow.client_conn.wfile, **{"mitmproxy.master": self}) diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py index fdc29df1b..38a6cb498 100644 --- a/libmproxy/protocol/http.py +++ b/libmproxy/protocol/http.py @@ -5,7 +5,7 @@ import threading from netlib import http, tcp, http_status import netlib.utils from netlib.odict import ODict, ODictCaseless -from .primitives import KILL, ProtocolHandler, TemporaryServerChangeMixin, Flow, Error +from .primitives import KILL, ProtocolHandler, LiveConnection, Flow, Error from ..proxy.connection import ServerConnection from .. import encoding, utils, controller, stateobject, proxy @@ -543,8 +543,8 @@ class HTTPRequest(HTTPMessage): self.path = path if host != self.get_host() or port != self.get_port(): - if self.flow.change_server: - self.flow.change_server((host, port), ssl=is_ssl) + if self.flow.live: + self.flow.live.change_server((host, port), ssl=is_ssl) else: # There's not live server connection, we're just changing the attributes here. self.flow.server_conn = ServerConnection((host, port), @@ -789,15 +789,15 @@ class HTTPFlow(Flow): The following additional attributes are exposed: intercepting: Is this flow currently being intercepted? + live: Does this flow have a live client connection? """ - def __init__(self, client_conn, server_conn, change_server=None): - super(HTTPFlow, self).__init__("http", client_conn, server_conn) + def __init__(self, client_conn, server_conn, live=None): + super(HTTPFlow, self).__init__("http", client_conn, server_conn, live) self.request = None """@type: HTTPRequest""" self.response = None """@type: HTTPResponse""" - self.change_server = change_server # Used by flow.request.set_url to change the server address self.intercepting = False # FIXME: Should that rather be an attribute of Flow? @@ -904,7 +904,7 @@ class HttpAuthenticationError(Exception): return "Proxy Authentication Required" -class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin): +class HTTPHandler(ProtocolHandler): def __init__(self, c): super(HTTPHandler, self).__init__(c) self.expected_form_in = c.config.http_form_in @@ -943,7 +943,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin): raise v def handle_flow(self): - flow = HTTPFlow(self.c.client_conn, self.c.server_conn, self.change_server) + flow = HTTPFlow(self.c.client_conn, self.c.server_conn, self.live) try: req = HTTPRequest.from_stream(self.c.client_conn.rfile, body_size_limit=self.c.config.body_size_limit) @@ -1038,7 +1038,8 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin): # If the user has changed the target server on this connection, # restore the original target server - self.restore_server() + flow.live.restore_server() + flow.live = None return True except (HttpAuthenticationError, http.HttpError, proxy.ProxyError, tcp.NetLibError), e: diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py index 7b936f7fb..a227d904d 100644 --- a/libmproxy/protocol/primitives.py +++ b/libmproxy/protocol/primitives.py @@ -71,12 +71,14 @@ class Error(stateobject.SimpleStateObject): class Flow(stateobject.SimpleStateObject, BackreferenceMixin): - def __init__(self, conntype, client_conn, server_conn): + def __init__(self, conntype, client_conn, server_conn, live=None): self.conntype = conntype self.client_conn = client_conn """@type: ClientConnection""" self.server_conn = server_conn """@type: ServerConnection""" + self.live = live # Used by flow.request.set_url to change the server address + """@type: LiveConnection""" self.error = None """@type: Error""" @@ -140,6 +142,8 @@ class ProtocolHandler(object): def __init__(self, c): self.c = c """@type: libmproxy.proxy.server.ConnectionHandler""" + self.live = LiveConnection(c) + """@type: LiveConnection""" def handle_messages(self): """ @@ -164,46 +168,50 @@ class ProtocolHandler(object): raise error # pragma: nocover -class TemporaryServerChangeMixin(object): +class LiveConnection(object): """ - This mixin allows safe modification of the target server, - without any need to expose the ConnectionHandler to the Flow. + This facade allows protocol handlers to interface with a live connection, + without requiring the expose the ConnectionHandler. """ - def change_server(self, address, ssl): + def __init__(self, c): + self._c = c + """@type: libmproxy.proxy.server.ConnectionHandler""" + + def change_server(self, address, ssl, persistent_change=False): address = netlib.tcp.Address.wrap(address) - if address == self.c.server_conn.address(): - return - priority = AddressPriority.MANUALLY_CHANGED + if address != self._c.server_conn.address: - self.c.log("Temporarily change server connection: %s:%s -> %s:%s" % ( - self.c.server_conn.address.host, - self.c.server_conn.address.port, - address.host, - address.port - ), "debug") + self._c.log("Change server connection: %s:%s -> %s:%s" % ( + self._c.server_conn.address.host, + self._c.server_conn.address.port, + address.host, + address.port + ), "debug") - if not hasattr(self, "_backup_server_conn"): - self._backup_server_conn = self.c.server_conn - self.c.server_conn = None - else: # This is at least the second temporary change. We can kill the current connection. - self.c.del_server_connection() + if not hasattr(self, "_backup_server_conn"): + self._backup_server_conn = self._c.server_conn + self._c.server_conn = None + else: # This is at least the second temporary change. We can kill the current connection. + 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) + self._c.set_server_address(address, AddressPriority.MANUALLY_CHANGED) + self._c.establish_server_connection(ask=False) + if ssl: + self._c.establish_ssl(server=True) + if hasattr(self, "_backup_server_conn") and persistent_change: + del self._backup_server_conn def restore_server(self): if not hasattr(self, "_backup_server_conn"): return - self.c.log("Restore original server connection: %s:%s -> %s:%s" % ( - self.c.server_conn.address.host, - self.c.server_conn.address.port, + self._c.log("Restore original server connection: %s:%s -> %s:%s" % ( + self._c.server_conn.address.host, + self._c.server_conn.address.port, self._backup_server_conn.address.host, self._backup_server_conn.address.port ), "debug") - self.c.del_server_connection() - self.c.server_conn = self._backup_server_conn + self._c.del_server_connection() + self._c.server_conn = self._backup_server_conn del self._backup_server_conn \ No newline at end of file diff --git a/test/test_server.py b/test/test_server.py index fea6791da..017faacb0 100644 --- a/test/test_server.py +++ b/test/test_server.py @@ -338,9 +338,9 @@ class MasterRedirectRequest(tservers.TestMaster): request.set_url(new) request.set_url(new) - request.flow.change_server(("127.0.0.1", self.redirect_port), False) + request.flow.live.change_server(("127.0.0.1", self.redirect_port), False) request.set_url(url) - tutils.raises("SSL handshake error", request.flow.change_server, ("127.0.0.1", self.redirect_port), True) + tutils.raises("SSL handshake error", request.flow.live.change_server, ("127.0.0.1", self.redirect_port), True) request.set_url(new) request.set_url(url) request.set_url(new)