Add WebSocketMessage.injected flag (#5105)

* Add WebSocketMessage.injected flag

* add flow format migration

Co-authored-by: Maximilian Hils <git@maximilianhils.com>

CHANGELOG.md

@@ -24,6 +24,7 @@
 * Fix random connection stalls (#5040, @EndUser509)
 * Add `n` new flow keybind to mitmweb (#5061, @ianklatzco)
 * Fix compatibility with BoringSSL (@pmoulton)
+* Added `WebSocketMessage.injected` flag (@Prinzhorn)
 * Add example addon for saving streamed data to individual files (@EndUser509)
 * Change connection event hooks to be blocking.
   Processing will only resume once the event hook has finished. (@Prinzhorn)

mitmproxy/io/compat.py

@@ -319,6 +319,17 @@ def convert_13_14(data):
     return data
 
 
+def convert_14_15(data):
+    data["version"] = 15
+    if data.get("websocket", None):
+        # Add "injected" attribute.
+        data["websocket"]["messages"] = [
+            msg + [False]
+            for msg in data["websocket"]["messages"]
+        ]
+    return data
+
+
 def _convert_dict_keys(o: Any) -> Any:
     if isinstance(o, dict):
         return {strutils.always_str(k): _convert_dict_keys(v) for k, v in o.items()}
@@ -380,6 +391,7 @@ converters = {
     11: convert_11_12,
     12: convert_12_13,
     13: convert_13_14,
+    14: convert_14_15,
 }
 
 

mitmproxy/proxy/layers/websocket.py

@@ -124,8 +124,10 @@ class WebsocketLayer(layer.Layer):
 
         if isinstance(event, events.ConnectionEvent):
             from_client = event.connection == self.context.client
+            injected = False
         elif isinstance(event, WebSocketMessageInjected):
             from_client = event.message.from_client
+            injected = True
         else:
             raise AssertionError(f"Unexpected event: {event}")
 
@@ -165,7 +167,7 @@ class WebsocketLayer(layer.Layer):
                     fragmentizer = Fragmentizer(src_ws.frame_buf, is_text)
                     src_ws.frame_buf = [b""]
 
-                    message = websocket.WebSocketMessage(typ, from_client, content)
+                    message = websocket.WebSocketMessage(typ, from_client, content, injected=injected)
                     self.flow.websocket.messages.append(message)
                     yield WebsocketMessageHook(self.flow)
 

mitmproxy/version.py

@@ -7,7 +7,7 @@ MITMPROXY = "mitmproxy " + VERSION
 
 # Serialization format version. This is displayed nowhere, it just needs to be incremented by one
 # for each change in the file format.
-FLOW_FORMAT_VERSION = 14
+FLOW_FORMAT_VERSION = 15
 
 
 def get_dev_version() -> str:

mitmproxy/websocket.py

@@ -14,7 +14,7 @@ from mitmproxy import stateobject
 from mitmproxy.coretypes import serializable
 from wsproto.frame_protocol import Opcode
 
-WebSocketMessageState = Tuple[int, bool, bytes, float, bool]
+WebSocketMessageState = Tuple[int, bool, bytes, float, bool, bool]
 
 
 class WebSocketMessage(serializable.Serializable):
@@ -47,6 +47,8 @@ class WebSocketMessage(serializable.Serializable):
     """Timestamp of when this message was received or created."""
     dropped: bool
     """True if the message has not been forwarded by mitmproxy, False otherwise."""
+    injected: bool
+    """True if the message was injected and did not originate from a client/server, False otherwise"""
 
     def __init__(
         self,
@@ -54,23 +56,25 @@ class WebSocketMessage(serializable.Serializable):
         from_client: bool,
         content: bytes,
         timestamp: Optional[float] = None,
-        killed: bool = False,
+        dropped: bool = False,
+        injected: bool = False,
     ) -> None:
         self.from_client = from_client
         self.type = Opcode(type)
         self.content = content
         self.timestamp: float = timestamp or time.time()
-        self.dropped = killed
+        self.dropped = dropped
+        self.injected = injected
 
     @classmethod
     def from_state(cls, state: WebSocketMessageState):
         return cls(*state)
 
     def get_state(self) -> WebSocketMessageState:
-        return int(self.type), self.from_client, self.content, self.timestamp, self.dropped
+        return int(self.type), self.from_client, self.content, self.timestamp, self.dropped, self.injected
 
     def set_state(self, state: WebSocketMessageState) -> None:
-        typ, self.from_client, self.content, self.timestamp, self.dropped = state
+        typ, self.from_client, self.content, self.timestamp, self.dropped, self.injected = state
         self.type = Opcode(typ)
 
     def __repr__(self):

test/mitmproxy/proxy/layers/test_websocket.py

@@ -438,6 +438,7 @@ def test_inject_message(ws_testdata):
     )
     assert flow.websocket.messages[-1].content == b"hello"
     assert flow.websocket.messages[-1].from_client is False
+    assert flow.websocket.messages[-1].injected is True
     assert (
             playbook
             >> reply()