Skip to content

Commit 1ec58b5

Browse files
Merge pull request #981 from taylorwilsdon/per-user-identity
feat(auth): trusted-gateway identity mode (verify proxy-signed assertion as principal)
2 parents d9d09fa + b28adc2 commit 1ec58b5

16 files changed

Lines changed: 1393 additions & 29 deletions

auth/auth_info_middleware.py

Lines changed: 87 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
Authentication middleware to populate context state with user information
33
"""
44

5+
import asyncio
56
import logging
67
import time
78

@@ -10,12 +11,23 @@
1011
from fastmcp.server.dependencies import get_http_headers
1112

1213
from auth.external_oauth_provider import get_session_time
14+
from auth.gateway_identity import GatewayIdentityError, extract_email_from_assertion
1315
from auth.oauth21_session_store import ensure_session_from_access_token
16+
from auth.oauth_config import get_oauth_config, is_trust_gateway_identity
1417
from auth.oauth_types import WorkspaceAccessToken
1518

1619
# Configure logging
1720
logger = logging.getLogger(__name__)
1821

22+
_AUTH_IDENTITY_STATE_KEYS = (
23+
"authenticated_user_email",
24+
"authenticated_via",
25+
"user_email",
26+
"username",
27+
"auth_provider_type",
28+
"token_type",
29+
)
30+
1931

2032
def _token_fingerprint(token: str) -> str:
2133
"""Return a safe, short fingerprint of a bearer token for logging."""
@@ -43,6 +55,71 @@ async def _process_request_for_auth(self, context: MiddlewareContext):
4355
authenticated_user = None
4456
auth_via = None
4557

58+
# Trusted-gateway identity: verify the SIGNED assertion the fronting proxy injects
59+
# and use the asserted email as the principal. This is the highest-priority and only
60+
# trusted source in this mode (MCP_ENABLE_OAUTH21 is off — the proxy owns the handshake).
61+
if is_trust_gateway_identity():
62+
try:
63+
# FastMCP state is session-scoped by default. Remove any identity left by
64+
# legacy authentication before installing this request's verified identity.
65+
for state_key in _AUTH_IDENTITY_STATE_KEYS:
66+
await context.fastmcp_context.delete_state(state_key)
67+
68+
header_name = get_oauth_config().gateway_identity_header
69+
hdrs = get_http_headers(include={header_name}) or {}
70+
assertion = hdrs.get(header_name)
71+
if not assertion:
72+
raise GatewayIdentityError(
73+
f"Missing trusted-gateway identity header '{header_name}'"
74+
)
75+
76+
# Offload the (synchronous) JWKS fetch/verify off the event loop —
77+
# PyJWKClient can do network I/O on cold start / key rotation.
78+
verified_email = await asyncio.to_thread(
79+
extract_email_from_assertion, assertion
80+
)
81+
if not verified_email:
82+
raise GatewayIdentityError(
83+
"Trusted-gateway identity assertion failed verification"
84+
)
85+
86+
# These values are authoritative only for this request. In particular, do
87+
# not persist them in FastMCP's session-scoped state store.
88+
await context.fastmcp_context.set_state(
89+
"authenticated_user_email",
90+
verified_email,
91+
serializable=False,
92+
)
93+
await context.fastmcp_context.set_state(
94+
"authenticated_via",
95+
"gateway_assertion",
96+
serializable=False,
97+
)
98+
await context.fastmcp_context.set_state(
99+
"user_email",
100+
verified_email,
101+
serializable=False,
102+
)
103+
await context.fastmcp_context.set_state(
104+
"username",
105+
verified_email,
106+
serializable=False,
107+
)
108+
logger.info("✓ Authenticated via gateway_assertion: %s", verified_email)
109+
return
110+
except GatewayIdentityError:
111+
logger.warning(
112+
"[AuthInfoMiddleware] Trusted-gateway authentication rejected"
113+
)
114+
raise
115+
except Exception as e:
116+
logger.error(
117+
f"[AuthInfoMiddleware] Error processing gateway identity assertion: {e}"
118+
)
119+
raise GatewayIdentityError(
120+
"Trusted-gateway identity verification failed"
121+
) from e
122+
46123
# First check if FastMCP has already validated an access token
47124
try:
48125
access_token = get_access_token()
@@ -383,9 +460,11 @@ async def on_call_tool(self, context: MiddlewareContext, call_next):
383460

384461
except Exception as e:
385462
# Check if this is an authentication error - don't log traceback for these
386-
if "GoogleAuthenticationError" in str(
387-
type(e)
388-
) or "Access denied: Cannot retrieve credentials" in str(e):
463+
if (
464+
isinstance(e, GatewayIdentityError)
465+
or "GoogleAuthenticationError" in str(type(e))
466+
or "Access denied: Cannot retrieve credentials" in str(e)
467+
):
389468
logger.info(f"Authentication check failed: {e}")
390469
else:
391470
logger.error(f"Error in on_call_tool middleware: {e}", exc_info=True)
@@ -405,9 +484,11 @@ async def on_get_prompt(self, context: MiddlewareContext, call_next):
405484

406485
except Exception as e:
407486
# Check if this is an authentication error - don't log traceback for these
408-
if "GoogleAuthenticationError" in str(
409-
type(e)
410-
) or "Access denied: Cannot retrieve credentials" in str(e):
487+
if (
488+
isinstance(e, GatewayIdentityError)
489+
or "GoogleAuthenticationError" in str(type(e))
490+
or "Access denied: Cannot retrieve credentials" in str(e)
491+
):
411492
logger.info(f"Authentication check failed in prompt: {e}")
412493
else:
413494
logger.error(f"Error in on_get_prompt middleware: {e}", exc_info=True)

auth/gateway_identity.py

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
"""
2+
Trusted-gateway identity verification.
3+
4+
When an MCP-aware reverse proxy fronts this server, it authenticates the user and
5+
attaches a SIGNED identity assertion (a JWT) to every upstream request. This module
6+
verifies that assertion against the proxy's JWKS and returns the verified claims
7+
(notably the user's email), so the asserted identity can be used as the per-request
8+
principal — without this server terminating MCP OAuth itself.
9+
10+
Provider-agnostic: works with any proxy that injects a JWKS-verifiable JWT identity
11+
header — e.g. Pomerium (x-pomerium-jwt-assertion, ES256), oauth2-proxy, Cloudflare
12+
Access (cf-access-jwt-assertion, RS256), Istio/Envoy, Traefik ForwardAuth. The header
13+
name, signing algorithm(s), JWKS URL, and optional issuer/audience are all configurable
14+
(see auth.oauth_config); defaults target Pomerium.
15+
16+
Security: the assertion is verified cryptographically (signature + expiry, and optional
17+
issuer/audience). An unverified or malformed assertion yields None — callers must treat
18+
that as "no identity" (fail closed), never as a trusted user.
19+
"""
20+
21+
import logging
22+
from typing import Any, Optional
23+
24+
import jwt
25+
from jwt import PyJWKClient
26+
from pydantic.networks import validate_email
27+
28+
from auth.oauth_config import get_oauth_config
29+
30+
logger = logging.getLogger(__name__)
31+
32+
33+
class GatewayIdentityError(Exception):
34+
"""Raised when a request lacks a valid trusted-gateway identity."""
35+
36+
37+
def normalize_principal_email(value: Any) -> Optional[str]:
38+
"""Return the canonical form used for gateway principals and credential keys."""
39+
if not isinstance(value, str):
40+
return None
41+
value = value.strip()
42+
if not value:
43+
return None
44+
try:
45+
_, canonical_email = validate_email(value)
46+
except ValueError:
47+
return None
48+
return canonical_email.lower()
49+
50+
51+
def require_gateway_principal(authenticated_user: Any, authenticated_via: Any) -> str:
52+
"""Return a canonical principal only when it came from gateway verification."""
53+
email = normalize_principal_email(authenticated_user)
54+
if authenticated_via != "gateway_assertion" or not email:
55+
raise GatewayIdentityError(
56+
"Trusted-gateway mode requires a request-scoped verified gateway principal"
57+
)
58+
return email
59+
60+
61+
async def get_verified_gateway_principal(context=None) -> str:
62+
"""Resolve the authoritative gateway principal from the active FastMCP request."""
63+
if context is None:
64+
try:
65+
from fastmcp.server.dependencies import get_context
66+
67+
context = get_context()
68+
except Exception as exc:
69+
raise GatewayIdentityError(
70+
"Trusted-gateway principal is unavailable outside an MCP request"
71+
) from exc
72+
73+
if context is None:
74+
raise GatewayIdentityError(
75+
"Trusted-gateway principal is unavailable outside an MCP request"
76+
)
77+
78+
authenticated_user = await context.get_state("authenticated_user_email")
79+
authenticated_via = await context.get_state("authenticated_via")
80+
return require_gateway_principal(authenticated_user, authenticated_via)
81+
82+
83+
# PyJWKClient caches fetched keys (and refreshes on unknown kid). Cache one client per
84+
# JWKS URL for the process lifetime.
85+
_jwks_clients: dict[str, PyJWKClient] = {}
86+
87+
88+
def _get_jwks_client(jwks_url: str) -> PyJWKClient:
89+
client = _jwks_clients.get(jwks_url)
90+
if client is None:
91+
# PyJWKClient keeps fetched signing keys in-memory and re-fetches when it sees an
92+
# unknown kid (key rotation), so this is safe to hold for the process lifetime.
93+
client = PyJWKClient(jwks_url, cache_keys=True)
94+
_jwks_clients[jwks_url] = client
95+
return client
96+
97+
98+
def verify_gateway_assertion(token: str) -> Optional[dict]:
99+
"""
100+
Verify a trusted-gateway identity-assertion JWT and return its claims.
101+
102+
Args:
103+
token: the raw JWT from the assertion header.
104+
105+
Returns:
106+
The verified claims dict (includes "email"/"sub") on success, else None.
107+
"""
108+
if not token:
109+
return None
110+
111+
config = get_oauth_config()
112+
jwks_url = config.gateway_identity_jwks_url
113+
if not jwks_url:
114+
logger.error(
115+
"verify_gateway_assertion called but GATEWAY_IDENTITY_JWKS_URL is unset"
116+
)
117+
return None
118+
audience = config.gateway_identity_audience
119+
if not audience:
120+
logger.error(
121+
"verify_gateway_assertion called but GATEWAY_IDENTITY_AUDIENCE is unset"
122+
)
123+
return None
124+
125+
try:
126+
signing_key = _get_jwks_client(jwks_url).get_signing_key_from_jwt(token)
127+
128+
decode_kwargs: dict = {
129+
# Pin to the configured algorithm(s) so a malicious token can't downgrade to
130+
# "alg: none" or trigger an HMAC/asymmetric confusion attack.
131+
"algorithms": config.gateway_identity_algorithms,
132+
# Require expiry and audience; issuer is additionally enforced when configured.
133+
"options": {
134+
"require": ["exp"],
135+
"verify_aud": True,
136+
},
137+
"audience": audience,
138+
}
139+
if config.gateway_identity_issuer:
140+
decode_kwargs["issuer"] = config.gateway_identity_issuer
141+
142+
claims = jwt.decode(token, signing_key.key, **decode_kwargs)
143+
return claims
144+
145+
except jwt.PyJWTError as e:
146+
# Invalid signature / expired / wrong aud-iss / unknown kid, etc.
147+
logger.warning(
148+
"SECURITY: rejected gateway identity assertion (%s: %s)",
149+
type(e).__name__,
150+
e,
151+
)
152+
return None
153+
except Exception as e: # noqa: BLE001 - JWKS fetch / network / unexpected
154+
logger.error(
155+
"Error verifying gateway identity assertion (%s: %s)",
156+
type(e).__name__,
157+
e,
158+
)
159+
return None
160+
161+
162+
def extract_email_from_assertion(token: str) -> Optional[str]:
163+
"""Verify the assertion and return the lowercased email claim, or None."""
164+
claims = verify_gateway_assertion(token)
165+
if not claims:
166+
return None
167+
email = normalize_principal_email(claims.get("email"))
168+
if not email:
169+
logger.warning(
170+
"SECURITY: verified gateway assertion has no usable 'email' claim (sub=%s)",
171+
claims.get("sub"),
172+
)
173+
return None
174+
return email

0 commit comments

Comments
 (0)