Move exception handling logic to Route

starlette/middleware/exceptions/__init__.py

@@ -0,0 +1,76 @@
+import typing
+
+from starlette.exceptions import HTTPException, WebSocketException
+from starlette.middleware.exceptions._wrapper import (
+    ExceptionHandlers,
+    StatusHandlers,
+    wrap_app_handling_exceptions,
+)
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse, Response
+from starlette.types import ASGIApp, Receive, Scope, Send
+from starlette.websockets import WebSocket
+
+
+class ExceptionMiddleware:
+    def __init__(
+        self,
+        app: ASGIApp,
+        handlers: typing.Optional[
+            typing.Mapping[typing.Any, typing.Callable[[Request, Exception], Response]]
+        ] = None,
+        debug: bool = False,
+    ) -> None:
+        self.app = app
+        self.debug = debug  # TODO: We ought to handle 404 cases if debug is set.
+        self._status_handlers: StatusHandlers = {}
+        self._exception_handlers: ExceptionHandlers = {
+            HTTPException: self.http_exception,
+            WebSocketException: self.websocket_exception,  # type: ignore[dict-item]
+        }
+        if handlers is not None:
+            for key, value in handlers.items():
+                self.add_exception_handler(key, value)
+
+    def add_exception_handler(
+        self,
+        exc_class_or_status_code: typing.Union[int, typing.Type[Exception]],
+        handler: typing.Callable[[Request, Exception], Response],
+    ) -> None:
+        if isinstance(exc_class_or_status_code, int):
+            self._status_handlers[exc_class_or_status_code] = handler
+        else:
+            assert issubclass(exc_class_or_status_code, Exception)
+            self._exception_handlers[exc_class_or_status_code] = handler
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] not in ("http", "websocket"):
+            await self.app(scope, receive, send)
+            return
+
+        scope["starlette.exception_handlers"] = (
+            self._exception_handlers,
+            self._status_handlers,
+        )
+
+        conn: typing.Union[Request, WebSocket]
+        if scope["type"] == "http":
+            conn = Request(scope, receive, send)
+        else:
+            conn = WebSocket(scope, receive, send)
+
+        await wrap_app_handling_exceptions(
+            self.app, self._exception_handlers, self._status_handlers, conn
+        )(scope, receive, send)
+
+    def http_exception(self, request: Request, exc: Exception) -> Response:
+        assert isinstance(exc, HTTPException)
+        if exc.status_code in {204, 304}:
+            return Response(status_code=exc.status_code, headers=exc.headers)
+        return PlainTextResponse(
+            exc.detail, status_code=exc.status_code, headers=exc.headers
+        )
+
+    async def websocket_exception(self, websocket: WebSocket, exc: Exception) -> None:
+        assert isinstance(exc, WebSocketException)
+        await websocket.close(code=exc.code, reason=exc.reason)  # pragma: no cover

starlette/middleware/exceptions/_wrapper.py

@@ -0,0 +1,72 @@
+import typing
+
+from starlette._utils import is_async_callable
+from starlette.concurrency import run_in_threadpool
+from starlette.exceptions import HTTPException
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+from starlette.websockets import WebSocket
+
+Handler = typing.Callable[..., typing.Any]
+ExceptionHandlers = typing.Dict[typing.Any, Handler]
+StatusHandlers = typing.Dict[int, Handler]
+
+
+def _lookup_exception_handler(
+    exc_handlers: ExceptionHandlers, exc: Exception
+) -> typing.Optional[Handler]:
+    for cls in type(exc).__mro__:
+        if cls in exc_handlers:
+            return exc_handlers[cls]
+    return None
+
+
+def wrap_app_handling_exceptions(
+    app: ASGIApp,
+    exc_handlers: ExceptionHandlers,
+    status_handlers: StatusHandlers,
+    conn: typing.Union[Request, WebSocket],
+) -> ASGIApp:
+    async def wrapped_app(scope: Scope, receive: Receive, send: Send) -> None:
+        response_started = False
+
+        async def sender(message: Message) -> None:
+            nonlocal response_started
+
+            if message["type"] == "http.response.start":
+                response_started = True
+            await send(message)
+
+        try:
+            await app(scope, receive, sender)
+        except Exception as exc:
+            handler = None
+
+            if isinstance(exc, HTTPException):
+                handler = status_handlers.get(exc.status_code)
+
+            if handler is None:
+                handler = _lookup_exception_handler(exc_handlers, exc)
+
+            if handler is None:
+                raise exc
+
+            if response_started:
+                msg = "Caught handled exception, but response already started."
+                raise RuntimeError(msg) from exc
+
+            if scope["type"] == "http":
+                response: Response
+                if is_async_callable(handler):
+                    response = await handler(conn, exc)
+                else:
+                    response = await run_in_threadpool(handler, conn, exc)
+                await response(scope, receive, sender)
+            elif scope["type"] == "websocket":
+                if is_async_callable(handler):
+                    await handler(conn, exc)
+                else:
+                    await run_in_threadpool(handler, conn, exc)
+
+    return wrapped_app

starlette/routing.py

@@ -15,6 +15,7 @@
 from starlette.datastructures import URL, Headers, URLPath
 from starlette.exceptions import HTTPException
 from starlette.middleware import Middleware
+from starlette.middleware.exceptions._wrapper import wrap_app_handling_exceptions
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse, RedirectResponse
 from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send
@@ -61,12 +62,23 @@ def request_response(func: typing.Callable) -> ASGIApp:
     is_coroutine = is_async_callable(func)
 
     async def app(scope: Scope, receive: Receive, send: Send) -> None:
-        request = Request(scope, receive=receive, send=send)
-        if is_coroutine:
-            response = await func(request)
-        else:
-            response = await run_in_threadpool(func, request)
-        await response(scope, receive, send)
+        request = Request(scope, receive, send)
+
+        async def app(scope: Scope, receive: Receive, send: Send) -> None:
+            if is_coroutine:
+                response = await func(request)
+            else:
+                response = await run_in_threadpool(func, request)
+            await response(scope, receive, send)
+
+        try:
+            exception_handlers, status_handlers = scope["starlette.exception_handlers"]
+        except KeyError:
+            exception_handlers, status_handlers = {}, {}
+
+        await wrap_app_handling_exceptions(
+            app, exception_handlers, status_handlers, request
+        )(scope, receive, send)
 
     return app
 
@@ -79,7 +91,21 @@ def websocket_session(func: typing.Callable) -> ASGIApp:
 
     async def app(scope: Scope, receive: Receive, send: Send) -> None:
         session = WebSocket(scope, receive=receive, send=send)
-        await func(session)
+
+        async def app(scope: Scope, receive: Receive, send: Send) -> None:
+            await func(session)
+
+        try:
+            exception_handlers, status_handlers = scope["starlette.exception_handlers"]
+        except KeyError:
+            exception_handlers, status_handlers = {}, {}
+
+        await wrap_app_handling_exceptions(
+            app,
+            exception_handlers,
+            status_handlers,
+            session,
+        )(scope, receive, send)
 
     return app
 

tests/test_exceptions.py

@@ -4,7 +4,8 @@
 
 from starlette.exceptions import HTTPException, WebSocketException
 from starlette.middleware.exceptions import ExceptionMiddleware
-from starlette.responses import PlainTextResponse
+from starlette.requests import Request
+from starlette.responses import JSONResponse, PlainTextResponse
 from starlette.routing import Route, Router, WebSocketRoute
 
 
@@ -28,6 +29,22 @@ def with_headers(request):
     raise HTTPException(status_code=200, headers={"x-potato": "always"})
 
 
+class BadBodyException(HTTPException):
+    pass
+
+
+async def read_body_and_raise_exc(request: Request):
+    await request.body()
+    raise BadBodyException(422)
+
+
+async def handler_that_reads_body(
+    request: Request, exc: BadBodyException
+) -> JSONResponse:
+    body = await request.body()
+    return JSONResponse(status_code=422, content={"body": body.decode()})
+
+
 class HandledExcAfterResponse:
     async def __call__(self, scope, receive, send):
         response = PlainTextResponse("OK", status_code=200)
@@ -44,11 +61,19 @@ async def __call__(self, scope, receive, send):
         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"],
+        ),
     ]
 )
 
 
-app = ExceptionMiddleware(router)
+app = ExceptionMiddleware(
+    router,
+    handlers={BadBodyException: handler_that_reads_body},  # type: ignore[dict-item]
+)
 
 
 @pytest.fixture
@@ -160,3 +185,9 @@ def test_exception_middleware_deprecation() -> None:
 
     with pytest.warns(DeprecationWarning):
         starlette.exceptions.ExceptionMiddleware
+
+
+def test_request_in_app_and_handler_is_the_same_object(client) -> None:
+    response = client.post("/consume_body_in_endpoint_and_handler", content=b"Hello!")
+    assert response.status_code == 422
+    assert response.json() == {"body": "Hello!"}

tests/test_routing.py

@@ -1033,13 +1033,9 @@ async def modified_send(msg: Message) -> None:
     assert resp.status_code == 200, resp.content
     assert "X-Mounted" in resp.headers
 
-    # this is the "surprising" behavior bit
-    # the middleware on the mount never runs because there
-    # is nothing to catch the HTTPException
-    # since Mount middlweare is not wrapped by ExceptionMiddleware
     resp = client.get("/mount/err")
     assert resp.status_code == 403, resp.content
-    assert "X-Mounted" not in resp.headers
+    assert "X-Mounted" in resp.headers
 
 
 def test_route_repr() -> None: