mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 00:01:36 +00:00
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:
parent
59c277effd
commit
6232622774
@ -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()
|
||||
|
@ -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
|
||||
|
@ -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):
|
||||
|
@ -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):
|
||||
|
Loading…
Reference in New Issue
Block a user