Repository: Kludex/uvicorn (Python)

Issue:
Bug: calling `WebSocketProtocol.asgi_receive` returns close frame even if there are data messages before close frame in read queue
### Checklist

- [x] The bug is reproducible against the latest release and/or `master`.
- [x] There are no similar issues or pull requests to fix it yet.

### Describe the bug

Once a client sends a close frame, calling `WebSocketProtocol.asgi_receive` returns `{"type": "websocket.disconnect", "code": exc.code}`, even if there are unread messages in the server's read queue that we sent **before** the close frame.

### To reproduce

The easiest way for me to repro is to use an ASGI application framework written on top of `uvicorn`, e.g. FastAPI.

- Install FastAPI and create a module `main.py`
- Run the websocket server: `uvicorn main:app --reload --host 0.0.0.0 --port 8001`
- Open the browser and create a websocket connection to this test endpoint

`main.py`

```py
import asyncio
from fastapi import FastAPI
from starlette.websockets import WebSocket


app = FastAPI()


@app.websocket("/echo")
async def echo_ws(ws: WebSocket):
    await ws.accept()
    await asyncio.sleep(1)

    data = await ws.receive_bytes()
    print(data)
```

Paste the following into the browser console:

```js
var socket = new WebSocket('ws://localhost:8001/atom_ws/echo')
socket.onopen = () => {
  console.log('socket opened')
  socket.send('first')
  socket.send('second')
  socket.close()
}
```

### Expected behavior

The call to `print(data)` should print `first`, and not raise a `starlette.websockets.WebSocketDisconnect` exception.

### Actual behavior

`starlette.websockets.WebSocketDisconnect` is raised on the first read, even though messages were successfully sent to the server before the close frame, and these messages are in the connection's read queue.

### Debugging material

Logs when running the server code from above:

<img width="839" alt="Screen Shot 2021-11-01 at 9 24 46 PM" src="https://user-images.githubusercontent.com/1524088/139782096-84024344-6044-453e-9b8f-dff193aef53a.png">

If you instead run a simple websocket server written with code directly from the `websockets` library, [as suggested in their docs](https://websockets.readthedocs.io/en/3.0/intro.html#basic-example), you don't have this problem:

```py
import asyncio
import websockets


async def echo(ws, path):
    print(f"accepted connection to path {path}")
    await asyncio.sleep(1)
    # By this time client has already closed connection

    # In uvicorn, `await ws.ensure_open()` is called before recv; this is the bug
    data = await ws.recv()
    print(data)  # Prints first
    data = await ws.recv()
    print(data)  # Prints second

    # The next `recv` call raises `websockets.exceptions.ConnectionClosedError`, because it reads the close frame
    data = await ws.recv()

start_server = websockets.serve(echo, 'localhost', 8001)

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()
```

We can see that the `first` and `second` messages are received, even though the client sent the close frame before the server tried to read any messages:

<img width="933" alt="Screen Shot 2021-11-01 at 9 30 37 PM" src="https://user-images.githubusercontent.com/1524088/139782584-0e436870-0c25-4d6b-b48e-11fdd36bec35.png">

### Environment

- MacOS 10.14.6 / Python 3.8.5 / uvicorn 0.15.0 (this bug is still present in 0.16.0, and in the latest commit in `master`)
- Can also repro on Ubuntu 20.04

- The exact command you're running uvicorn with, all flags you passed included: `uvicorn main:app --reload --host 0.0.0.0 --port 8001`

### Additional context

The bug is caused by this line of code: https://github.com/encode/uvicorn/blob/48edc940522a3d0d7529922a23ac019eeb53f629/uvicorn/protocols/websockets/websockets_impl.py#L286-L286

```py
await self.ensure_open()
data = await self.recv()
```

This change was made in this commit: https://github.com/encode/uvicorn/commit/9a3040c9cd56844631b28631acd5862b5a4eafdd

`uvicorn` depends on `websockets` under the hood, but it shouldn't be calling `ensure_open` before calling `recv`, because `ensure_open` raises an exception if a close frame has been sent by the client **even if there are earlier unread messages in the read queue**.

I'm not sure what the intent of that line of code is, but the server shouldn't raise a "connection closed" exception on read until the actual close frame is read. Otherwise, neither the client nor the server has any way of knowing that data that was successfully sent from the client to the server was ignored by the server.


Discussion:
@Kludex @florimondmanca @tomchristie @euri10 

If we can revert the `await self.ensure_open()` change in `asgi_receive` without breaking other things in `uvicorn`, that would be an easy solve.

Or, if we could expose the underlying `WebSocketCommonProtocol` instance that client code in websocket views could avoid using `asgi_receive` and just call `send` and `recv` directly, that would work as well, but I'm guessing this isn't possible.

This bug is pretty serious, but I hope fixing it isn't that difficult.



Patch:
diff --git a/uvicorn/protocols/websockets/websockets_impl.py b/uvicorn/protocols/websockets/websockets_impl.py
index d829f49a6..b69e31535 100644
--- a/uvicorn/protocols/websockets/websockets_impl.py
+++ b/uvicorn/protocols/websockets/websockets_impl.py
@@ -282,10 +282,17 @@ async def asgi_receive(self):
             return {"type": "websocket.connect"}
 
         await self.handshake_completed_event.wait()
+
+        if self.closed_event.is_set():
+            # If client disconnected, use WebSocketServerProtocol.close_code property.
+            # If the handshake failed or the app closed before handshake completion,
+            # use 1006 Abnormal Closure.
+            return {"type": "websocket.disconnect", "code": self.close_code or 1006}
+
         try:
-            await self.ensure_open()
             data = await self.recv()
         except websockets.ConnectionClosed as exc:
+            self.closed_event.set()
             return {"type": "websocket.disconnect", "code": exc.code}
 
         msg = {"type": "websocket.receive"}


Test patch:
diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py
index d21be54de..94b625a36 100644
--- a/tests/protocols/test_websocket.py
+++ b/tests/protocols/test_websocket.py
@@ -1,3 +1,5 @@
+import asyncio
+
 import httpx
 import pytest
 
@@ -539,3 +541,67 @@ async def send_text(url):
             with pytest.raises(websockets.ConnectionClosedError) as e:
                 data = await send_text("ws://127.0.0.1:8000")
             assert e.value.code == expected_result
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_server_reject_connection(ws_protocol_cls, http_protocol_cls):
+    async def app(scope, receive, send):
+        assert scope["type"] == "websocket"
+
+        # Pull up first recv message.
+        message = await receive()
+        assert message["type"] == "websocket.connect"
+
+        # Reject the connection.
+        await send({"type": "websocket.close"})
+        # -- At this point websockets' recv() is unusable. --
+
+        # This doesn't raise `TypeError`:
+        # See https://github.com/encode/uvicorn/issues/244
+        message = await receive()
+        assert message["type"] == "websocket.disconnect"
+
+    async def websocket_session(url):
+        try:
+            async with websockets.connect(url):
+                pass
+        except Exception:
+            pass
+
+    config = Config(app=app, ws=ws_protocol_cls, http=http_protocol_cls, lifespan="off")
+    async with run_server(config):
+        await websocket_session("ws://127.0.0.1:8000")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_server_can_read_messages_in_buffer_after_close(
+    ws_protocol_cls, http_protocol_cls
+):
+    frames = []
+
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send({"type": "websocket.accept"})
+            # Ensure server doesn't start reading frames from read buffer until
+            # after client has sent close frame, but server is still able to
+            # read these frames
+            await asyncio.sleep(0.2)
+
+        async def websocket_receive(self, message):
+            frames.append(message.get("bytes"))
+
+    async def send_text(url):
+        async with websockets.connect(url) as websocket:
+            await websocket.send(b"abc")
+            await websocket.send(b"abc")
+            await websocket.send(b"abc")
+
+    config = Config(app=App, ws=ws_protocol_cls, http=http_protocol_cls, lifespan="off")
+    async with run_server(config):
+        await send_text("ws://127.0.0.1:8000")
+
+    assert frames == [b"abc", b"abc", b"abc"]