Acknowledge HTTP/2 data received for an already closed stream

src/hypercorn/protocol/h2.py

@@ -256,9 +256,14 @@ async def _handle_events(self, events: list[h2.events.Event]) -> None:
                 if self.keep_alive_requests > self.config.keep_alive_max_requests:
                     self.connection.close_connection()
             elif isinstance(event, h2.events.DataReceived):
-                await self.streams[event.stream_id].handle(
-                    Body(stream_id=event.stream_id, data=event.data)
-                )
+                try:
+                    await self.streams[event.stream_id].handle(
+                        Body(stream_id=event.stream_id, data=event.data)
+                    )
+                except KeyError:
+                    # Response sent before full request received, the
+                    # data is dropped but must still be acknowledged.
+                    pass
                 self.connection.acknowledge_received_data(
                     event.flow_controlled_length, event.stream_id
                 )

tests/protocol/test_h2.py

@@ -5,7 +5,7 @@
 
 import pytest
 from h2.connection import H2Connection
-from h2.events import ConnectionTerminated
+from h2.events import ConnectionTerminated, WindowUpdated
 
 from hypercorn.asyncio.worker_context import EventWrapper, WorkerContext
 from hypercorn.config import Config
@@ -116,3 +116,44 @@ async def test_protocol_keep_alive_max_requests() -> None:
     protocol.send.assert_awaited()  # type: ignore
     events = client.receive_data(protocol.send.call_args_list[1].args[0].data)  # type: ignore
     assert isinstance(events[-1], ConnectionTerminated)
+
+
+@pytest.mark.asyncio
+async def test_protocol_data_received_after_stream_closed() -> None:
+    protocol = H2Protocol(
+        Mock(),
+        Config(),
+        WorkerContext(None),
+        AsyncMock(),
+        ConnectionState({}),
+        False,
+        None,
+        None,
+        AsyncMock(),
+    )
+    client = H2Connection()
+    client.initiate_connection()
+    headers = [
+        (":method", "POST"),
+        (":path", "/"),
+        (":authority", "hypercorn"),
+        (":scheme", "https"),
+    ]
+    client.send_headers(1, headers)
+    await protocol.handle(RawData(data=client.data_to_send()))
+    # The app responded before reading the request body, so the stream is gone
+    await protocol._close_stream(1)
+    protocol.send.reset_mock()  # type: ignore
+
+    for _ in range(3):
+        client.send_data(1, b"a" * client.max_outbound_frame_size)
+    await protocol.handle(RawData(data=client.data_to_send()))
+
+    # The dropped data is still acknowledged, so the connection window is restored
+    data = b"".join(
+        call.args[0].data
+        for call in protocol.send.call_args_list  # type: ignore
+        if isinstance(call.args[0], RawData)
+    )
+    events = client.receive_data(data)
+    assert any(isinstance(event, WindowUpdated) and event.stream_id == 0 for event in events)