Improve guard security integration

Author: ItayTheDar

docs/guards.md

@@ -1,6 +1,6 @@
 # Guards and Authentication
 
-PyNest now supports route guards similar to NestJS. Guards are classes that implement custom authorization logic. Use the `UseGuards` decorator to attach one or more guards to a controller or to specific routes.
+PyNest now supports route guards similar to NestJS. Guards are classes that implement custom authorization logic. Use the `UseGuards` decorator to attach one or more guards to a controller or to specific routes.  If a guard defines a FastAPI security scheme via the ``security_scheme`` attribute, the generated OpenAPI schema will mark the route as protected and the interactive docs will allow entering credentials.
 
 ```python
 from fastapi import Request
@@ -23,26 +23,31 @@ When the guard returns `False`, a `403 Forbidden` response is sent automatically
 
 ## JWT Authentication Example
 
-You can use third-party libraries like `pyjwt` to validate tokens inside a guard.
+You can use third-party libraries like `pyjwt` together with FastAPI's security utilities.
 
 ```python
 import jwt
 from fastapi import Request
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from nest.core import BaseGuard
 
 class JWTGuard(BaseGuard):
-    def can_activate(self, request: Request) -> bool:
-        auth = request.headers.get("Authorization", "")
-        token = auth.replace("Bearer ", "")
+    security_scheme = HTTPBearer()
+
+    def can_activate(
+        self, request: Request, credentials: HTTPAuthorizationCredentials
+    ) -> bool:
         try:
-            payload = jwt.decode(token, "your-secret", algorithms=["HS256"])
+            payload = jwt.decode(
+                credentials.credentials, "your-secret", algorithms=["HS256"]
+            )
         except jwt.PyJWTError:
             return False
         request.state.user = payload.get("sub")
         return True
 ```
 
-Attach the guard with `@UseGuards(JWTGuard)` on controllers or routes to secure them.
+Attach the guard with `@UseGuards(JWTGuard)` on controllers or routes to secure them. Because ``JWTGuard`` specifies a ``security_scheme`` the route will display a lock icon in the docs and allow entering a token.
 
 ## Controller vs. Route Guards
 
@@ -102,3 +107,10 @@ class AsyncGuard(BaseGuard):
 
 PyNest awaits the result automatically.
 
+## OpenAPI Integration
+
+When a guard sets the ``security_scheme`` attribute, the generated OpenAPI schema
+includes the corresponding security requirement. The docs page will show a lock
+icon next to the route and present an input box for the token or credentials.
+This works with any ``fastapi.security`` scheme (e.g. ``HTTPBearer``, ``OAuth2PasswordBearer``).
+

nest/core/decorators/controller.py

@@ -138,7 +138,9 @@ def _collect_guards(cls: Type, method: callable) -> List[BaseGuard]:
     return guards
 
 
-def add_route_to_router(router: APIRouter, method_function: callable, cls: Type) -> None:
+def add_route_to_router(
+    router: APIRouter, method_function: callable, cls: Type
+) -> None:
     """Add the configured route to the router."""
     route_kwargs = {
         "path": method_function.__route_path__,
@@ -154,7 +156,7 @@ def add_route_to_router(router: APIRouter, method_function: callable, cls: Type)
     if guards:
         dependencies = route_kwargs.get("dependencies", [])
         for guard in guards:
-            dependencies.append(Depends(guard()))
+            dependencies.append(Depends(guard.as_dependency()))
         route_kwargs["dependencies"] = dependencies
 
     router.add_api_route(**route_kwargs)

nest/core/guards.py

@@ -1,21 +1,51 @@
-from fastapi import Request, HTTPException, status
+from fastapi import Request, HTTPException, status, Security
+from fastapi.security.base import SecurityBase
 import inspect
 
 
 class BaseGuard:
-    """Base class for creating route guards."""
+    """Base class for creating route guards.
 
-    def can_activate(self, request: Request) -> bool:
+    If ``security_scheme`` is set to an instance of ``fastapi.security.SecurityBase``
+    the guard will be injected with the credentials from that scheme and the
+    corresponding security requirement will appear in the generated OpenAPI
+    schema.
+    """
+
+    security_scheme: SecurityBase | None = None
+
+    def can_activate(self, request: Request, credentials=None) -> bool:
         """Override this method with your authorization logic."""
         raise NotImplementedError
 
-    async def __call__(self, request: Request):
-        result = self.can_activate(request)
+    async def __call__(self, request: Request, credentials=None):
+        result = self.can_activate(request, credentials)
         if inspect.isawaitable(result):
             result = await result
         if not result:
-            raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Forbidden")
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN, detail="Forbidden"
+            )
+
+    @classmethod
+    def as_dependency(cls):
+        """Return a dependency callable for FastAPI routes."""
+
+        if cls.security_scheme is None:
+
+            async def dependency(request: Request):
+                guard = cls()
+                await guard(request)
+
+            return dependency
+
+        security_scheme = cls.security_scheme
+
+        async def dependency(request: Request, credentials=Security(security_scheme)):
+            guard = cls()
+            await guard(request, credentials)
 
+        return dependency
 
 
 def UseGuards(*guards):

tests/test_core/test_decorators/test_guard.py

@@ -2,6 +2,8 @@
 
 from fastapi import Request
 
+from fastapi.security import HTTPBearer
+
 from nest.core import Controller, Get, UseGuards, BaseGuard
 
 
@@ -14,6 +16,13 @@ def can_activate(self, request: Request) -> bool:
         return True
 
 
+class BearerGuard(BaseGuard):
+    security_scheme = HTTPBearer()
+
+    def can_activate(self, request: Request, credentials) -> bool:
+        return True
+
+
 @Controller("/guard")
 class GuardController:
     @Get("/")
@@ -32,4 +41,17 @@ def test_guard_added_to_route_dependencies():
     route = router.routes[0]
     deps = route.dependencies
     assert len(deps) == 1
-    assert isinstance(deps[0].dependency, SimpleGuard)
+    assert callable(deps[0].dependency)
+
+
+def test_openapi_security_requirement():
+    @Controller("/bearer")
+    class BearerController:
+        @Get("/")
+        @UseGuards(BearerGuard)
+        def root(self):
+            return {"ok": True}
+
+    router = BearerController.get_router()
+    route = router.routes[0]
+    assert route.dependant.security_requirements