Allow private JWT headers (#4290)

Author: jlowin

fastmcp_slim/fastmcp/server/auth/providers/jwt.py

@@ -13,6 +13,8 @@
 from cryptography.hazmat.primitives.asymmetric import rsa
 from joserfc import jwk, jwt
 from joserfc.errors import JoseError
+from joserfc.jws import JWSRegistry
+from joserfc.registry import JWS_HEADER_REGISTRY
 from pydantic import AnyHttpUrl, SecretStr
 from typing_extensions import TypedDict
 
@@ -24,6 +26,7 @@
 logger = get_logger(__name__)
 
 JWKKeyData: TypeAlias = dict[str, str | list[str]]
+SUPPORTED_JWS_HEADER_FIELDS = frozenset(JWS_HEADER_REGISTRY)
 
 
 def _import_key_for_algorithm(key: str | bytes | JWKKeyData, algorithm: str):
@@ -45,6 +48,21 @@ def _jwk_to_pem(key_data: JWKKeyData) -> str:
     raise ValueError(f"Unsupported JWK key type: {key_type!r}")
 
 
+def _has_unsupported_critical_headers(header: dict[str, Any]) -> bool:
+    crit = header.get("crit")
+    if crit is None:
+        return False
+    if not isinstance(crit, list):
+        return True
+
+    return any(
+        not isinstance(header_name, str)
+        or header_name not in header
+        or header_name not in SUPPORTED_JWS_HEADER_FIELDS
+        for header_name in crit
+    )
+
+
 class JWKData(TypedDict, total=False):
     """JSON Web Key data structure."""
 
@@ -429,7 +447,22 @@ async def load_access_token(self, token: str) -> AccessToken | None:
 
             # Decode and verify the JWT token
             key = _import_key_for_algorithm(verification_key, self.algorithm)
-            claims = jwt.decode(token, key, algorithms=[self.algorithm]).claims
+            header = decode_jwt_header(token)
+            if _has_unsupported_critical_headers(header):
+                self.logger.debug(
+                    "Token validation failed: unsupported critical JWT header"
+                )
+                return None
+
+            claims = jwt.decode(
+                token,
+                key,
+                algorithms=[self.algorithm],
+                registry=JWSRegistry(
+                    algorithms=[self.algorithm],
+                    strict_check_header=False,
+                ),
+            ).claims
 
             # Extract client ID early for logging
             client_id = (

tests/server/auth/test_jwt_provider.py

@@ -6,6 +6,8 @@
 import pytest
 from joserfc import jwk as jose_jwk
 from joserfc import jwt
+from joserfc.jws import JWSRegistry
+from joserfc.registry import HeaderParameter
 from pytest_httpx import HTTPXMock
 
 from fastmcp import FastMCP
@@ -182,6 +184,73 @@ def test_create_token_with_scopes(self, rsa_key_pair: RSAKeyPair):
         # We'll validate the scopes in the BearerToken tests
 
 
+class TestJWTVerifierHeaders:
+    async def test_non_critical_private_header_is_allowed(
+        self, rsa_key_pair: RSAKeyPair
+    ):
+        signing_key = jose_jwk.import_key(
+            rsa_key_pair.private_key.get_secret_value(),
+            "RSA",
+        )
+        token = jwt.encode(
+            {
+                "alg": "RS256",
+                "cat": "cl_example",
+            },
+            {
+                "sub": "test-user",
+                "iss": "https://test.example.com",
+                "exp": int(time.time()) + 3600,
+            },
+            signing_key,
+            algorithms=["RS256"],
+            registry=JWSRegistry(strict_check_header=False),
+        )
+        verifier = JWTVerifier(
+            public_key=rsa_key_pair.public_key,
+            issuer="https://test.example.com",
+        )
+
+        access_token = await verifier.verify_token(token)
+
+        assert access_token is not None
+        assert access_token.client_id == "test-user"
+
+    async def test_critical_private_header_is_rejected(self, rsa_key_pair: RSAKeyPair):
+        signing_key = jose_jwk.import_key(
+            rsa_key_pair.private_key.get_secret_value(),
+            "RSA",
+        )
+        token = jwt.encode(
+            {
+                "alg": "RS256",
+                "crit": ["cat"],
+                "cat": "cl_example",
+            },
+            {
+                "sub": "test-user",
+                "iss": "https://test.example.com",
+                "exp": int(time.time()) + 3600,
+            },
+            signing_key,
+            algorithms=["RS256"],
+            registry=JWSRegistry(
+                header_registry={
+                    "cat": HeaderParameter("Custom private header", "str")
+                },
+                strict_check_header=False,
+            ),
+        )
+        verifier = JWTVerifier(
+            public_key=rsa_key_pair.public_key,
+            issuer="https://test.example.com",
+        )
+
+        access_token = await verifier.verify_token(token)
+
+        assert access_token is None
+
+
 class TestSymmetricKeyJWT:
     """Tests for JWT verification using symmetric keys (HMAC algorithms)."""