Skip to content

Commit

Permalink
Check if handshake is completed before sending frame on wsproto shutdown
Browse files Browse the repository at this point in the history
  • Loading branch information
Kludex committed Oct 28, 2022
1 parent 5ba3f13 commit f9fe643
Showing 1 changed file with 16 additions and 13 deletions.
29 changes: 16 additions & 13 deletions uvicorn/protocols/websockets/wsproto_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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})
Expand Down

0 comments on commit f9fe643

Please sign in to comment.