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.

For now, the module expects the JWT to contain specific fields for user
ID and username, which will cause an error if they are not present.
These fields will be configurable in the future but for now, they match
the fields used in Red Hat SSO JWTs.

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,13 @@ 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
+
+    if configuration.authentication_configuration is None:
+        err_msg = "Authentication configuration is not set"
+        logger.error(err_msg)
+        raise ValueError(err_msg)
+
+    module = configuration.authentication_configuration.module
 
     logger.debug(
         "Initializing authentication dependency: module='%s', virtual_path='%s'",
@@ -32,6 +38,15 @@ def get_auth_dependency(
             )
         case constants.AUTH_MOD_K8S:
             return k8s.K8SAuthDependency(virtual_path=virtual_path)
+        case constants.AUTH_MOD_JWK_TOKEN:
+            if configuration.authentication_configuration.jwk_url is None:
+                err_msg = "JWK URL must be provided for JWK Token authentication"
+                logger.error(err_msg)
+                raise ValueError(err_msg)
+            return jwk_token.JwkTokenAuthDependency(
+                configuration.authentication_configuration.jwk_url,
+                virtual_path=virtual_path,
+            )
         case _:
             err_msg = f"Unsupported authentication module '{module}'"
             logger.error(err_msg)

src/auth/jwk_token.py

@@ -0,0 +1,96 @@
+"""Manage authentication flow for FastAPI endpoints with no-op auth."""
+
+import logging
+from asyncio import Lock
+
+from fastapi import Request
+from authlib.jose import JsonWebKey, KeySet, jwt, JoseError
+from cachetools import TTLCache
+import aiohttp
+from pydantic import AnyHttpUrl
+
+from constants import (
+    DEFAULT_USER_NAME,
+    DEFAULT_USER_UID,
+    DEFAULT_VIRTUAL_PATH,
+)
+from auth.interface import AuthInterface
+from auth.utils import extract_user_token
+
+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 a RWLock, but it would require adding a dependency on
+# aiorwlock
+_jwk_cache_lock = Lock()
+
+# TODO(omertuc): These default field names currently match the values in the Red Hat SSO JWTs
+# (e.g. `ocm token` command). Ideally they would be configurable so users of
+# lightspeed-core could set them to match their own JWT claims.
+DEFAULT_JWT_FIELD_USER_UID = "user_id"
+DEFAULT_JWT_FIELD_USER_NAME = "username"
+
+
+async def get_jwk_set(jwk_url):
+    """
+    Fetch the JWK set from the cache, or fetch it from the URL if not cached.
+    """
+    async with _jwk_cache_lock:
+        if jwk_url not in _jwk_cache:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(jwk_url) as resp:
+                    resp.raise_for_status()
+                    data = await resp.json()
+                    _jwk_cache[jwk_url] = JsonWebKey.import_key_set(data)
+        return _jwk_cache[jwk_url]
+
+
+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, jwks_url: AnyHttpUrl, virtual_path: str = DEFAULT_VIRTUAL_PATH
+    ) -> None:
+        """Initialize the required allowed paths for authorization checks."""
+        self.virtual_path = virtual_path
+        self.jwks_url: AnyHttpUrl = jwks_url
+
+    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.jwks_url))
+
+        if not user_token:
+            logger.error("No authorization token found in request")
+            return DEFAULT_USER_UID, DEFAULT_USER_NAME, self.virtual_path
+
+        try:
+            claims = jwt.decode(user_token, jwk_set)
+
+            try:
+                user_id = claims[DEFAULT_JWT_FIELD_USER_UID]
+            except KeyError as exc:
+                raise JoseError(
+                    f"JWT token does not contain required field '{DEFAULT_JWT_FIELD_USER_UID}'"
+                ) from exc
+
+            try:
+                username = claims[DEFAULT_JWT_FIELD_USER_NAME]
+            except KeyError as exc:
+                raise JoseError(
+                    f"JWT token does not contain required field '{DEFAULT_JWT_FIELD_USER_NAME}'"
+                ) from exc
+
+            logger.info(
+                "Successfully authenticated user %s (ID: %s)", username, user_id
+            )
+            return user_id, username, user_token
+        except JoseError:
+            logger.error("JWT token validation failed")
+            raise
+        except Exception:
+            logger.error("Unexpected error during token validation")
+            raise

src/constants.py

@@ -33,12 +33,14 @@
 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

src/models/config.py

@@ -112,6 +112,7 @@ class AuthenticationConfiguration(BaseModel):
     skip_tls_verification: bool = False
     k8s_cluster_api: Optional[AnyHttpUrl] = None
     k8s_ca_cert_path: Optional[FilePath] = None
+    jwk_url: Optional[AnyHttpUrl] = None
 
     @model_validator(mode="after")
     def check_authentication_model(self) -> Self: