Auth

Author: omertuc

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "starlette>=0.47.1",
     "aiohttp>=3.12.14",
     "authlib>=1.6.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."""
     # ensure that configuration is loaded

src/app/endpoints/conversations.py

@@ -11,6 +11,8 @@
 from configuration import configuration
 from models.responses import ConversationResponse, ConversationDeleteResponse
 from auth import get_auth_dependency
+from authorization.middleware import authorize
+from authorization.models import Action
 from utils.endpoints import check_configuration_loaded
 from utils.suid import check_suid
 
@@ -110,6 +112,7 @@ def simplify_session_data(session_data: dict) -> list[dict[str, Any]]:
 
 
 @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),
@@ -179,6 +182,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 (
     FeedbackResponse,
@@ -64,6 +66,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."""
     # Setup the model metrics if not already done. This is a one-time setup

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,8 +50,15 @@
 
 
 @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."""
+
+    # Used only by the middleware
+    _ = auth
+
     check_configuration_loaded(configuration)
 
     llama_stack_configuration = configuration.llama_stack_configuration

src/app/endpoints/query.py

@@ -26,6 +26,8 @@
 from models.responses import QueryResponse, UnauthorizedResponse, ForbiddenResponse
 from models.requests import QueryRequest, Attachment
 import constants
+from authorization.middleware import authorize
+from authorization.models import Action
 from utils.endpoints import check_configuration_loaded, get_agent, get_system_prompt
 from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups
 from utils.suid import get_suid
@@ -66,6 +68,7 @@ def is_transcripts_enabled() -> bool:
 
 
 @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
@@ -380,6 +382,7 @@ def _handle_heartbeat_event(chunk_id: int) -> Iterator[str]:
 
 
 @router.post("/streaming_query")
+@authorize(Action.STREAMING_QUERY)
 async def streaming_query_endpoint_handler(
     _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."""