This commit is contained in:
Maximilian Hils 2015-02-07 16:26:19 +01:00
parent 0d608c1fb3
commit 53fadd05f4
4 changed files with 35 additions and 22 deletions

View File

@ -824,12 +824,10 @@ class FlowMaster(controller.Master):
if self.stickycookie_state: if self.stickycookie_state:
self.stickycookie_state.handle_response(f) self.stickycookie_state.handle_response(f)
def replay_request(self, f, block=False): def replay_request(self, f, block=False, run_scripthooks=True):
""" """
Returns None if successful, or error message if not. Returns None if successful, or error message if not.
""" """
if f.live:
return "Can't replay request which is still live..."
if f.intercepted: if f.intercepted:
return "Can't replay while intercepting..." return "Can't replay while intercepting..."
if f.request.content == http.CONTENT_MISSING: if f.request.content == http.CONTENT_MISSING:
@ -845,7 +843,7 @@ class FlowMaster(controller.Master):
rt = http.RequestReplayThread( rt = http.RequestReplayThread(
self.server.config, self.server.config,
f, f,
self.masterq, self.masterq if run_scripthooks else False,
self.should_exit self.should_exit
) )
rt.start() # pragma: no cover rt.start() # pragma: no cover

View File

@ -1416,20 +1416,31 @@ class RequestReplayThread(threading.Thread):
name = "RequestReplayThread" name = "RequestReplayThread"
def __init__(self, config, flow, masterq, should_exit): def __init__(self, config, flow, masterq, should_exit):
self.config, self.flow, self.channel = config, flow, controller.Channel(masterq, should_exit) """
threading.Thread.__init__(self) masterqueue can be a queue or None, if no scripthooks should be processed.
"""
self.config, self.flow = config, flow
if masterq:
self.channel = controller.Channel(masterq, should_exit)
else:
self.channel = None
super(RequestReplayThread, self).__init__()
def run(self): def run(self):
r = self.flow.request r = self.flow.request
form_out_backup = r.form_out form_out_backup = r.form_out
try: try:
self.flow.response = None self.flow.response = None
# If we have a channel, run script hooks.
if self.channel:
request_reply = self.channel.ask("request", self.flow) request_reply = self.channel.ask("request", self.flow)
if request_reply is None or request_reply == KILL: if request_reply is None or request_reply == KILL:
raise KillSignal() raise KillSignal()
elif isinstance(request_reply, HTTPResponse): elif isinstance(request_reply, HTTPResponse):
self.flow.response = request_reply self.flow.response = request_reply
else:
if not self.flow.response:
# In all modes, we directly connect to the server displayed # In all modes, we directly connect to the server displayed
if self.config.mode == "upstream": if self.config.mode == "upstream":
server_address = self.config.mode.get_upstream_server(self.flow.client_conn)[2:] server_address = self.config.mode.get_upstream_server(self.flow.client_conn)[2:]
@ -1453,13 +1464,16 @@ class RequestReplayThread(threading.Thread):
self.flow.server_conn = server self.flow.server_conn = server
self.flow.response = HTTPResponse.from_stream(server.rfile, r.method, self.flow.response = HTTPResponse.from_stream(server.rfile, r.method,
body_size_limit=self.config.body_size_limit) body_size_limit=self.config.body_size_limit)
if self.channel:
response_reply = self.channel.ask("response", self.flow) response_reply = self.channel.ask("response", self.flow)
if response_reply is None or response_reply == KILL: if response_reply is None or response_reply == KILL:
raise KillSignal() raise KillSignal()
except (proxy.ProxyError, http.HttpError, tcp.NetLibError) as v: except (proxy.ProxyError, http.HttpError, tcp.NetLibError) as v:
self.flow.error = Error(repr(v)) self.flow.error = Error(repr(v))
if self.channel:
self.channel.ask("error", self.flow) self.channel.ask("error", self.flow)
except KillSignal: except KillSignal:
# KillSignal should only be raised if there's a channel in the first place.
self.channel.tell("log", proxy.Log("Connection killed", "info")) self.channel.tell("log", proxy.Log("Connection killed", "info"))
finally: finally:
r.form_out = form_out_backup r.form_out = form_out_backup

View File

@ -36,7 +36,7 @@ class ScriptContext:
Replay the request on the current flow. The response will be added Replay the request on the current flow. The response will be added
to the flow object. to the flow object.
""" """
self._master.replay_request(f) return self._master.replay_request(f, block=True, run_scripthooks=False)
@property @property
def app_registry(self): def app_registry(self):
@ -139,8 +139,12 @@ def _handle_concurrent_reply(fn, o, *args, **kwargs):
def run(): def run():
fn(*args, **kwargs) fn(*args, **kwargs)
o.reply() # If the script did not call .reply(), we have to do it now. reply_proxy() # If the script did not call .reply(), we have to do it now.
threading.Thread(target=run, name="ScriptThread").start() ScriptThread(target=run).start()
class ScriptThread(threading.Thread):
name = "ScriptThread"
def concurrent(fn): def concurrent(fn):

View File

@ -666,9 +666,6 @@ class TestFlowMaster:
f.intercepted = True f.intercepted = True
assert "intercepting" in fm.replay_request(f) assert "intercepting" in fm.replay_request(f)
f.live = True
assert "live" in fm.replay_request(f)
def test_script_reqerr(self): def test_script_reqerr(self):
s = flow.State() s = flow.State()
fm = flow.FlowMaster(None, s) fm = flow.FlowMaster(None, s)