Check if handshake is completed before sending frame on wsproto shutdown

encode · Oct 28, 2022 · f9fe643 · f9fe643
1 parent 5ba3f13
commit f9fe643
Showing 1 changed file with 16 additions and 13 deletions.
diff --git a/uvicorn/protocols/websockets/wsproto_impl.py b/uvicorn/protocols/websockets/wsproto_impl.py
@@ -125,9 +125,12 @@ def resume_writing(self):
         self.writable.set()
 
     def shutdown(self):
-        self.queue.put_nowait({"type": "websocket.disconnect", "code": 1012})
-        output = self.conn.send(wsproto.events.CloseConnection(code=1012))
-        self.transport.write(output)
+        if self.handshake_complete:
+            self.queue.put_nowait({"type": "websocket.disconnect", "code": 1012})
+            output = self.conn.send(wsproto.events.CloseConnection(code=1012))
+            self.transport.write(output)
+        else:
+            self.send_500_response()
         self.transport.close()
 
     def on_task_complete(self, task):
@@ -222,9 +225,8 @@ def send_500_response(self):
     async def run_asgi(self):
         try:
             result = await self.app(self.scope, self.receive, self.send)
-        except BaseException as exc:
-            msg = "Exception in ASGI application\n"
-            self.logger.error(msg, exc_info=exc)
+        except BaseException:
+            self.logger.exception("Exception in ASGI application\n")
             if not self.handshake_complete:
                 self.send_500_response()
             self.transport.close()
@@ -257,14 +259,15 @@ async def send(self, message):
                 extensions = []
                 if self.config.ws_per_message_deflate:
                     extensions.append(PerMessageDeflate())
-                output = self.conn.send(
-                    wsproto.events.AcceptConnection(
-                        subprotocol=subprotocol,
-                        extensions=extensions,
-                        extra_headers=extra_headers,
+                if not self.transport.is_closing():
+                    output = self.conn.send(
+                        wsproto.events.AcceptConnection(
+                            subprotocol=subprotocol,
+                            extensions=extensions,
+                            extra_headers=extra_headers,
+                        )
                     )
-                )
-                self.transport.write(output)
+                    self.transport.write(output)
 
             elif message_type == "websocket.close":
                 self.queue.put_nowait({"type": "websocket.disconnect", "code": None})