fix Flow.kill behaviour

This now just sets a kill reply instead of committing directly.
First, this seems like the more sane thing to do.
Second, we have an iffy race condition where we call Reply.commit()
before the addonmanager finishes its invocation, the proxy thread then progresses
and sets a new flow.reply attribute, and the addonmanager then gets confused
when finishing. This commit doesn't fix that, but mitigates it for Flow.kill
which is now committed by the addonmanager.
This commit is contained in:
Maximilian Hils 2017-12-29 22:44:18 +01:00 committed by Thomas Kriechbaumer
parent 59c277effd
commit 6232622774
4 changed files with 15 additions and 18 deletions

View File

@ -230,7 +230,7 @@ class AddonManager:
self.trigger(name, message)
if message.reply.state != "taken":
if message.reply.state == "start":
message.reply.take()
if not message.reply.has_message:
message.reply.ack()

View File

@ -105,16 +105,16 @@ class Reply:
self.q.put(self.value)
def ack(self, force=False):
if self.state not in {"start", "taken"}:
raise exceptions.ControlException(
"Reply is {}, but expected it to be start or taken.".format(self.state)
)
self.send(self.obj, force)
def kill(self, force=False):
self.send(exceptions.Kill, force)
def send(self, msg, force=False):
if self.state not in {"start", "taken"}:
raise exceptions.ControlException(
"Reply is {}, but expected it to be start or taken.".format(self.state)
)
if self.has_message and not force:
raise exceptions.ControlException("There is already a reply message.")
self.value = msg

View File

@ -1,13 +1,12 @@
import time
import typing # noqa
import uuid
from mitmproxy import controller # noqa
from mitmproxy import stateobject
from mitmproxy import connections
from mitmproxy import controller, exceptions # noqa
from mitmproxy import stateobject
from mitmproxy import version
import typing # noqa
class Error(stateobject.StateObject):
@ -145,7 +144,11 @@ class Flow(stateobject.StateObject):
@property
def killable(self):
return self.reply and self.reply.state == "taken"
return (
self.reply and
self.reply.state in {"start", "taken"} and
self.reply.value != exceptions.Kill
)
def kill(self):
"""
@ -153,13 +156,7 @@ class Flow(stateobject.StateObject):
"""
self.error = Error("Connection killed")
self.intercepted = False
# reply.state should be "taken" here, or .take() will raise an
# exception.
if self.reply.state != "taken":
self.reply.take()
self.reply.kill(force=True)
self.reply.commit()
self.live = False
def intercept(self):

View File

@ -236,8 +236,8 @@ class TestKillFlow(_WebSocketTest):
self.master.addons.add(KillFlow())
self.setup_connection()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'foo'
with pytest.raises(exceptions.TcpDisconnect):
websockets.Frame.from_file(self.client.rfile)
class TestSimpleTLS(_WebSocketTest):