Auth

Author: omertuc

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
     "openai==1.99.1",
     "sqlalchemy>=2.0.42",
     "email-validator>=2.2.0",
+    "jsonpath-ng>=1.6.1",
 ]
 
 [tool.pyright]

src/app/endpoints/config.py

@@ -7,6 +7,8 @@
 
 from models.config import Configuration
 from configuration import configuration
+from authorization.middleware import authorize
+from authorization.models import Action
 from utils.endpoints import check_configuration_loaded
 
 logger = logging.getLogger(__name__)
@@ -56,6 +58,7 @@
 
 
 @router.get("/config", responses=get_config_responses)
+@authorize(Action.GET_CONFIG)
 def config_endpoint_handler(_request: Request) -> Configuration:
     """
     Handle requests to the /config endpoint.

src/app/endpoints/conversations.py

@@ -19,6 +19,8 @@
 from auth import get_auth_dependency
 from app.database import get_session
 from utils.endpoints import check_configuration_loaded, validate_conversation_ownership
+from authorization.middleware import authorize
+from authorization.models import Action
 from utils.suid import check_suid
 
 logger = logging.getLogger("app.endpoints.handlers")
@@ -200,6 +202,7 @@ def get_conversations_list_endpoint_handler(
 
 
 @router.get("/conversations/{conversation_id}", responses=conversation_responses)
+@authorize(Action.GET_CONVERSATION)
 async def get_conversation_endpoint_handler(
     conversation_id: str,
     auth: Any = Depends(auth_dependency),
@@ -309,6 +312,7 @@ async def get_conversation_endpoint_handler(
 @router.delete(
     "/conversations/{conversation_id}", responses=conversation_delete_responses
 )
+@authorize(Action.DELETE_CONVERSATION)
 async def delete_conversation_endpoint_handler(
     conversation_id: str,
     auth: Any = Depends(auth_dependency),

src/app/endpoints/feedback.py

@@ -9,6 +9,8 @@
 
 from auth import get_auth_dependency
 from auth.interface import AuthTuple
+from authorization.middleware import authorize
+from authorization.models import Action
 from configuration import configuration
 from models.responses import (
     ErrorResponse,
@@ -79,6 +81,7 @@ async def assert_feedback_enabled(_request: Request) -> None:
 
 
 @router.post("", responses=feedback_response)
+@authorize(Action.FEEDBACK)
 def feedback_endpoint_handler(
     feedback_request: FeedbackRequest,
     auth: Annotated[AuthTuple, Depends(auth_dependency)],

src/app/endpoints/metrics.py

@@ -7,12 +7,15 @@
     CONTENT_TYPE_LATEST,
 )
 
+from authorization.middleware import authorize
+from authorization.models import Action
 from metrics.utils import setup_model_metrics
 
 router = APIRouter(tags=["metrics"])
 
 
 @router.get("/metrics", response_class=PlainTextResponse)
+@authorize(Action.GET_METRICS)
 async def metrics_endpoint_handler(_request: Request) -> PlainTextResponse:
     """
     Handle request to the /metrics endpoint.

src/app/endpoints/models.py

@@ -3,18 +3,25 @@
 import logging
 from typing import Any
 
+from fastapi.params import Depends
 from llama_stack_client import APIConnectionError
 from fastapi import APIRouter, HTTPException, Request, status
 
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
+from authorization.middleware import authorize
+from authorization.models import Action
 from models.responses import ModelsResponse
 from utils.endpoints import check_configuration_loaded
+from auth import get_auth_dependency
 
 logger = logging.getLogger(__name__)
 router = APIRouter(tags=["models"])
 
 
+auth_dependency = get_auth_dependency()
+
+
 models_responses: dict[int | str, dict[str, Any]] = {
     200: {
         "models": [
@@ -43,7 +50,8 @@
 
 
 @router.get("/models", responses=models_responses)
-async def models_endpoint_handler(_request: Request) -> ModelsResponse:
+@authorize(Action.GET_MODELS)
+async def models_endpoint_handler(_request: Request, auth: Any = Depends(get_auth_dependency())) -> ModelsResponse:
     """
     Handle requests to the /models endpoint.
 
@@ -57,6 +65,10 @@ async def models_endpoint_handler(_request: Request) -> ModelsResponse:
     Returns:
         ModelsResponse: An object containing the list of available models.
     """
+    
+    # Used only by the middleware
+    _ = auth
+
     check_configuration_loaded(configuration)
 
     llama_stack_configuration = configuration.llama_stack_configuration

src/app/endpoints/query.py

@@ -34,6 +34,8 @@
     get_system_prompt,
     validate_conversation_ownership,
 )
+from authorization.middleware import authorize
+from authorization.models import Action
 from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups
 from utils.suid import get_suid
 
@@ -147,6 +149,7 @@ def evaluate_model_hints(
 
 
 @router.post("/query", responses=query_response)
+@authorize(Action.QUERY)
 async def query_endpoint_handler(
     query_request: QueryRequest,
     auth: Annotated[AuthTuple, Depends(auth_dependency)],

src/app/endpoints/streaming_query.py

@@ -19,6 +19,8 @@
 
 from auth import get_auth_dependency
 from auth.interface import AuthTuple
+from authorization.middleware import authorize
+from authorization.models import Action
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
 import metrics
@@ -384,6 +386,7 @@ def _handle_heartbeat_event(chunk_id: int) -> Iterator[str]:
 
 
 @router.post("/streaming_query")
+@authorize(Action.STREAMING_QUERY)
 async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
     _request: Request,
     query_request: QueryRequest,

src/auth/jwk_token.py

@@ -3,6 +3,7 @@
 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
@@ -188,4 +189,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, user_token
+        return user_id, username, json.dumps(claims)

src/authorization/__init__.py

@@ -0,0 +1 @@
+"""Authorization module for role-based access control."""