Support sending trailing headers on h2/h3

This adds support for the ASGI http-trailers extension which allows
for trailing headers to be sent after a response.

Many thanks to @jeffsawatzky from whom parts of this implementation
are based.

src/hypercorn/protocol/events.py

@@ -27,6 +27,11 @@ class EndBody(Event):
     pass
 
 
+@dataclass(frozen=True)
+class Trailers(Event):
+    headers: List[Tuple[bytes, bytes]]
+
+
 @dataclass(frozen=True)
 class Data(Event):
     data: bytes

src/hypercorn/protocol/h2.py

@@ -18,6 +18,7 @@
     Request,
     Response,
     StreamClosed,
+    Trailers,
 )
 from .http_stream import HTTPStream
 from .ws_stream import WSStream
@@ -213,6 +214,9 @@ async def stream_send(self, event: StreamEvent) -> None:
                 self.priority.unblock(event.stream_id)
                 await self.has_data.set()
                 await self.stream_buffers[event.stream_id].drain()
+            elif isinstance(event, Trailers):
+                self.connection.send_headers(event.stream_id, event.headers)
+                await self._flush()
             elif isinstance(event, StreamClosed):
                 await self._close_stream(event.stream_id)
                 idle = len(self.streams) == 0 or all(

src/hypercorn/protocol/h3.py

@@ -18,6 +18,7 @@
     Request,
     Response,
     StreamClosed,
+    Trailers,
 )
 from .http_stream import HTTPStream
 from .ws_stream import WSStream
@@ -79,6 +80,9 @@ async def stream_send(self, event: StreamEvent) -> None:
         elif isinstance(event, (EndBody, EndData)):
             self.connection.send_data(event.stream_id, b"", True)
             await self.send()
+        elif isinstance(event, Trailers):
+            self.connection.send_headers(event.stream_id, event.headers)
+            await self.send()
         elif isinstance(event, StreamClosed):
             pass  # ??
         elif isinstance(event, Request):

src/hypercorn/protocol/http_stream.py

@@ -5,7 +5,16 @@
 from typing import Awaitable, Callable, Optional, Tuple
 from urllib.parse import unquote
 
-from .events import Body, EndBody, Event, InformationalResponse, Request, Response, StreamClosed
+from .events import (
+    Body,
+    EndBody,
+    Event,
+    InformationalResponse,
+    Request,
+    Response,
+    StreamClosed,
+    Trailers,
+)
 from ..config import Config
 from ..typing import (
     AppWrapper,
@@ -22,6 +31,7 @@
     valid_server_name,
 )
 
+TRAILERS_VERSIONS = {"2", "3"}
 PUSH_VERSIONS = {"2", "3"}
 EARLY_HINTS_VERSIONS = {"2", "3"}
 
@@ -32,6 +42,7 @@ class ASGIHTTPState(Enum):
     # state tracking is required.
     REQUEST = auto()
     RESPONSE = auto()
+    TRAILERS = auto()
     CLOSED = auto()
 
 
@@ -88,6 +99,10 @@ async def handle(self, event: Event) -> None:
                 "server": self.server,
                 "extensions": {},
             }
+
+            if event.http_version in TRAILERS_VERSIONS:
+                self.scope["extensions"]["http.response.trailers"] = {}
+
             if event.http_version in PUSH_VERSIONS:
                 self.scope["extensions"]["http.response.push"] = {}
 
@@ -182,13 +197,33 @@ async def app_send(self, message: Optional[ASGISendEvent]) -> None:
                     )
 
                 if not message.get("more_body", False):
-                    if self.state != ASGIHTTPState.CLOSED:
+                    await self.send(EndBody(stream_id=self.stream_id))
+
+                    if self.response.get("trailers", False):
+                        self.state = ASGIHTTPState.TRAILERS
+                    else:
                         self.state = ASGIHTTPState.CLOSED
                         await self.config.log.access(
                             self.scope, self.response, time() - self.start_time
                         )
-                        await self.send(EndBody(stream_id=self.stream_id))
                         await self.send(StreamClosed(stream_id=self.stream_id))
+            elif (
+                message["type"] == "http.response.trailers"
+                and self.scope["http_version"] in TRAILERS_VERSIONS
+                and self.state == ASGIHTTPState.TRAILERS
+            ):
+                for name, value in self.scope["headers"]:
+                    if name == b"te" and value == b"trailers":
+                        headers = build_and_validate_headers(message["headers"])
+                        await self.send(Trailers(stream_id=self.stream_id, headers=headers))
+                        break
+
+                if not message.get("more_trailers", False):
+                    self.state = ASGIHTTPState.CLOSED
+                    await self.config.log.access(
+                        self.scope, self.response, time() - self.start_time
+                    )
+                    await self.send(StreamClosed(stream_id=self.stream_id))
             else:
                 raise UnexpectedMessageError(self.state, message["type"])
 

src/hypercorn/typing.py

@@ -22,6 +22,11 @@
 
 from .config import Config, Sockets
 
+try:
+    from typing import NotRequired
+except ImportError:
+    from typing_extensions import NotRequired
+
 H11SendableEvent = Union[h11.Data, h11.EndOfMessage, h11.InformationalResponse, h11.Response]
 
 WorkerFunc = Callable[[Config, Optional[Sockets], Optional[EventType]], None]
@@ -83,6 +88,7 @@ class HTTPResponseStartEvent(TypedDict):
     type: Literal["http.response.start"]
     status: int
     headers: Iterable[Tuple[bytes, bytes]]
+    trailers: NotRequired[bool]
 
 
 class HTTPResponseBodyEvent(TypedDict):
@@ -91,6 +97,12 @@ class HTTPResponseBodyEvent(TypedDict):
     more_body: bool
 
 
+class HTTPResponseTrailersEvent(TypedDict):
+    type: Literal["http.response.trailers"]
+    headers: Iterable[Tuple[bytes, bytes]]
+    more_trailers: NotRequired[bool]
+
+
 class HTTPServerPushEvent(TypedDict):
     type: Literal["http.response.push"]
     path: str
@@ -191,6 +203,7 @@ class LifespanShutdownFailedEvent(TypedDict):
 ASGISendEvent = Union[
     HTTPResponseStartEvent,
     HTTPResponseBodyEvent,
+    HTTPResponseTrailersEvent,
     HTTPServerPushEvent,
     HTTPEarlyHintEvent,
     HTTPDisconnectEvent,

tests/protocol/test_http_stream.py

@@ -17,6 +17,7 @@
     Request,
     Response,
     StreamClosed,
+    Trailers,
 )
 from hypercorn.protocol.http_stream import ASGIHTTPState, HTTPStream
 from hypercorn.typing import HTTPResponseBodyEvent, HTTPResponseStartEvent, HTTPScope
@@ -84,7 +85,11 @@ async def test_handle_request_http_2(stream: HTTPStream) -> None:
         "headers": [],
         "client": None,
         "server": None,
-        "extensions": {"http.response.early_hint": {}, "http.response.push": {}},
+        "extensions": {
+            "http.response.trailers": {},
+            "http.response.early_hint": {},
+            "http.response.push": {},
+        },
     }
 
 
@@ -204,6 +209,65 @@ async def test_send_early_hint(stream: HTTPStream, http_scope: HTTPScope) -> Non
     ]
 
 
+@pytest.mark.asyncio
+async def test_send_trailers(stream: HTTPStream) -> None:
+    await stream.handle(
+        Request(
+            stream_id=1,
+            http_version="2",
+            headers=[(b"te", b"trailers")],
+            raw_path=b"/?a=b",
+            method="GET",
+        )
+    )
+    await stream.app_send(
+        cast(
+            HTTPResponseStartEvent,
+            {"type": "http.response.start", "status": 200, "trailers": True},
+        )
+    )
+    await stream.app_send(
+        cast(HTTPResponseBodyEvent, {"type": "http.response.body", "body": b"Body"})
+    )
+    await stream.app_send({"type": "http.response.trailers", "headers": [(b"X", b"V")]})
+    assert stream.send.call_args_list == [  # type: ignore
+        call(Response(stream_id=1, headers=[], status_code=200)),
+        call(Body(stream_id=1, data=b"Body")),
+        call(EndBody(stream_id=1)),
+        call(Trailers(stream_id=1, headers=[(b"X", b"V")])),
+        call(StreamClosed(stream_id=1)),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_send_trailers_ignored(stream: HTTPStream) -> None:
+    await stream.handle(
+        Request(
+            stream_id=1,
+            http_version="2",
+            headers=[],  # no TE: trailers header
+            raw_path=b"/?a=b",
+            method="GET",
+        )
+    )
+    await stream.app_send(
+        cast(
+            HTTPResponseStartEvent,
+            {"type": "http.response.start", "status": 200, "trailers": True},
+        )
+    )
+    await stream.app_send(
+        cast(HTTPResponseBodyEvent, {"type": "http.response.body", "body": b"Body"})
+    )
+    await stream.app_send({"type": "http.response.trailers", "headers": [(b"X", b"V")]})
+    assert stream.send.call_args_list == [  # type: ignore
+        call(Response(stream_id=1, headers=[], status_code=200)),
+        call(Body(stream_id=1, data=b"Body")),
+        call(EndBody(stream_id=1)),
+        call(StreamClosed(stream_id=1)),
+    ]
+
+
 @pytest.mark.asyncio
 async def test_send_app_error(stream: HTTPStream) -> None:
     await stream.handle(
@@ -229,16 +293,20 @@ async def test_send_app_error(stream: HTTPStream) -> None:
     "state, message_type",
     [
         (ASGIHTTPState.REQUEST, "not_a_real_type"),
+        (ASGIHTTPState.REQUEST, "http.response.trailers"),
         (ASGIHTTPState.RESPONSE, "http.response.start"),
+        (ASGIHTTPState.TRAILERS, "http.response.start"),
         (ASGIHTTPState.CLOSED, "http.response.start"),
         (ASGIHTTPState.CLOSED, "http.response.body"),
+        (ASGIHTTPState.CLOSED, "http.response.trailers"),
     ],
 )
 @pytest.mark.asyncio
 async def test_send_invalid_message_given_state(
-    stream: HTTPStream, state: ASGIHTTPState, message_type: str
+    stream: HTTPStream, state: ASGIHTTPState, http_scope: HTTPScope, message_type: str
 ) -> None:
     stream.state = state
+    stream.scope = http_scope
     with pytest.raises(UnexpectedMessageError):
         await stream.app_send({"type": message_type})  # type: ignore