diff options
| -rw-r--r-- | examples/nonblocking.py | 4 | ||||
| -rw-r--r-- | mitmproxy/script/concurrent.py | 44 | ||||
| -rw-r--r-- | test/mitmproxy/script/test_concurrent.py | 26 |
3 files changed, 19 insertions, 55 deletions
diff --git a/examples/nonblocking.py b/examples/nonblocking.py index 41674b2a..4609f389 100644 --- a/examples/nonblocking.py +++ b/examples/nonblocking.py @@ -4,6 +4,6 @@ from mitmproxy.script import concurrent @concurrent # Remove this and see what happens def request(context, flow): - print("handle request: %s%s" % (flow.request.host, flow.request.path)) + context.log("handle request: %s%s" % (flow.request.host, flow.request.path)) time.sleep(5) - print("start request: %s%s" % (flow.request.host, flow.request.path)) + context.log("start request: %s%s" % (flow.request.host, flow.request.path)) diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py index b81f2ab1..89c835f6 100644 --- a/mitmproxy/script/concurrent.py +++ b/mitmproxy/script/concurrent.py @@ -8,43 +8,6 @@ from mitmproxy import controller import threading -class ReplyProxy(object): - - def __init__(self, reply_func, script_thread): - self.reply_func = reply_func - self.script_thread = script_thread - self.master_reply = None - - def send(self, message): - if self.master_reply is None: - self.master_reply = message - self.script_thread.start() - return - self.reply_func(message) - - def done(self): - self.reply_func.send(self.master_reply) - - def __getattr__(self, k): - return getattr(self.reply_func, k) - - -def _handle_concurrent_reply(fn, o, *args, **kwargs): - # Make first call to o.reply a no op and start the script thread. - # We must not start the script thread before, as this may lead to a nasty race condition - # where the script thread replies a different response before the normal reply, which then gets swallowed. - - def run(): - fn(*args, **kwargs) - # If the script did not call .reply(), we have to do it now. - reply_proxy.done() - - script_thread = ScriptThread(target=run) - - reply_proxy = ReplyProxy(o.reply, script_thread) - o.reply = reply_proxy - - class ScriptThread(threading.Thread): name = "ScriptThread" @@ -56,5 +19,10 @@ def concurrent(fn): ) def _concurrent(ctx, obj): - _handle_concurrent_reply(fn, obj, ctx, obj) + def run(): + fn(ctx, obj) + if not obj.reply.acked: + obj.reply.ack() + obj.reply.take() + ScriptThread(target=run).start() return _concurrent diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index c2f169ad..62541f3f 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -1,29 +1,25 @@ -from threading import Event - from mitmproxy.script import Script from test.mitmproxy import tutils +from mitmproxy import controller +import time -class Dummy: - def __init__(self, reply): - self.reply = reply +class Thing: + def __init__(self): + self.reply = controller.DummyReply() @tutils.skip_appveyor def test_concurrent(): with Script(tutils.test_data.path("data/scripts/concurrent_decorator.py"), None) as s: - def reply(): - reply.acked.set() - reply.acked = Event() - - f1, f2 = Dummy(reply), Dummy(reply) + f1, f2 = Thing(), Thing() s.run("request", f1) - f1.reply() s.run("request", f2) - f2.reply() - assert f1.reply.acked == reply.acked - assert not reply.acked.is_set() - assert reply.acked.wait(10) + start = time.time() + while time.time() - start < 5: + if f1.reply.acked and f2.reply.acked: + return + raise ValueError("Script never acked") def test_concurrent_err(): |