Merge cdda96f881 into cbe89d2bdd

llama_stack/distribution/server/auth_providers.py

@@ -18,6 +18,7 @@ from llama_stack.distribution.datatypes import (
     AuthenticationConfig,
     CustomAuthConfig,
     GitHubTokenAuthConfig,
+    KubernetesAuthProviderConfig,
     OAuth2TokenAuthConfig,
     User,
 )
@@ -176,7 +177,7 @@ class OAuth2TokenAuthProvider(AuthProvider):
                     attributes=access_attributes,
                 )
         except httpx.TimeoutException:
-            logger.exception("Token introspection request timed out")
+            logger.warning("Token introspection request timed out")
             raise
         except ValueError:
             # Re-raise ValueError exceptions to preserve their message
@@ -374,6 +375,90 @@ async def _get_github_user_info(access_token: str, github_api_base_url: str) ->
         }
 
 
+class KubernetesAuthProvider(AuthProvider):
+    """
+    Kubernetes authentication provider that validates tokens using the Kubernetes SelfSubjectReview API.
+    This provider integrates with Kubernetes API server by using the
+    /apis/authentication.k8s.io/v1/selfsubjectreviews endpoint to validate tokens and extract user information.
+    """
+
+    def __init__(self, config: KubernetesAuthProviderConfig):
+        self.config = config
+
+    async def validate_token(self, token: str, scope: dict | None = None) -> User:
+        """Validate a token using Kubernetes SelfSubjectReview API endpoint."""
+
+        # Configure SSL context
+        ssl_ctxt = None
+        if self.config.tls_cafile:
+            ssl_ctxt = ssl.create_default_context(cafile=self.config.tls_cafile.as_posix())
+        elif not self.config.verify_tls:
+            ssl_ctxt = ssl.create_default_context()
+            ssl_ctxt.check_hostname = False
+            ssl_ctxt.verify_mode = ssl.CERT_NONE
+
+        # Build the Kubernetes SelfSubjectReview API endpoint URL
+        review_api_url = f"{self.config.api_server_url.rstrip('/')}/apis/authentication.k8s.io/v1/selfsubjectreviews"
+
+        # Create SelfSubjectReview request body
+        review_request = {"apiVersion": "authentication.k8s.io/v1", "kind": "SelfSubjectReview"}
+
+        try:
+            async with httpx.AsyncClient(verify=ssl_ctxt if ssl_ctxt else self.config.verify_tls) as client:
+                response = await client.post(
+                    review_api_url,
+                    json=review_request,
+                    headers={
+                        "Authorization": f"Bearer {token}",
+                        "Content-Type": "application/json",
+                    },
+                    timeout=10.0,
+                )
+
+                if response.status_code == 401:
+                    raise ValueError("Invalid token")
+                elif response.status_code != 201:
+                    logger.warning(f"Kubernetes SelfSubjectReview API failed with status code: {response.status_code}")
+                    raise ValueError(f"Token validation failed: {response.status_code}")
+
+                review_response = response.json()
+
+                # Extract user information from SelfSubjectReview response
+                status = review_response.get("status", {})
+                if not status:
+                    raise ValueError("No status found in SelfSubjectReview response")
+
+                user_info = status.get("userInfo", {})
+                if not user_info:
+                    raise ValueError("No userInfo found in SelfSubjectReview response")
+
+                username = user_info.get("username")
+                if not username:
+                    raise ValueError("No username found in SelfSubjectReview response")
+
+                # Build user attributes from Kubernetes user info
+                user_attributes = get_attributes_from_claims(user_info, self.config.claims_mapping)
+
+                return User(
+                    principal=username,
+                    attributes=user_attributes,
+                )
+
+        except httpx.TimeoutException:
+            logger.warning("Kubernetes SelfSubjectReview API request timed out")
+            raise ValueError("Token validation timeout") from None
+        except ValueError:
+            # Re-raise ValueError exceptions to preserve their message
+            raise
+        except Exception as e:
+            logger.warning(f"Error during token validation: {str(e)}")
+            raise ValueError(f"Token validation error: {str(e)}") from e
+
+    async def close(self):
+        """Close any resources."""
+        pass
+
+
 def create_auth_provider(config: AuthenticationConfig) -> AuthProvider:
     """Factory function to create the appropriate auth provider."""
     provider_config = config.provider_config
@@ -384,5 +469,7 @@ def create_auth_provider(config: AuthenticationConfig) -> AuthProvider:
         return OAuth2TokenAuthProvider(provider_config)
     elif isinstance(provider_config, GitHubTokenAuthConfig):
         return GitHubTokenAuthProvider(provider_config)
+    elif isinstance(provider_config, KubernetesAuthProviderConfig):
+        return KubernetesAuthProvider(provider_config)
     else:
         raise ValueError(f"Unknown authentication provider config type: {type(provider_config)}")