Skip to content

Commit 2aa70ab

Browse files
yezz123yezz123
authored
🐛 Add custom exception handlers and enhance OpenAPI security in AuthX (#873)
* ✨ Add custom exception handlers for AuthX exceptions in FastAPI - Implemented methods to set request exception handlers for various AuthX exceptions, ensuring proper error responses. - Added a method to ensure default exception handlers are installed for active requests, preventing 500 responses for auth errors. - Included specific handlers for exceptions like MissingCSRFTokenError, LoginTypeMismatchError, and RateLimitExceeded, among others. * ✨ Enhance OpenAPI security schemes in AuthX - Introduced new methods for handling OpenAPI security dependencies, including support for header, cookie, and query token locations. - Added descriptions for JWT usage in OpenAPI documentation to improve clarity for users. - Refactored token validation dependencies in `AuthX` and `AuthManager` to utilize the new security schemes, ensuring a more flexible and robust authentication process. * 🔧 Add comprehensive tests for OpenAPI security in AuthX - Introduced a new test suite for OpenAPI security features, validating the integration of access token requirements and error handling. - Implemented tests to ensure proper OpenAPI metadata is generated for various token locations (headers, cookies, query). - Enhanced error handling tests for revoked, missing, and invalid tokens, ensuring consistent responses across different scenarios. * ⬆️ Update dependencies in `pyproject.toml` and `uv.lock` - Added `uvicorn` version `0.39.0` to the dependencies. - Updated `anyio` from `4.7.0` to `4.12.1` and added a new version `4.14.0` with updated resolution markers. - Introduced `httpx2` version `2.0.0` to the test dependencies. - Adjusted Python version resolution markers for compatibility. * 🐛 Add tests for request exception and rate limit handlers in error_handler - Implemented tests to verify the functionality of request exception handlers for MissingTokenError and RateLimitExceeded. - Ensured that existing handlers are preserved when setting new handlers. - Added checks for requests without starlette scope to confirm proper behavior. - Enhanced OpenAPI security metadata tests to include custom header token locations.
2 parents 7b2a373 + b870fa9 commit 2aa70ab

11 files changed

Lines changed: 2756 additions & 1135 deletions

‎authx/_internal/_error.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,141 @@ async def exception_handler_wrapper(request: Request, exc: exceptions.AuthXExcep
9090
# The exception handler will return a JSONResponse with the specified status code and message
9191
app.exception_handler(exception)(exception_handler_wrapper)
9292

93+
def _has_request_exception_handler(
94+
self,
95+
request: Request,
96+
exception: type[exceptions.AuthXException],
97+
) -> bool:
98+
handlers_scope = request.scope.get("starlette.exception_handlers")
99+
if handlers_scope is None:
100+
return False
101+
exception_handlers, _ = handlers_scope
102+
return any(cls in exception_handlers for cls in exception.__mro__)
103+
104+
def _set_request_exception_handler(
105+
self,
106+
request: Request,
107+
exception: type[exceptions.AuthXException],
108+
status_code: int,
109+
message: Optional[str],
110+
) -> None:
111+
handlers_scope = request.scope.get("starlette.exception_handlers")
112+
if handlers_scope is None or self._has_request_exception_handler(request, exception):
113+
return
114+
exception_handlers, _ = handlers_scope
115+
116+
async def exception_handler_wrapper(request: Request, exc: exceptions.AuthXException) -> JSONResponse:
117+
return await self._error_handler(request, exc, status_code, message)
118+
119+
exception_handlers[exception] = exception_handler_wrapper
120+
121+
def _set_request_rate_limit_handler(self, request: Request) -> None:
122+
handlers_scope = request.scope.get("starlette.exception_handlers")
123+
if handlers_scope is None or self._has_request_exception_handler(request, exceptions.RateLimitExceeded):
124+
return
125+
exception_handlers, _ = handlers_scope
126+
127+
async def rate_limit_wrapper(request: Request, exc: exceptions.RateLimitExceeded) -> JSONResponse:
128+
return await self._rate_limit_handler(request, exc)
129+
130+
exception_handlers[exceptions.RateLimitExceeded] = rate_limit_wrapper
131+
132+
def ensure_request_exception_handlers(self, request: Request) -> None:
133+
"""Install default AuthX handlers for the active request if the app has none.
134+
135+
FastAPI exception handlers are normally registered with ``handle_errors(app)``.
136+
Minimal apps that omit it should still return HTTP auth errors instead of
137+
leaking AuthX exceptions as 500 responses.
138+
"""
139+
# Register subclasses before their parent classes so detailed handlers win
140+
# during Starlette's MRO-based exception lookup.
141+
self._set_request_exception_handler(
142+
request,
143+
exception=exceptions.MissingCSRFTokenError,
144+
status_code=401,
145+
message=None,
146+
)
147+
self._set_request_exception_handler(
148+
request,
149+
exception=exceptions.LoginTypeMismatchError,
150+
status_code=401,
151+
message=None,
152+
)
153+
self._set_request_exception_handler(request, exception=exceptions.JWTDecodeError, status_code=422, message=None)
154+
self._set_request_exception_handler(
155+
request,
156+
exception=exceptions.MissingTokenError,
157+
status_code=401,
158+
message=self.MSG_TokenError,
159+
)
160+
self._set_request_exception_handler(
161+
request,
162+
exception=exceptions.RevokedTokenError,
163+
status_code=401,
164+
message=self.MSG_RevokedTokenError,
165+
)
166+
self._set_request_exception_handler(
167+
request,
168+
exception=exceptions.TokenRequiredError,
169+
status_code=401,
170+
message=self.MSG_TokenRequiredError,
171+
)
172+
self._set_request_exception_handler(
173+
request,
174+
exception=exceptions.FreshTokenRequiredError,
175+
status_code=401,
176+
message=self.MSG_FreshTokenRequiredError,
177+
)
178+
self._set_request_exception_handler(
179+
request,
180+
exception=exceptions.AccessTokenRequiredError,
181+
status_code=401,
182+
message=self.MSG_AccessTokenRequiredError,
183+
)
184+
self._set_request_exception_handler(
185+
request,
186+
exception=exceptions.RefreshTokenRequiredError,
187+
status_code=401,
188+
message=self.MSG_RefreshTokenRequiredError,
189+
)
190+
self._set_request_exception_handler(
191+
request,
192+
exception=exceptions.TokenTypeError,
193+
status_code=401,
194+
message=self.MSG_TokenTypeError,
195+
)
196+
self._set_request_exception_handler(
197+
request,
198+
exception=exceptions.CSRFError,
199+
status_code=401,
200+
message=self.MSG_CSRFError,
201+
)
202+
self._set_request_exception_handler(
203+
request,
204+
exception=exceptions.InsufficientScopeError,
205+
status_code=403,
206+
message=None,
207+
)
208+
self._set_request_exception_handler(
209+
request,
210+
exception=exceptions.PolicyDeniedError,
211+
status_code=403,
212+
message=None,
213+
)
214+
self._set_request_exception_handler(
215+
request,
216+
exception=exceptions.PolicyEvaluationError,
217+
status_code=500,
218+
message=self.MSG_PolicyEvaluationError,
219+
)
220+
self._set_request_exception_handler(
221+
request,
222+
exception=exceptions.SessionRevoked,
223+
status_code=401,
224+
message="Session has been revoked",
225+
)
226+
self._set_request_rate_limit_handler(request)
227+
93228
def handle_errors(self, app: FastAPI) -> None:
94229
"""Add the `FastAPI.exception_handlers` relative to AuthX exceptions.
95230

‎authx/main.py‎

Lines changed: 97 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
)
1313

1414
from fastapi import Depends, Request, Response, WebSocket
15+
from fastapi.security import APIKeyCookie, APIKeyHeader, APIKeyQuery, HTTPBearer
1516

1617
from authx._internal._callback import _CallbackHandler
1718
from authx._internal._error import _ErrorHandler
@@ -39,6 +40,21 @@
3940
)
4041

4142

43+
def _noop_openapi_security() -> None:
44+
"""Placeholder dependency for token locations that OpenAPI cannot represent."""
45+
return None
46+
47+
48+
_OPENAPI_BEARER_DESCRIPTION = (
49+
"Paste an AuthX JWT from create_access_token/create_refresh_token, not JWT_SECRET_KEY. "
50+
"The Bearer prefix is optional in Swagger UI."
51+
)
52+
_OPENAPI_HEADER_DESCRIPTION = "Provide an AuthX JWT in this header, not JWT_SECRET_KEY."
53+
_OPENAPI_ACCESS_COOKIE_DESCRIPTION = "Provide an AuthX access token in this cookie."
54+
_OPENAPI_REFRESH_COOKIE_DESCRIPTION = "Provide an AuthX refresh token in this cookie."
55+
_OPENAPI_QUERY_DESCRIPTION = "Provide an AuthX JWT in this query parameter."
56+
57+
4258
class AuthX(_CallbackHandler[T], _ErrorHandler):
4359
"""The base class for AuthX.
4460
@@ -763,6 +779,56 @@ def get_dependency(self, request: Request, response: Response) -> AuthXDependenc
763779
"""
764780
return AuthXDependency(self, request=request, response=response)
765781

782+
def _openapi_header_security_scheme(self) -> Callable[..., Any]:
783+
if self.config.JWT_HEADER_NAME.lower() == "authorization" and self.config.JWT_HEADER_TYPE.lower() == "bearer":
784+
return HTTPBearer(
785+
scheme_name="AuthXBearer",
786+
bearerFormat="JWT",
787+
description=_OPENAPI_BEARER_DESCRIPTION,
788+
auto_error=False,
789+
)
790+
return APIKeyHeader(
791+
name=self.config.JWT_HEADER_NAME,
792+
scheme_name="AuthXHeader",
793+
description=_OPENAPI_HEADER_DESCRIPTION,
794+
auto_error=False,
795+
)
796+
797+
def _openapi_cookie_security_scheme(self, type: str) -> Callable[..., Any]:
798+
if type == "refresh":
799+
return APIKeyCookie(
800+
name=self.config.JWT_REFRESH_COOKIE_NAME,
801+
scheme_name="AuthXRefreshCookie",
802+
description=_OPENAPI_REFRESH_COOKIE_DESCRIPTION,
803+
auto_error=False,
804+
)
805+
return APIKeyCookie(
806+
name=self.config.JWT_ACCESS_COOKIE_NAME,
807+
scheme_name="AuthXAccessCookie",
808+
description=_OPENAPI_ACCESS_COOKIE_DESCRIPTION,
809+
auto_error=False,
810+
)
811+
812+
def _openapi_query_security_scheme(self) -> Callable[..., Any]:
813+
return APIKeyQuery(
814+
name=self.config.JWT_QUERY_STRING_NAME,
815+
scheme_name="AuthXQuery",
816+
description=_OPENAPI_QUERY_DESCRIPTION,
817+
auto_error=False,
818+
)
819+
820+
def _openapi_security_dependencies(
821+
self,
822+
type: str = "access",
823+
locations: Optional[TokenLocations] = None,
824+
) -> tuple[Callable[..., Any], Callable[..., Any], Callable[..., Any]]:
825+
effective_locations = locations if locations is not None else self.config.JWT_TOKEN_LOCATION
826+
return (
827+
self._openapi_header_security_scheme() if "headers" in effective_locations else _noop_openapi_security,
828+
self._openapi_cookie_security_scheme(type) if "cookies" in effective_locations else _noop_openapi_security,
829+
self._openapi_query_security_scheme() if "query" in effective_locations else _noop_openapi_security,
830+
)
831+
766832
def token_required(
767833
self,
768834
type: str = "access",
@@ -783,8 +849,21 @@ def token_required(
783849
Returns:
784850
Callable[[Request], TokenPayload]: Dependency for Valid token Payload retrieval
785851
"""
786-
787-
async def _auth_required(request: Request) -> Any:
852+
header_scheme, cookie_scheme, query_scheme = self._openapi_security_dependencies(
853+
type=type,
854+
locations=locations,
855+
)
856+
header_dependency = Depends(header_scheme)
857+
cookie_dependency = Depends(cookie_scheme)
858+
query_dependency = Depends(query_scheme)
859+
860+
async def _auth_required(
861+
request: Request,
862+
_authx_openapi_header: Any = header_dependency,
863+
_authx_openapi_cookie: Any = cookie_dependency,
864+
_authx_openapi_query: Any = query_dependency,
865+
) -> Any:
866+
self.ensure_request_exception_handlers(request)
788867
return await self._auth_required(
789868
request=request,
790869
type=type,
@@ -876,8 +955,21 @@ async def admin_users(): ...
876955
```
877956
"""
878957
required_scopes = list(scopes)
879-
880-
async def _scopes_required(request: Request) -> TokenPayload:
958+
header_scheme, cookie_scheme, query_scheme = self._openapi_security_dependencies(
959+
type="access",
960+
locations=locations,
961+
)
962+
header_dependency = Depends(header_scheme)
963+
cookie_dependency = Depends(cookie_scheme)
964+
query_dependency = Depends(query_scheme)
965+
966+
async def _scopes_required(
967+
request: Request,
968+
_authx_openapi_header: Any = header_dependency,
969+
_authx_openapi_cookie: Any = cookie_dependency,
970+
_authx_openapi_query: Any = query_dependency,
971+
) -> TokenPayload:
972+
self.ensure_request_exception_handlers(request)
881973
payload = await self._auth_required(
882974
request=request,
883975
type="access",
@@ -909,6 +1001,7 @@ async def get_current_subject(self, request: Request) -> Optional[T]:
9091001
Returns:
9101002
The authenticated subject if present, otherwise None.
9111003
"""
1004+
self.ensure_request_exception_handlers(request)
9121005
token: TokenPayload = await self._auth_required(request=request)
9131006
uid = token.sub
9141007
return await self._get_current_subject(uid=uid)

‎authx/manager.py‎

Lines changed: 51 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@
33
from collections.abc import Awaitable, Callable, Mapping, Sequence
44
from typing import Any, Optional
55

6-
from fastapi import Request
6+
from fastapi import Depends, Request
7+
from fastapi.security import HTTPBearer
78

89
from authx._internal._error import _ErrorHandler
910
from authx.exceptions import (
@@ -13,7 +14,7 @@
1314
PolicyDeniedError,
1415
RevokedTokenError,
1516
)
16-
from authx.main import AuthX
17+
from authx.main import _OPENAPI_BEARER_DESCRIPTION, AuthX, _noop_openapi_security
1718
from authx.policy import (
1819
PolicyContext,
1920
PolicyEngine,
@@ -140,8 +141,22 @@ def token_required(
140141
locations: Optional[TokenLocations] = None,
141142
) -> Callable[[Request], Awaitable[TokenPayload]]:
142143
"""Dependency factory requiring a token for a specific login type."""
143-
144-
async def _auth_required(request: Request) -> TokenPayload:
144+
header_scheme, cookie_scheme, query_scheme = self._openapi_security_dependencies(
145+
login_type=login_type,
146+
type=type,
147+
locations=locations,
148+
)
149+
header_dependency = Depends(header_scheme)
150+
cookie_dependency = Depends(cookie_scheme)
151+
query_dependency = Depends(query_scheme)
152+
153+
async def _auth_required(
154+
request: Request,
155+
_authx_openapi_header: Any = header_dependency,
156+
_authx_openapi_cookie: Any = cookie_dependency,
157+
_authx_openapi_query: Any = query_dependency,
158+
) -> TokenPayload:
159+
self.ensure_request_exception_handlers(request)
145160
return await self._auth_required(
146161
login_type=login_type,
147162
request=request,
@@ -154,6 +169,26 @@ async def _auth_required(request: Request) -> TokenPayload:
154169

155170
return _auth_required
156171

172+
def _openapi_security_dependencies(
173+
self,
174+
login_type: str,
175+
type: str = "access",
176+
locations: Optional[TokenLocations] = None,
177+
) -> tuple[Callable[..., Any], Callable[..., Any], Callable[..., Any]]:
178+
try:
179+
return self.get(login_type)._openapi_security_dependencies(type=type, locations=locations)
180+
except BadConfigurationError:
181+
return (
182+
HTTPBearer(
183+
scheme_name="AuthXBearer",
184+
bearerFormat="JWT",
185+
description=_OPENAPI_BEARER_DESCRIPTION,
186+
auto_error=False,
187+
),
188+
_noop_openapi_security,
189+
_noop_openapi_security,
190+
)
191+
157192
def access_token_required(self, login_type: str) -> Callable[[Request], Awaitable[TokenPayload]]:
158193
"""Dependency factory requiring an access token for a login type."""
159194
return self.token_required(login_type=login_type, type="access")
@@ -292,8 +327,18 @@ def policy_required(
292327
env: Optional[Mapping[str, Any]] = None,
293328
) -> Callable[[Request], Awaitable[TokenPayload]]:
294329
"""Dependency factory requiring policy authorization."""
295-
296-
async def _policy_required(request: Request) -> TokenPayload:
330+
header_scheme, cookie_scheme, query_scheme = self._openapi_security_dependencies(login_type=login_type)
331+
header_dependency = Depends(header_scheme)
332+
cookie_dependency = Depends(cookie_scheme)
333+
query_dependency = Depends(query_scheme)
334+
335+
async def _policy_required(
336+
request: Request,
337+
_authx_openapi_header: Any = header_dependency,
338+
_authx_openapi_cookie: Any = cookie_dependency,
339+
_authx_openapi_query: Any = query_dependency,
340+
) -> TokenPayload:
341+
self.ensure_request_exception_handlers(request)
297342
return await self.authorize(
298343
login_type=login_type,
299344
action=action,

‎pyproject.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ dependencies = [
5454
"python-dateutil>=2.8,<3.0.0",
5555
"pytz>=2023.3,<2027.0",
5656
"itsdangerous>=2.2.0,<3.0.0",
57+
"uvicorn>=0.39.0",
5758
]
5859

5960
dynamic = ["version"]
@@ -76,6 +77,7 @@ test = [
7677
"websockets",
7778
"freezegun",
7879
"async-asgi-testclient",
80+
"httpx2>=2.0.0",
7981
]
8082

8183
lint = [

0 commit comments

Comments
 (0)