mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-16 18:49:14 +02:00
52feff6878* Transparently refresh upstream token in OAuthProxy.load_access_token() When upstream token validation fails during load_access_token, attempt to refresh using the stored refresh token before returning None. This prevents premature 401s that force clients into expensive full re-auth flows when the upstream token expires. Co-authored-by: Claude <noreply@anthropic.com> * Gate transparent refresh on token expiry, add advisory lock Only attempt upstream refresh when the token is actually expired, not on any validation failure (scope mismatch, revocation, etc.). Add per-token advisory lock to prevent concurrent async tasks from racing to refresh the same upstream token. * Re-check expiry inside lock, reload from storage after refresh failure --------- Co-authored-by: Claude <noreply@anthropic.com>
857 lines
33 KiB
Python
857 lines
33 KiB
Python
"""Tests for OAuth proxy token endpoint and handling."""
|
|
|
|
import time
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
import pytest
|
|
from key_value.aio.stores.memory import MemoryStore
|
|
from mcp.server.auth.handlers.token import TokenErrorResponse
|
|
from mcp.server.auth.handlers.token import TokenHandler as SDKTokenHandler
|
|
from mcp.server.auth.provider import AccessToken, AuthorizationCode
|
|
from mcp.shared.auth import OAuthClientInformationFull
|
|
from pydantic import AnyUrl
|
|
|
|
from fastmcp.server.auth.auth import RefreshToken, TokenHandler, TokenVerifier
|
|
from fastmcp.server.auth.oauth_proxy import OAuthProxy
|
|
from fastmcp.server.auth.oauth_proxy.models import (
|
|
DEFAULT_ACCESS_TOKEN_EXPIRY_NO_REFRESH_SECONDS,
|
|
DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS,
|
|
ClientCode,
|
|
JTIMapping,
|
|
UpstreamTokenSet,
|
|
_hash_token,
|
|
)
|
|
from fastmcp.server.auth.providers.jwt import JWTVerifier
|
|
|
|
|
|
class TestOAuthProxyTokenEndpointAuth:
|
|
"""Tests for token endpoint authentication methods."""
|
|
|
|
def test_token_auth_method_initialization(self, jwt_verifier):
|
|
"""Test different token endpoint auth methods."""
|
|
# client_secret_post
|
|
proxy_post = OAuthProxy(
|
|
upstream_authorization_endpoint="https://oauth.example.com/authorize",
|
|
upstream_token_endpoint="https://oauth.example.com/token",
|
|
upstream_client_id="client",
|
|
upstream_client_secret="secret",
|
|
token_verifier=jwt_verifier,
|
|
base_url="https://proxy.example.com",
|
|
token_endpoint_auth_method="client_secret_post",
|
|
jwt_signing_key="test-secret",
|
|
client_storage=MemoryStore(),
|
|
)
|
|
assert proxy_post._token_endpoint_auth_method == "client_secret_post"
|
|
|
|
# client_secret_basic (default)
|
|
proxy_basic = OAuthProxy(
|
|
upstream_authorization_endpoint="https://oauth.example.com/authorize",
|
|
upstream_token_endpoint="https://oauth.example.com/token",
|
|
upstream_client_id="client",
|
|
upstream_client_secret="secret",
|
|
token_verifier=jwt_verifier,
|
|
base_url="https://proxy.example.com",
|
|
token_endpoint_auth_method="client_secret_basic",
|
|
jwt_signing_key="test-secret",
|
|
client_storage=MemoryStore(),
|
|
)
|
|
assert proxy_basic._token_endpoint_auth_method == "client_secret_basic"
|
|
|
|
# None (use authlib default)
|
|
proxy_default = OAuthProxy(
|
|
upstream_authorization_endpoint="https://oauth.example.com/authorize",
|
|
upstream_token_endpoint="https://oauth.example.com/token",
|
|
upstream_client_id="client",
|
|
upstream_client_secret="secret",
|
|
token_verifier=jwt_verifier,
|
|
base_url="https://proxy.example.com",
|
|
jwt_signing_key="test-secret",
|
|
client_storage=MemoryStore(),
|
|
)
|
|
assert proxy_default._token_endpoint_auth_method is None
|
|
|
|
async def test_token_auth_method_passed_to_client(self, jwt_verifier):
|
|
"""Test that auth method is passed to AsyncOAuth2Client."""
|
|
proxy = OAuthProxy(
|
|
upstream_authorization_endpoint="https://oauth.example.com/authorize",
|
|
upstream_token_endpoint="https://oauth.example.com/token",
|
|
upstream_client_id="client-id",
|
|
upstream_client_secret="client-secret",
|
|
token_verifier=jwt_verifier,
|
|
base_url="https://proxy.example.com",
|
|
token_endpoint_auth_method="client_secret_post",
|
|
jwt_signing_key="test-secret",
|
|
client_storage=MemoryStore(),
|
|
)
|
|
|
|
# Initialize JWT issuer before token operations
|
|
proxy.set_mcp_path("/mcp")
|
|
|
|
# First, create a valid FastMCP token via full OAuth flow
|
|
client = OAuthClientInformationFull(
|
|
client_id="test-client",
|
|
client_secret="test-secret",
|
|
redirect_uris=[AnyUrl("http://localhost:12345/callback")],
|
|
)
|
|
|
|
# Mock the upstream OAuth provider response
|
|
with patch(
|
|
"fastmcp.server.auth.oauth_proxy.proxy.AsyncOAuth2Client"
|
|
) as MockClient:
|
|
mock_client = AsyncMock()
|
|
|
|
# Mock initial token exchange (authorization code flow)
|
|
mock_client.fetch_token = AsyncMock(
|
|
return_value={
|
|
"access_token": "upstream-access-token",
|
|
"refresh_token": "upstream-refresh-token",
|
|
"expires_in": 3600,
|
|
"token_type": "Bearer",
|
|
}
|
|
)
|
|
|
|
# Mock token refresh
|
|
mock_client.refresh_token = AsyncMock(
|
|
return_value={
|
|
"access_token": "new-upstream-token",
|
|
"refresh_token": "new-upstream-refresh",
|
|
"expires_in": 3600,
|
|
"token_type": "Bearer",
|
|
}
|
|
)
|
|
MockClient.return_value = mock_client
|
|
|
|
# Register client and do initial OAuth flow to get valid FastMCP tokens
|
|
await proxy.register_client(client)
|
|
|
|
# Store client code that would be created during OAuth callback
|
|
client_code = ClientCode(
|
|
code="test-auth-code",
|
|
client_id="test-client",
|
|
redirect_uri="http://localhost:12345/callback",
|
|
code_challenge="",
|
|
code_challenge_method="S256",
|
|
scopes=["read"],
|
|
idp_tokens={
|
|
"access_token": "upstream-access-token",
|
|
"refresh_token": "upstream-refresh-token",
|
|
"expires_in": 3600,
|
|
"token_type": "Bearer",
|
|
},
|
|
expires_at=time.time() + 300,
|
|
created_at=time.time(),
|
|
)
|
|
await proxy._code_store.put(key=client_code.code, value=client_code)
|
|
|
|
# Exchange authorization code to get FastMCP tokens
|
|
auth_code = AuthorizationCode(
|
|
code="test-auth-code",
|
|
scopes=["read"],
|
|
expires_at=time.time() + 300,
|
|
client_id="test-client",
|
|
code_challenge="",
|
|
redirect_uri=AnyUrl("http://localhost:12345/callback"),
|
|
redirect_uri_provided_explicitly=True,
|
|
)
|
|
result = await proxy.exchange_authorization_code(
|
|
client=client,
|
|
authorization_code=auth_code,
|
|
)
|
|
|
|
# Now test refresh with the valid FastMCP refresh token
|
|
assert result.refresh_token is not None
|
|
fastmcp_refresh = RefreshToken(
|
|
token=result.refresh_token,
|
|
client_id="test-client",
|
|
scopes=["read"],
|
|
expires_at=None,
|
|
)
|
|
|
|
# Reset mock to check refresh call
|
|
MockClient.reset_mock()
|
|
mock_client.refresh_token = AsyncMock(
|
|
return_value={
|
|
"access_token": "new-upstream-token-2",
|
|
"refresh_token": "new-upstream-refresh-2",
|
|
"expires_in": 3600,
|
|
"token_type": "Bearer",
|
|
}
|
|
)
|
|
MockClient.return_value = mock_client
|
|
|
|
await proxy.exchange_refresh_token(client, fastmcp_refresh, ["read"])
|
|
|
|
# Verify auth method was passed to OAuth client
|
|
MockClient.assert_called_with(
|
|
client_id="client-id",
|
|
client_secret="client-secret",
|
|
token_endpoint_auth_method="client_secret_post",
|
|
timeout=30.0,
|
|
)
|
|
|
|
|
|
class TestTokenHandlerErrorTransformation:
|
|
"""Tests for TokenHandler's OAuth 2.1 compliant error transformation."""
|
|
|
|
async def test_transforms_client_auth_failure_to_invalid_client_401(self):
|
|
"""Test that client authentication failures return invalid_client with 401."""
|
|
handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
|
|
|
|
# Create a mock 401 response like the SDK returns for auth failures
|
|
mock_response = Mock()
|
|
mock_response.status_code = 401
|
|
mock_response.body = (
|
|
b'{"error":"unauthorized_client","error_description":"Invalid client_id"}'
|
|
)
|
|
|
|
# Patch the parent class's handle() to return our mock response
|
|
with patch.object(
|
|
SDKTokenHandler,
|
|
"handle",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_response,
|
|
):
|
|
response = await handler.handle(Mock())
|
|
|
|
# Should transform to OAuth 2.1 compliant response
|
|
assert response.status_code == 401
|
|
assert b'"error":"invalid_client"' in response.body
|
|
assert b'"error_description":"Invalid client_id"' in response.body
|
|
|
|
def test_does_not_transform_grant_type_unauthorized_to_invalid_client(self):
|
|
"""Test that grant type authorization errors stay as unauthorized_client with 400."""
|
|
handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
|
|
|
|
# Simulate error from grant_type not in client_info.grant_types
|
|
error_response = TokenErrorResponse(
|
|
error="unauthorized_client",
|
|
error_description="Client not authorized for this grant type",
|
|
)
|
|
|
|
response = handler.response(error_response)
|
|
|
|
# Should NOT transform - keep as 400 unauthorized_client
|
|
assert response.status_code == 400
|
|
assert b'"error":"unauthorized_client"' in response.body
|
|
|
|
async def test_transforms_invalid_grant_to_401(self):
|
|
"""Test that invalid_grant errors return 401 per MCP spec.
|
|
|
|
Per MCP spec: "Invalid or expired tokens MUST receive a HTTP 401 response."
|
|
The SDK incorrectly returns 400 for all TokenErrorResponse including invalid_grant.
|
|
"""
|
|
handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
|
|
|
|
# Create a mock 400 response like the SDK returns for invalid_grant
|
|
mock_response = Mock()
|
|
mock_response.status_code = 400
|
|
mock_response.body = (
|
|
b'{"error":"invalid_grant","error_description":"refresh token has expired"}'
|
|
)
|
|
|
|
# Patch the parent class's handle() to return our mock response
|
|
with patch.object(
|
|
SDKTokenHandler,
|
|
"handle",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_response,
|
|
):
|
|
response = await handler.handle(Mock())
|
|
|
|
# Should transform to MCP-compliant 401 response
|
|
assert response.status_code == 401
|
|
assert b'"error":"invalid_grant"' in response.body
|
|
assert b'"error_description":"refresh token has expired"' in response.body
|
|
|
|
def test_does_not_transform_other_400_errors(self):
|
|
"""Test that non-invalid_grant 400 errors pass through unchanged."""
|
|
handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
|
|
|
|
# Test with invalid_request error (should stay 400)
|
|
error_response = TokenErrorResponse(
|
|
error="invalid_request",
|
|
error_description="Missing required parameter",
|
|
)
|
|
|
|
response = handler.response(error_response)
|
|
|
|
# Should pass through unchanged as 400
|
|
assert response.status_code == 400
|
|
assert b'"error":"invalid_request"' in response.body
|
|
|
|
|
|
class TestFallbackAccessTokenExpiry:
|
|
"""Test fallback access token expiry constants and configuration."""
|
|
|
|
def test_default_constants(self):
|
|
"""Verify the default expiry constants are set correctly."""
|
|
assert DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS == 60 * 60 # 1 hour
|
|
assert (
|
|
DEFAULT_ACCESS_TOKEN_EXPIRY_NO_REFRESH_SECONDS == 60 * 60 * 24 * 365
|
|
) # 1 year
|
|
|
|
def test_fallback_parameter_stored(self):
|
|
"""Verify fallback_access_token_expiry_seconds is stored on provider."""
|
|
provider = OAuthProxy(
|
|
upstream_authorization_endpoint="https://idp.example.com/authorize",
|
|
upstream_token_endpoint="https://idp.example.com/token",
|
|
upstream_client_id="test-client",
|
|
upstream_client_secret="test-secret",
|
|
token_verifier=JWTVerifier(
|
|
jwks_uri="https://idp.example.com/.well-known/jwks.json",
|
|
issuer="https://idp.example.com",
|
|
),
|
|
base_url="http://localhost:8000",
|
|
jwt_signing_key="test-signing-key",
|
|
fallback_access_token_expiry_seconds=86400,
|
|
client_storage=MemoryStore(),
|
|
)
|
|
|
|
assert provider._fallback_access_token_expiry_seconds == 86400
|
|
|
|
def test_fallback_parameter_defaults_to_none(self):
|
|
"""Verify fallback defaults to None (enabling smart defaults)."""
|
|
provider = OAuthProxy(
|
|
upstream_authorization_endpoint="https://idp.example.com/authorize",
|
|
upstream_token_endpoint="https://idp.example.com/token",
|
|
upstream_client_id="test-client",
|
|
upstream_client_secret="test-secret",
|
|
token_verifier=JWTVerifier(
|
|
jwks_uri="https://idp.example.com/.well-known/jwks.json",
|
|
issuer="https://idp.example.com",
|
|
),
|
|
base_url="http://localhost:8000",
|
|
jwt_signing_key="test-signing-key",
|
|
client_storage=MemoryStore(),
|
|
)
|
|
|
|
assert provider._fallback_access_token_expiry_seconds is None
|
|
|
|
|
|
class TestUpstreamTokenStorageTTL:
|
|
"""Tests for upstream token storage TTL calculation (issue #2670).
|
|
|
|
The TTL should use max(refresh_expires_in, expires_in) to handle cases where
|
|
the refresh token has a shorter lifetime than the access token (e.g., Keycloak
|
|
with sliding session windows).
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def jwt_verifier(self):
|
|
"""Create a mock JWT verifier."""
|
|
verifier = Mock(spec=TokenVerifier)
|
|
verifier.required_scopes = ["read", "write"]
|
|
verifier.verify_token = AsyncMock(return_value=None)
|
|
return verifier
|
|
|
|
@pytest.fixture
|
|
def proxy(self, jwt_verifier):
|
|
"""Create an OAuth proxy for testing."""
|
|
proxy = OAuthProxy(
|
|
upstream_authorization_endpoint="https://idp.example.com/authorize",
|
|
upstream_token_endpoint="https://idp.example.com/token",
|
|
upstream_client_id="test-client",
|
|
upstream_client_secret="test-secret",
|
|
token_verifier=jwt_verifier,
|
|
base_url="https://proxy.example.com",
|
|
jwt_signing_key="test-secret-key",
|
|
client_storage=MemoryStore(),
|
|
)
|
|
proxy.set_mcp_path("/mcp")
|
|
return proxy
|
|
|
|
async def test_ttl_uses_max_when_refresh_shorter_than_access(self, proxy):
|
|
"""TTL should use access token expiry when refresh is shorter.
|
|
|
|
This is the xsreality case: Keycloak returns refresh_expires_in=120 (2 min)
|
|
but expires_in=28800 (8 hours). The upstream tokens should persist for
|
|
8 hours (the access token lifetime), not 2 minutes.
|
|
"""
|
|
# Register client
|
|
client = OAuthClientInformationFull(
|
|
client_id="test-client",
|
|
client_secret="test-secret",
|
|
redirect_uris=[AnyUrl("http://localhost:12345/callback")],
|
|
)
|
|
await proxy.register_client(client)
|
|
|
|
# Simulate xsreality's Keycloak setup: short refresh, long access
|
|
client_code = ClientCode(
|
|
code="test-auth-code",
|
|
client_id="test-client",
|
|
redirect_uri="http://localhost:12345/callback",
|
|
code_challenge="test-challenge",
|
|
code_challenge_method="S256",
|
|
scopes=["read", "write"],
|
|
idp_tokens={
|
|
"access_token": "upstream-access-token",
|
|
"refresh_token": "upstream-refresh-token",
|
|
"expires_in": 28800, # 8 hours (access token)
|
|
"refresh_expires_in": 120, # 2 minutes (refresh token) - SHORTER!
|
|
"token_type": "Bearer",
|
|
},
|
|
expires_at=time.time() + 300,
|
|
created_at=time.time(),
|
|
)
|
|
await proxy._code_store.put(key=client_code.code, value=client_code)
|
|
|
|
# Exchange the code
|
|
auth_code = AuthorizationCode(
|
|
code="test-auth-code",
|
|
scopes=["read", "write"],
|
|
expires_at=time.time() + 300,
|
|
client_id="test-client",
|
|
code_challenge="test-challenge",
|
|
redirect_uri=AnyUrl("http://localhost:12345/callback"),
|
|
redirect_uri_provided_explicitly=True,
|
|
)
|
|
|
|
result = await proxy.exchange_authorization_code(
|
|
client=client,
|
|
authorization_code=auth_code,
|
|
)
|
|
|
|
# Verify tokens were issued
|
|
assert result.access_token is not None
|
|
assert result.refresh_token is not None
|
|
|
|
# The key test: verify upstream tokens are stored with TTL=max(120, 28800)=28800
|
|
# We can verify this by checking the tokens are still accessible after 2 minutes
|
|
# would have passed (if TTL was incorrectly set to 120)
|
|
#
|
|
# Since we can't easily time-travel in tests, we verify the storage directly
|
|
# by checking that we can still look up the tokens for refresh purposes.
|
|
#
|
|
# Extract the JTI from the refresh token to look up the mapping
|
|
refresh_payload = proxy.jwt_issuer.verify_token(
|
|
result.refresh_token, expected_token_use="refresh"
|
|
)
|
|
refresh_jti = refresh_payload["jti"]
|
|
|
|
# The JTI mapping should exist
|
|
jti_mapping = await proxy._jti_mapping_store.get(key=refresh_jti)
|
|
assert jti_mapping is not None
|
|
|
|
# The upstream tokens should exist
|
|
upstream_tokens = await proxy._upstream_token_store.get(
|
|
key=jti_mapping.upstream_token_id
|
|
)
|
|
assert upstream_tokens is not None
|
|
assert upstream_tokens.access_token == "upstream-access-token"
|
|
assert upstream_tokens.refresh_token == "upstream-refresh-token"
|
|
|
|
async def test_ttl_uses_refresh_when_refresh_longer_than_access(self, proxy):
|
|
"""TTL should use refresh token expiry when refresh is longer.
|
|
|
|
This is the ianw case: IdP returns expires_in=300 (5 min) but
|
|
refresh_expires_in=32318 (9 hours). The upstream tokens should persist
|
|
for 9 hours (the refresh token lifetime).
|
|
"""
|
|
# Register client
|
|
client = OAuthClientInformationFull(
|
|
client_id="test-client",
|
|
client_secret="test-secret",
|
|
redirect_uris=[AnyUrl("http://localhost:12345/callback")],
|
|
)
|
|
await proxy.register_client(client)
|
|
|
|
# Simulate ianw's setup: short access, long refresh (typical)
|
|
client_code = ClientCode(
|
|
code="test-auth-code-2",
|
|
client_id="test-client",
|
|
redirect_uri="http://localhost:12345/callback",
|
|
code_challenge="test-challenge",
|
|
code_challenge_method="S256",
|
|
scopes=["read", "write"],
|
|
idp_tokens={
|
|
"access_token": "upstream-access-token-2",
|
|
"refresh_token": "upstream-refresh-token-2",
|
|
"expires_in": 300, # 5 minutes (access token)
|
|
"refresh_expires_in": 32318, # 9 hours (refresh token) - LONGER
|
|
"token_type": "Bearer",
|
|
},
|
|
expires_at=time.time() + 300,
|
|
created_at=time.time(),
|
|
)
|
|
await proxy._code_store.put(key=client_code.code, value=client_code)
|
|
|
|
# Exchange the code
|
|
auth_code = AuthorizationCode(
|
|
code="test-auth-code-2",
|
|
scopes=["read", "write"],
|
|
expires_at=time.time() + 300,
|
|
client_id="test-client",
|
|
code_challenge="test-challenge",
|
|
redirect_uri=AnyUrl("http://localhost:12345/callback"),
|
|
redirect_uri_provided_explicitly=True,
|
|
)
|
|
|
|
result = await proxy.exchange_authorization_code(
|
|
client=client,
|
|
authorization_code=auth_code,
|
|
)
|
|
|
|
# Verify tokens were issued
|
|
assert result.access_token is not None
|
|
assert result.refresh_token is not None
|
|
|
|
# Verify upstream tokens are accessible
|
|
refresh_payload = proxy.jwt_issuer.verify_token(
|
|
result.refresh_token, expected_token_use="refresh"
|
|
)
|
|
refresh_jti = refresh_payload["jti"]
|
|
|
|
jti_mapping = await proxy._jti_mapping_store.get(key=refresh_jti)
|
|
assert jti_mapping is not None
|
|
|
|
upstream_tokens = await proxy._upstream_token_store.get(
|
|
key=jti_mapping.upstream_token_id
|
|
)
|
|
assert upstream_tokens is not None
|
|
|
|
async def test_refresh_expires_in_zero_issues_refresh_token(self, proxy):
|
|
"""refresh_expires_in=0 should fall back to 30-day default.
|
|
|
|
Keycloak returns refresh_expires_in=0 for offline tokens (offline_access scope),
|
|
meaning "no fixed time-based expiry". The proxy should still issue a PROXY_RT.
|
|
"""
|
|
client = OAuthClientInformationFull(
|
|
client_id="test-client",
|
|
client_secret="test-secret",
|
|
redirect_uris=[AnyUrl("http://localhost:12345/callback")],
|
|
)
|
|
await proxy.register_client(client)
|
|
|
|
client_code = ClientCode(
|
|
code="test-auth-code-keycloak-offline",
|
|
client_id="test-client",
|
|
redirect_uri="http://localhost:12345/callback",
|
|
code_challenge="test-challenge",
|
|
code_challenge_method="S256",
|
|
scopes=["read", "write"],
|
|
idp_tokens={
|
|
"access_token": "upstream-access-token-kc",
|
|
"refresh_token": "upstream-refresh-token-kc",
|
|
"expires_in": 3600,
|
|
"refresh_expires_in": 0, # Keycloak offline token convention
|
|
"token_type": "Bearer",
|
|
},
|
|
expires_at=time.time() + 300,
|
|
created_at=time.time(),
|
|
)
|
|
await proxy._code_store.put(key=client_code.code, value=client_code)
|
|
|
|
auth_code = AuthorizationCode(
|
|
code="test-auth-code-keycloak-offline",
|
|
scopes=["read", "write"],
|
|
expires_at=time.time() + 300,
|
|
client_id="test-client",
|
|
code_challenge="test-challenge",
|
|
redirect_uri=AnyUrl("http://localhost:12345/callback"),
|
|
redirect_uri_provided_explicitly=True,
|
|
)
|
|
|
|
result = await proxy.exchange_authorization_code(
|
|
client=client,
|
|
authorization_code=auth_code,
|
|
)
|
|
|
|
# refresh_expires_in=0 must NOT prevent refresh token issuance
|
|
assert result.access_token is not None
|
|
assert result.refresh_token is not None
|
|
|
|
# Verify refresh token metadata was stored
|
|
refresh_meta = await proxy._refresh_token_store.get(
|
|
key=_hash_token(result.refresh_token)
|
|
)
|
|
assert refresh_meta is not None
|
|
|
|
|
|
class TestTransparentUpstreamRefresh:
|
|
"""Tests for transparent upstream token refresh in load_access_token.
|
|
|
|
When the upstream token expires but a refresh token is available, the proxy
|
|
should transparently refresh the upstream token rather than returning a 401
|
|
that forces the client into a full re-authentication flow.
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def mock_verifier(self):
|
|
"""Token verifier that rejects expired tokens, accepts refreshed ones."""
|
|
verifier = Mock(spec=TokenVerifier)
|
|
verifier.required_scopes = ["read"]
|
|
|
|
async def verify(token: str) -> AccessToken | None:
|
|
if token.startswith("refreshed-"):
|
|
return AccessToken(
|
|
token=token,
|
|
client_id="test-client",
|
|
scopes=["read"],
|
|
expires_at=int(time.time() + 3600),
|
|
)
|
|
return None
|
|
|
|
verifier.verify_token = AsyncMock(side_effect=verify)
|
|
return verifier
|
|
|
|
@pytest.fixture
|
|
def proxy(self, mock_verifier):
|
|
proxy = OAuthProxy(
|
|
upstream_authorization_endpoint="https://idp.example.com/authorize",
|
|
upstream_token_endpoint="https://idp.example.com/token",
|
|
upstream_client_id="test-client",
|
|
upstream_client_secret="test-secret",
|
|
token_verifier=mock_verifier,
|
|
base_url="https://proxy.example.com",
|
|
jwt_signing_key="test-secret-key",
|
|
client_storage=MemoryStore(),
|
|
)
|
|
proxy.set_mcp_path("/mcp")
|
|
return proxy
|
|
|
|
async def _setup_expired_session(
|
|
self,
|
|
proxy: OAuthProxy,
|
|
*,
|
|
upstream_refresh_token: str | None = "upstream-refresh-tok",
|
|
) -> str:
|
|
"""Set up a proxy JWT pointing at an expired upstream token.
|
|
|
|
Returns the proxy JWT (access token) that can be passed to
|
|
load_access_token.
|
|
"""
|
|
upstream_token_id = "upstream-tok-id"
|
|
access_jti = "test-access-jti"
|
|
|
|
upstream_token_set = UpstreamTokenSet(
|
|
upstream_token_id=upstream_token_id,
|
|
access_token="expired-upstream-access",
|
|
refresh_token=upstream_refresh_token,
|
|
refresh_token_expires_at=time.time() + 86400
|
|
if upstream_refresh_token
|
|
else None,
|
|
expires_at=time.time() - 60, # expired 1 minute ago
|
|
token_type="Bearer",
|
|
scope="read",
|
|
client_id="test-client",
|
|
created_at=time.time() - 3600,
|
|
)
|
|
await proxy._upstream_token_store.put(
|
|
key=upstream_token_id,
|
|
value=upstream_token_set,
|
|
ttl=86400,
|
|
)
|
|
|
|
await proxy._jti_mapping_store.put(
|
|
key=access_jti,
|
|
value=JTIMapping(
|
|
jti=access_jti,
|
|
upstream_token_id=upstream_token_id,
|
|
created_at=time.time(),
|
|
),
|
|
ttl=3600,
|
|
)
|
|
|
|
fastmcp_jwt = proxy.jwt_issuer.issue_access_token(
|
|
client_id="test-client",
|
|
scopes=["read"],
|
|
jti=access_jti,
|
|
expires_in=3600,
|
|
)
|
|
return fastmcp_jwt
|
|
|
|
async def test_transparent_refresh_on_expired_upstream(self, proxy):
|
|
"""load_access_token refreshes upstream token when validation fails."""
|
|
fastmcp_jwt = await self._setup_expired_session(proxy)
|
|
|
|
mock_oauth_client = AsyncMock()
|
|
mock_oauth_client.refresh_token = AsyncMock(
|
|
return_value={
|
|
"access_token": "refreshed-upstream-access",
|
|
"token_type": "Bearer",
|
|
"expires_in": 3600,
|
|
"refresh_token": "upstream-refresh-tok",
|
|
"scope": "read",
|
|
}
|
|
)
|
|
|
|
with patch.object(
|
|
proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
|
|
):
|
|
result = await proxy.load_access_token(fastmcp_jwt)
|
|
|
|
assert result is not None
|
|
assert result.token == "refreshed-upstream-access"
|
|
mock_oauth_client.refresh_token.assert_called_once()
|
|
|
|
async def test_transparent_refresh_updates_stored_token(self, proxy):
|
|
"""After transparent refresh, the stored upstream token is updated."""
|
|
fastmcp_jwt = await self._setup_expired_session(proxy)
|
|
|
|
mock_oauth_client = AsyncMock()
|
|
mock_oauth_client.refresh_token = AsyncMock(
|
|
return_value={
|
|
"access_token": "refreshed-upstream-access",
|
|
"token_type": "Bearer",
|
|
"expires_in": 7200,
|
|
"refresh_token": "upstream-refresh-tok",
|
|
"scope": "read",
|
|
}
|
|
)
|
|
|
|
with patch.object(
|
|
proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
|
|
):
|
|
await proxy.load_access_token(fastmcp_jwt)
|
|
|
|
stored = await proxy._upstream_token_store.get(key="upstream-tok-id")
|
|
assert stored is not None
|
|
assert stored.access_token == "refreshed-upstream-access"
|
|
assert stored.expires_at > time.time()
|
|
|
|
async def test_no_refresh_without_refresh_token(self, proxy):
|
|
"""Without a refresh token, load_access_token returns None immediately."""
|
|
fastmcp_jwt = await self._setup_expired_session(
|
|
proxy, upstream_refresh_token=None
|
|
)
|
|
|
|
result = await proxy.load_access_token(fastmcp_jwt)
|
|
|
|
assert result is None
|
|
|
|
async def test_returns_none_when_refresh_fails(self, proxy):
|
|
"""If the upstream refresh call fails, load_access_token returns None."""
|
|
fastmcp_jwt = await self._setup_expired_session(proxy)
|
|
|
|
mock_oauth_client = AsyncMock()
|
|
mock_oauth_client.refresh_token = AsyncMock(
|
|
side_effect=Exception("upstream refused refresh")
|
|
)
|
|
|
|
with patch.object(
|
|
proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
|
|
):
|
|
result = await proxy.load_access_token(fastmcp_jwt)
|
|
|
|
assert result is None
|
|
|
|
async def test_returns_none_when_refreshed_token_still_invalid(self, proxy):
|
|
"""If the refreshed token also fails validation, returns None."""
|
|
fastmcp_jwt = await self._setup_expired_session(proxy)
|
|
|
|
mock_oauth_client = AsyncMock()
|
|
mock_oauth_client.refresh_token = AsyncMock(
|
|
return_value={
|
|
"access_token": "still-bad-token",
|
|
"token_type": "Bearer",
|
|
"expires_in": 3600,
|
|
"scope": "read",
|
|
}
|
|
)
|
|
|
|
with patch.object(
|
|
proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
|
|
):
|
|
result = await proxy.load_access_token(fastmcp_jwt)
|
|
|
|
# "still-bad-token" doesn't start with "refreshed-" so verifier rejects it
|
|
assert result is None
|
|
|
|
async def test_no_refresh_when_token_not_expired(self, proxy):
|
|
"""Non-expiry validation failures (e.g. revocation) should not trigger refresh."""
|
|
upstream_token_id = "upstream-tok-id"
|
|
access_jti = "test-access-jti"
|
|
|
|
# Token is NOT expired — verification failure is for another reason
|
|
upstream_token_set = UpstreamTokenSet(
|
|
upstream_token_id=upstream_token_id,
|
|
access_token="revoked-upstream-access",
|
|
refresh_token="upstream-refresh-tok",
|
|
refresh_token_expires_at=time.time() + 86400,
|
|
expires_at=time.time() + 3600, # still valid for 1 hour
|
|
token_type="Bearer",
|
|
scope="read",
|
|
client_id="test-client",
|
|
created_at=time.time() - 60,
|
|
)
|
|
await proxy._upstream_token_store.put(
|
|
key=upstream_token_id,
|
|
value=upstream_token_set,
|
|
ttl=86400,
|
|
)
|
|
await proxy._jti_mapping_store.put(
|
|
key=access_jti,
|
|
value=JTIMapping(
|
|
jti=access_jti,
|
|
upstream_token_id=upstream_token_id,
|
|
created_at=time.time(),
|
|
),
|
|
ttl=3600,
|
|
)
|
|
fastmcp_jwt = proxy.jwt_issuer.issue_access_token(
|
|
client_id="test-client",
|
|
scopes=["read"],
|
|
jti=access_jti,
|
|
expires_in=3600,
|
|
)
|
|
|
|
mock_oauth_client = AsyncMock()
|
|
mock_oauth_client.refresh_token = AsyncMock()
|
|
|
|
with patch.object(
|
|
proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
|
|
):
|
|
result = await proxy.load_access_token(fastmcp_jwt)
|
|
|
|
assert result is None
|
|
# Refresh should NOT have been attempted
|
|
mock_oauth_client.refresh_token.assert_not_called()
|
|
|
|
async def test_reload_from_storage_after_refresh_failure(self, proxy):
|
|
"""If refresh fails, re-read from storage in case another worker refreshed."""
|
|
fastmcp_jwt = await self._setup_expired_session(proxy)
|
|
|
|
# Simulate: refresh fails (stale refresh token), but another worker
|
|
# already wrote a fresh upstream token to storage.
|
|
refreshed_token_set = UpstreamTokenSet(
|
|
upstream_token_id="upstream-tok-id",
|
|
access_token="refreshed-upstream-access",
|
|
refresh_token="new-refresh-tok",
|
|
refresh_token_expires_at=time.time() + 86400,
|
|
expires_at=time.time() + 3600,
|
|
token_type="Bearer",
|
|
scope="read",
|
|
client_id="test-client",
|
|
created_at=time.time(),
|
|
)
|
|
|
|
# The store is read three times:
|
|
# 1. Initial lookup in load_access_token
|
|
# 2. Re-read inside the advisory lock
|
|
# 3. Re-read after refresh failure (recovery path)
|
|
original_get = proxy._upstream_token_store.get
|
|
call_count = 0
|
|
|
|
async def mock_get(key: str) -> UpstreamTokenSet | None:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count <= 2:
|
|
return await original_get(key)
|
|
# After refresh failure, return the "other worker's" refreshed token
|
|
return refreshed_token_set
|
|
|
|
mock_oauth_client = AsyncMock()
|
|
mock_oauth_client.refresh_token = AsyncMock(
|
|
side_effect=Exception("refresh token rotated by another worker")
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
|
|
),
|
|
patch.object(proxy._upstream_token_store, "get", side_effect=mock_get),
|
|
):
|
|
result = await proxy.load_access_token(fastmcp_jwt)
|
|
|
|
assert result is not None
|
|
assert result.token == "refreshed-upstream-access"
|