Skip to content

Commit

Permalink
Improve the trailing headers support
Browse files Browse the repository at this point in the history
The EndBody event should only be sent after the headers as this
results in the stream being closed.

It is acceptable to send no response and only trailing headers, in
which case a default 200 status code response is sent with the
headers.
  • Loading branch information
pgjones committed May 27, 2024
1 parent e47757d commit d16b503
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 13 deletions.
44 changes: 33 additions & 11 deletions src/hypercorn/protocol/http_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,16 +194,36 @@ async def app_send(self, message: Optional[ASGISendEvent]) -> None:
)

if not message.get("more_body", False):
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_closed()
elif (
message["type"] == "http.response.trailers"
and self.scope["http_version"] in TRAILERS_VERSIONS
and self.state == ASGIHTTPState.REQUEST
):
for name, value in self.scope["headers"]:
if name == b"te" and value == b"trailers":
headers = build_and_validate_headers(message["headers"])
self.response = {
"type": "http.response.start",
"status": 200,
"headers": headers,
}
await self.send(
Response(
stream_id=self.stream_id,
headers=headers,
status_code=200,
)
)
await self.send(StreamClosed(stream_id=self.stream_id))
self.state = ASGIHTTPState.TRAILERS
break

if not message.get("more_trailers", False):
await self._send_closed()

elif (
message["type"] == "http.response.trailers"
and self.scope["http_version"] in TRAILERS_VERSIONS
Expand All @@ -216,14 +236,16 @@ async def app_send(self, message: Optional[ASGISendEvent]) -> None:
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))
await self._send_closed()
else:
raise UnexpectedMessageError(self.state, message["type"])

async def _send_closed(self) -> None:
await self.send(EndBody(stream_id=self.stream_id))
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))

async def _send_error_response(self, status_code: int) -> None:
await self.send(
Response(
Expand Down
3 changes: 1 addition & 2 deletions tests/protocol/test_http_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,8 +269,8 @@ async def test_send_trailers(stream: HTTPStream) -> None:
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(EndBody(stream_id=1)),
call(StreamClosed(stream_id=1)),
]

Expand Down Expand Up @@ -337,7 +337,6 @@ 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"),
Expand Down

0 comments on commit d16b503

Please sign in to comment.