Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Send close frame after ASGI application returned #2332

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 26 additions & 3 deletions tests/protocols/test_websocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -555,9 +555,9 @@ async def connect(url: str):

@pytest.mark.anyio
async def test_asgi_return_value(ws_protocol_cls: WSProtocol, http_protocol_cls: HTTPProtocol, unused_tcp_port: int):
"""
The ASGI callable should return 'None'. If it doesn't, make sure that
the connection is closed with an error condition.
"""The ASGI callable should return 'None'.

If it doesn't, make sure that the connection is closed with an error condition.
"""

async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):
Expand All @@ -581,6 +581,29 @@ async def connect(url: str):
assert exc_info.value.code == 1006


@pytest.mark.anyio
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test would hang before the changes.

async def test_close_transport_on_asgi_return(
ws_protocol_cls: WSProtocol, http_protocol_cls: HTTPProtocol, unused_tcp_port: int
):
"""The ASGI callable should call the `websocket.close` event.

If it doesn't, the server should still send a close frame to the client.
"""

async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):
await send({"type": "websocket.accept"})

async def connect(url: str):
async with websockets.client.connect(url) as websocket:
_ = await websocket.recv()

config = Config(app=app, ws=ws_protocol_cls, http=http_protocol_cls, lifespan="off", port=unused_tcp_port)
async with run_server(config):
with pytest.raises(websockets.exceptions.ConnectionClosed) as exc_info:
await connect(f"ws://127.0.0.1:{unused_tcp_port}")
assert exc_info.value.code == 1006


@pytest.mark.anyio
@pytest.mark.parametrize("code", [None, 1000, 1001])
@pytest.mark.parametrize(
Expand Down
18 changes: 6 additions & 12 deletions uvicorn/protocols/websockets/websockets_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,28 +240,22 @@ async def run_asgi(self) -> None:
result = await self.app(self.scope, self.asgi_receive, self.asgi_send)
except ClientDisconnected:
self.closed_event.set()
self.transport.close()
except BaseException as exc:
except BaseException:
self.closed_event.set()
msg = "Exception in ASGI application\n"
self.logger.error(msg, exc_info=exc)
self.logger.exception("Exception in ASGI application\n")
if not self.handshake_started_event.is_set():
self.send_500_response()
else:
await self.handshake_completed_event.wait()
self.transport.close()
else:
self.closed_event.set()
if not self.handshake_started_event.is_set():
msg = "ASGI callable returned without sending handshake."
self.logger.error(msg)
self.logger.error("ASGI callable returned without sending handshake.")
self.send_500_response()
self.transport.close()
elif result is not None:
msg = "ASGI callable should return None, but returned '%s'."
self.logger.error(msg, result)
await self.handshake_completed_event.wait()
self.transport.close()
self.logger.error("ASGI callable should return None, but returned '%s'.", result)
await self.handshake_completed_event.wait()
self.transport.close()

async def asgi_send(self, message: ASGISendEvent) -> None:
message_type = message["type"]
Expand Down
12 changes: 4 additions & 8 deletions uvicorn/protocols/websockets/wsproto_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,21 +232,17 @@ async def run_asgi(self) -> None:
try:
result = await self.app(self.scope, self.receive, self.send)
except ClientDisconnected:
self.transport.close()
pass
except BaseException:
self.logger.exception("Exception in ASGI application\n")
self.send_500_response()
self.transport.close()
else:
if not self.handshake_complete:
msg = "ASGI callable returned without completing handshake."
self.logger.error(msg)
self.logger.error("ASGI callable returned without completing handshake.")
self.send_500_response()
self.transport.close()
elif result is not None:
msg = "ASGI callable should return None, but returned '%s'."
self.logger.error(msg, result)
self.transport.close()
self.logger.error("ASGI callable should return None, but returned '%s'.", result)
self.transport.close()

async def send(self, message: ASGISendEvent) -> None:
await self.writable.wait()
Expand Down
Loading