auth5

Author: omertuc

src/app/endpoints/config.py

@@ -1,10 +1,12 @@
 """Handler for REST API call to retrieve service configuration."""
 
 import logging
-from typing import Any
+from typing import Annotated, Any
 
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, Request, Depends
 
+from auth.interface import AuthTuple
+from auth import get_auth_dependency
 from authorization.middleware import authorize
 from configuration import configuration
 from models.config import Action, Configuration
@@ -13,6 +15,8 @@
 logger = logging.getLogger(__name__)
 router = APIRouter(tags=["config"])
 
+auth_dependency = get_auth_dependency()
+
 
 get_config_responses: dict[int | str, dict[str, Any]] = {
     200: {
@@ -58,7 +62,10 @@
 
 @router.get("/config", responses=get_config_responses)
 @authorize(Action.GET_CONFIG)
-def config_endpoint_handler(_request: Request) -> Configuration:
+def config_endpoint_handler(
+    auth: Annotated[AuthTuple, Depends(auth_dependency)],
+    _request: Request,
+) -> Configuration:
     """
     Handle requests to the /config endpoint.
 
@@ -68,6 +75,9 @@ def config_endpoint_handler(_request: Request) -> Configuration:
     Returns:
         Configuration: The loaded service configuration object.
     """
+    # Used only for authorization
+    _ = auth
+
     # ensure that configuration is loaded
     check_configuration_loaded(configuration)
 

src/app/endpoints/metrics.py

@@ -1,22 +1,30 @@
 """Handler for REST API call to provide metrics."""
 
+from typing import Annotated
 from fastapi.responses import PlainTextResponse
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, Request, Depends
 from prometheus_client import (
     generate_latest,
     CONTENT_TYPE_LATEST,
 )
 
+from auth.interface import AuthTuple
+from auth import get_auth_dependency
 from authorization.middleware import authorize
 from models.config import Action
 from metrics.utils import setup_model_metrics
 
 router = APIRouter(tags=["metrics"])
 
+auth_dependency = get_auth_dependency()
+
 
 @router.get("/metrics", response_class=PlainTextResponse)
 @authorize(Action.GET_METRICS)
-async def metrics_endpoint_handler(_request: Request) -> PlainTextResponse:
+async def metrics_endpoint_handler(
+    auth: Annotated[AuthTuple, Depends(auth_dependency)],
+    _request: Request,
+) -> PlainTextResponse:
     """
     Handle request to the /metrics endpoint.
 
@@ -27,6 +35,9 @@ async def metrics_endpoint_handler(_request: Request) -> PlainTextResponse:
     set up, then responds with the current metrics snapshot in
     Prometheus format.
     """
+    # Used only for authorization
+    _ = auth
+
     # Setup the model metrics if not already done. This is a one-time setup
     # and will not be run again on subsequent calls to this endpoint
     await setup_model_metrics()

src/auth/jwk_token.py

@@ -3,7 +3,6 @@
 import logging
 from asyncio import Lock
 from typing import Any, Callable
-import json
 
 from fastapi import Request, HTTPException, status
 from authlib.jose import JsonWebKey, KeySet, jwt, Key
@@ -189,4 +188,4 @@ async def __call__(self, request: Request) -> tuple[str, str, str]:
 
         logger.info("Successfully authenticated user %s (ID: %s)", username, user_id)
 
-        return user_id, username, json.dumps(claims)
+        return user_id, username, user_token

src/authorization/engine.py

@@ -145,7 +145,7 @@ def __init__(self, access_rules: list[AccessRule]):
 
     async def check_access(self, action: Action, user_roles: UserRoles) -> bool:
         """Check if the user has access to the specified action based on their roles."""
-        if action != Action.ADMIN and self.check_access(action.ADMIN, user_roles):
+        if action != Action.ADMIN and await self.check_access(Action.ADMIN, user_roles):
             # Recurse to check if the roles allow the user to perform the admin action,
             # if they do, then we allow any action
             return True

src/authorization/middleware.py

@@ -73,7 +73,7 @@ async def _perform_authorization_check(action: Action, kwargs: dict[str, Any]) -
             detail="Internal server error",
         ) from exc
 
-    if not access_resolver.check_access(
+    if not await access_resolver.check_access(
         action, await role_resolver.resolve_roles(auth)
     ):
         raise HTTPException(
@@ -84,12 +84,18 @@ async def _perform_authorization_check(action: Action, kwargs: dict[str, Any]) -
 
 def authorize(action: Action) -> Callable:
     """Check authorization for an endpoint (async version)."""
+    import asyncio
 
     def decorator(func: Callable) -> Callable:
         @wraps(func)
         async def wrapper(*args: Any, **kwargs: Any) -> Any:
             await _perform_authorization_check(action, kwargs)
-            return await func(*args, **kwargs)
+
+            # Handle both sync and async functions
+            result = func(*args, **kwargs)
+            if asyncio.iscoroutine(result):
+                return await result
+            return result
 
         return wrapper
 

tests/unit/app/endpoints/test_config.py

@@ -7,8 +7,19 @@
 from configuration import AppConfig
 
 
-def test_config_endpoint_handler_configuration_not_loaded(mocker):
+@pytest.mark.asyncio
+async def test_config_endpoint_handler_configuration_not_loaded(mocker):
     """Test the config endpoint handler."""
+    # Mock authorization resolvers
+    mock_resolvers = mocker.patch(
+        "authorization.middleware.get_authorization_resolvers"
+    )
+    mock_role_resolver = mocker.AsyncMock()
+    mock_access_resolver = mocker.AsyncMock()
+    mock_role_resolver.resolve_roles.return_value = []
+    mock_access_resolver.check_access.return_value = True
+    mock_resolvers.return_value = (mock_role_resolver, mock_access_resolver)
+
     mocker.patch(
         "app.endpoints.config.configuration._configuration",
         new=None,
@@ -20,14 +31,27 @@ def test_config_endpoint_handler_configuration_not_loaded(mocker):
             "type": "http",
         }
     )
-    with pytest.raises(HTTPException) as e:
-        config_endpoint_handler(request)
-        assert e.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
-        assert e.detail["response"] == "Configuration is not loaded"
+    auth = ("test_user", "token", {})
+    with pytest.raises(HTTPException) as exc_info:
+        await config_endpoint_handler(auth=auth, _request=request)
 
+    assert exc_info.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
+    assert exc_info.value.detail["response"] == "Configuration is not loaded"
 
-def test_config_endpoint_handler_configuration_loaded():
+
+@pytest.mark.asyncio
+async def test_config_endpoint_handler_configuration_loaded(mocker):
     """Test the config endpoint handler."""
+    # Mock authorization resolvers
+    mock_resolvers = mocker.patch(
+        "authorization.middleware.get_authorization_resolvers"
+    )
+    mock_role_resolver = mocker.AsyncMock()
+    mock_access_resolver = mocker.AsyncMock()
+    mock_role_resolver.resolve_roles.return_value = []
+    mock_access_resolver.check_access.return_value = True
+    mock_resolvers.return_value = (mock_role_resolver, mock_access_resolver)
+
     config_dict = {
         "name": "foo",
         "service": {
@@ -49,15 +73,21 @@ def test_config_endpoint_handler_configuration_loaded():
         "authentication": {
             "module": "noop",
         },
+        "authorization": {"access_rules": []},
         "customization": None,
     }
     cfg = AppConfig()
     cfg.init_from_dict(config_dict)
+
+    # Mock configuration
+    mocker.patch("app.endpoints.config.configuration", cfg)
+
     request = Request(
         scope={
             "type": "http",
         }
     )
-    response = config_endpoint_handler(request)
+    auth = ("test_user", "token", {})
+    response = await config_endpoint_handler(auth=auth, _request=request)
     assert response is not None
     assert response == cfg.configuration

tests/unit/app/endpoints/test_feedback.py

@@ -62,9 +62,20 @@ async def test_assert_feedback_enabled(mocker):
     ],
     ids=["no_categories", "with_negative_categories"],
 )
-def test_feedback_endpoint_handler(mocker, feedback_request_data):
+@pytest.mark.asyncio
+async def test_feedback_endpoint_handler(mocker, feedback_request_data):
     """Test that feedback_endpoint_handler processes feedback for different payloads."""
 
+    # Mock authorization resolvers
+    mock_resolvers = mocker.patch(
+        "authorization.middleware.get_authorization_resolvers"
+    )
+    mock_role_resolver = mocker.AsyncMock()
+    mock_access_resolver = mocker.AsyncMock()
+    mock_role_resolver.resolve_roles.return_value = []
+    mock_access_resolver.check_access.return_value = True
+    mock_resolvers.return_value = (mock_role_resolver, mock_access_resolver)
+
     # Mock the dependencies
     mocker.patch("app.endpoints.feedback.assert_feedback_enabled", return_value=None)
     mocker.patch("app.endpoints.feedback.store_feedback", return_value=None)
@@ -74,7 +85,7 @@ def test_feedback_endpoint_handler(mocker, feedback_request_data):
     feedback_request.model_dump.return_value = feedback_request_data
 
     # Call the endpoint handler
-    result = feedback_endpoint_handler(
+    result = await feedback_endpoint_handler(
         feedback_request=feedback_request,
         _ensure_feedback_enabled=assert_feedback_enabled,
         auth=("test_user_id", "test_username", "test_token"),
@@ -84,8 +95,19 @@ def test_feedback_endpoint_handler(mocker, feedback_request_data):
     assert result.response == "feedback received"
 
 
-def test_feedback_endpoint_handler_error(mocker):
+@pytest.mark.asyncio
+async def test_feedback_endpoint_handler_error(mocker):
     """Test that feedback_endpoint_handler raises an HTTPException on error."""
+    # Mock authorization resolvers
+    mock_resolvers = mocker.patch(
+        "authorization.middleware.get_authorization_resolvers"
+    )
+    mock_role_resolver = mocker.AsyncMock()
+    mock_access_resolver = mocker.AsyncMock()
+    mock_role_resolver.resolve_roles.return_value = []
+    mock_access_resolver.check_access.return_value = True
+    mock_resolvers.return_value = (mock_role_resolver, mock_access_resolver)
+
     # Mock the dependencies
     mocker.patch("app.endpoints.feedback.assert_feedback_enabled", return_value=None)
     mocker.patch(
@@ -98,7 +120,7 @@ def test_feedback_endpoint_handler_error(mocker):
 
     # Call the endpoint handler and assert it raises an exception
     with pytest.raises(HTTPException) as exc_info:
-        feedback_endpoint_handler(
+        await feedback_endpoint_handler(
             feedback_request=feedback_request,
             _ensure_feedback_enabled=assert_feedback_enabled,
             auth=("test_user_id", "test_username", "test_token"),

tests/unit/app/endpoints/test_health.py

@@ -1,5 +1,6 @@
 """Unit tests for the /health REST API endpoint."""
 
+import pytest
 from unittest.mock import Mock
 
 from llama_stack.providers.datatypes import HealthStatus
@@ -12,8 +13,19 @@
 from models.responses import ProviderHealthStatus, ReadinessResponse
 
 
+@pytest.mark.asyncio
 async def test_readiness_probe_fails_due_to_unhealthy_providers(mocker):
     """Test the readiness endpoint handler fails when providers are unhealthy."""
+    # Mock authorization resolvers
+    mock_resolvers = mocker.patch(
+        "authorization.middleware.get_authorization_resolvers"
+    )
+    mock_role_resolver = mocker.AsyncMock()
+    mock_access_resolver = mocker.AsyncMock()
+    mock_role_resolver.resolve_roles.return_value = []
+    mock_access_resolver.check_access.return_value = True
+    mock_resolvers.return_value = (mock_role_resolver, mock_access_resolver)
+
     # Mock get_providers_health_statuses to return an unhealthy provider
     mock_get_providers_health_statuses = mocker.patch(
         "app.endpoints.health.get_providers_health_statuses"
@@ -38,8 +50,19 @@ async def test_readiness_probe_fails_due_to_unhealthy_providers(mocker):
     assert mock_response.status_code == 503
 
 
+@pytest.mark.asyncio
 async def test_readiness_probe_success_when_all_providers_healthy(mocker):
     """Test the readiness endpoint handler succeeds when all providers are healthy."""
+    # Mock authorization resolvers
+    mock_resolvers = mocker.patch(
+        "authorization.middleware.get_authorization_resolvers"
+    )
+    mock_role_resolver = mocker.AsyncMock()
+    mock_access_resolver = mocker.AsyncMock()
+    mock_role_resolver.resolve_roles.return_value = []
+    mock_access_resolver.check_access.return_value = True
+    mock_resolvers.return_value = (mock_role_resolver, mock_access_resolver)
+
     # Mock get_providers_health_statuses to return healthy providers
     mock_get_providers_health_statuses = mocker.patch(
         "app.endpoints.health.get_providers_health_statuses"
@@ -70,10 +93,21 @@ async def test_readiness_probe_success_when_all_providers_healthy(mocker):
     assert len(response.providers) == 0
 
 
-def test_liveness_probe():
+@pytest.mark.asyncio
+async def test_liveness_probe(mocker):
     """Test the liveness endpoint handler."""
+    # Mock authorization resolvers
+    mock_resolvers = mocker.patch(
+        "authorization.middleware.get_authorization_resolvers"
+    )
+    mock_role_resolver = mocker.AsyncMock()
+    mock_access_resolver = mocker.AsyncMock()
+    mock_role_resolver.resolve_roles.return_value = []
+    mock_access_resolver.check_access.return_value = True
+    mock_resolvers.return_value = (mock_role_resolver, mock_access_resolver)
+
     auth = ("test_user", "token", {})
-    response = liveness_probe_get_method(auth=auth)
+    response = await liveness_probe_get_method(auth=auth)
     assert response is not None
     assert response.alive is True
 

tests/unit/app/endpoints/test_info.py

@@ -1,12 +1,14 @@
 """Unit tests for the /info REST API endpoint."""
 
+import pytest
 from fastapi import Request
 
 from app.endpoints.info import info_endpoint_handler
 from configuration import AppConfig
 
 
-def test_info_endpoint():
+@pytest.mark.asyncio
+async def test_info_endpoint(mocker):
     """Test the info endpoint handler."""
     config_dict = {
         "name": "foo",
@@ -27,16 +29,32 @@ def test_info_endpoint():
             "feedback_enabled": False,
         },
         "customization": None,
+        "authorization": {"access_rules": []},
+        "authentication": {"module": "noop"},
     }
     cfg = AppConfig()
     cfg.init_from_dict(config_dict)
+
+    # Mock configuration
+    mocker.patch("configuration.configuration", cfg)
+
+    # Mock authorization resolvers
+    mock_resolvers = mocker.patch(
+        "authorization.middleware.get_authorization_resolvers"
+    )
+    mock_role_resolver = mocker.AsyncMock()
+    mock_access_resolver = mocker.AsyncMock()
+    mock_role_resolver.resolve_roles.return_value = []
+    mock_access_resolver.check_access.return_value = True
+    mock_resolvers.return_value = (mock_role_resolver, mock_access_resolver)
+
     request = Request(
         scope={
             "type": "http",
         }
     )
     auth = ("test_user", "token", {})
-    response = info_endpoint_handler(auth, request)
+    response = await info_endpoint_handler(auth=auth, _request=request)
     assert response is not None
     assert response.name is not None
     assert response.version is not None