fix(auth): Use token in AuthTuple

While implementing authz I've made the authentication module put the
claims in the `AuthTuple` instead of putting the original token because
I assumed the original token was unnecessary and I didn't want to make
the authz module re-parse the token to get the claims.

That was a mistake, because we need to pass on the user token to the MCP
server for it to validate the user's access. So this commit changes the
`AuthTuple` to contain the original token instead of the claims.

The authz module has been updated to parse the claims from the token
instead of receiving them as a JSON string. Somewhat hackily, since we
don't want to re-verify the token signature, so we assume that the token
has already been verified during authentication and just decode the
claims from the middle section of the JWT token.

Author: omertuc

src/auth/jwk_token.py

@@ -1,7 +1,6 @@
 """Manage authentication flow for FastAPI endpoints with JWK based JWT auth."""
 
 import logging
-import json
 from asyncio import Lock
 from typing import Any, Callable
 
@@ -191,4 +190,4 @@ async def __call__(self, request: Request) -> AuthTuple:
 
         logger.info("Successfully authenticated user %s (ID: %s)", username, user_id)
 
-        return user_id, username, json.dumps(claims)
+        return user_id, username, user_token

src/authorization/resolvers.py

@@ -1,8 +1,9 @@
 """Authorization resolvers for role evaluation and access control."""
 
 from abc import ABC, abstractmethod
-import json
 import logging
+import base64
+import json
 from typing import Any
 
 from jsonpath_ng import parse
@@ -38,6 +39,17 @@ async def resolve_roles(self, auth: AuthTuple) -> UserRoles:
         return set()
 
 
+def unsafe_get_claims(token: str) -> dict[str, Any]:
+    """Get claims from a token without validating the signature.
+
+    A somewhat hacky way to get JWT claims without verifying the signature.
+    We assume verification has already been done during authentication.
+    """
+    payload = token.split(".")[1]
+    padded = payload + "=" * (-len(payload) % 4)
+    return json.loads(base64.urlsafe_b64decode(padded))
+
+
 class JwtRolesResolver(RolesResolver):  # pylint: disable=too-few-public-methods
     """Processes JWT claims with the given JSONPath rules to get roles."""
 
@@ -76,7 +88,7 @@ def _get_claims(auth: AuthTuple) -> dict[str, Any]:
             # No claims for guests
             return {}
 
-        jwt_claims = json.loads(token)
+        jwt_claims = unsafe_get_claims(token)
 
         if not jwt_claims:
             raise RoleResolutionError(

tests/unit/auth/test_jwk_token.py

@@ -3,8 +3,6 @@
 """Unit tests for functions defined in auth/jwk_token.py"""
 
 import time
-import json
-import base64
 
 import pytest
 from fastapi import HTTPException, Request
@@ -174,19 +172,12 @@ def set_auth_header(request: Request, token: str):
     request.scope["headers"] = new_headers
 
 
-def get_claims(token: str):
-    """Extract claims from a JWT token without validating it."""
-    payload = token.split(".")[1]
-    padded = payload + "=" * (-len(payload) % 4)
-    return json.loads(base64.urlsafe_b64decode(padded))
-
-
-def ensure_test_user_id_and_name(auth_tuple, token):
+def ensure_test_user_id_and_name(auth_tuple, expected_token):
     """Utility to ensure that the values in the auth tuple match the test values."""
-    user_id, username, token_claims = auth_tuple
+    user_id, username, token = auth_tuple
     assert user_id == TEST_USER_ID
     assert username == TEST_USER_NAME
-    assert json.loads(token_claims) == get_claims(token)
+    assert token == expected_token
 
 
 async def test_valid(

tests/unit/authorization/test_resolvers.py

@@ -1,9 +1,23 @@
 """Unit tests for the authorization resolvers."""
 
+import json
+import base64
+
 from authorization.resolvers import JwtRolesResolver, GenericAccessResolver
 from models.config import JwtRoleRule, AccessRule, JsonPathOperator, Action
 
 
+def claims_to_token(claims: dict) -> str:
+    """Convert JWT claims dictionary to a JSON string token."""
+
+    string_claims = json.dumps(claims)
+    b64_encoded_claims = (
+        base64.urlsafe_b64encode(string_claims.encode()).decode().rstrip("=")
+    )
+
+    return f"foo_header.{b64_encoded_claims}.foo_signature"
+
+
 class TestJwtRolesResolver:
     """Test cases for JwtRolesResolver."""
 
@@ -33,7 +47,7 @@ async def test_resolve_roles_redhat_employee(self):
         }
 
         # Mock auth tuple with JWT claims as third element
-        auth = ("user", "token", str(jwt_claims).replace("'", '"'))
+        auth = ("user", "token", claims_to_token(jwt_claims))
         roles = await jwt_resolver.resolve_roles(auth)
         assert "employee" in roles
 
@@ -57,7 +71,7 @@ async def test_resolve_roles_no_match(self):
         }
 
         # Mock auth tuple with JWT claims as third element
-        auth = ("user", "token", str(jwt_claims).replace("'", '"'))
+        auth = ("user", "token", claims_to_token(jwt_claims))
         roles = await jwt_resolver.resolve_roles(auth)
         assert len(roles) == 0