Bug: calling WebSocketProtocol.asgi_receive returns close frame even if there are data messages before close frame in read queue

[kylebebak]

Checklist

• The bug is reproducible against the latest release and/or master.
• 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

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:

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:

If you instead run a simple websocket server written with code directly from the websockets library, as suggested in their docs, you don't have this problem:

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:

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:

uvicorn/uvicorn/protocols/websockets/websockets_impl.py

Line 286 in 48edc94

await self.ensure_open()

await self.ensure_open()
data = await self.recv()

This change was made in this commit: 9a3040c

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.

[kylebebak]

@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.