Add support for JWT Token authentication with JWK

This commit introduces a new authentication module `jwk-token` that
allows the application to authenticate users using JSON Web Tokens (JWT)
signed with JSON Web Keys (JWK). The new module fetches JWKs from a
specified URL and validates incoming JWTs against these keys.

Example config:

```yaml
authentication:
  module: jwk-token
  jwk_config:
    url: https://sso.redhat.com/auth/realms/redhat-external/protocol/openid-connect/certs
    jwt_configuration:
      user_id_claim: user_id
      username_claim: username
```

Author: omertuc

pyproject.toml

@@ -13,6 +13,8 @@ dependencies = [
     "llama-stack>=0.2.13",
     "rich>=14.0.0",
     "cachetools>=6.1.0",
+    "aiohttp>=3.12.14",
+    "authlib>=1.6.0",
 ]
 
 [tool.pyright]

src/auth/__init__.py

@@ -3,7 +3,7 @@
 import logging
 
 from auth.interface import AuthInterface
-from auth import noop, noop_with_token, k8s
+from auth import noop, noop_with_token, k8s, jwk_token
 from configuration import configuration
 import constants
 
@@ -15,7 +15,7 @@ def get_auth_dependency(
     virtual_path: str = constants.DEFAULT_VIRTUAL_PATH,
 ) -> AuthInterface:
     """Select the configured authentication dependency interface."""
-    module = configuration.authentication_configuration.module  # pyright: ignore
+    module = configuration.authentication_configuration.module
 
     logger.debug(
         "Initializing authentication dependency: module='%s', virtual_path='%s'",
@@ -32,6 +32,11 @@ def get_auth_dependency(
             )
         case constants.AUTH_MOD_K8S:
             return k8s.K8SAuthDependency(virtual_path=virtual_path)
+        case constants.AUTH_MOD_JWK_TOKEN:
+            return jwk_token.JwkTokenAuthDependency(
+                configuration.authentication_configuration.jwk_configuration,
+                virtual_path=virtual_path,
+            )
         case _:
             err_msg = f"Unsupported authentication module '{module}'"
             logger.error(err_msg)

src/auth/jwk_token.py

@@ -0,0 +1,149 @@
+"""Manage authentication flow for FastAPI endpoints with no-op auth."""
+
+import logging
+from asyncio import Lock
+
+from fastapi import Request, HTTPException, status
+from authlib.jose import JsonWebKey, KeySet, jwt
+from authlib.jose.errors import (
+    BadSignatureError,
+    DecodeError,
+    ExpiredTokenError,
+    JoseError,
+)
+from cachetools import TTLCache
+import aiohttp
+
+from constants import (
+    DEFAULT_VIRTUAL_PATH,
+)
+from auth.interface import AuthInterface
+from auth.utils import extract_user_token
+from models.config import JwkConfiguration
+
+logger = logging.getLogger(__name__)
+
+# Global JWK registry to avoid re-fetching JWKs for each request. Cached for 1
+# hour, keys are unlikely to change frequently.
+_jwk_cache: TTLCache[str, KeySet] = TTLCache(maxsize=3, ttl=3600)
+# Ideally this would be an RWLock, but it would require adding a dependency on
+# aiorwlock
+_jwk_cache_lock = Lock()
+
+
+async def get_jwk_set(url: str):
+    """
+    Fetch the JWK set from the cache, or fetch it from the URL if not cached.
+    """
+    async with _jwk_cache_lock:
+        if url not in _jwk_cache:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url) as resp:
+                    resp.raise_for_status()
+                    _jwk_cache[url] = JsonWebKey.import_key_set(await resp.json())
+        return _jwk_cache[url]
+
+
+class KeyNotFoundError(Exception):
+    """
+    Exception raised when a key is not found in the JWK set based on kid/alg.
+    """
+
+
+class JwkTokenAuthDependency(AuthInterface):  # pylint: disable=too-few-public-methods
+    """JWK AuthDependency class that checks that the given bearer token is
+    valid based on a given JWKS URL."""
+
+    def __init__(
+        self, config: JwkConfiguration, virtual_path: str = DEFAULT_VIRTUAL_PATH
+    ) -> None:
+        """Initialize the required allowed paths for authorization checks."""
+        self.virtual_path: str = virtual_path
+        self.config: JwkConfiguration = config
+
+    async def __call__(self, request: Request) -> tuple[str, str, str]:
+        user_token = extract_user_token(request.headers)
+
+        jwk_set = await get_jwk_set(str(self.config.url))
+
+        def resolve_key(header, _payload):
+            """Resolve the key from the JWK set based on the JWT header. Also
+            match the algorithm to make sure the algorithm stated by the user
+            is the same algorithm the key itself expects."""
+            key = jwk_set.find_by_kid(header["kid"])
+
+            if header["alg"] != key["alg"]:
+                raise KeyNotFoundError
+
+            return key
+
+        if not user_token:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Authorization header is missing or invalid",
+            )
+
+        try:
+            claims = jwt.decode(user_token, key=resolve_key)
+        except KeyNotFoundError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid token: signed by unknown key or algorithm mismatch",
+            ) from exc
+        except BadSignatureError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid token: bad signature",
+            ) from exc
+        except DecodeError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Invalid token: decode error",
+            ) from exc
+        except JoseError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Invalid token: unknown error",
+            ) from exc
+        except Exception as exc:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Internal server error",
+            ) from exc
+
+        try:
+            claims.validate()
+        except ExpiredTokenError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED, detail="Token has expired"
+            ) from exc
+        except JoseError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Error validating token",
+            ) from exc
+        except Exception as exc:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Internal server error during token validation",
+            ) from exc
+
+        try:
+            user_id: str = claims[self.config.jwt_configuration.user_id_claim]
+        except KeyError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail=f"Token missing required claim: {self.config.jwt_configuration.user_id_claim}",
+            ) from exc
+
+        try:
+            username: str = claims[self.config.jwt_configuration.username_claim]
+        except KeyError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail=f"Token missing required claim: {self.config.jwt_configuration.username_claim}",
+            ) from exc
+
+        logger.info("Successfully authenticated user %s (ID: %s)", username, user_id)
+
+        return user_id, username, user_token

src/configuration.py

@@ -84,11 +84,16 @@ def mcp_servers(self) -> list[ModelContextProtocolServer]:
         return self._configuration.mcp_servers
 
     @property
-    def authentication_configuration(self) -> Optional[AuthenticationConfiguration]:
+    def authentication_configuration(self) -> AuthenticationConfiguration:
         """Return authentication configuration."""
         assert (
             self._configuration is not None
         ), "logic error: configuration is not loaded"
+
+        assert (
+            self._configuration.authentication is not None
+        ), "logic error: authentication configuration is not loaded"
+
         return self._configuration.authentication
 
     @property

src/constants.py

@@ -33,15 +33,19 @@
 AUTH_MOD_K8S = "k8s"
 AUTH_MOD_NOOP = "noop"
 AUTH_MOD_NOOP_WITH_TOKEN = "noop-with-token"
+AUTH_MOD_JWK_TOKEN = "jwk-token"
 # Supported authentication modules
 SUPPORTED_AUTHENTICATION_MODULES = frozenset(
     {
         AUTH_MOD_K8S,
         AUTH_MOD_NOOP,
         AUTH_MOD_NOOP_WITH_TOKEN,
+        AUTH_MOD_JWK_TOKEN,
     }
 )
 DEFAULT_AUTHENTICATION_MODULE = AUTH_MOD_NOOP
+DEFAULT_JWT_UID_CLAIM = "user_id"
+DEFAULT_JWT_USER_NAME_CLAIM = "username"
 
 # Data collector constants
 DATA_COLLECTOR_COLLECTION_INTERVAL = 7200  # 2 hours in seconds

src/models/config.py

@@ -131,13 +131,40 @@ def check_storage_location_is_set_when_needed(self) -> Self:
         return self
 
 
+class JwtConfiguration(BaseModel):
+    """JWT configuration."""
+
+    user_id_claim: str = constants.DEFAULT_JWT_UID_CLAIM
+    username_claim: str = constants.DEFAULT_JWT_USER_NAME_CLAIM
+
+    @model_validator(mode="after")
+    def check_jwt_configuration(self) -> Self:
+        """Validate JWT configuration."""
+        return self
+
+
+class JwkConfiguration(BaseModel):
+    """JWK configuration."""
+
+    url: Optional[AnyHttpUrl] = None
+    jwt_configuration: JwtConfiguration = JwtConfiguration()
+
+    @model_validator(mode="after")
+    def check_jwk_configuration(self) -> Self:
+        """Validate JWK configuration."""
+        if self.url is None:
+            raise ValueError("JWK URL must be specified")
+        return self
+
+
 class AuthenticationConfiguration(BaseModel):
     """Authentication configuration."""
 
     module: str = constants.DEFAULT_AUTHENTICATION_MODULE
     skip_tls_verification: bool = False
     k8s_cluster_api: Optional[AnyHttpUrl] = None
     k8s_ca_cert_path: Optional[FilePath] = None
+    jwk_config: Optional[JwkConfiguration] = None
 
     @model_validator(mode="after")
     def check_authentication_model(self) -> Self:
@@ -148,8 +175,26 @@ def check_authentication_model(self) -> Self:
                 f"Unsupported authentication module '{self.module}'. "
                 f"Supported modules: {supported_modules}"
             )
+
+        if self.module == constants.AUTH_MOD_JWK_TOKEN:
+            if self.jwk_config is None:
+                raise ValueError(
+                    "JWK configuration must be specified when using JWK token authentication"
+                )
+
         return self
 
+    @property
+    def jwk_configuration(self) -> JwkConfiguration:
+        """Return JWK configuration if the module is JWK token."""
+
+        if self.module != constants.AUTH_MOD_JWK_TOKEN:
+            raise ValueError(
+                "JWK configuration is only available for JWK token authentication module"
+            )
+        assert self.jwk_config is not None, "JWK configuration should not be None"
+        return self.jwk_config
+
 
 class Customization(BaseModel):
     """Service customization."""