Allow to raise HTTPException before websocket.accept() (#2725)

* Allow to raise `HTTPException` before `websocket.accept()`

* move <<

* Add documentation

docs/exceptions.md

@@ -115,14 +115,30 @@ In order to deal with this behaviour correctly, the middleware stack of a
 
 ## HTTPException
 
-The `HTTPException` class provides a base class that you can use for any
-handled exceptions. The `ExceptionMiddleware` implementation defaults to
-returning plain-text HTTP responses for any `HTTPException`.
+The `HTTPException` class provides a base class that you can use for any handled exceptions.
+The `ExceptionMiddleware` implementation defaults to returning plain-text HTTP responses for any `HTTPException`.
 
 * `HTTPException(status_code, detail=None, headers=None)`
 
-You should only raise `HTTPException` inside routing or endpoints. Middleware
-classes should instead just return appropriate responses directly.
+You should only raise `HTTPException` inside routing or endpoints.
+Middleware classes should instead just return appropriate responses directly.
+
+You can use an `HTTPException` on a WebSocket endpoint in case it's raised before `websocket.accept()`.
+The connection is not upgraded to a WebSocket connection, and the proper HTTP response is returned.
+
+```python
+from starlette.applications import Starlette
+from starlette.exceptions import HTTPException
+from starlette.routing import WebSocketRoute
+from starlette.websockets import WebSocket
+
+
+async def websocket_endpoint(websocket: WebSocket):
+ raise HTTPException(status_code=400, detail="Bad request")
+
+
+app = Starlette(routes=[WebSocketRoute("/ws", websocket_endpoint)])
+```
 
 ## WebSocketException
 

starlette/_exception_handler.py

@@ -6,16 +6,7 @@
 from starlette.concurrency import run_in_threadpool
 from starlette.exceptions import HTTPException
 from starlette.requests import Request
-from starlette.types import (
- ASGIApp,
- ExceptionHandler,
- HTTPExceptionHandler,
- Message,
- Receive,
- Scope,
- Send,
- WebSocketExceptionHandler,
-)
+from starlette.types import ASGIApp, ExceptionHandler, Message, Receive, Scope, Send
 from starlette.websockets import WebSocket
 
 ExceptionHandlers = typing.Dict[typing.Any, ExceptionHandler]
@@ -62,24 +53,13 @@ async def sender(message: Message) -> None:
  raise exc
 
  if response_started:
- msg = "Caught handled exception, but response already started."
- raise RuntimeError(msg) from exc
-
- if scope["type"] == "http":
- nonlocal conn
- handler = typing.cast(HTTPExceptionHandler, handler)
- conn = typing.cast(Request, conn)
- if is_async_callable(handler):
- response = await handler(conn, exc)
- else:
- response = await run_in_threadpool(handler, conn, exc)
+ raise RuntimeError("Caught handled exception, but response already started.") from exc
+
+ if is_async_callable(handler):
+ response = await handler(conn, exc)
+ else:
+ response = await run_in_threadpool(handler, conn, exc) # type: ignore
+ if response is not None:
  await response(scope, receive, sender)
- elif scope["type"] == "websocket":
- handler = typing.cast(WebSocketExceptionHandler, handler)
- conn = typing.cast(WebSocket, conn)
- if is_async_callable(handler):
- await handler(conn, exc)
- else:
- await run_in_threadpool(handler, conn, exc)
 
  return wrapped_app

starlette/testclient.py

@@ -178,11 +178,7 @@ def _raise_on_close(self, message: Message) -> None:
  body.append(message["body"])
  if not message.get("more_body", False):
  break
- raise WebSocketDenialResponse(
- status_code=status_code,
- headers=headers,
- content=b"".join(body),
- )
+ raise WebSocketDenialResponse(status_code=status_code, headers=headers, content=b"".join(body))
 
  def send(self, message: Message) -> None:
  self._receive_queue.put(message)

tests/test_applications.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import AsyncGenerator, AsyncIterator, Generator
 
-import anyio
+import anyio.from_thread
 import pytest
 
 from starlette import status
@@ -17,7 +17,7 @@
 from starlette.responses import JSONResponse, PlainTextResponse
 from starlette.routing import Host, Mount, Route, Router, WebSocketRoute
 from starlette.staticfiles import StaticFiles
-from starlette.testclient import TestClient
+from starlette.testclient import TestClient, WebSocketDenialResponse
 from starlette.types import ASGIApp, Receive, Scope, Send
 from starlette.websockets import WebSocket
 from tests.types import TestClientFactory
@@ -71,11 +71,15 @@ async def websocket_endpoint(session: WebSocket) -> None:
  await session.close()
 
 
-async def websocket_raise_websocket(websocket: WebSocket) -> None:
+async def websocket_raise_websocket_exception(websocket: WebSocket) -> None:
  await websocket.accept()
  raise WebSocketException(code=status.WS_1003_UNSUPPORTED_DATA)
 
 
+async def websocket_raise_http_exception(websocket: WebSocket) -> None:
+ raise HTTPException(status_code=401, detail="Unauthorized")
+
+
 class CustomWSException(Exception):
  pass
 
@@ -118,7 +122,8 @@ def custom_ws_exception_handler(websocket: WebSocket, exc: CustomWSException) ->
  Route("/class", endpoint=Homepage),
  Route("/500", endpoint=runtime_error),
  WebSocketRoute("/ws", endpoint=websocket_endpoint),
- WebSocketRoute("/ws-raise-websocket", endpoint=websocket_raise_websocket),
+ WebSocketRoute("/ws-raise-websocket", endpoint=websocket_raise_websocket_exception),
+ WebSocketRoute("/ws-raise-http", endpoint=websocket_raise_http_exception),
  WebSocketRoute("/ws-raise-custom", endpoint=websocket_raise_custom),
  Mount("/users", app=users),
  Host("{subdomain}.example.org", app=subdomain),
@@ -219,6 +224,14 @@ def test_websocket_raise_websocket_exception(client: TestClient) -> None:
  }
 
 
+def test_websocket_raise_http_exception(client: TestClient) -> None:
+ with pytest.raises(WebSocketDenialResponse) as exc:
+ with client.websocket_connect("/ws-raise-http"):
+ pass # pragma: no cover
+ assert exc.value.status_code == 401
+ assert exc.value.content == b'{"detail":"Unauthorized"}'
+
+
 def test_websocket_raise_custom_exception(client: TestClient) -> None:
  with client.websocket_connect("/ws-raise-custom") as session:
  response = session.receive()
@@ -243,7 +256,8 @@ def test_routes() -> None:
  Route("/class", endpoint=Homepage),
  Route("/500", endpoint=runtime_error, methods=["GET"]),
  WebSocketRoute("/ws", endpoint=websocket_endpoint),
- WebSocketRoute("/ws-raise-websocket", endpoint=websocket_raise_websocket),
+ WebSocketRoute("/ws-raise-websocket", endpoint=websocket_raise_websocket_exception),
+ WebSocketRoute("/ws-raise-http", endpoint=websocket_raise_http_exception),
  WebSocketRoute("/ws-raise-custom", endpoint=websocket_raise_custom),
  Mount(
  "/users",

tests/test_exceptions.py

@@ -63,11 +63,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  Route("/with_headers", endpoint=with_headers),
  Route("/handled_exc_after_response", endpoint=HandledExcAfterResponse()),
  WebSocketRoute("/runtime_error", endpoint=raise_runtime_error),
- Route(
- "/consume_body_in_endpoint_and_handler",
- endpoint=read_body_and_raise_exc,
- methods=["POST"],
- ),
+ Route("/consume_body_in_endpoint_and_handler", endpoint=read_body_and_raise_exc, methods=["POST"]),
  ]
 )
 
@@ -114,13 +110,10 @@ def test_websockets_should_raise(client: TestClient) -> None:
  pass # pragma: no cover
 
 
-def test_handled_exc_after_response(
- test_client_factory: TestClientFactory,
- client: TestClient,
-) -> None:
+def test_handled_exc_after_response(test_client_factory: TestClientFactory, client: TestClient) -> None:
  # A 406 HttpException is raised *after* the response has already been sent.
  # The exception middleware should raise a RuntimeError.
- with pytest.raises(RuntimeError):
+ with pytest.raises(RuntimeError, match="Caught handled exception, but response already started."):
  client.get("/handled_exc_after_response")
 
  # If `raise_server_exceptions=False` then the test client will still allow
@@ -132,7 +125,7 @@ def test_handled_exc_after_response(
 
 
 def test_force_500_response(test_client_factory: TestClientFactory) -> None:
- # use a sentinal variable to make sure we actually
+ # use a sentinel variable to make sure we actually
  # make it into the endpoint and don't get a 500
  # from an incorrect ASGI app signature or something
  called = False