Repository: mitmproxy/mitmproxy
Pull request: fix huge (>65kb) http2 responses corrupted
#### Description

BufferedH2Connection is able to buffer huge response and send it to client, properly following flow control specs.
But trailers are not subject to flow control and were mistankenly sent immediate.
That led to the situation that mitmproxy worked fine on small responses (that fit in 65k bytes - initial size of flow control window) and broke on huge responses.

With this changes I delay trailers until all data gets sent.
I also altered `end_stream` method because it is hard to expose proper stream state to `Http2Connection.is_closed` now.
Without the change an empty data frame with `end_stream` flag was sent before trailers and it also violated protocol.

Fixes #5428

I will remove excessive logging and add tests/changelog in the next iterations.
@mhils are you ok with the solution?

#### Checklist

 - [x] I have updated tests where applicable.
 - [x] I have added an entry to the CHANGELOG.


Problem statement:
[grpc] mitmproxy corrupts huge (>65kb) responses
#### Problem Description
Grpc requests with huge responses fails with error `io.grpc.StatusException: INTERNAL: Encountered end-of-stream mid-frame` if mitmproxy is used. Requests succeed without mitmproxy.
In my case response was 200k bytes.


#### Steps to reproduce the behavior:
1. Just make any grpc request that produces huge response


#### System Information
Mitmproxy: 8.0.0
Python:    3.10.2
OpenSSL:   OpenSSL 1.1.1n  15 Mar 2022
Platform:  macOS-10.16-x86_64-i386-64bit

Mitmproxy: 9.0.0.dev (+13, commit 3eab9f0)
Python:    3.9.12
OpenSSL:   OpenSSL 3.0.3 3 May 2022
Platform:  macOS-10.16-x86_64-i386-64bit

Hints from issue discussion:

Commit message: add changelog for #5428

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d54e35a0fe..35078673c1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,8 @@
   ([#5372](https://github.com/mitmproxy/mitmproxy/discussions/5372), @jixunmoe)
 * Drop pkg_resources dependency.
   ([#5401](https://github.com/mitmproxy/mitmproxy/issues/5401), @PavelICS)
+* Fix huge (>65kb) http2 responses corrupted
+  ([#5428](https://github.com/mitmproxy/mitmproxy/issues/5428), @dhabensky)
 
 ## 15 May 2022: mitmproxy 8.1.0
 
diff --git a/mitmproxy/proxy/layers/http/_http2.py b/mitmproxy/proxy/layers/http/_http2.py
index 3e34c325da..a014c3efbb 100644
--- a/mitmproxy/proxy/layers/http/_http2.py
+++ b/mitmproxy/proxy/layers/http/_http2.py
@@ -117,9 +117,7 @@ def _handle_event(self, event: Event) -> CommandGenerator[None]:
             elif isinstance(event, (RequestTrailers, ResponseTrailers)):
                 if self.is_open_for_us(event.stream_id):
                     trailers = [*event.trailers.fields]
-                    self.h2_conn.send_headers(
-                        event.stream_id, trailers, end_stream=True
-                    )
+                    self.h2_conn.send_trailers(event.stream_id, trailers)
             elif isinstance(event, (RequestEndOfMessage, ResponseEndOfMessage)):
                 if self.is_open_for_us(event.stream_id):
                     self.h2_conn.end_stream(event.stream_id)
diff --git a/mitmproxy/proxy/layers/http/_http_h2.py b/mitmproxy/proxy/layers/http/_http_h2.py
index 894de1496e..ff8f589f48 100644
--- a/mitmproxy/proxy/layers/http/_http_h2.py
+++ b/mitmproxy/proxy/layers/http/_http_h2.py
@@ -1,5 +1,5 @@
 import collections
-from typing import NamedTuple
+from typing import Any, Dict, NamedTuple
 
 import h2.config
 import h2.connection
@@ -34,10 +34,12 @@ class BufferedH2Connection(h2.connection.H2Connection):
     """
 
     stream_buffers: collections.defaultdict[int, collections.deque[SendH2Data]]
+    stream_trailers: Dict[int, Any]
 
     def __init__(self, config: h2.config.H2Configuration):
         super().__init__(config)
         self.stream_buffers = collections.defaultdict(collections.deque)
+        self.stream_trailers = {}
 
     def send_data(
         self,
@@ -77,8 +79,16 @@ def send_data(
                 # We can't send right now, so we buffer.
                 self.stream_buffers[stream_id].append(SendH2Data(data, end_stream))
 
+    def send_trailers(self, stream_id, trailers):
+        if self.stream_buffers.get(stream_id, None):
+            # Though trailers are not subject to flow control, we need to queue them and send strictly after data frames
+            self.stream_trailers[stream_id] = trailers
+        else:
+            self.send_headers(stream_id, trailers, end_stream=True)
+
     def end_stream(self, stream_id) -> None:
-        self.send_data(stream_id, b"", end_stream=True)
+        if not self.stream_trailers.get(stream_id):
+            self.send_data(stream_id, b"", end_stream=True)
 
     def reset_stream(self, stream_id: int, error_code: int = 0) -> None:
         self.stream_buffers.pop(stream_id, None)
@@ -147,6 +157,8 @@ def stream_window_updated(self, stream_id: int) -> bool:
             available_window -= len(chunk.data)
             if not self.stream_buffers[stream_id]:
                 del self.stream_buffers[stream_id]
+                if self.stream_trailers.get(stream_id):
+                    self.send_headers(stream_id, self.stream_trailers.pop(stream_id), end_stream=True)
             sent_any_data = True
 
         return sent_any_data
diff --git a/test/mitmproxy/proxy/layers/http/test_http2.py b/test/mitmproxy/proxy/layers/http/test_http2.py
index 3bace49ac6..6e8297af49 100644
--- a/test/mitmproxy/proxy/layers/http/test_http2.py
+++ b/test/mitmproxy/proxy/layers/http/test_http2.py
@@ -288,6 +288,163 @@ def test_upstream_error(tctx):
     assert b"server &lt;&gt; error" in d.data
 
 
+@pytest.mark.parametrize("trailers", ["trailers", ""])
+def test_long_response(tctx: Context, trailers):
+    playbook, cff = start_h2_client(tctx)
+    flow = Placeholder(HTTPFlow)
+    server = Placeholder(Server)
+    initial = Placeholder(bytes)
+    assert (
+        playbook
+        >> DataReceived(
+            tctx.client,
+            cff.build_headers_frame(
+                example_request_headers, flags=["END_STREAM"]
+            ).serialize(),
+        )
+        << http.HttpRequestHeadersHook(flow)
+        >> reply()
+        << http.HttpRequestHook(flow)
+        >> reply()
+        << OpenConnection(server)
+        >> reply(None, side_effect=make_h2)
+        << SendData(server, initial)
+    )
+    frames = decode_frames(initial())
+    assert [type(x) for x in frames] == [
+        hyperframe.frame.SettingsFrame,
+        hyperframe.frame.HeadersFrame,
+    ]
+    sff = FrameFactory()
+    assert (
+        playbook
+        # a conforming h2 server would send settings first, we disregard this for now.
+        >> DataReceived(
+            server, sff.build_headers_frame(example_response_headers).serialize()
+        )
+        << http.HttpResponseHeadersHook(flow)
+        >> reply()
+        >> DataReceived(
+            server,
+            sff.build_data_frame(b"a" * 10000, flags=[]).serialize()
+        )
+        >> DataReceived(
+            server,
+            sff.build_data_frame(b"a" * 10000, flags=[]).serialize(),
+        )
+        >> DataReceived(
+            server,
+            sff.build_data_frame(b"a" * 10000, flags=[]).serialize(),
+        )
+        >> DataReceived(
+            server,
+            sff.build_data_frame(b"a" * 10000, flags=[]).serialize(),
+        )
+        << SendData(
+            server,
+            sff.build_window_update_frame(0, 40000).serialize()
+            + sff.build_window_update_frame(1, 40000).serialize(),
+        )
+        >> DataReceived(
+            server,
+            sff.build_data_frame(b"a" * 10000, flags=[]).serialize(),
+        )
+        >> DataReceived(
+            server,
+            sff.build_data_frame(b"a" * 10000, flags=[]).serialize(),
+        )
+        >> DataReceived(
+            server,
+            sff.build_data_frame(b"a" * 10000, flags=[]).serialize(),
+        )
+    )
+    if trailers:
+        (
+            playbook
+            >> DataReceived(
+                server,
+                sff.build_headers_frame(
+                    example_response_trailers, flags=["END_STREAM"]
+                ).serialize(),
+            )
+        )
+    else:
+        (
+            playbook
+            >> DataReceived(
+                server,
+                sff.build_data_frame(
+                    b'', flags=["END_STREAM"]
+                ).serialize(),
+            )
+        )
+    (
+        playbook
+        << http.HttpResponseHook(flow)
+        >> reply()
+        << SendData(
+            tctx.client,
+            cff.build_headers_frame(example_response_headers).serialize()
+            + cff.build_data_frame(b"a" * 16384).serialize(),
+        )
+        << SendData(
+            tctx.client,
+            cff.build_data_frame(b"a" * 16384).serialize(),
+        )
+        << SendData(
+            tctx.client,
+            cff.build_data_frame(b"a" * 16384).serialize(),
+        )
+        << SendData(
+            tctx.client,
+            cff.build_data_frame(b"a" * 16383).serialize(),
+        )
+        >> DataReceived(
+            tctx.client,
+            cff.build_window_update_frame(0, 65535).serialize()
+            + cff.build_window_update_frame(1, 65535).serialize(),
+        )
+    )
+    if trailers:
+        assert (
+            playbook
+            << SendData(
+                tctx.client,
+                cff.build_data_frame(b"a" * 1).serialize(),
+            )
+            << SendData(
+                tctx.client,
+                cff.build_data_frame(b"a" * 4464).serialize()
+            )
+            << SendData(
+                tctx.client,
+                cff.build_headers_frame(
+                    example_response_trailers, flags=["END_STREAM"]
+                ).serialize(),
+            )
+        )
+    else:
+        assert (
+            playbook
+            << SendData(
+                tctx.client,
+                cff.build_data_frame(b"a" * 1).serialize(),
+            )
+            << SendData(
+                tctx.client,
+                cff.build_data_frame(b"a" * 4464).serialize()
+            )
+            << SendData(
+                tctx.client,
+                cff.build_data_frame(
+                    b"", flags=["END_STREAM"]
+                ).serialize(),
+            )
+        )
+    assert flow().request.url == "http://example.com/"
+    assert flow().response.text == "a" * 70000
+
+
 @pytest.mark.parametrize("stream", ["stream", ""])
 @pytest.mark.parametrize("when", ["request", "response"])
 @pytest.mark.parametrize("how", ["RST", "disconnect", "RST+disconnect"])


Code review comments:
diff_hunk: @@ -77,8 +79,16 @@ def send_data(
                 # We can't send right now, so we buffer.
                 self.stream_buffers[stream_id].append(SendH2Data(data, end_stream))
 
+    def send_trailers(self, stream_id, trailers):
+        if self.stream_buffers.get(stream_id, None):
+            # Though trailers are not subject to flow control, we need to queue them and send strictly after data frames
+            self.stream_trailers[stream_id] = trailers
+        else:
+            self.send_headers(stream_id, trailers, end_stream=True)
+
     def end_stream(self, stream_id) -> None:
-        self.send_data(stream_id, b"", end_stream=True)
+        if not self.stream_trailers.get(stream_id):
+            self.send_data(stream_id, b"", end_stream=True)
path: mitmproxy/proxy/layers/http/_http_h2.py
text: @user1:
This path here wasn't immediately clear to me, let's add a short explainer.
```suggestion
        if stream_id in self.stream_trailers:
            return  # we already have trailers queued up that will end the stream.
        self.send_data(stream_id, b"", end_stream=True)
```

diff_hunk: @@ -34,10 +34,12 @@ class BufferedH2Connection(h2.connection.H2Connection):
     """
 
     stream_buffers: collections.defaultdict[int, collections.deque[SendH2Data]]
+    stream_trailers: Dict[int, Any]
path: mitmproxy/proxy/layers/http/_http_h2.py
text: @user1:
While we're here, could we make this `List[Tuple[bytes, bytes]]` instead of `Any` please (and add type annotations below)? :)

diff_hunk: @@ -147,6 +157,8 @@ def stream_window_updated(self, stream_id: int) -> bool:
             available_window -= len(chunk.data)
             if not self.stream_buffers[stream_id]:
                 del self.stream_buffers[stream_id]
+                if self.stream_trailers.get(stream_id):
path: mitmproxy/proxy/layers/http/_http_h2.py
text: @user1:
```suggestion
                if stream_id in self.stream_trailers:
```
(We may have an empty trailers frame, which would otherwise evaluate to False)