kylefang/mitmproxy
1
0
Fork 0
mirror of https://github.com/zhigang1992/mitmproxy.git synced 2026-04-24 04:14:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

If a message has been acked, all other processors are skipped

Browse Source 
This applies the constraint, but does to clumsily. When we've unified modules
and processors it will be much nicer. We also make some exceptions for the
master processors that we may want to re-evaluate down the track.
This commit is contained in:
Aldo Cortesi
2016-06-08 12:58:58 +12:00
parent b3bf754e53
commit a5cb241c7c
3 changed files with 40 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
mitmproxy/flow/master.py
View File

@@ -103,9 +103,10 @@ class FlowMaster(controller.Master):
except script.ScriptException as e:
self.add_event("Script error:\n{}".format(e), "error")
def run_script_hook(self, name, *args, **kwargs):
def run_scripts(self, name, msg):
for script_obj in self.scripts:
self._run_single_script_hook(script_obj, name, *args, **kwargs)
if not msg.reply.acked:
self._run_single_script_hook(script_obj, name, msg)
def get_ignore_filter(self):
return self.server.config.check_ignore.patterns
@@ -373,28 +374,28 @@ class FlowMaster(controller.Master):
@controller.handler
def clientconnect(self, root_layer):
self.run_script_hook("clientconnect", root_layer)
self.run_scripts("clientconnect", root_layer)
@controller.handler
def clientdisconnect(self, root_layer):
self.run_script_hook("clientdisconnect", root_layer)
self.run_scripts("clientdisconnect", root_layer)
@controller.handler
def serverconnect(self, server_conn):
self.run_script_hook("serverconnect", server_conn)
self.run_scripts("serverconnect", server_conn)
@controller.handler
def serverdisconnect(self, server_conn):
self.run_script_hook("serverdisconnect", server_conn)
self.run_scripts("serverdisconnect", server_conn)
@controller.handler
def next_layer(self, top_layer):
self.run_script_hook("next_layer", top_layer)
self.run_scripts("next_layer", top_layer)
@controller.handler
def error(self, f):
self.state.update_flow(f)
self.run_script_hook("error", f)
self.run_scripts("error", f)
if self.client_playback:
self.client_playback.clear(f)
return f
@@ -416,10 +417,14 @@ class FlowMaster(controller.Master):
if f not in self.state.flows: # don't add again on replay
self.state.add_flow(f)
self.active_flows.add(f)
self.replacehooks.run(f)
self.setheaders.run(f)
self.process_new_request(f)
self.run_script_hook("request", f)
if not f.reply.acked:
self.replacehooks.run(f)
if not f.reply.acked:
self.setheaders.run(f)
if not f.reply.acked:
self.process_new_request(f)
if not f.reply.acked:
self.run_scripts("request", f)
return f
@controller.handler
@@ -430,18 +435,21 @@ class FlowMaster(controller.Master):
except netlib.exceptions.HttpException:
f.reply.kill()
return
self.run_script_hook("responseheaders", f)
self.run_scripts("responseheaders", f)
return f
@controller.handler
def response(self, f):
self.active_flows.discard(f)
self.state.update_flow(f)
self.replacehooks.run(f)
self.setheaders.run(f)
self.run_script_hook("response", f)
if self.client_playback:
self.client_playback.clear(f)
if not f.reply.acked:
self.replacehooks.run(f)
if not f.reply.acked:
self.setheaders.run(f)
self.run_scripts("response", f)
if not f.reply.acked:
if self.client_playback:
self.client_playback.clear(f)
self.process_new_response(f)
if self.stream:
self.stream.add(f)
@@ -487,11 +495,11 @@ class FlowMaster(controller.Master):
# TODO: This would break mitmproxy currently.
# self.state.add_flow(flow)
self.active_flows.add(flow)
self.run_script_hook("tcp_open", flow)
self.run_scripts("tcp_open", flow)
@controller.handler
def tcp_message(self, flow):
self.run_script_hook("tcp_message", flow)
self.run_scripts("tcp_message", flow)
message = flow.messages[-1]
direction = "->" if message.from_client else "<-"
self.add_event("{client} {direction} tcp {direction} {server}".format(
@@ -507,14 +515,14 @@ class FlowMaster(controller.Master):
repr(flow.server_conn.address),
flow.error
), "info")
self.run_script_hook("tcp_error", flow)
self.run_scripts("tcp_error", flow)
@controller.handler
def tcp_close(self, flow):
self.active_flows.discard(flow)
if self.stream:
self.stream.add(flow)
self.run_script_hook("tcp_close", flow)
self.run_scripts("tcp_close", flow)
def shutdown(self):
super(FlowMaster, self).shutdown()

2
test/mitmproxy/mastertest.py
View File

@@ -16,7 +16,9 @@ class MasterTest:
master.request(f)
if not f.error:
f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content))
f.reply.acked = False
f = master.response(f)
f.client_conn.reply.acked = False
master.clientdisconnect(f.client_conn)
return f

8
test/mitmproxy/test_flow.py
View File

@@ -807,17 +807,22 @@ class TestFlowMaster:
fm.load_script(tutils.test_data.path("data/scripts/all.py"))
f = tutils.tflow(resp=True)
f.client_conn.acked = False
fm.clientconnect(f.client_conn)
assert fm.scripts[0].ns["log"][-1] == "clientconnect"
f.server_conn.acked = False
fm.serverconnect(f.server_conn)
assert fm.scripts[0].ns["log"][-1] == "serverconnect"
f.reply.acked = False
fm.request(f)
assert fm.scripts[0].ns["log"][-1] == "request"
f.reply.acked = False
fm.response(f)
assert fm.scripts[0].ns["log"][-1] == "response"
# load second script
fm.load_script(tutils.test_data.path("data/scripts/all.py"))
assert len(fm.scripts) == 2
f.server_conn.reply.acked = False
fm.clientdisconnect(f.server_conn)
assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
assert fm.scripts[1].ns["log"][-1] == "clientdisconnect"
@@ -828,6 +833,7 @@ class TestFlowMaster:
fm.load_script(tutils.test_data.path("data/scripts/all.py"))
f.error = tutils.terr()
f.reply.acked = False
fm.error(f)
assert fm.scripts[0].ns["log"][-1] == "error"
@@ -977,10 +983,12 @@ class TestFlowMaster:
f = tutils.tflow(resp=True)
f.response.headers["set-cookie"] = "foo=bar"
fm.request(f)
f.reply.acked = False
fm.response(f)
assert fm.stickycookie_state.jar
assert "cookie" not in f.request.headers
f = f.copy()
f.reply.acked = False
fm.request(f)
assert f.request.headers["cookie"] == "foo=bar"