"""Tests for OAuth proxy token endpoint and handling.""" import logging import time from unittest.mock import AsyncMock, Mock, patch from urllib.parse import parse_qs, urlparse 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 AuthorizationCode from mcp.shared.auth import OAuthClientInformationFull from pydantic import AnyUrl from fastmcp.server.auth.auth import ( AccessToken, 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, DEFAULT_REFRESH_TOKEN_EXPIRY_SECONDS, ClientCode, JTIMapping, OAuthTransaction, RefreshTokenMetadata, 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, ) mock_client.aclose.assert_awaited_once() async def test_callback_closes_upstream_oauth_client(self, jwt_verifier): 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", require_authorization_consent=False, jwt_signing_key="test-secret", client_storage=MemoryStore(), ) await proxy._transaction_store.put( key="txn-id", value=OAuthTransaction( txn_id="txn-id", client_id="test-client", client_redirect_uri="http://localhost:12345/callback", client_state="client-state", code_challenge="", code_challenge_method="S256", scopes=["read"], created_at=time.time(), ), ) mock_request = Mock() mock_request.query_params = {"code": "idp-code", "state": "txn-id"} mock_request.cookies = {} mock_client = AsyncMock() mock_client.fetch_token = AsyncMock( return_value={ "access_token": "upstream-access-token", "refresh_token": "upstream-refresh-token", "expires_in": 3600, "token_type": "Bearer", } ) with patch.object( proxy, "_create_upstream_oauth_client", return_value=mock_client ): response = await proxy._handle_idp_callback(mock_request) assert response.status_code == 302 mock_client.fetch_token.assert_awaited_once() mock_client.aclose.assert_awaited_once() async def test_callback_redirect_includes_proxy_issuer(self, jwt_verifier): 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", require_authorization_consent=False, jwt_signing_key="test-secret", client_storage=MemoryStore(), ) await proxy._transaction_store.put( key="txn-id", value=OAuthTransaction( txn_id="txn-id", client_id="test-client", client_redirect_uri="http://localhost:12345/callback", client_state="client-state", code_challenge="", code_challenge_method="S256", scopes=["read"], created_at=time.time(), ), ) mock_request = Mock() mock_request.query_params = {"code": "idp-code", "state": "txn-id"} mock_request.cookies = {} mock_client = AsyncMock() mock_client.fetch_token = AsyncMock( return_value={ "access_token": "upstream-access-token", "refresh_token": "upstream-refresh-token", "expires_in": 3600, "token_type": "Bearer", } ) with patch.object( proxy, "_create_upstream_oauth_client", return_value=mock_client ): response = await proxy._handle_idp_callback(mock_request) assert response.status_code == 302 location = response.headers["location"] query_params = parse_qs(urlparse(location).query) assert "code" in query_params assert query_params["state"] == ["client-state"] assert query_params["iss"] == ["https://proxy.example.com/"] mock_client.aclose.assert_awaited_once() async def test_callback_rejects_unsafe_transaction_redirect(self, jwt_verifier): 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", require_authorization_consent=False, jwt_signing_key="test-secret", client_storage=MemoryStore(), ) await proxy._transaction_store.put( key="txn-id", value=OAuthTransaction( txn_id="txn-id", client_id="test-client", client_redirect_uri="javascript:alert(document.cookie)//", client_state="client-state", code_challenge="", code_challenge_method="S256", scopes=["read"], created_at=time.time(), ), ) mock_request = Mock() mock_request.query_params = {"code": "idp-code", "state": "txn-id"} mock_request.cookies = {} mock_client = AsyncMock() mock_client.fetch_token = AsyncMock() with patch.object( proxy, "_create_upstream_oauth_client", return_value=mock_client ) as create_upstream_oauth_client: response = await proxy._handle_idp_callback(mock_request) assert response.status_code == 400 assert "location" not in response.headers assert "Invalid redirect URI" in bytes(response.body).decode() create_upstream_oauth_client.assert_not_called() mock_client.fetch_token.assert_not_called() 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 TestFallbackRefreshTokenExpiry: """Tests for fallback_refresh_token_expiry_seconds (issue #3987). When upstream omits `refresh_expires_in`, the FastMCP refresh JWT and the JTI mapping must use the configured fallback (default 1 year), not a hardcoded 30-day value. The FastMCP refresh JWT is a signed pointer; the real upstream refresh remains the source of truth. """ @pytest.fixture def jwt_verifier(self): verifier = Mock(spec=TokenVerifier) verifier.required_scopes = ["read", "write"] verifier.verify_token = AsyncMock(return_value=None) return verifier def test_default_constant(self): assert DEFAULT_REFRESH_TOKEN_EXPIRY_SECONDS == 60 * 60 * 24 * 365 # 1 year def test_fallback_defaults_to_one_year(self, jwt_verifier): 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=jwt_verifier, base_url="http://localhost:8000", jwt_signing_key="test-signing-key", client_storage=MemoryStore(), ) assert ( provider._fallback_refresh_token_expiry_seconds == DEFAULT_REFRESH_TOKEN_EXPIRY_SECONDS ) def test_fallback_parameter_stored(self, jwt_verifier): 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=jwt_verifier, base_url="http://localhost:8000", jwt_signing_key="test-signing-key", fallback_refresh_token_expiry_seconds=7 * 24 * 60 * 60, # 7 days client_storage=MemoryStore(), ) assert provider._fallback_refresh_token_expiry_seconds == 7 * 24 * 60 * 60 async def test_initial_exchange_jti_mapping_aligned_with_refresh_expiry( self, jwt_verifier ): """Bug 1: JTI mapping TTL must align with refresh_expires_in. Previously hardcoded to 30 days, which silently expired the mapping early when upstream returned a longer refresh lifetime. """ 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") client = OAuthClientInformationFull( client_id="test-client", client_secret="test-secret", redirect_uris=[AnyUrl("http://localhost:12345/callback")], ) await proxy.register_client(client) # Upstream returns 1-year refresh expiry — well past the old 30d default one_year = 60 * 60 * 24 * 365 client_code = ClientCode( code="long-refresh-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", "refresh_token": "upstream-refresh", "expires_in": 3600, "refresh_expires_in": one_year, "token_type": "Bearer", }, expires_at=time.time() + 300, created_at=time.time(), ) await proxy._code_store.put(key=client_code.code, value=client_code) result = await proxy.exchange_authorization_code( client=client, authorization_code=AuthorizationCode( code="long-refresh-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, ), ) assert result.refresh_token is not None # The refresh JWT exp claim should reflect the 1-year upstream lifetime refresh_payload = proxy.jwt_issuer.verify_token( result.refresh_token, expected_token_use="refresh" ) # exp - iat should be approximately one_year (allow small clock skew) assert refresh_payload["exp"] - refresh_payload["iat"] == pytest.approx( one_year, abs=5 ) async def test_initial_exchange_uses_fallback_when_upstream_silent( self, jwt_verifier ): """When upstream omits refresh_expires_in, fall back to configured value. Default is 1 year (was previously 30 days, ignoring user config). """ custom_fallback = 60 * 60 * 24 * 90 # 90 days 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", fallback_refresh_token_expiry_seconds=custom_fallback, client_storage=MemoryStore(), ) proxy.set_mcp_path("/mcp") 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="silent-refresh-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", "refresh_token": "upstream-refresh", "expires_in": 3600, # no refresh_expires_in — upstream is silent "token_type": "Bearer", }, expires_at=time.time() + 300, created_at=time.time(), ) await proxy._code_store.put(key=client_code.code, value=client_code) result = await proxy.exchange_authorization_code( client=client, authorization_code=AuthorizationCode( code="silent-refresh-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, ), ) assert result.refresh_token is not None refresh_payload = proxy.jwt_issuer.verify_token( result.refresh_token, expected_token_use="refresh" ) assert refresh_payload["exp"] - refresh_payload["iat"] == pytest.approx( custom_fallback, abs=5 ) async def test_refresh_jwt_exp_aligned_when_upstream_omits_refresh_token( self, jwt_verifier ): """Bug 2: JWT exp must align with storage TTL when upstream omits refresh_token. Cognito does NOT return a refresh_token on refresh. Previously the new FastMCP refresh JWT was issued with hardcoded 30-day exp, while the storage TTL correctly used the existing upstream refresh lifetime — forcing re-login at day 30 even when the upstream refresh was valid for much longer. """ 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") client = OAuthClientInformationFull( client_id="test-client", client_secret="test-secret", redirect_uris=[AnyUrl("http://localhost:12345/callback")], ) await proxy.register_client(client) # Seed an upstream session whose refresh token has 6 months left now = time.time() upstream_refresh_expires_at = now + (60 * 60 * 24 * 180) upstream_token_id = "upstream-id-cognito" upstream_token_set = UpstreamTokenSet( upstream_token_id=upstream_token_id, access_token="upstream-access-old", refresh_token="upstream-refresh-tok", refresh_token_expires_at=upstream_refresh_expires_at, expires_at=now + 60, # access nearly expired (irrelevant here) token_type="Bearer", scope="read write", client_id="test-client", created_at=now, raw_token_data={}, ) await proxy._upstream_token_store.put( key=upstream_token_id, value=upstream_token_set, ttl=int(upstream_refresh_expires_at - now), ) # Issue a FastMCP refresh JWT pointing at this session old_refresh_jti = "old-refresh-jti" old_refresh_jwt = proxy.jwt_issuer.issue_refresh_token( client_id="test-client", scopes=["read", "write"], jti=old_refresh_jti, expires_in=int(upstream_refresh_expires_at - now), ) await proxy._jti_mapping_store.put( key=old_refresh_jti, value=JTIMapping( jti=old_refresh_jti, upstream_token_id=upstream_token_id, created_at=now, ), ttl=int(upstream_refresh_expires_at - now), ) await proxy._refresh_token_store.put( key=_hash_token(old_refresh_jwt), value=RefreshTokenMetadata( client_id="test-client", scopes=["read", "write"], expires_at=int(upstream_refresh_expires_at), created_at=now, ), ttl=int(upstream_refresh_expires_at - now), ) # Mock the upstream refresh response — Cognito-style: NO refresh_token async def fake_refresh(url, refresh_token, scope=None, **kwargs): return { "access_token": "upstream-access-new", "expires_in": 3600, "token_type": "Bearer", # NO refresh_token, NO refresh_expires_in } oauth_client_mock = Mock() oauth_client_mock.refresh_token = AsyncMock(side_effect=fake_refresh) oauth_client_mock.aclose = AsyncMock() with patch.object( proxy, "_create_upstream_oauth_client", return_value=oauth_client_mock, ): new_token = await proxy.exchange_refresh_token( client=client, refresh_token=RefreshToken( token=old_refresh_jwt, client_id="test-client", scopes=["read", "write"], expires_at=int(upstream_refresh_expires_at), ), scopes=["read", "write"], ) assert new_token.refresh_token is not None # New refresh JWT exp must reflect remaining upstream refresh lifetime # (~6 months), not a hardcoded 30 days new_refresh_payload = proxy.jwt_issuer.verify_token( new_token.refresh_token, expected_token_use="refresh" ) ttl_remaining = new_refresh_payload["exp"] - new_refresh_payload["iat"] # Should be close to 180 days, well above the old 30-day cliff assert ttl_remaining > 60 * 60 * 24 * 90 # at least 90 days class TestFastMCPAccessTokenExpiry: """Tests for fastmcp_access_token_expiry_seconds (issue #4252). The FastMCP-issued access token is a reference into FastMCP storage; its lifetime can be decoupled from the upstream provider's short `expires_in` so MCP clients that don't refresh gracefully (e.g. mcp-remote) aren't forced through a full re-auth on every idle period. The upstream token's real expiry is preserved internally to drive transparent refresh. """ @pytest.fixture def jwt_verifier(self): verifier = Mock(spec=TokenVerifier) verifier.required_scopes = ["read", "write"] verifier.verify_token = AsyncMock(return_value=None) return verifier def _make_proxy(self, jwt_verifier, **kwargs): return 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(), **kwargs, ) def test_parameter_stored(self, jwt_verifier): proxy = self._make_proxy( jwt_verifier, fastmcp_access_token_expiry_seconds=86400 ) assert proxy._fastmcp_access_token_expiry_seconds == 86400 def test_parameter_defaults_to_none(self, jwt_verifier): proxy = self._make_proxy(jwt_verifier) assert proxy._fastmcp_access_token_expiry_seconds is None async def _exchange(self, proxy, code="test-code", **idp_token_overrides): proxy.set_mcp_path("/mcp") client = OAuthClientInformationFull( client_id="test-client", client_secret="test-secret", redirect_uris=[AnyUrl("http://localhost:12345/callback")], ) await proxy.register_client(client) idp_tokens = { "access_token": "upstream-access", "refresh_token": "upstream-refresh", "expires_in": 3600, "token_type": "Bearer", **idp_token_overrides, } # Allow callers to drop a default key by overriding it with None # (e.g. refresh_token=None to simulate a provider that issues none). idp_tokens = {k: v for k, v in idp_tokens.items() if v is not None} client_code = ClientCode( code=code, client_id="test-client", redirect_uri="http://localhost:12345/callback", code_challenge="test-challenge", code_challenge_method="S256", scopes=["read", "write"], idp_tokens=idp_tokens, expires_at=time.time() + 300, created_at=time.time(), ) await proxy._code_store.put(key=client_code.code, value=client_code) return client, await proxy.exchange_authorization_code( client=client, authorization_code=AuthorizationCode( code=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, ), ) async def test_initial_exchange_decouples_access_token_from_upstream( self, jwt_verifier ): """A long FastMCP TTL applies even though upstream returns expires_in=3600.""" one_day = 60 * 60 * 24 proxy = self._make_proxy( jwt_verifier, fastmcp_access_token_expiry_seconds=one_day ) _, result = await self._exchange(proxy) # Response and JWT exp reflect the configured FastMCP lifetime, not 3600 assert result.expires_in == one_day access_payload = proxy.jwt_issuer.verify_token(result.access_token) assert access_payload["exp"] - access_payload["iat"] == pytest.approx( one_day, abs=5 ) async def test_upstream_token_expiry_preserved_for_transparent_refresh( self, jwt_verifier ): """Decoupling must not corrupt the upstream token's real expiry. The stored upstream token still expires at ~upstream expires_in so that transparent refresh fires; only the FastMCP-issued token lives longer. """ one_day = 60 * 60 * 24 proxy = self._make_proxy( jwt_verifier, fastmcp_access_token_expiry_seconds=one_day ) _, result = await self._exchange(proxy) access_jti = proxy.jwt_issuer.verify_token(result.access_token)["jti"] jti_mapping = await proxy._jti_mapping_store.get(key=access_jti) assert jti_mapping is not None stored = await proxy._upstream_token_store.get( key=jti_mapping.upstream_token_id ) assert stored is not None # Upstream access token expiry tracks the upstream lifetime (~3600s), # NOT the 1-day FastMCP token lifetime. assert stored.expires_at - time.time() == pytest.approx(3600, abs=30) async def test_access_jti_mapping_ttl_matches_configured_lifetime( self, jwt_verifier ): """The access JTI mapping must outlive the upstream access token. The JWT exp and the JTI mapping TTL are set from the same value, so a drift between them would let the JWT verify while its storage lookup has already expired — silently breaking long-idle sessions. Guard the TTL passed to storage directly, since wall-clock expiry can't be exercised in a fast unit test. """ one_week = 60 * 60 * 24 * 7 proxy = self._make_proxy( jwt_verifier, fastmcp_access_token_expiry_seconds=one_week ) original_put = proxy._jti_mapping_store.put calls: list[dict] = [] async def spy(**kwargs): calls.append(kwargs) return await original_put(**kwargs) with patch.object(proxy._jti_mapping_store, "put", side_effect=spy): _, result = await self._exchange(proxy) access_jti = proxy.jwt_issuer.verify_token(result.access_token)["jti"] access_call = next(c for c in calls if c["value"].jti == access_jti) assert access_call["ttl"] == one_week @pytest.mark.parametrize( "configured, expected", [ (60 * 60 * 24 * 7, 3600), # configured > upstream -> capped at upstream (600, 600), # configured < upstream -> honored (still <= upstream) ], ) async def test_no_refresh_token_does_not_extend_past_upstream( self, jwt_verifier, configured, expected ): """Without an upstream refresh token, the FastMCP token can't be renewed. Issuing a token that claims to outlive the upstream access token would be a lie — there's no way to transparently refresh it — so the lifetime is capped at the upstream `expires_in` when no refresh token is present. """ proxy = self._make_proxy( jwt_verifier, fastmcp_access_token_expiry_seconds=configured ) _, result = await self._exchange(proxy, refresh_token=None) assert result.expires_in == expected access_payload = proxy.jwt_issuer.verify_token(result.access_token) assert access_payload["exp"] - access_payload["iat"] == pytest.approx( expected, abs=5 ) async def test_extended_token_survives_upstream_expiry_via_refresh(self): """End-to-end: a long-lived FastMCP token keeps working after the upstream access token expires, by transparently refreshing underneath. Proves the pieces integrate: the long-exp JWT still verifies, its JTI mapping still resolves, and an expired upstream token triggers transparent refresh rather than a 401. """ one_week = 60 * 60 * 24 * 7 verifier = Mock(spec=TokenVerifier) verifier.required_scopes = ["read", "write"] async def verify(token: str) -> AccessToken | None: if token.startswith("refreshed-"): return AccessToken( token=token, client_id="test-client", scopes=["read", "write"], expires_at=int(time.time() + 3600), ) return None # original upstream token is treated as invalid/expired verifier.verify_token = AsyncMock(side_effect=verify) proxy = self._make_proxy(verifier, fastmcp_access_token_expiry_seconds=one_week) _, result = await self._exchange(proxy) # Force the stored upstream access token to be expired. access_jti = proxy.jwt_issuer.verify_token(result.access_token)["jti"] jti_mapping = await proxy._jti_mapping_store.get(key=access_jti) assert jti_mapping is not None stored = await proxy._upstream_token_store.get( key=jti_mapping.upstream_token_id ) assert stored is not None stored.expires_at = time.time() - 60 await proxy._upstream_token_store.put( key=jti_mapping.upstream_token_id, value=stored, ttl=one_week ) 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", "scope": "read write", } ) with patch.object( proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client ): loaded = await proxy.load_access_token(result.access_token) assert loaded is not None assert loaded.token == "refreshed-upstream-access" mock_oauth_client.refresh_token.assert_called_once() async def test_initial_exchange_default_mirrors_upstream(self, jwt_verifier): """With the param unset, the FastMCP access token mirrors upstream.""" proxy = self._make_proxy(jwt_verifier) _, result = await self._exchange(proxy) assert result.expires_in == 3600 access_payload = proxy.jwt_issuer.verify_token(result.access_token) assert access_payload["exp"] - access_payload["iat"] == pytest.approx( 3600, abs=5 ) async def test_refresh_exchange_decouples_access_token(self, jwt_verifier): """Re-issued access tokens on refresh also honor the configured lifetime.""" one_day = 60 * 60 * 24 proxy = self._make_proxy( jwt_verifier, fastmcp_access_token_expiry_seconds=one_day ) client, result = await self._exchange(proxy) assert result.refresh_token is not None 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", "scope": "read write", } ) with patch.object( proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client ): refreshed = await proxy.exchange_refresh_token( client=client, refresh_token=RefreshToken( token=result.refresh_token, client_id="test-client", scopes=["read", "write"], ), scopes=["read", "write"], ) assert refreshed.expires_in == one_day access_payload = proxy.jwt_issuer.verify_token(refreshed.access_token) assert access_payload["exp"] - access_payload["iat"] == pytest.approx( one_day, abs=5 ) 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 the configured 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"] # Cache tokens to test mutation of returned objects cache: dict[str, AccessToken] = {} async def verify(token: str) -> AccessToken | None: if token in cache: return cache[token] if token.startswith("refreshed-") or token.startswith("valid-"): t = AccessToken( token=token, client_id="test-client", scopes=["read"], expires_at=int(time.time() + 3600), ) cache[token] = t return t 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 _setup_session_with_claims(self, proxy, *, upstream_claims=None): """Set up a proxy JWT pointing at a valid upstream token, with optional upstream_claims.""" upstream_token_id = "upstream-tok-id" access_jti = "test-claims-jti" upstream_token_set = UpstreamTokenSet( upstream_token_id=upstream_token_id, access_token="valid-upstream-access", refresh_token=None, refresh_token_expires_at=None, expires_at=time.time() + 3600, token_type="Bearer", scope="read", client_id="test-client", created_at=time.time(), ) await proxy._upstream_token_store.put( key=upstream_token_id, value=upstream_token_set, ttl=3600, ) 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, ) return proxy.jwt_issuer.issue_access_token( client_id="test-client", scopes=["read"], jti=access_jti, expires_in=3600, upstream_claims=upstream_claims, ) 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() mock_oauth_client.aclose.assert_awaited_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" def test_refresh_lock_cache_bounded(self, proxy, monkeypatch): """_get_refresh_lock never grows beyond _REFRESH_LOCK_CACHE_SIZE.""" monkeypatch.setattr( "fastmcp.server.auth.oauth_proxy.proxy._REFRESH_LOCK_CACHE_SIZE", 3 ) for i in range(10): proxy._get_refresh_lock(f"token-{i}") assert len(proxy._refresh_locks) == 3 def test_refresh_lock_lru_evicts_least_recently_used(self, proxy, monkeypatch): """Touching an entry promotes it; eviction removes the oldest untouched entry.""" monkeypatch.setattr( "fastmcp.server.auth.oauth_proxy.proxy._REFRESH_LOCK_CACHE_SIZE", 3 ) proxy._get_refresh_lock("a") proxy._get_refresh_lock("b") proxy._get_refresh_lock("c") # Touch "a" to move it to MRU position proxy._get_refresh_lock("a") # Adding "d" should evict "b" (oldest untouched) proxy._get_refresh_lock("d") assert "b" not in proxy._refresh_locks assert "a" in proxy._refresh_locks assert "c" in proxy._refresh_locks assert "d" in proxy._refresh_locks def test_refresh_lock_same_token_returns_same_lock(self, proxy): """Requesting the same token ID twice returns the same lock object.""" lock1 = proxy._get_refresh_lock("tok") lock2 = proxy._get_refresh_lock("tok") assert lock1 is lock2 async def test_upstream_claims_propagated(self, proxy): jwt = await self._setup_session_with_claims( proxy, upstream_claims={"sub": "user-123"} ) result = await proxy.load_access_token(jwt) assert result is not None assert result.claims["upstream_claims"] == {"sub": "user-123"} async def test_upstream_claims_not_mutated_on_cached_token( self, proxy, mock_verifier ): jwt = await self._setup_session_with_claims( proxy, upstream_claims={"sub": "user-123"} ) result = await proxy.load_access_token(jwt) assert result is not None assert result.claims["upstream_claims"] == {"sub": "user-123"} # Original verifier result must not be mutated for call in list(mock_verifier.verify_token.call_args_list): returned = await mock_verifier.verify_token(call.args[0]) if returned: assert "upstream_claims" not in returned.claims async def test_transparent_refresh_triggered_by_threshold(self, mock_verifier): """Token within expiry threshold is treated as expired and refreshed.""" 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(), token_expiry_threshold_seconds=60, ) proxy.set_mcp_path("/mcp") upstream_token_id = "upstream-tok-threshold" access_jti = "test-threshold-jti" upstream_token_set = UpstreamTokenSet( upstream_token_id=upstream_token_id, access_token="almost-expired-upstream-access", refresh_token="upstream-refresh-tok", refresh_token_expires_at=time.time() + 86400, expires_at=time.time() + 30, # expires in 30s, within 60s threshold 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, ) 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_no_refresh_when_within_threshold_but_no_refresh_token( self, mock_verifier ): """Token within threshold but without refresh token is not refreshed.""" 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(), token_expiry_threshold_seconds=60, ) proxy.set_mcp_path("/mcp") upstream_token_id = "upstream-tok-no-refresh" access_jti = "test-no-refresh-jti" upstream_token_set = UpstreamTokenSet( upstream_token_id=upstream_token_id, access_token="valid-upstream-access", refresh_token=None, refresh_token_expires_at=None, expires_at=time.time() + 30, # expires in 30s, within 60s threshold 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, ) result = await proxy.load_access_token(fastmcp_jwt) # Token validates via the verifier (starts with "valid-"), so even though # it's within threshold, no refresh is attempted without a refresh token. assert result is not None assert result.token == "valid-upstream-access" async def test_zero_threshold_only_refreshes_on_actual_expiry(self, mock_verifier): """With threshold=0, token that's not yet expired is not refreshed.""" 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(), token_expiry_threshold_seconds=0, ) proxy.set_mcp_path("/mcp") upstream_token_id = "upstream-tok-zero-threshold" access_jti = "test-zero-threshold-jti" upstream_token_set = UpstreamTokenSet( upstream_token_id=upstream_token_id, access_token="valid-upstream-access", refresh_token="upstream-refresh-tok", refresh_token_expires_at=time.time() + 86400, expires_at=time.time() + 30, # expires in 30s, but threshold is 0 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, ) result = await proxy.load_access_token(fastmcp_jwt) # With threshold=0, token with 30s remaining is still valid assert result is not None assert result.token == "valid-upstream-access" async def test_proactive_refresh_when_validated_but_within_threshold( self, mock_verifier ): """Token that passes validation but is within threshold gets proactively refreshed.""" 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(), token_expiry_threshold_seconds=60, ) proxy.set_mcp_path("/mcp") upstream_token_id = "upstream-tok-proactive" access_jti = "test-proactive-jti" # Token starts with "valid-" so the mock verifier accepts it, # but it expires in 30s which is within the 60s threshold. upstream_token_set = UpstreamTokenSet( upstream_token_id=upstream_token_id, access_token="valid-but-near-expiry", refresh_token="upstream-refresh-tok", refresh_token_expires_at=time.time() + 86400, expires_at=time.time() + 30, 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, ) 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() class TestRefreshTokenMissLogging: """A refresh-token miss forces a user-visible reconnect, so it must not be silent.""" async def test_unknown_refresh_token_logs_warning(self, oauth_proxy, caplog): client = OAuthClientInformationFull( client_id="test-client", client_secret="test-secret", redirect_uris=[AnyUrl("http://localhost:12345/callback")], ) proxy_logger = logging.getLogger("fastmcp.server.auth.oauth_proxy.proxy") caplog.set_level(logging.WARNING) proxy_logger.addHandler(caplog.handler) try: result = await oauth_proxy.load_refresh_token(client, "nonexistent-token") finally: proxy_logger.removeHandler(caplog.handler) assert result is None assert any( record.levelno == logging.WARNING and "Refresh token not found" in record.getMessage() and "test-client" in record.getMessage() for record in caplog.records )