Timeout close handshake WebSocket connections

This ensures that WebSocket connections will close even if the client
doesn't respond with a close frame (after a duration equal to the keep
alive timeout).

hypercorn/protocol/h11.py

@@ -94,7 +94,7 @@ def __init__(
 
     @property
     def idle(self) -> bool:
-        return self.stream is None
+        return self.stream is None or self.stream.idle
 
     async def initiate(self) -> None:
         pass

hypercorn/protocol/h2.py

@@ -59,7 +59,7 @@ def __init__(
 
     @property
     def idle(self) -> bool:
-        return len(self.streams) == 0
+        return len(self.streams) == 0 or all(stream.idle for stream in self.streams.values())
 
     async def initiate(
         self, headers: Optional[List[Tuple[bytes, bytes]]] = None, settings: Optional[str] = None

hypercorn/protocol/http_stream.py

@@ -41,6 +41,10 @@ def __init__(
         self.state = ASGIHTTPState.REQUEST
         self.stream_id = stream_id
 
+    @property
+    def idle(self) -> bool:
+        return False
+
     async def handle(self, event: Event) -> None:
         if isinstance(event, Request):
             path, _, query_string = event.raw_path.partition(b"?")

hypercorn/protocol/ws_stream.py

@@ -162,6 +162,10 @@ def __init__(
         self.connection: Connection
         self.handshake: Handshake
 
+    @property
+    def idle(self) -> bool:
+        return self.state in {ASGIWebsocketState.CLOSED, ASGIWebsocketState.HTTPCLOSED}
+
     async def handle(self, event: Event) -> None:
         if isinstance(event, Request):
             self.handshake = Handshake(event.headers, event.http_version)
@@ -247,14 +251,14 @@ async def app_send(self, message: Optional[dict]) -> None:
             elif (
                 message["type"] == "websocket.close" and self.state == ASGIWebsocketState.HANDSHAKE
             ):
-                await self._send_error_response(403)
                 self.state = ASGIWebsocketState.HTTPCLOSED
+                await self._send_error_response(403)
             elif message["type"] == "websocket.close":
+                self.state = ASGIWebsocketState.CLOSED
                 await self._send_wsproto_event(
                     CloseConnection(code=int(message.get("code", CloseReason.NORMAL_CLOSURE)))
                 )
                 await self.send(EndData(stream_id=self.stream_id))
-                self.state = ASGIWebsocketState.CLOSED
             else:
                 raise UnexpectedMessage(self.state, message["type"])
 
@@ -311,6 +315,6 @@ async def _send_rejection(self, message: dict) -> None:
         if not body_suppressed:
             await self.send(Body(stream_id=self.stream_id, data=bytes(message.get("body", b""))))
         if not message.get("more_body", False):
-            await self.send(EndBody(stream_id=self.stream_id))
             self.state = ASGIWebsocketState.HTTPCLOSED
+            await self.send(EndBody(stream_id=self.stream_id))
             await self.config.log.access(self.scope, self.response, time() - self.start_time)

tests/protocol/test_http_stream.py

@@ -144,3 +144,7 @@ async def test_send_invalid_message(
     with pytest.raises((TypeError, ValueError)):
         await stream.app_send({"type": "http.response.start", "headers": headers, "status": status})
         await stream.app_send({"type": "http.response.body", "body": body})
+
+
+def test_stream_idle(stream: HTTPStream) -> None:
+    assert stream.idle is False

tests/protocol/test_ws_stream.py

@@ -347,3 +347,17 @@ async def test_send_invalid_http_message(
             {"type": "websocket.http.response.start", "headers": headers, "status": status}
         )
         await stream.app_send({"type": "websocket.http.response.body", "body": body})
+
+
+@pytest.mark.parametrize(
+    "state, idle",
+    [
+        (state, False)
+        for state in ASGIWebsocketState
+        if state not in {ASGIWebsocketState.CLOSED, ASGIWebsocketState.HTTPCLOSED}
+    ]
+    + [(ASGIWebsocketState.CLOSED, True), (ASGIWebsocketState.HTTPCLOSED, True)],
+)
+def test_stream_idle(stream: WSStream, state: ASGIWebsocketState, idle: bool) -> None:
+    stream.state = state
+    assert stream.idle is idle