Unblock ASGI ws receive on close without accept (#819)

Author: gi0baro

src/asgi/io.rs

@@ -476,16 +476,18 @@ impl ASGIWebsocketProtocol {
 
     #[inline(always)]
     fn close<'p>(&self, py: Python<'p>, frame: Option<wsframe::CloseFrame>) -> PyResult<Bound<'p, PyAny>> {
-        let closed = self.closed.clone();
+        let init_ev = self.init_event.clone();
         let ws_rx = self.ws_rx.clone();
         let ws_tx = self.ws_tx.clone();
+        self.closed.store(true, atomic::Ordering::Release);
 
         future_into_py_futlike(self.rt.clone(), py, async move {
             if let Some(tx) = ws_tx.lock().await.take() {
-                closed.store(true, atomic::Ordering::Release);
                 WebsocketDetachedTransport::new(true, ws_rx.lock().await.take(), Some(tx), frame)
                     .close()
                     .await;
+            } else {
+                init_ev.notify_one();
             }
             FutureResultToPy::None
         })

tests/apps/asgi.py

@@ -70,6 +70,12 @@ async def ws_reject(scope, receive, send):
     return
 
 
+async def ws_reject_explicit(scope, receive, send):
+    await receive()
+    await send({'type': 'websocket.close'})
+    await receive()
+
+
 async def ws_reject_custom(scope, receive, send):
     await receive()
     await send(
@@ -216,6 +222,7 @@ def app(scope, receive, send):
         '/echo': echo,
         '/file': pathsend,
         '/ws_reject': ws_reject,
+        '/ws_rejecte': ws_reject_explicit,
         '/ws_rejectc': ws_reject_custom,
         '/ws_info': ws_info,
         '/ws_echo': ws_echo,

tests/test_ws.py

@@ -4,6 +4,7 @@
 
 import pytest
 import websockets
+import websockets.exceptions
 
 
 @pytest.mark.asyncio
@@ -52,6 +53,17 @@ async def test_asgi_server_close(asgi_server, runtime_mode, tmp_path):
     assert target.exists()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize('runtime_mode', ['mt', 'st'])
+async def test_asgi_reject_explicit(asgi_server, runtime_mode):
+    async with asgi_server(runtime_mode) as port:
+        with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
+            async with websockets.connect(f'ws://localhost:{port}/ws_rejecte'):
+                pass
+
+    assert exc.value.response.status_code == 403
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize('runtime_mode', ['mt', 'st'])
 async def test_asgi_reject_custom(asgi_server, runtime_mode):