Method relay_messages

mitmproxy/proxy/layers/websocket.py:137–220

(self, event: events.Event)

    @expect(events.DataReceived, events.ConnectionClosed, WebSocketMessageInjected)
    def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:
        assert self.flow.websocket  # satisfy type checker

        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}")

        from_str = "client" if from_client else "server"
        if from_client:
            src_ws = self.client_ws
            dst_ws = self.server_ws
        else:
            src_ws = self.server_ws
            dst_ws = self.client_ws

        if isinstance(event, events.DataReceived):
            src_ws.receive_data(event.data)
        elif isinstance(event, events.ConnectionClosed):
            src_ws.receive_data(None)
        elif isinstance(event, WebSocketMessageInjected):
            fragmentizer = Fragmentizer([], event.message.type == Opcode.TEXT)
            src_ws._events.extend(fragmentizer(event.message.content))
        else:  # pragma: no cover
            raise AssertionError(f"Unexpected event: {event}")

        for ws_event in src_ws.events():
            if isinstance(ws_event, wsproto.events.Message):
                is_text = isinstance(ws_event.data, str)
                if is_text:
                    typ = Opcode.TEXT
                    src_ws.frame_buf[-1] += ws_event.data.encode()
                else:
                    typ = Opcode.BINARY
                    src_ws.frame_buf[-1] += ws_event.data

                if ws_event.message_finished:
                    content = b"".join(src_ws.frame_buf)

                    fragmentizer = Fragmentizer(src_ws.frame_buf, is_text)
                    src_ws.frame_buf = [b""]

                    message = websocket.WebSocketMessage(
                        typ, from_client, content, injected=injected
                    )
                    self.flow.websocket.messages.append(message)
                    yield WebsocketMessageHook(self.flow)

                    if not message.dropped:
                        for msg in fragmentizer(message.content):
                            yield dst_ws.send2(msg)

                elif ws_event.frame_finished:
                    src_ws.frame_buf.append(b"")