Repository: mitmproxy/mitmproxy

Issue:
Proxy crashes on flow.kill()
##### Steps to reproduce the problem:

0. (Sorry if this is my own mistake, I just started poking around with mitmproxy)
1. Load the proxy with the scenario which is supposed to kill every websocket request with 50% chance
```
from mitmproxy import ctx, websocket
import random


def websocket_message(flow: websocket.WebSocketFlow) -> None:
    ctx.log.info(flow.messages[-1].content.strip()[:20])
    if random.random() < 50:
        flow.kill()
```
2. Start the test which opens a websocket connection through the proxy and sends several binary chunks of data.
=> Proxy crashes

```
Traceback (most recent call last):
  File "/Users/eugene/gitrepo/mitmproxy/mitmproxy/tools/console/master.py", line 216, in run
    self.loop.run()
  File "/Users/eugene/gitrepo/mitmproxy/venv/lib/python3.6/site-packages/urwid/main_loop.py", line 278, in run
    self._run()
  File "/Users/eugene/gitrepo/mitmproxy/venv/lib/python3.6/site-packages/urwid/main_loop.py", line 376, in _run
    self.event_loop.run()
  File "/Users/eugene/gitrepo/mitmproxy/venv/lib/python3.6/site-packages/urwid/main_loop.py", line 682, in run
    self._loop()
  File "/Users/eugene/gitrepo/mitmproxy/venv/lib/python3.6/site-packages/urwid/main_loop.py", line 715, in _loop
    alarm_callback()
  File "/Users/eugene/gitrepo/mitmproxy/venv/lib/python3.6/site-packages/urwid/main_loop.py", line 164, in cb
    callback(self, user_data)
  File "/Users/eugene/gitrepo/mitmproxy/mitmproxy/tools/console/master.py", line 175, in ticker
    changed = self.tick(timeout=0)
  File "/Users/eugene/gitrepo/mitmproxy/mitmproxy/master.py", line 109, in tick
    self.addons.handle_lifecycle(mtype, obj)
  File "/Users/eugene/gitrepo/mitmproxy/mitmproxy/addonmanager.py", line 234, in handle_lifecycle
    message.reply.take()
  File "/Users/eugene/gitrepo/mitmproxy/mitmproxy/controller.py", line 88, in take
    "Reply is {}, but expected it to be start.".format(self.state)
mitmproxy.exceptions.ControlException: Reply is committed, but expected it to be start.

mitmproxy has crashed!
```

```
Mitmproxy: 3.0.0dev1075 (a6a4b1c3) 
Python:    3.6.3
OpenSSL:   OpenSSL 1.1.0g  2 Nov 2017
Platform:  Darwin-17.2.0-x86_64-i386-64bit
```


Discussion:
Do you want to kill a websocket connection with 50% probability or prevent exchanging messages with 50%?
`flow.kill` kills the whole connection - and we might have a bug here because websocket is a bit different than regular HTTP.

@mhils: how should this behave? Can we kill individual websocket messages?
What is the expected behaviour for TCP flows/streams?
I would like to prevent message exchange at 50%. The API documentation is absent on `websocket.WebSocketFlow` so I took it for granted it could be more or less the same as `HTTPFlow` and `HTTPFlow.kill()` says it kills the request.

I take it `kill_flow()` should have aborted the connection.
@Kriechi I'm sorry to trod on the issue, but the original reported problem about the crash on `flow.kill()` is still present in the branch where you changed the websocket API/docs.

@mhils suggests to use `flow.kill()` on the [forums](https://discourse.mitmproxy.org/t/killing-tcp-connections/775/3), but I refer to this issue as a blocker.

Should there be a new API for TCP killing as well? If not, then this issue is still relevant.
I tried to reproduce it - but unfortunately I don't see the traceback you mentioned.
I'm not sure about TCP flows, @mhils?
WebSocketFlows should be killable with `flow.kill()`, and it should not result in a traceback. Killing a WebSocketFlow will terminate the entire connection (websocket session & underlying socket).

Are you testing with a specific website / domain we could use to debug this?
This issue reproduces on a local environment, which I am not free to share, but I can provide with both the message emitter scenario (Python 2.7):

```
from websocket import create_connection


def getResponse():
    ws = create_connection('ws://localhost:9000/listen',
                           http_proxy_host="localhost",
                           http_proxy_port=8080)
    listen = """{
  "type":"LISTEN",
  "id": "42",
  "data":{
    "lang": "en-US",
    "hotphrase": false,
    "rules": []
  }
}"""
    ws.send(listen)

    ws.close()
    print("done")


if __name__ == "__main__":
    getResponse()
```

and the MITM script:
```
from mitmproxy import ctx, websocket
import random

global i


def websocket_message(flow: websocket.WebSocketFlow) -> None:
    global i
    i += 1
    ctx.log.info(i)
    flow.kill()

def websocket_end(flow: websocket.WebSocketFlow) -> None:
    global i
    ctx.log.info("WebSocket messages sent in total: " + str(i))
```

From the stacktrace, I gather that the message is in committed state too early, when it should be in the start state (if I am reading this right).

Is there an entrypoint to launch MITM through an IDE? I might be able to debug it.

Resolving patch:
diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
--- a/mitmproxy/addonmanager.py
+++ b/mitmproxy/addonmanager.py
@@ -230,7 +230,7 @@ def handle_lifecycle(self, name, message):
 
         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
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -105,16 +105,16 @@ def commit(self):
         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
--- 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 @@ def revert(self):
 
     @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 @@ def kill(self):
         """
         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):


Test patch:
diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py
--- a/test/mitmproxy/proxy/protocol/test_websocket.py
+++ b/test/mitmproxy/proxy/protocol/test_websocket.py
@@ -221,6 +221,25 @@ def websocket_message(self, f):
         assert frame.payload == b'foo'
 
 
+class TestKillFlow(_WebSocketTest):
+
+    @classmethod
+    def handle_websockets(cls, rfile, wfile):
+        wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'server-foobar')))
+        wfile.flush()
+
+    def test_kill(self):
+        class KillFlow:
+            def websocket_message(self, f):
+                f.kill()
+
+        self.master.addons.add(KillFlow())
+        self.setup_connection()
+
+        with pytest.raises(exceptions.TcpDisconnect):
+            websockets.Frame.from_file(self.client.rfile)
+
+
 class TestSimpleTLS(_WebSocketTest):
     ssl = True