Merge pull request #2732 from mitmproxy/websocket-flow-kill

fix flow.kill

3 files changed, 13 insertions, 16 deletions

diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
index 70cfda30..37c501ee 100644
--- a/mitmproxy/addonmanager.py
+++ b/mitmproxy/addonmanager.py
@@ -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()
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index 63117ef0..f39c1b24 100644
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -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
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 111566b8..944c032d 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -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):