diff options
| -rw-r--r-- | doc-src/scripting/inlinescripts.html | 5 | ||||
| -rw-r--r-- | examples/stub.py | 7 | ||||
| -rw-r--r-- | libmproxy/proxy.py | 20 | ||||
| -rw-r--r-- | libmproxy/script.py | 2 | ||||
| -rw-r--r-- | test/scripts/all.py | 4 | ||||
| -rw-r--r-- | test/scripts/concurrent_decorator.py | 23 | ||||
| -rw-r--r-- | test/test_dump.py | 4 | ||||
| -rw-r--r-- | test/test_flow.py | 6 | ||||
| -rw-r--r-- | test/test_script.py | 10 |
9 files changed, 60 insertions, 21 deletions
diff --git a/doc-src/scripting/inlinescripts.html b/doc-src/scripting/inlinescripts.html index 19ae89a1..7ab1c101 100644 --- a/doc-src/scripting/inlinescripts.html +++ b/doc-src/scripting/inlinescripts.html @@ -36,6 +36,11 @@ Called when a client initiates a connection to the proxy. Note that a connection can correspond to multiple HTTP requests. +### serverconnect(ScriptContext, ServerConnection) + +Called when the proxy initiates a connection to the target server. Note that +a connection can correspond to multiple HTTP requests. + ### request(ScriptContext, Flow) Called when a client request has been received. The __Flow__ object is diff --git a/examples/stub.py b/examples/stub.py index 42b2935a..78cbfcf2 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -14,6 +14,13 @@ def clientconnect(ctx, client_connect): """ ctx.log("clientconnect") +def serverconnect(ctx, server_connection): + """ + Called when the proxy initiates a connection to the target server. Note that a + connection can correspond to multiple HTTP requests + """ + ctx.log("serverconnect") + def request(ctx, flow): """ Called when a client request has been received. diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index 3098aff4..609ffb62 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -133,18 +133,25 @@ class ProxyHandler(tcp.BaseHandler): self.server_conn = None tcp.BaseHandler.__init__(self, connection, client_address, server) - def get_server_connection(self, cc, scheme, host, port, sni): + def get_server_connection(self, cc, scheme, host, port, sni, request=None): """ When SNI is in play, this means we have an SSL-encrypted connection, which means that the entire handler is dedicated to a single server connection - no multiplexing. If this assumption ever breaks, we'll have to do something different with the SNI host variable on the handler object. + + `conn_info` holds the initial connection's parameters, as the + hook might change them. Also, the hook might require an initial + request to figure out connection settings; in this case it can + set require_request, which will cause the connection to be + re-opened after the client's request arrives. """ sc = self.server_conn if not sni: sni = host - if sc and (scheme, host, port, sni) != (sc.scheme, sc.host, sc.port, sc.sni): + conn_info = (scheme, host, port, sni) + if sc and (conn_info != sc.conn_info or (request and sc.require_request)): sc.terminate() self.server_conn = None self.log( @@ -159,6 +166,13 @@ class ProxyHandler(tcp.BaseHandler): if not self.server_conn: try: self.server_conn = ServerConnection(self.config, scheme, host, port, sni) + + # Additional attributes, used if the server_connect hook + # needs to change parameters + self.server_conn.request = request + self.server_conn.require_request = False + + self.server_conn.conn_info = conn_info self.channel.ask(self.server_conn) self.server_conn.connect() except tcp.NetLibError, v: @@ -223,7 +237,7 @@ class ProxyHandler(tcp.BaseHandler): # the case, we want to reconnect without sending an error # to the client. while 1: - sc = self.get_server_connection(cc, scheme, host, port, self.sni) + sc = self.get_server_connection(cc, scheme, host, port, self.sni, request=request) sc.send(request) if sc.requestcount == 1: # add timestamps only for first request (others are not directly affected) request.tcp_setup_timestamp = sc.tcp_setup_timestamp diff --git a/libmproxy/script.py b/libmproxy/script.py index d1b714db..f8a0d085 100644 --- a/libmproxy/script.py +++ b/libmproxy/script.py @@ -78,7 +78,7 @@ def concurrent(fn): r = getattr(flow, fn.func_name) _handle_concurrent_reply(fn, r, [ctx, flow]) return _concurrent - elif fn.func_name in ["clientconnect", "clientdisconnect", "serverconnect"]: + elif fn.func_name in ["clientconnect", "serverconnect", "clientdisconnect"]: def _concurrent(ctx, conn): _handle_concurrent_reply(fn, conn, [ctx, conn]) return _concurrent diff --git a/test/scripts/all.py b/test/scripts/all.py index e6da7e51..7d30d757 100644 --- a/test/scripts/all.py +++ b/test/scripts/all.py @@ -3,6 +3,10 @@ def clientconnect(ctx, cc): ctx.log("XCLIENTCONNECT") log.append("clientconnect") +def serverconnect(ctx, cc): + ctx.log("XSERVERCONNECT") + log.append("serverconnect") + def request(ctx, r): ctx.log("XREQUEST") log.append("request") diff --git a/test/scripts/concurrent_decorator.py b/test/scripts/concurrent_decorator.py index c1c2651e..8e132006 100644 --- a/test/scripts/concurrent_decorator.py +++ b/test/scripts/concurrent_decorator.py @@ -1,6 +1,17 @@ import time from libmproxy.script import concurrent + +@concurrent +def clientconnect(context, cc): + context.log("clientconnect") + + +@concurrent +def serverconnect(context, sc): + context.log("serverconnect") + + @concurrent def request(context, flow): time.sleep(0.1) @@ -17,15 +28,5 @@ def error(context, err): @concurrent -def clientconnect(context, cc): - context.log("clientconnect") - - -@concurrent def clientdisconnect(context, dc): - context.log("clientdisconnect") - - -@concurrent -def serverconnect(context, sc): - context.log("serverconnect")
\ No newline at end of file + context.log("clientdisconnect")
\ No newline at end of file diff --git a/test/test_dump.py b/test/test_dump.py index 3b79c721..3d375f16 100644 --- a/test/test_dump.py +++ b/test/test_dump.py @@ -30,6 +30,9 @@ class TestDumpMaster: resp = tutils.tresp(req) resp.content = content m.handle_clientconnect(cc) + sc = proxy.ServerConnection(m.o, req.scheme, req.host, req.port, None) + sc.reply = mock.MagicMock() + m.handle_serverconnection(sc) m.handle_request(req) f = m.handle_response(resp) cd = flow.ClientDisconnect(cc) @@ -153,6 +156,7 @@ class TestDumpMaster: scripts=[[tutils.test_data.path("scripts/all.py")]], verbosity=0, eventlog=True ) assert "XCLIENTCONNECT" in ret + assert "XSERVERCONNECT" in ret assert "XREQUEST" in ret assert "XRESPONSE" in ret assert "XCLIENTDISCONNECT" in ret diff --git a/test/test_flow.py b/test/test_flow.py index 9844e0fd..c614960b 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -1,7 +1,7 @@ import Queue, time, os.path from cStringIO import StringIO import email.utils -from libmproxy import filt, flow, controller, utils, tnetstring +from libmproxy import filt, flow, controller, utils, tnetstring, proxy import tutils @@ -575,6 +575,10 @@ class TestFlowMaster: req = tutils.treq() fm.handle_clientconnect(req.client_conn) assert fm.scripts[0].ns["log"][-1] == "clientconnect" + sc = proxy.ServerConnection(None, req.scheme, req.host, req.port, None) + sc.reply = controller.DummyReply() + fm.handle_serverconnection(sc) + assert fm.scripts[0].ns["log"][-1] == "serverconnect" f = fm.handle_request(req) assert fm.scripts[0].ns["log"][-1] == "request" resp = tutils.tresp(req) diff --git a/test/test_script.py b/test/test_script.py index 766cfb8c..ad2296ef 100644 --- a/test/test_script.py +++ b/test/test_script.py @@ -103,11 +103,11 @@ class TestScript: f.error = tutils.terr(f.request) f.reply = f.request.reply - print s.run("response", f) - print s.run("error", f) - print s.run("clientconnect", f) - print s.run("clientdisconnect", f) - print s.run("serverconnect", f) + s.run("clientconnect", f) + s.run("serverconnect", f) + s.run("response", f) + s.run("error", f) + s.run("clientdisconnect", f) time.sleep(0.1) assert ctx.count == 5 |