Restrict the IOError

Author: Kludex

starlette/websockets.py

@@ -61,36 +61,34 @@ async def send(self, message: Message) -> None:
         """
         Send ASGI websocket messages, ensuring valid state transitions.
         """
-        try:
-            if self.application_state == WebSocketState.CONNECTING:
-                message_type = message["type"]
-                if message_type not in {"websocket.accept", "websocket.close"}:
-                    raise RuntimeError(
-                        'Expected ASGI message "websocket.accept" or '
-                        f'"websocket.close", but got {message_type!r}'
-                    )
-                if message_type == "websocket.close":
-                    self.application_state = WebSocketState.DISCONNECTED
-                else:
-                    self.application_state = WebSocketState.CONNECTED
-                await self._send(message)
-            elif self.application_state == WebSocketState.CONNECTED:
-                message_type = message["type"]
-                if message_type not in {"websocket.send", "websocket.close"}:
-                    raise RuntimeError(
-                        'Expected ASGI message "websocket.send" or "websocket.close", '
-                        f"but got {message_type!r}"
-                    )
-                if message_type == "websocket.close":
-                    self.application_state = WebSocketState.DISCONNECTED
-                await self._send(message)
+        if self.application_state == WebSocketState.CONNECTING:
+            message_type = message["type"]
+            if message_type not in {"websocket.accept", "websocket.close"}:
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.accept" or '
+                    f'"websocket.close", but got {message_type!r}'
+                )
+            if message_type == "websocket.close":
+                self.application_state = WebSocketState.DISCONNECTED
             else:
+                self.application_state = WebSocketState.CONNECTED
+            await self._send(message)
+        elif self.application_state == WebSocketState.CONNECTED:
+            message_type = message["type"]
+            if message_type not in {"websocket.send", "websocket.close"}:
                 raise RuntimeError(
-                    'Cannot call "send" once a close message has been sent.'
+                    'Expected ASGI message "websocket.send" or "websocket.close", '
+                    f"but got {message_type!r}"
                 )
-        except IOError:  # pragma: no cover
-            self.application_state = WebSocketState.DISCONNECTED
-            raise WebSocketDisconnect(code=1006)
+            if message_type == "websocket.close":
+                self.application_state = WebSocketState.DISCONNECTED
+            try:
+                await self._send(message)
+            except IOError:
+                self.application_state = WebSocketState.DISCONNECTED
+                raise WebSocketDisconnect(code=1006)
+        else:
+            raise RuntimeError('Cannot call "send" once a close message has been sent.')
 
     async def accept(
         self,

tests/test_websockets.py

@@ -255,6 +255,28 @@ async def app(scope: Scope, receive: Receive, send: Send) -> None:
     assert close_reason == "Going Away"
 
 
+@pytest.mark.anyio
+async def test_client_disconnect_on_send():
+    async def app(scope: Scope, receive: Receive, send: Send) -> None:
+        websocket = WebSocket(scope, receive=receive, send=send)
+        await websocket.accept()
+        await websocket.send_text("Hello, world!")
+
+    async def receive() -> Message:
+        return {"type": "websocket.connect"}
+
+    async def send(message: Message) -> None:
+        if message["type"] == "websocket.accept":
+            return
+        # Simulate the exception the server would send to the application when the
+        # client disconnects.
+        raise IOError
+
+    with pytest.raises(WebSocketDisconnect) as ctx:
+        await app({"type": "websocket", "path": "/"}, receive, send)
+    assert ctx.value.code == status.WS_1006_ABNORMAL_CLOSURE
+
+
 def test_application_close(test_client_factory: Callable[..., TestClient]):
     async def app(scope: Scope, receive: Receive, send: Send) -> None:
         websocket = WebSocket(scope, receive=receive, send=send)