Send code 1012 on shutdown for websockets #1816

Kludex · 2022-12-26T14:09:12Z

Closes Incorrect websocket close code on server shutdown #1344

Changes

Send 500 or 1012 to the client on server shutdown, instead of the previous 1006.
Improve test suite to check what the client receives, instead of only the ASGI application.
Increase coverage from 98.5% to 98.8%.

uvicorn/protocols/websockets/websockets_impl.py

Kludex · 2022-12-26T14:10:10Z

uvicorn/protocols/websockets/websockets_impl.py

@@ -90,7 +90,6 @@ def __init__(
        self.connect_sent = False
        self.lost_connection_before_handshake = False
        self.accepted_subprotocol: Optional[Subprotocol] = None
-        self.transfer_data_task: asyncio.Task = None  # type: ignore[assignment]


This is necessary to call fail_connection(), because internally there's a check like if hasattr(self, "transfer_data_task"), which should succeed.

Kludex · 2022-12-26T20:47:18Z

uvicorn/protocols/websockets/wsproto_impl.py

-        if self.conn.connection is None:
-            output = self.conn.send(wsproto.events.RejectConnection(status_code=500))
-        else:
-            msg = h11.Response(
-                status_code=500, headers=headers, reason="Internal Server Error"
+        output = self.conn.send(
+            wsproto.events.RejectConnection(
+                status_code=500, headers=headers, has_body=True
            )
-            output = self.conn.send(msg)
-            msg = h11.Data(data=b"Internal Server Error")
-            output += self.conn.send(msg)
-            msg = h11.EndOfMessage()
-            output += self.conn.send(msg)
+        )
+        output += self.conn.send(
+            wsproto.events.RejectData(data=b"Internal Server Error")
+        )


Don't be afraid here. I'll explain what's happening.

We were sending RejectConnection(status_code=500) without a body, but on websockets implementation we were sending the "Internal Server Error" body on the analogous behavior. The RejectData matches the behavior.

We are removing the conditional because it's never reached, and the reason for it is that we only call send_500_response is we didn't complete the handshake - which makes a lot of sense to remove it.

Hope it's clear.

Am I missing EndOfMessage? I need to check this.

No, I'm not missing it. wsproto internally adds the h11.EndOfMessage. 🙏

Kludex · 2022-12-26T20:48:38Z

uvicorn/protocols/websockets/websockets_impl.py

+        if self.handshake_completed_event.is_set():
+            self.fail_connection(1012)
+        else:
+            self.send_500_response()


This is the change that motivated this PR.

We were only sending 1006 to the client, even when the handshake was not completed.

Why self.fail_connection(1012) instead of just calling self.close(1012)?

I've been look at websockets server.close() and procotol.close() to see why uvicorn isn't calling `close(), but haven't gotten terribly far yet.

Because close() is async.

Ahh, missed the lack of async here.

Looking deeper at the difference, it seems like the primary difference is fail_connection will proactively cancel the data transfer task but close will simply send the close frame and then wait.

Is there any consequence to that distinction here? My very, very basic testing seems to suggest with uvicorn tasks are given a chance to shutdown cleanly, but I'm still testing.

Yeah, this code still waits for it to finish. 👍

Kludex · 2022-12-26T20:49:03Z

tests/protocols/test_websocket.py

+    assert response is not None
+    assert response.status_code == 500, response.text
+    assert response.text == "Internal Server Error"


We are now testing what the client receives, instead of only the ASGI application.

Kludex · 2022-12-26T21:21:51Z

tests/protocols/test_websocket.py

+    task.cancel()
+    assert response is not None
+    assert response.status_code == 500, response.text
+    assert response.text == "Internal Server Error"
    assert disconnect_message == {"type": "websocket.disconnect", "code": 1006}


Hmm... On a second thought... This doesn't make sense, does it? Why are we even sending a websocket.disconnect when the handshake was not even completed? 🤔

I think this was on purpose because then the application could receive a websocket event, but thinking about it again, does it make sense?

I need to check the issues/PRs about this decision. It shouldn't be a blocker for this PR tho.

According to django/asgiref#364 it's fine.

jalaziz · 2022-12-26T23:19:06Z

Can confirm that wscat, Chrome, and Firefox all show the correct close codes. I also checked that the client receives the correct close code if an app sends {"type": "websocket.close", "code": code} and that works too.

~~Just trying to figure out what's going on with the websockets client.~~
~~Figured it out, websockets is fine, it's just the __main__.py script doesn't show the correct code when the server initiates the disconnect.~~
I'm an idiot and was accidentally using an testing with an older version of websockets. The latest version (10.4) shows the correct codes.

tomchristie

So PRs like this do highlight to me how awkward it is to try to follow the flow when using protocols rather than streams.

Anyways.

Kludex · 2023-01-06T12:07:14Z

So PRs like this do highlight to me how awkward it is to try to follow the flow when using protocols rather than streams.

Anyways.

I agree. 🤷

Send code 1012 on shutdown for websockets

782f396

Kludex commented Dec 26, 2022

View reviewed changes

uvicorn/protocols/websockets/websockets_impl.py Outdated Show resolved Hide resolved

Kludex commented Dec 26, 2022

View reviewed changes

Kludex mentioned this pull request Dec 26, 2022

Incorrect websocket close code on server shutdown #1344

Closed

2 tasks

Kludex force-pushed the send-1012-on-shutdown branch from 4aeb55c to eb21ef0 Compare December 26, 2022 20:35

Improve tests

ed7d669

Kludex force-pushed the send-1012-on-shutdown branch from eb21ef0 to ed7d669 Compare December 26, 2022 20:39

Remove unused code on send_500_response on wsproto

e4f5d2a

Kludex commented Dec 26, 2022

View reviewed changes

Kludex added 2 commits December 26, 2022 21:49

Remove unused import

01f1a8f

Decrease coverage

abc9db3

Kludex requested review from tomchristie and a team December 26, 2022 20:59

Kludex commented Dec 26, 2022

View reviewed changes

tomchristie approved these changes Jan 6, 2023

View reviewed changes

Kludex merged commit 23b9f05 into master Jan 6, 2023

Kludex deleted the send-1012-on-shutdown branch January 6, 2023 12:07

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Send code 1012 on shutdown for websockets #1816

Send code 1012 on shutdown for websockets #1816

Kludex commented Dec 26, 2022 •

edited

Kludex Dec 26, 2022 •

edited

Kludex Dec 26, 2022

Kludex Dec 27, 2022

Kludex Dec 29, 2022

Kludex Dec 26, 2022

jalaziz Dec 26, 2022

Kludex Dec 26, 2022 •

edited

jalaziz Dec 26, 2022

Kludex Dec 29, 2022

Kludex Dec 26, 2022

Kludex Dec 26, 2022 •

edited

Kludex Dec 26, 2022

Kludex Dec 27, 2022

jalaziz commented Dec 26, 2022 •

edited

tomchristie left a comment

Kludex commented Jan 6, 2023

Send code 1012 on shutdown for websockets #1816

Send code 1012 on shutdown for websockets #1816

Conversation

Kludex commented Dec 26, 2022 • edited

Changes

Kludex Dec 26, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Kludex Dec 26, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Kludex Dec 26, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jalaziz commented Dec 26, 2022 • edited

tomchristie left a comment

Choose a reason for hiding this comment

Kludex commented Jan 6, 2023

Kludex commented Dec 26, 2022 •

edited

Kludex Dec 26, 2022 •

edited

Kludex Dec 26, 2022 •

edited

Kludex Dec 26, 2022 •

edited

jalaziz commented Dec 26, 2022 •

edited