Repository: mitmproxy/mitmproxy

Issue:
"AssertionError: Unexpected event type at HttpStream.state_done" in presence of transformation function
#### Problem Description

```
127.0.0.1:50894: GET https://alive.github.com/_sockets/u/<redacted>/ws?session=<redacted>
              << 101 Switching Protocols (content missing)
127.0.0.1:50894: mitmproxy has crashed!
Traceback (most recent call last):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/server.py", line 279, in server_event
    for command in layer_commands:
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 255, in handle_event
    yield from self._handle(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 699, in _handle_event
    yield from self.event_to_child(stream, event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 749, in event_to_child
    for command in child.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 617, in passthrough
    for command in self.child_layer.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/tunnel.py", line 85, in _handle_event
    yield from self.event_to_child(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/tls.py", line 320, in event_to_child
    yield from super().event_to_child(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/tunnel.py", line 104, in event_to_child
    for command in self.child_layer.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/tunnel.py", line 85, in _handle_event
    yield from self.event_to_child(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/tunnel.py", line 104, in event_to_child
    for command in self.child_layer.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 699, in _handle_event
    yield from self.event_to_child(stream, event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 764, in event_to_child
    yield from self.event_to_child(conn, command.event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 761, in event_to_child
    yield from self.event_to_child(stream, command.event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 749, in event_to_child
    for command in child.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 136, in _handle_event
    yield from self.client_state(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/utils.py", line 23, in _check_event_type
    raise AssertionError(
AssertionError: Unexpected event type at HttpStream.state_done: Expected no events, got RequestData(stream_id=1, data=bytearray(b'<redacted>')).

127.0.0.1:50874: mitmproxy has crashed!
Traceback (most recent call last):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/server.py", line 279, in server_event
    for command in layer_commands:
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 255, in handle_event
    yield from self._handle(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 740, in _handle_event
    yield from self.event_to_child(handler, event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 749, in event_to_child
    for command in child.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 617, in passthrough
    for command in self.child_layer.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/tunnel.py", line 72, in _handle_event
    yield from self.receive_data(event.data)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/tls.py", line 261, in receive_data
    yield from self.event_to_child(
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/tls.py", line 320, in event_to_child
    yield from super().event_to_child(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/tunnel.py", line 104, in event_to_child
    for command in self.child_layer.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/tunnel.py", line 85, in _handle_event
    yield from self.event_to_child(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/tunnel.py", line 104, in event_to_child
    for command in self.child_layer.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 740, in _handle_event
    yield from self.event_to_child(handler, event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 761, in event_to_child
    yield from self.event_to_child(stream, command.event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 749, in event_to_child
    for command in child.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 138, in _handle_event
    yield from self.server_state(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/utils.py", line 23, in _check_event_type
    raise AssertionError(
AssertionError: Unexpected event type at HttpStream.state_done: Expected no events, got ResponseEndOfMessage(stream_id=1).
```

Similar exceptions also pop up for, among other things, [https://push.services.mozilla.com/]() and [https://qa.sockets.stackexchange.com/]().

#### Steps to reproduce the behavior:

1.
```python
from time import sleep

class Test:
        def _process_chunk(self, data):
                sleep(0.01) # not reproducible without this
                return data # noop
        def responseheaders(self, flow):
                flow.response.stream = self._process_chunk

addons = [
        Test(),
]
```
2. Visit GitHub, StackExchange, or wait for Firefox to phone home.

#### System Information

```
Mitmproxy: 7.0.4
Python:    3.9.1
OpenSSL:   OpenSSL 1.1.1g  21 Apr 2020
Platform:  Linux-5.10.13-gentoo-x86_64-AMD_Ryzen_5_PRO_4650G_with_Radeon_Graphics-with-glibc2.32
```


Discussion:
Reproducible locally:

```sh
while :; do date; sleep .5; done | websocat -b ws-l:0.0.0.0:12345 stdio:
```

(Not reproducible without binary server messages, hence `-b`)

```html
<link rel="icon" type="image/gif" href="data:image/gif;base64,R0lGODlhAQABAPAAAP///wAAACH5BAAAAAAALAAAAAABAAEAAAICRAEAOw=="> <!-- don't spam logs with missing favicons -->
<script>
let s = new WebSocket('ws://192.168.1.2:12345/ws');

s.onopen = function(e) {
        console.log('opened');
        s.send('asdf');
};

s.onmessage = function(ev) {
        console.log('rcvd', ev.data);
};

s.onclose = function(event) {
        console.log('closed');
};

s.onerror = function(err) {
        console.log('error', err.message);
};
</script>
```

(Not reproducible without initial `send()`)

JS console:

```
opened
closed
```

(note lack of rcvd messages)

```
127.0.0.1:38420: server connect 192.168.1.2:12345
127.0.0.1:38420: Streaming response from 192.168.1.2.
127.0.0.1:38420: GET http://192.168.1.2:12345/ws
    Host: 192.168.1.2:12345
    <...snip...>
    Upgrade: websocket
 << 101 Switching Protocols (content missing)
    Sec-WebSocket-Accept: 5jIIDxDZEKGf9M1XZ1XIHtUDZzY=
    Connection: Upgrade
    Upgrade: websocket
127.0.0.1:38420: mitmproxy has crashed!
Traceback (most recent call last):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/server.py", line 279, in server_event
    for command in layer_commands:
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 255, in handle_event
    yield from self._handle(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 699, in _handle_event
    yield from self.event_to_child(stream, event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 749, in event_to_child
    for command in child.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 617, in passthrough
    for command in self.child_layer.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 699, in _handle_event
    yield from self.event_to_child(stream, event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 764, in event_to_child
    yield from self.event_to_child(conn, command.event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 761, in event_to_child
    yield from self.event_to_child(stream, command.event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 749, in event_to_child
    for command in child.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 136, in _handle_event
    yield from self.client_state(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/utils.py", line 23, in _check_event_type
    raise AssertionError(
AssertionError: Unexpected event type at HttpStream.state_done: Expected no events, got RequestData(stream_id=1, data=bytearray(b'\x81\x84r\xb8\xbfG\x13\xcb\xdb!')).

127.0.0.1:38420: mitmproxy has crashed!
Traceback (most recent call last):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/server.py", line 279, in server_event
    for command in layer_commands:
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 255, in handle_event
    yield from self._handle(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 740, in _handle_event
    yield from self.event_to_child(handler, event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 749, in event_to_child
    for command in child.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 617, in passthrough
    for command in self.child_layer.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 740, in _handle_event
    yield from self.event_to_child(handler, event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 761, in event_to_child
    yield from self.event_to_child(stream, command.event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 749, in event_to_child
    for command in child.handle_event(event):
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layer.py", line 144, in handle_event
    command = command_generator.send(send)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/layers/http/__init__.py", line 138, in _handle_event
    yield from self.server_state(event)
  File "/usr/lib/python3.9/site-packages/mitmproxy/proxy/utils.py", line 23, in _check_event_type
    raise AssertionError(
AssertionError: Unexpected event type at HttpStream.state_done: Expected no events, got ResponseData(stream_id=1, data=b'\x82 Mon Dec 27 06:28:29 PM MSK 2021\n').
```
It's worth noting that I also see random connection hangs and glitches that don't seem to be related to websockets at all, but those are way harder to reproduce.

Resolving patch:
diff --git a/mitmproxy/proxy/layers/http/__init__.py b/mitmproxy/proxy/layers/http/__init__.py
--- a/mitmproxy/proxy/layers/http/__init__.py
+++ b/mitmproxy/proxy/layers/http/__init__.py
@@ -389,7 +389,6 @@ def send_response(self, already_streamed: bool = False):
 
         if self.flow.response.trailers:
             yield SendHttp(ResponseTrailers(self.stream_id, self.flow.response.trailers), self.context.client)
-        yield SendHttp(ResponseEndOfMessage(self.stream_id), self.context.client)
 
         if self.flow.response.status_code == 101:
             if is_websocket:
@@ -405,7 +404,10 @@ def send_response(self, already_streamed: bool = False):
                 yield commands.Log(f"{self.debug}[http] upgrading to {self.child_layer}", "debug")
             yield from self.child_layer.handle_event(events.Start())
             self._handle_event = self.passthrough
-            return
+
+        # delay sending EOM until the child layer is set up,
+        # we may get data immediately and need to be prepared to handle it.
+        yield SendHttp(ResponseEndOfMessage(self.stream_id), self.context.client)
 
     def check_body_size(self, request: bool) -> layer.CommandGenerator[bool]:
         """


Test patch:
diff --git a/test/mitmproxy/proxy/layers/test_websocket.py b/test/mitmproxy/proxy/layers/test_websocket.py
--- a/test/mitmproxy/proxy/layers/test_websocket.py
+++ b/test/mitmproxy/proxy/layers/test_websocket.py
@@ -109,6 +109,61 @@ def test_upgrade(tctx):
     assert flow().websocket.messages[1].type == Opcode.BINARY
 
 
+def test_upgrade_streamed(tctx):
+    """If the HTTP response is streamed, we may get early data from the client."""
+    tctx.server.address = ("example.com", 80)
+    tctx.server.state = ConnectionState.OPEN
+    flow = Placeholder(HTTPFlow)
+
+    def enable_streaming(flow: HTTPFlow):
+        flow.response.stream = True
+
+    assert (
+            Playbook(http.HttpLayer(tctx, HTTPMode.transparent))
+            >> DataReceived(tctx.client,
+                            b"GET / HTTP/1.1\r\n"
+                            b"Connection: upgrade\r\n"
+                            b"Upgrade: websocket\r\n"
+                            b"Sec-WebSocket-Version: 13\r\n"
+                            b"\r\n")
+            << http.HttpRequestHeadersHook(flow)
+            >> reply()
+            << http.HttpRequestHook(flow)
+            >> reply()
+            << SendData(tctx.server, b"GET / HTTP/1.1\r\n"
+                                     b"Connection: upgrade\r\n"
+                                     b"Upgrade: websocket\r\n"
+                                     b"Sec-WebSocket-Version: 13\r\n"
+                                     b"\r\n")
+            >> DataReceived(tctx.server, b"HTTP/1.1 101 Switching Protocols\r\n"
+                                         b"Upgrade: websocket\r\n"
+                                         b"Connection: Upgrade\r\n"
+                                         b"\r\n")
+            << http.HttpResponseHeadersHook(flow)
+            >> reply(side_effect=enable_streaming)
+            << SendData(tctx.client, b"HTTP/1.1 101 Switching Protocols\r\n"
+                                     b"Upgrade: websocket\r\n"
+                                     b"Connection: Upgrade\r\n"
+                                     b"\r\n")
+            << http.HttpResponseHook(flow)
+            >> DataReceived(tctx.client, masked_bytes(b"\x81\x0bhello world"))  # early !!
+            >> reply(to=-2)
+            << websocket.WebsocketStartHook(flow)
+            >> reply()
+            << websocket.WebsocketMessageHook(flow)
+            >> reply()
+            << SendData(tctx.server, masked(b"\x81\x0bhello world"))
+            >> DataReceived(tctx.server, b"\x82\nhello back")
+            << websocket.WebsocketMessageHook(flow)
+            >> reply()
+            << SendData(tctx.client, b"\x82\nhello back")
+            >> DataReceived(tctx.client, masked_bytes(b"\x81\x0bhello again"))
+            << websocket.WebsocketMessageHook(flow)
+            >> reply()
+            << SendData(tctx.server, masked(b"\x81\x0bhello again"))
+    )
+
+
 @pytest.fixture()
 def ws_testdata(tctx):
     tctx.server.address = ("example.com", 80)