Updated exception handling in internal logic

Author: Andrej Simurka

src/authentication/jwk_token.py

@@ -1,26 +1,28 @@
 """Manage authentication flow for FastAPI endpoints with JWK based JWT auth."""
 
+import json
 import logging
 from asyncio import Lock
 from typing import Any, Callable
 
-from fastapi import Request, HTTPException, status
-from authlib.jose import JsonWebKey, KeySet, jwt, Key
+import aiohttp
+from authlib.jose import JsonWebKey, Key, KeySet, jwt
 from authlib.jose.errors import (
     BadSignatureError,
     DecodeError,
     ExpiredTokenError,
     JoseError,
 )
 from cachetools import TTLCache
-import aiohttp
+from fastapi import HTTPException, Request
 
+from authentication.interface import NO_AUTH_TUPLE, AuthInterface, AuthTuple
+from authentication.utils import extract_user_token
 from constants import (
     DEFAULT_VIRTUAL_PATH,
 )
-from authentication.interface import NO_AUTH_TUPLE, AuthInterface, AuthTuple
-from authentication.utils import extract_user_token
 from models.config import JwkConfiguration
+from models.responses import UnauthorizedResponse
 
 logger = logging.getLogger(__name__)
 
@@ -126,68 +128,67 @@ async def __call__(self, request: Request) -> AuthTuple:
             return NO_AUTH_TUPLE
 
         user_token = extract_user_token(request.headers)
-        jwk_set = await get_jwk_set(str(self.config.url))
 
         try:
-            claims = jwt.decode(user_token, key=key_resolver_func(jwk_set))
-        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
+            jwk_set = await get_jwk_set(str(self.config.url))
+        except aiohttp.ClientError as exc:
+            logger.error("Failed to fetch JWK set: %s", exc)
+            response = UnauthorizedResponse(
+                cause="Unable to reach authentication key server"
+            )
+            raise HTTPException(**response.model_dump()) from exc
+        except json.JSONDecodeError as exc:
+            logger.error("Invalid JSON in JWK set response: %s", exc)
+            response = UnauthorizedResponse(
+                cause="Authentication key server returned invalid data"
+            )
+            raise HTTPException(**response.model_dump()) 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
+            logger.error("Invalid JWK set format: %s", exc)
+            response = UnauthorizedResponse(cause="Authentication keys are malformed")
+            raise HTTPException(**response.model_dump()) from exc
+
+        try:
+            claims = jwt.decode(user_token, key=key_resolver_func(jwk_set))
+        except (KeyNotFoundError, BadSignatureError, DecodeError, JoseError) as exc:
+            logger.warning("Token decode error: %s", exc)
+            cause_map = {
+                KeyNotFoundError: "Token signed by unknown key",
+                BadSignatureError: "Invalid token signature",
+                DecodeError: "Token could not be decoded",
+                JoseError: "Token format error",
+            }
+            response = UnauthorizedResponse(
+                cause=cause_map.get(type(exc), "Unknown token error")
+            )
+            raise HTTPException(**response.model_dump()) from exc
 
         try:
             claims.validate()
         except ExpiredTokenError as exc:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED, detail="Token has expired"
-            ) from exc
+            response = UnauthorizedResponse(cause="Token has expired")
+            raise HTTPException(**response.model_dump()) 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
+            response = UnauthorizedResponse(cause="Token validation failed")
+            raise HTTPException(**response.model_dump()) 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
+            missing_claim = self.config.jwt_configuration.user_id_claim
+            response = UnauthorizedResponse(
+                cause=f"Token missing claim: {missing_claim}"
+            )
+            raise HTTPException(**response.model_dump()) 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
+            missing_claim = self.config.jwt_configuration.username_claim
+            response = UnauthorizedResponse(
+                cause=f"Token missing claim: {missing_claim}"
+            )
+            raise HTTPException(**response.model_dump()) from exc
 
         logger.info("Successfully authenticated user %s (ID: %s)", username, user_id)
 

src/authentication/k8s.py

@@ -4,15 +4,22 @@
 import os
 from pathlib import Path
 from typing import Optional, Self
-from fastapi import Request, HTTPException
 
 import kubernetes.client
+from fastapi import HTTPException, Request
 from kubernetes.client.rest import ApiException
 from kubernetes.config import ConfigException
 
-from configuration import configuration
 from authentication.interface import AuthInterface
+from authentication.utils import extract_user_token
+from configuration import configuration
 from constants import DEFAULT_VIRTUAL_PATH
+from models.responses import (
+    ForbiddenResponse,
+    InternalServerErrorResponse,
+    ServiceUnavailableResponse,
+    UnauthorizedResponse,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -172,8 +179,20 @@ def get_user_info(token: str) -> Optional[kubernetes.client.V1TokenReview]:
 
     Returns:
         The user information if the token is valid, None otherwise.
+
+    Raises:
+        HTTPException: If unable to connect to Kubernetes API or unexpected error occurs.
     """
-    auth_api = K8sClientSingleton.get_authn_api()
+    try:
+        auth_api = K8sClientSingleton.get_authn_api()
+    except Exception as e:
+        logger.error("Failed to get Kubernetes authentication API: %s", e)
+        response = ServiceUnavailableResponse(
+            backend_name="Kubernetes API",
+            cause="Unable to initialize Kubernetes client",
+        )
+        raise HTTPException(**response.model_dump()) from e
+
     token_review = kubernetes.client.V1TokenReview(
         spec=kubernetes.client.V1TokenReviewSpec(token=token)
     )
@@ -182,31 +201,9 @@ def get_user_info(token: str) -> Optional[kubernetes.client.V1TokenReview]:
         if response.status.authenticated:
             return response.status
         return None
-    except ApiException as e:
+    except Exception as e:  # pylint: disable=broad-exception-caught
         logger.error("API exception during TokenReview: %s", e)
         return None
-    except Exception as e:
-        logger.error("Unexpected error during TokenReview - Unauthorized: %s", e)
-        raise HTTPException(
-            status_code=500,
-            detail={"response": "Forbidden: Unable to Review Token", "cause": str(e)},
-        ) from e
-
-
-def _extract_bearer_token(header: str) -> str:
-    """Extract the bearer token from an HTTP authorization header.
-
-    Args:
-        header: The authorization header containing the token.
-
-    Returns:
-        The extracted token if present, else an empty string.
-    """
-    try:
-        scheme, token = header.split(" ", 1)
-        return token if scheme.lower() == "bearer" else ""
-    except ValueError:
-        return ""
 
 
 class K8SAuthDependency(AuthInterface):  # pylint: disable=too-few-public-methods
@@ -239,47 +236,51 @@ async def __call__(self, request: Request) -> tuple[str, str, bool, str]:
             user_id check should never be skipped with K8s authentication
             If user_id check should be skipped - always return False for k8s
             User's token
-        """
-        authorization_header = request.headers.get("Authorization")
-        if not authorization_header:
-            raise HTTPException(
-                status_code=401, detail="Unauthorized: No auth header found"
-            )
-
-        token = _extract_bearer_token(authorization_header)
-        if not token:
-            raise HTTPException(
-                status_code=401,
-                detail="Unauthorized: Bearer token not found or invalid",
-            )
 
+        Raises:
+            HTTPException: If authentication or authorization fails.
+        """
+        token = extract_user_token(request.headers)
         user_info = get_user_info(token)
+
         if user_info is None:
-            raise HTTPException(
-                status_code=403, detail="Forbidden: Invalid or expired token"
-            )
+            response = UnauthorizedResponse(cause="Invalid or expired Kubernetes token")
+            raise HTTPException(**response.model_dump())
+
         if user_info.user.username == "kube:admin":
-            user_info.user.uid = K8sClientSingleton.get_cluster_id()
-        authorization_api = K8sClientSingleton.get_authz_api()
-
-        sar = kubernetes.client.V1SubjectAccessReview(
-            spec=kubernetes.client.V1SubjectAccessReviewSpec(
-                user=user_info.user.username,
-                groups=user_info.user.groups,
-                non_resource_attributes=kubernetes.client.V1NonResourceAttributes(
-                    path=self.virtual_path, verb="get"
-                ),
-            )
-        )
+            try:
+                user_info.user.uid = K8sClientSingleton.get_cluster_id()
+            except ClusterIDUnavailableError as e:
+                logger.error("Failed to get cluster ID: %s", e)
+                response = InternalServerErrorResponse(
+                    response="Internal server error",
+                    cause="Unable to retrieve cluster ID",
+                )
+                raise HTTPException(**response.model_dump()) from e
+
         try:
+            authorization_api = K8sClientSingleton.get_authz_api()
+            sar = kubernetes.client.V1SubjectAccessReview(
+                spec=kubernetes.client.V1SubjectAccessReviewSpec(
+                    user=user_info.user.username,
+                    groups=user_info.user.groups,
+                    non_resource_attributes=kubernetes.client.V1NonResourceAttributes(
+                        path=self.virtual_path, verb="get"
+                    ),
+                )
+            )
             response = authorization_api.create_subject_access_review(sar)
+
             if not response.status.allowed:
-                raise HTTPException(
-                    status_code=403, detail="Forbidden: User does not have access"
-                )
-        except ApiException as e:
+                response = ForbiddenResponse.endpoint(user_id=user_info.user.uid)
+                raise HTTPException(**response.model_dump())
+        except Exception as e:
             logger.error("API exception during SubjectAccessReview: %s", e)
-            raise HTTPException(status_code=403, detail="Internal server error") from e
+            response = ServiceUnavailableResponse(
+                backend_name="Kubernetes API",
+                cause="Unable to perform authorization check",
+            )
+            raise HTTPException(**response.model_dump()) from e
 
         return (
             user_info.user.uid,

src/authentication/utils.py

@@ -2,25 +2,25 @@
 
 from fastapi import HTTPException
 from starlette.datastructures import Headers
+from models.responses import UnauthorizedResponse
 
 
 def extract_user_token(headers: Headers) -> str:
     """Extract the bearer token from an HTTP authorization header.
 
     Args:
         header: The authorization header containing the token.
-
     Returns:
         The extracted token if present, else an empty string.
     """
     authorization_header = headers.get("Authorization")
     if not authorization_header:
-        raise HTTPException(status_code=400, detail="No Authorization header found")
+        response = UnauthorizedResponse(cause="No Authorization header found")
+        raise HTTPException(**response.model_dump())
 
     scheme_and_token = authorization_header.strip().split()
     if len(scheme_and_token) != 2 or scheme_and_token[0].lower() != "bearer":
-        raise HTTPException(
-            status_code=400, detail="No token found in Authorization header"
-        )
+        response = UnauthorizedResponse(cause="No token found in Authorization header")
+        raise HTTPException(**response.model_dump())
 
     return scheme_and_token[1]