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:
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.
"""
if f.live:
return "Can't replay request which is still live..."
if f.intercepted:
return "Can't replay while intercepting..."
if f.request.content == http.CONTENT_MISSING:
@ -845,7 +843,7 @@ class FlowMaster(controller.Master):
rt = http.RequestReplayThread(
self.server.config,
f,
self.masterq,
self.masterq if run_scripthooks else False,
self.should_exit
)
rt.start() # pragma: no cover

View File

@ -1416,20 +1416,31 @@ class RequestReplayThread(threading.Thread):
name = "RequestReplayThread"
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):
r = self.flow.request
form_out_backup = r.form_out
try:
self.flow.response = None
# If we have a channel, run script hooks.
if self.channel:
request_reply = self.channel.ask("request", self.flow)
if request_reply is None or request_reply == KILL:
raise KillSignal()
elif isinstance(request_reply, HTTPResponse):
self.flow.response = request_reply
else:
if not self.flow.response:
# In all modes, we directly connect to the server displayed
if self.config.mode == "upstream":
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.response = HTTPResponse.from_stream(server.rfile, r.method,
body_size_limit=self.config.body_size_limit)
if self.channel:
response_reply = self.channel.ask("response", self.flow)
if response_reply is None or response_reply == KILL:
raise KillSignal()
except (proxy.ProxyError, http.HttpError, tcp.NetLibError) as v:
self.flow.error = Error(repr(v))
if self.channel:
self.channel.ask("error", self.flow)
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"))
finally:
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
to the flow object.
"""
self._master.replay_request(f)
return self._master.replay_request(f, block=True, run_scripthooks=False)
@property
def app_registry(self):
@ -139,8 +139,12 @@ def _handle_concurrent_reply(fn, o, *args, **kwargs):
def run():
fn(*args, **kwargs)
o.reply() # If the script did not call .reply(), we have to do it now.
threading.Thread(target=run, name="ScriptThread").start()
reply_proxy() # If the script did not call .reply(), we have to do it now.
ScriptThread(target=run).start()
class ScriptThread(threading.Thread):
name = "ScriptThread"
def concurrent(fn):

View File

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