Ensure that ws connection is open before receiving data
* Revert "Fix recv() errors when closed during handshake (#704)"
This reverts commit
* Fix 204 by using ensure_open that raises the disconnect
@@ -216,7 +216,6 @@ class WebSocketProtocol(websockets.WebSocketServerProtocol):
elif message_type == "websocket.close":
code = message.get("code", 1000)
- self.close_code = code # for WebSocketServerProtocol
await self.close(code)
self.closed_event.set()
@@ -237,15 +236,8 @@ class WebSocketProtocol(websockets.WebSocketServerProtocol):
return {"type": "websocket.connect"}
await self.handshake_completed_event.wait()
-
- if self.closed_event.is_set():
- # If the client disconnected: WebSocketServerProtocol set self.close_code.
- # If the handshake failed or the app closed before handshake completion,
- # use 1006 Abnormal Closure.
- code = getattr(self, "close_code", 1006)
- return {"type": "websocket.disconnect", "code": code}
-
try:
+ await self.ensure_open()
data = await self.recv()
except websockets.ConnectionClosed as exc:
return {"type": "websocket.disconnect", "code": exc.code}
