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.

Resolving patch:
diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py
--- a/mitmproxy/proxy/protocol/websocket.py
+++ b/mitmproxy/proxy/protocol/websocket.py
@@ -109,7 +109,7 @@ def _handle_data_received(self, event, source_conn, other_conn, is_server):
             self.flow.messages.append(websocket_message)
             self.channel.ask("websocket_message", self.flow)
 
-            if not self.flow.stream:
+            if not self.flow.stream and not websocket_message.killed:
                 def get_chunk(payload):
                     if len(payload) == length:
                         # message has the same length, we can reuse the same sizes
@@ -129,14 +129,9 @@ def get_chunk(payload):
                     self.connections[other_conn].send_data(chunk, final)
                     other_conn.send(self.connections[other_conn].bytes_to_send())
 
-            else:
-                self.connections[other_conn].send_data(event.data, event.message_finished)
-                other_conn.send(self.connections[other_conn].bytes_to_send())
-
-        elif self.flow.stream:
+        if self.flow.stream:
             self.connections[other_conn].send_data(event.data, event.message_finished)
             other_conn.send(self.connections[other_conn].bytes_to_send())
-
         return True
 
     def _handle_ping_received(self, event, source_conn, other_conn, is_server):
diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -10,23 +10,33 @@
 
 
 class WebSocketMessage(serializable.Serializable):
+    """
+    A WebSocket message sent from one endpoint to the other.
+    """
+
     def __init__(
-        self, type: int, from_client: bool, content: bytes, timestamp: Optional[int]=None
+        self, type: int, from_client: bool, content: bytes, timestamp: Optional[int]=None, killed: bool=False
     ) -> None:
         self.type = wsproto.frame_protocol.Opcode(type)  # type: ignore
+        """indicates either TEXT or BINARY (from wsproto.frame_protocol.Opcode)."""
         self.from_client = from_client
+        """True if this messages was sent by the client."""
         self.content = content
+        """A byte-string representing the content of this message."""
         self.timestamp = timestamp or int(time.time())  # type: int
+        """Timestamp of when this message was received or created."""
+        self.killed = killed
+        """True if this messages was killed and should not be sent to the other endpoint."""
 
     @classmethod
     def from_state(cls, state):
         return cls(*state)
 
     def get_state(self):
-        return int(self.type), self.from_client, self.content, self.timestamp
+        return int(self.type), self.from_client, self.content, self.timestamp, self.killed
 
     def set_state(self, state):
-        self.type, self.from_client, self.content, self.timestamp = state
+        self.type, self.from_client, self.content, self.timestamp, self.killed = state
         self.type = wsproto.frame_protocol.Opcode(self.type)  # replace enum with bare int
 
     def __repr__(self):
@@ -35,20 +45,37 @@ def __repr__(self):
         else:
             return "binary message: {}".format(strutils.bytes_to_escaped_str(self.content))
 
+    def kill(self):
+        """
+        Kill this message.
+
+        It will not be sent to the other endpoint. This has no effect in streaming mode.
+        """
+        self.killed = True
+
 
 class WebSocketFlow(flow.Flow):
     """
-    A WebsocketFlow is a simplified representation of a Websocket session.
+    A WebsocketFlow is a simplified representation of a Websocket connection.
     """
 
     def __init__(self, client_conn, server_conn, handshake_flow, live=None):
         super().__init__("websocket", client_conn, server_conn, live)
+
         self.messages = []  # type: List[WebSocketMessage]
+        """A list containing all WebSocketMessage's."""
         self.close_sender = 'client'
+        """'client' if the client initiated connection closing."""
         self.close_code = wsproto.frame_protocol.CloseReason.NORMAL_CLOSURE
+        """WebSocket close code."""
         self.close_message = '(message missing)'
+        """WebSocket close message."""
         self.close_reason = 'unknown status code'
+        """WebSocket close reason."""
         self.stream = False
+        """True of this connection is streaming directly to the other endpoint."""
+        self.handshake_flow = handshake_flow
+        """The HTTP flow containing the initial WebSocket handshake."""
 
         if handshake_flow:
             self.client_key = websockets.get_client_key(handshake_flow.request.headers)
@@ -65,8 +92,6 @@ def __init__(self, client_conn, server_conn, handshake_flow, live=None):
             self.server_protocol = ''
             self.server_extensions = ''
 
-        self.handshake_flow = handshake_flow
-
     _stateobject_attributes = flow.Flow._stateobject_attributes.copy()
     # mypy doesn't support update with kwargs
     _stateobject_attributes.update(dict(


Test patch:
diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py
--- a/test/mitmproxy/test_websocket.py
+++ b/test/mitmproxy/test_websocket.py
@@ -3,6 +3,7 @@
 
 from mitmproxy.io import tnetstring
 from mitmproxy import flowfilter
+from mitmproxy.exceptions import Kill, ControlException
 from mitmproxy.test import tflow
 
 
@@ -42,6 +43,20 @@ def test_copy(self):
         assert f.error.get_state() == f2.error.get_state()
         assert f.error is not f2.error
 
+    def test_kill(self):
+        f = tflow.twebsocketflow()
+        with pytest.raises(ControlException):
+            f.intercept()
+            f.resume()
+            f.kill()
+
+        f = tflow.twebsocketflow()
+        f.intercept()
+        assert f.killable
+        f.kill()
+        assert not f.killable
+        assert f.reply.value == Kill
+
     def test_match(self):
         f = tflow.twebsocketflow()
         assert not flowfilter.match("~b nonexistent", f)
@@ -71,3 +86,9 @@ def test_serialize(self):
         d = tflow.twebsocketflow().handshake_flow.get_state()
         tnetstring.dump(d, b)
         assert b.getvalue()
+
+    def test_message_kill(self):
+        f = tflow.twebsocketflow()
+        assert not f.messages[-1].killed
+        f.messages[-1].kill()
+        assert f.messages[-1].killed