mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-24 06:24:18 +02:00
Fix OAuth proxy refresh token storage for multi-instance deployments (#2483)
* Fix OAuth proxy refresh token storage for multi-instance deployments - Use pluggable client_storage instead of local dict for refresh tokens - Store refresh tokens by SHA-256 hash for defense in depth - Remove unused access token and relationship mapping stores - Simplify revocation logic * Address review feedback for refresh token storage - Use calculated refresh_expires_in for TTL instead of hardcoded 30 days - Populate expires_at field with actual expiry timestamp - Add client_id validation in load_refresh_token to prevent cross-client token usage
This commit is contained in:
Jeremiah Lowin
GitHub
parent
3341c0c893
commit
01ecc91807
1 changed files with 92 additions and 69 deletions
|
|
@ -179,6 +179,28 @@ class JTIMapping(BaseModel):
|
|||
created_at: float # Unix timestamp
|
||||
|
||||
|
||||
class RefreshTokenMetadata(BaseModel):
|
||||
"""Metadata for a refresh token, stored keyed by token hash.
|
||||
|
||||
We store only metadata (not the token itself) for security - if storage
|
||||
is compromised, attackers get hashes they can't reverse into usable tokens.
|
||||
"""
|
||||
|
||||
client_id: str
|
||||
scopes: list[str]
|
||||
expires_at: int | None = None
|
||||
created_at: float
|
||||
|
||||
|
||||
def _hash_token(token: str) -> str:
|
||||
"""Hash a token for secure storage lookup.
|
||||
|
||||
Uses SHA-256 to create a one-way hash. The original token cannot be
|
||||
recovered from the hash, providing defense in depth if storage is compromised.
|
||||
"""
|
||||
return hashlib.sha256(token.encode()).hexdigest()
|
||||
|
||||
|
||||
class ProxyDCRClient(OAuthClientInformationFull):
|
||||
"""Client for DCR proxy with configurable redirect URI validation.
|
||||
|
||||
|
|
@ -624,14 +646,18 @@ class OAuthProxy(OAuthProvider):
|
|||
|
||||
State Management
|
||||
---------------
|
||||
The proxy maintains minimal but crucial state:
|
||||
The proxy maintains minimal but crucial state via pluggable storage (client_storage):
|
||||
- _oauth_transactions: Active authorization flows with client context
|
||||
- _client_codes: Authorization codes with PKCE challenges and upstream tokens
|
||||
- _access_tokens, _refresh_tokens: Token storage for revocation
|
||||
- Token relationship mappings for cleanup and rotation
|
||||
- _jti_mapping_store: Maps FastMCP token JTIs to upstream token IDs
|
||||
- _refresh_token_store: Refresh token metadata (keyed by token hash)
|
||||
|
||||
All state is stored in the configured client_storage backend (Redis, disk, etc.)
|
||||
enabling horizontal scaling across multiple instances.
|
||||
|
||||
Security Considerations
|
||||
----------------------
|
||||
- Refresh tokens stored by hash only (defense in depth if storage compromised)
|
||||
- PKCE enforced end-to-end (client to proxy, proxy to upstream)
|
||||
- Authorization codes are single-use with short expiry
|
||||
- Transaction IDs are cryptographically random
|
||||
|
|
@ -895,13 +921,17 @@ class OAuthProxy(OAuthProvider):
|
|||
raise_on_validation_error=True,
|
||||
)
|
||||
|
||||
# Local state for token bookkeeping only (no client caching)
|
||||
self._access_tokens: dict[str, AccessToken] = {}
|
||||
self._refresh_tokens: dict[str, RefreshToken] = {}
|
||||
|
||||
# Token relation mappings for cleanup
|
||||
self._access_to_refresh: dict[str, str] = {}
|
||||
self._refresh_to_access: dict[str, str] = {}
|
||||
# Refresh token metadata storage, keyed by token hash for security.
|
||||
# We only store metadata (not the token itself) - if storage is compromised,
|
||||
# attackers get hashes they can't reverse into usable tokens.
|
||||
self._refresh_token_store: PydanticAdapter[RefreshTokenMetadata] = (
|
||||
PydanticAdapter[RefreshTokenMetadata](
|
||||
key_value=self._client_storage,
|
||||
pydantic_model=RefreshTokenMetadata,
|
||||
default_collection="mcp-refresh-tokens",
|
||||
raise_on_validation_error=True,
|
||||
)
|
||||
)
|
||||
|
||||
# Use the provided token validator
|
||||
self._token_validator: TokenVerifier = token_verifier
|
||||
|
|
@ -1254,25 +1284,18 @@ class OAuthProxy(OAuthProvider):
|
|||
ttl=60 * 60 * 24 * 30, # Auto-expire with refresh token (30 days)
|
||||
)
|
||||
|
||||
# Store FastMCP access token for MCP framework validation
|
||||
self._access_tokens[fastmcp_access_token] = AccessToken(
|
||||
token=fastmcp_access_token,
|
||||
client_id=client.client_id,
|
||||
scopes=authorization_code.scopes,
|
||||
expires_at=int(time.time() + expires_in),
|
||||
)
|
||||
|
||||
# Store FastMCP refresh token if provided
|
||||
if fastmcp_refresh_token:
|
||||
self._refresh_tokens[fastmcp_refresh_token] = RefreshToken(
|
||||
token=fastmcp_refresh_token,
|
||||
client_id=client.client_id,
|
||||
scopes=authorization_code.scopes,
|
||||
expires_at=None,
|
||||
# Store refresh token metadata (keyed by hash for security)
|
||||
if fastmcp_refresh_token and refresh_expires_in:
|
||||
await self._refresh_token_store.put(
|
||||
key=_hash_token(fastmcp_refresh_token),
|
||||
value=RefreshTokenMetadata(
|
||||
client_id=client.client_id,
|
||||
scopes=authorization_code.scopes,
|
||||
expires_at=int(time.time()) + refresh_expires_in,
|
||||
created_at=time.time(),
|
||||
),
|
||||
ttl=refresh_expires_in,
|
||||
)
|
||||
# Maintain token relationships for cleanup
|
||||
self._access_to_refresh[fastmcp_access_token] = fastmcp_refresh_token
|
||||
self._refresh_to_access[fastmcp_refresh_token] = fastmcp_access_token
|
||||
|
||||
logger.debug(
|
||||
"Issued FastMCP tokens for client=%s (access_jti=%s, refresh_jti=%s)",
|
||||
|
|
@ -1316,8 +1339,29 @@ class OAuthProxy(OAuthProvider):
|
|||
client: OAuthClientInformationFull,
|
||||
refresh_token: str,
|
||||
) -> RefreshToken | None:
|
||||
"""Load refresh token from local storage."""
|
||||
return self._refresh_tokens.get(refresh_token)
|
||||
"""Load refresh token metadata from distributed storage.
|
||||
|
||||
Looks up by token hash and reconstructs the RefreshToken object.
|
||||
Validates that the token belongs to the requesting client.
|
||||
"""
|
||||
token_hash = _hash_token(refresh_token)
|
||||
metadata = await self._refresh_token_store.get(key=token_hash)
|
||||
if not metadata:
|
||||
return None
|
||||
# Verify token belongs to this client (prevents cross-client token usage)
|
||||
if metadata.client_id != client.client_id:
|
||||
logger.warning(
|
||||
"Refresh token client_id mismatch: expected %s, got %s",
|
||||
client.client_id,
|
||||
metadata.client_id,
|
||||
)
|
||||
return None
|
||||
return RefreshToken(
|
||||
token=refresh_token,
|
||||
client_id=metadata.client_id,
|
||||
scopes=metadata.scopes,
|
||||
expires_at=metadata.expires_at,
|
||||
)
|
||||
|
||||
async def exchange_refresh_token(
|
||||
self,
|
||||
|
|
@ -1488,30 +1532,20 @@ class OAuthProxy(OAuthProvider):
|
|||
"Rotated refresh token (old JTI invalidated - one-time use enforced)"
|
||||
)
|
||||
|
||||
# Update local token tracking
|
||||
self._access_tokens[new_fastmcp_access] = AccessToken(
|
||||
token=new_fastmcp_access,
|
||||
client_id=client.client_id,
|
||||
scopes=scopes,
|
||||
expires_at=int(time.time() + new_expires_in),
|
||||
)
|
||||
self._refresh_tokens[new_fastmcp_refresh] = RefreshToken(
|
||||
token=new_fastmcp_refresh,
|
||||
client_id=client.client_id,
|
||||
scopes=scopes,
|
||||
expires_at=None,
|
||||
# Store new refresh token metadata (keyed by hash)
|
||||
await self._refresh_token_store.put(
|
||||
key=_hash_token(new_fastmcp_refresh),
|
||||
value=RefreshTokenMetadata(
|
||||
client_id=client.client_id,
|
||||
scopes=scopes,
|
||||
expires_at=int(time.time()) + refresh_ttl,
|
||||
created_at=time.time(),
|
||||
),
|
||||
ttl=refresh_ttl,
|
||||
)
|
||||
|
||||
# Update token relationship mappings
|
||||
self._access_to_refresh[new_fastmcp_access] = new_fastmcp_refresh
|
||||
self._refresh_to_access[new_fastmcp_refresh] = new_fastmcp_access
|
||||
|
||||
# Clean up old token from in-memory tracking
|
||||
self._refresh_tokens.pop(refresh_token.token, None)
|
||||
old_access = self._refresh_to_access.pop(refresh_token.token, None)
|
||||
if old_access:
|
||||
self._access_tokens.pop(old_access, None)
|
||||
self._access_to_refresh.pop(old_access, None)
|
||||
# Delete old refresh token (by hash)
|
||||
await self._refresh_token_store.delete(key=_hash_token(refresh_token.token))
|
||||
|
||||
logger.info(
|
||||
"Issued new FastMCP tokens (rotated refresh) for client=%s (access_jti=%s, refresh_jti=%s)",
|
||||
|
|
@ -1592,24 +1626,13 @@ class OAuthProxy(OAuthProvider):
|
|||
async def revoke_token(self, token: AccessToken | RefreshToken) -> None:
|
||||
"""Revoke token locally and with upstream server if supported.
|
||||
|
||||
Removes tokens from local storage and attempts to revoke them with
|
||||
the upstream server if a revocation endpoint is configured.
|
||||
For refresh tokens, removes from local storage by hash.
|
||||
For all tokens, attempts upstream revocation if endpoint is configured.
|
||||
Access token JTI mappings expire via TTL.
|
||||
"""
|
||||
# Clean up local token storage
|
||||
if isinstance(token, AccessToken):
|
||||
self._access_tokens.pop(token.token, None)
|
||||
# Also remove associated refresh token
|
||||
paired_refresh = self._access_to_refresh.pop(token.token, None)
|
||||
if paired_refresh:
|
||||
self._refresh_tokens.pop(paired_refresh, None)
|
||||
self._refresh_to_access.pop(paired_refresh, None)
|
||||
else: # RefreshToken
|
||||
self._refresh_tokens.pop(token.token, None)
|
||||
# Also remove associated access token
|
||||
paired_access = self._refresh_to_access.pop(token.token, None)
|
||||
if paired_access:
|
||||
self._access_tokens.pop(paired_access, None)
|
||||
self._access_to_refresh.pop(paired_access, None)
|
||||
# For refresh tokens, delete from local storage by hash
|
||||
if isinstance(token, RefreshToken):
|
||||
await self._refresh_token_store.delete(key=_hash_token(token.token))
|
||||
|
||||
# Attempt upstream revocation if endpoint is configured
|
||||
if self._upstream_revocation_endpoint:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue