Fix bug in `WebSocketProtocol.asgi_receive`

[kylebebak]

@Kludex @euri10

Note: this PR obsoletes this old one, which was closed anyway

There were 2 things to fix in https://github.com/encode/uvicorn/issues/1230:

• ensure the server can read frames in the read queue after the client has sent the close frame
• ensure the server can read frames in the read queue after the server has sent the close frame

The first item was fixed was fixed by this PR: https://github.com/encode/uvicorn/pull/1252. But I realized after this PR was merged that the second item hadn't been fixed.

This PR fixes the second item. It ensures that the server can "drain" unread frames in the read queue even after the server has sent the close frame, as long as those frames were sent before the server sent the close frame.

As with the first item, the second item is part of the websockets spec. The following code, using just the websockets library, makes it clear that this is the intended behavior of this library.

I did not include code in this PR that fixes this bug in wsproto_impl.py, because I'm not familiar with wsproto. That said, I don't think we should gate this fix on the corresponding fix in wsproto.

client.py

import asyncio

import websockets


async def hello():
    async with websockets.connect("ws://localhost:8765") as websocket:
        await websocket.send("message 1")
        await websocket.send("message 2")
        await websocket.send("message 3")

        # Sleep until after server has sent close frame
        await asyncio.sleep(0.3)

        try:
            await websocket.send("message 4")
        except websockets.exceptions.ConnectionClosed as e:
            print("message 4 not sent, connection closed:", e)


asyncio.run(hello())

server.py

import asyncio

import websockets


async def handler(websocket, *args):
    # Give client time to send frames that go into read queue before server sends close frame
    await asyncio.sleep(0.2)
    await websocket.close(1000)

    # Server can still read these messages from read queue after sending close frame
    print(await websocket.recv())
    print(await websocket.recv())
    print(await websocket.recv())


async def main():
    async with websockets.serve(handler, "localhost", 8765):
        await asyncio.Future()  # Run forever


asyncio.run(main())

[Kludex]

That said, I don't think we should gate this fix on the corresponding fix in wsproto.

Why?

[kylebebak]

Because I can't make the fix in wsproto. There's a bug in uvicorn right now. Fixing it in just websockets_impl.py is better than not fixing it at all.

IMO it's much better, because by default FastAPI depends on the code in websockets_impl.py, not in wsproto. Fixing the bug in websockets_impl.py will fix it for the vast majority of FastAPI deployments.

If there's someone that can port the fix to wsproto that's even better, but no one has volunteered to do that.

[Kludex]

I understand the effects of fixing a bug...

I'm trying to understand your statement.

We have two implementations, and ideally we want to match the behavior. So when you say "I don't think we should gate the fix to wsproto", I need to understand why you said that. Is it because it cannot be done? Is it because there's some kind of blocking issue?

But now the follow up message was "I can't fix", which is different than the previous statement. It's not that you don't think we should, is that you don't want/can/need to dedicate time for it. Which is fine. 🙏

[iudeen]

I'll look into this, for wsproto ... if that's okay.

[Kludex]

I spent a lot of time on this PR trying to fix the wsproto part. Also, I come to a more conceptual understanding about this PR.

Let's go with some notes.

On the description, you mentioned that we should solve the following issues:

1. ensure the server can read frames in the read queue after the client has sent the close frame
2. ensure the server can read frames in the read queue after the server has sent the close frame

From what I understand, we really want 1, but I'm not sure about 2.

Let me see if I can express myself here. 1 is fine, if the client sent data, we can first read that data, and then pass it to the application, and then we can send the disconnect event.

On 2 there's a big difference. The application already stated that it wants to close the connection, the server will listen to that request. So... Even if more data can be read from the client, it doesn't matter. The client already stated that it doesn't care anymore.

[Kludex]

As a note, this PR doesn't send a disconnect event to the application - only the one from the server shutdown is sent.

[Kludex]

Ok. Given https://github.com/django/asgiref/issues/349, and my comments above, I'll be closing this PR.

Sorry for taking long here, and thanks for the PR @kylebebak . :pray:

[kylebebak]

Sorry for not responding to this until now @Kludex , thank you for reviewing this PR.

I didn't explain this as well as I should have, and I think the client/server/application terminology was a little confusing.

I still think the current uvicorn behavior isn't ideal, and the best way to summarize my concerns is the code I pasted in the first comment, that uses client/server terminology in example code in the websockets docs.

If you do this, e.g. with Python 3.11...

• pip install websockets==11.0.3
• Create client.py and server.py as in the first comment
• Run python server.py to start WS handler
• Run python client.py to connect to this handler and send it messages

...you'll see that the WS handler can read messages in its read queue with await websocket.recv() after it has sent the close frame with await websocket.close(1000), as long as those messages were sent to it before it sent the close frame.

I believe this is more than an implementation detail, and that it's an important difference in behavior between uvicorn and websockets for this use case. I think uvicorn should do the same thing as websockets here.