websocket: docs++ and kill messages

fixes #2724

5 files changed, 53 insertions, 15 deletions

diff --git a/docs/scripting/api.rst b/docs/scripting/api.rst
index e82afef4..368b9ba8 100644
--- a/docs/scripting/api.rst
+++ b/docs/scripting/api.rst
@@ -10,6 +10,9 @@ API
     - `mitmproxy.http.HTTPRequest <#mitmproxy.http.HTTPRequest>`_
     - `mitmproxy.http.HTTPResponse <#mitmproxy.http.HTTPResponse>`_
     - `mitmproxy.http.HTTPFlow <#mitmproxy.http.HTTPFlow>`_
+- WebSocket
+    - `mitmproxy.websocket.WebSocketFlow <#mitmproxy.websocket.WebSocketFlow>`_
+    - `mitmproxy.websocket.WebSocketMessage <#mitmproxy.websocket.WebSocketMessage>`_
 - Logging
     - `mitmproxy.log.Log <#mitmproxy.controller.Log>`_
     - `mitmproxy.log.LogEntry <#mitmproxy.controller.LogEntry>`_
@@ -33,6 +36,15 @@ HTTP
 .. autoclass:: mitmproxy.http.HTTPFlow
     :inherited-members:
 
+WebSocket
+---------
+
+.. autoclass:: mitmproxy.websocket.WebSocketFlow
+    :inherited-members:
+
+.. autoclass:: mitmproxy.websocket.WebSocketMessage
+    :inherited-members:
+
 Logging
 --------
 
diff --git a/docs/scripting/events.rst b/docs/scripting/events.rst
index 9e84dacf..4d74b220 100644
--- a/docs/scripting/events.rst
+++ b/docs/scripting/events.rst
@@ -187,8 +187,8 @@ are issued, only new WebSocket messages are called.
 
         - Called when a WebSocket message is received from the client or server. The
           sender and receiver are identifiable. The most recent message will be
-          ``flow.messages[-1]``. The message is user-modifiable. Currently there are
-          two types of messages, corresponding to the BINARY and TEXT frame types.
+          ``flow.messages[-1]``. The message is user-modifiable and is killable.
+          A message is either of TEXT or BINARY type.
 
           *flow*
             A ``models.WebSocketFlow`` object.
diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py
index 1bd5284d..92f99518 100644
--- a/mitmproxy/proxy/protocol/websocket.py
+++ b/mitmproxy/proxy/protocol/websocket.py
@@ -109,7 +109,7 @@ class WebSocketLayer(base.Layer):
             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 @@ class WebSocketLayer(base.Layer):
                     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
index 8efd4117..a37edb54 100644
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -10,23 +10,33 @@ from mitmproxy.utils import strutils, human
 
 
 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 @@ class WebSocketMessage(serializable.Serializable):
         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 @@ class WebSocketFlow(flow.Flow):
             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(
diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py
index 7c53a4b0..0a37f370 100644
--- a/test/mitmproxy/test_websocket.py
+++ b/test/mitmproxy/test_websocket.py
@@ -71,3 +71,9 @@ class TestWebSocketFlow:
         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