Add support for JWT Token authentication with JWK

omertuc · omertuc · commit 08d8ac1f741d · 2025-07-18T14:41:17.000+02:00
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 ```
diff --git a/pyproject.toml b/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]
diff --git a/src/auth/__init__.py b/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)
diff --git a/src/auth/jwk_token.py b/src/auth/jwk_token.py
@@ -0,0 +1,148 @@
+"""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 for JWK-based JWT authentication."""
+
+    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]:
+        """Authenticate the JWT in the headers against the keys from the JWK url."""
+        user_token = extract_user_token(request.headers)
+
+        jwk_set = await get_jwk_set(str(self.config.url))
+
+        def resolve_key(header, _payload):
+            """Match kid and alg from the JWT header to the JWK set.
+
+            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 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 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
diff --git a/src/configuration.py b/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
diff --git a/src/constants.py b/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
diff --git a/src/models/config.py b/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,25 @@ 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."""
diff --git a/tests/unit/auth/test_jwk_token.py b/tests/unit/auth/test_jwk_token.py
diff --git a/tests/unit/models/test_config.py b/tests/unit/models/test_config.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,8 @@ dependencies = [`
`13`	`13`	`"llama-stack>=0.2.13",`
`14`	`14`	`"rich>=14.0.0",`
`15`	`15`	`"cachetools>=6.1.0",`
	`16`	`+ "aiohttp>=3.12.14",`
	`17`	`+ "authlib>=1.6.0",`
`16`	`18`	`]`
`17`	`19`
`18`	`20`	`[tool.pyright]`
Original file line number	Diff line number	Diff line change
`@@ -33,15 +33,19 @@`
`33`	`33`	`AUTH_MOD_K8S = "k8s"`
`34`	`34`	`AUTH_MOD_NOOP = "noop"`
`35`	`35`	`AUTH_MOD_NOOP_WITH_TOKEN = "noop-with-token"`
	`36`	`+AUTH_MOD_JWK_TOKEN = "jwk-token"`
`36`	`37`	`# Supported authentication modules`
`37`	`38`	`SUPPORTED_AUTHENTICATION_MODULES = frozenset(`
`38`	`39`	`{`
`39`	`40`	`AUTH_MOD_K8S,`
`40`	`41`	`AUTH_MOD_NOOP,`
`41`	`42`	`AUTH_MOD_NOOP_WITH_TOKEN,`
	`43`	`+ AUTH_MOD_JWK_TOKEN,`
`42`	`44`	`}`
`43`	`45`	`)`
`44`	`46`	`DEFAULT_AUTHENTICATION_MODULE = AUTH_MOD_NOOP`
	`47`	`+DEFAULT_JWT_UID_CLAIM = "user_id"`
	`48`	`+DEFAULT_JWT_USER_NAME_CLAIM = "username"`
`45`	`49`
`46`	`50`	`# Data collector constants`
`47`	`51`	`DATA_COLLECTOR_COLLECTION_INTERVAL = 7200 # 2 hours in seconds`