d365fo-client 0.2.4__py3-none-any.whl → 0.3.1__py3-none-any.whl

d365fo_client/mcp/auth_server/auth/oauth_proxy.py

@@ -0,0 +1,989 @@
+"""OAuth Proxy Provider for FastMCP.
+
+This provider acts as a transparent proxy to an upstream OAuth Authorization Server,
+handling Dynamic Client Registration locally while forwarding all other OAuth flows.
+This enables authentication with upstream providers that don't support DCR or have
+restricted client registration policies.
+
+Key features:
+- Proxies authorization and token endpoints to upstream server
+- Implements local Dynamic Client Registration with fixed upstream credentials
+- Validates tokens using upstream JWKS
+- Maintains minimal local state for bookkeeping
+- Enhanced logging with request correlation
+
+This implementation is based on the OAuth 2.1 specification and is designed for
+production use with enterprise identity providers.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import secrets
+import time
+from base64 import urlsafe_b64encode
+from typing import TYPE_CHECKING, Any, Final
+from urllib.parse import urlencode
+
+import httpx
+from authlib.common.security import generate_token
+from authlib.integrations.httpx_client import AsyncOAuth2Client
+from mcp.server.auth.provider import (
+    AccessToken,
+    AuthorizationCode,
+    AuthorizationParams,
+    RefreshToken,
+    TokenError,
+)
+from mcp.server.auth.settings import (
+    ClientRegistrationOptions,
+    RevocationOptions,
+)
+from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
+from pydantic import AnyHttpUrl, AnyUrl, SecretStr
+from starlette.requests import Request
+from starlette.responses import RedirectResponse
+from starlette.routing import Route
+
+from .auth import OAuthProvider
+from .auth import TokenVerifier
+from .redirect_validation import validate_redirect_uri
+from d365fo_client.mcp.utilities.logging import get_logger
+
+if TYPE_CHECKING:
+    pass
+
+logger = get_logger(__name__)
+
+
+class ProxyDCRClient(OAuthClientInformationFull):
+    """Client for DCR proxy with configurable redirect URI validation.
+
+    This special client class is critical for the OAuth proxy to work correctly
+    with Dynamic Client Registration (DCR). Here's why it exists:
+
+    Problem:
+    --------
+    When MCP clients use OAuth, they dynamically register with random localhost
+    ports (e.g., http://localhost:55454/callback). The OAuth proxy needs to:
+    1. Accept these dynamic redirect URIs from clients based on configured patterns
+    2. Use its own fixed redirect URI with the upstream provider (Google, GitHub, etc.)
+    3. Forward the authorization code back to the client's dynamic URI
+
+    Solution:
+    ---------
+    This class validates redirect URIs against configurable patterns,
+    while the proxy internally uses its own fixed redirect URI with the upstream
+    provider. This allows the flow to work even when clients reconnect with
+    different ports or when tokens are cached.
+
+    Without proper validation, clients could get "Redirect URI not registered" errors
+    when trying to authenticate with cached tokens, or security vulnerabilities could
+    arise from accepting arbitrary redirect URIs.
+    """
+
+    def __init__(
+        self, *args, allowed_redirect_uri_patterns: list[str] | None = None, **kwargs
+    ):
+        """Initialize with allowed redirect URI patterns.
+
+        Args:
+            allowed_redirect_uri_patterns: List of allowed redirect URI patterns with wildcard support.
+                                          If None, defaults to localhost-only patterns.
+                                          If empty list, allows all redirect URIs.
+        """
+        super().__init__(*args, **kwargs)
+        self._allowed_redirect_uri_patterns = allowed_redirect_uri_patterns
+
+    def validate_redirect_uri(self, redirect_uri: AnyUrl | None) -> AnyUrl:
+        """Validate redirect URI against allowed patterns.
+
+        Since we're acting as a proxy and clients register dynamically,
+        we validate their redirect URIs against configurable patterns.
+        This is essential for cached token scenarios where the client may
+        reconnect with a different port.
+        """
+        if redirect_uri is not None:
+            # Validate against allowed patterns
+            if validate_redirect_uri(redirect_uri, self._allowed_redirect_uri_patterns):
+                return redirect_uri
+            # Fall back to normal validation if not in allowed patterns
+            return super().validate_redirect_uri(redirect_uri)
+        # If no redirect_uri provided, use default behavior
+        return super().validate_redirect_uri(redirect_uri)
+
+
+# Default token expiration times
+DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS: Final[int] = 60 * 60  # 1 hour
+DEFAULT_AUTH_CODE_EXPIRY_SECONDS: Final[int] = 5 * 60  # 5 minutes
+
+# HTTP client timeout
+HTTP_TIMEOUT_SECONDS: Final[int] = 30
+
+
+class OAuthProxy(OAuthProvider):
+    """OAuth provider that presents a DCR-compliant interface while proxying to non-DCR IDPs.
+
+    Purpose
+    -------
+    MCP clients expect OAuth providers to support Dynamic Client Registration (DCR),
+    where clients can register themselves dynamically and receive unique credentials.
+    Most enterprise IDPs (Google, GitHub, Azure AD, etc.) don't support DCR and require
+    pre-registered OAuth applications with fixed credentials.
+
+    This proxy bridges that gap by:
+    - Presenting a full DCR-compliant OAuth interface to MCP clients
+    - Translating DCR registration requests to use pre-configured upstream credentials
+    - Proxying all OAuth flows to the upstream IDP with appropriate translations
+    - Managing the state and security requirements of both protocols
+
+    Architecture Overview
+    --------------------
+    The proxy maintains a single OAuth app registration with the upstream provider
+    while allowing unlimited MCP clients to register and authenticate dynamically.
+    It implements the complete OAuth 2.1 + DCR specification for clients while
+    translating to whatever OAuth variant the upstream provider requires.
+
+    Key Translation Challenges Solved
+    ---------------------------------
+    1. Dynamic Client Registration:
+       - MCP clients expect to register dynamically and get unique credentials
+       - Upstream IDPs require pre-registered apps with fixed credentials
+       - Solution: Accept DCR requests, return shared upstream credentials
+
+    2. Dynamic Redirect URIs:
+       - MCP clients use random localhost ports that change between sessions
+       - Upstream IDPs require fixed, pre-registered redirect URIs
+       - Solution: Use proxy's fixed callback URL with upstream, forward to client's dynamic URI
+
+    3. Authorization Code Mapping:
+       - Upstream returns codes for the proxy's redirect URI
+       - Clients expect codes for their own redirect URIs
+       - Solution: Exchange upstream code server-side, issue new code to client
+
+    4. State Parameter Collision:
+       - Both client and proxy need to maintain state through the flow
+       - Only one state parameter available in OAuth
+       - Solution: Use transaction ID as state with upstream, preserve client's state
+
+    5. Token Management:
+       - Clients may expect different token formats/claims than upstream provides
+       - Need to track tokens for revocation and refresh
+       - Solution: Store token relationships, forward upstream tokens transparently
+
+    OAuth Flow Implementation
+    ------------------------
+    1. Client Registration (DCR):
+       - Accept any client registration request
+       - Store ProxyDCRClient that accepts dynamic redirect URIs
+
+    2. Authorization:
+       - Store transaction mapping client details to proxy flow
+       - Redirect to upstream with proxy's fixed redirect URI
+       - Use transaction ID as state parameter with upstream
+
+    3. Upstream Callback:
+       - Exchange upstream authorization code for tokens (server-side)
+       - Generate new authorization code bound to client's PKCE challenge
+       - Redirect to client's original dynamic redirect URI
+
+    4. Token Exchange:
+       - Validate client's code and PKCE verifier
+       - Return previously obtained upstream tokens
+       - Clean up one-time use authorization code
+
+    5. Token Refresh:
+       - Forward refresh requests to upstream using authlib
+       - Handle token rotation if upstream issues new refresh token
+       - Update local token mappings
+
+    State Management
+    ---------------
+    The proxy maintains minimal but crucial state:
+    - _clients: DCR registrations (all use ProxyDCRClient for flexibility)
+    - _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
+
+    Security Considerations
+    ----------------------
+    - PKCE enforced end-to-end (client to proxy, proxy to upstream)
+    - Authorization codes are single-use with short expiry
+    - Transaction IDs are cryptographically random
+    - All state is cleaned up after use to prevent replay
+    - Token validation delegates to upstream provider
+
+    Provider Compatibility
+    ---------------------
+    Works with any OAuth 2.0 provider that supports:
+    - Authorization code flow
+    - Fixed redirect URI (configured in provider's app settings)
+    - Standard token endpoint
+
+    Handles provider-specific requirements:
+    - Google: Ensures minimum scope requirements
+    - GitHub: Compatible with OAuth Apps and GitHub Apps
+    - Azure AD: Handles tenant-specific endpoints
+    - Generic: Works with any spec-compliant provider
+    """
+
+    def __init__(
+        self,
+        *,
+        # Upstream server configuration
+        upstream_authorization_endpoint: str,
+        upstream_token_endpoint: str,
+        upstream_client_id: str,
+        upstream_client_secret: str,
+        upstream_revocation_endpoint: str | None = None,
+        # Token validation
+        token_verifier: TokenVerifier,
+        # FastMCP server configuration
+        base_url: AnyHttpUrl | str,
+        redirect_path: str = "/auth/callback",
+        issuer_url: AnyHttpUrl | str | None = None,
+        service_documentation_url: AnyHttpUrl | str | None = None,
+        # Client redirect URI validation
+        allowed_client_redirect_uris: list[str] | None = None,
+        valid_scopes: list[str] | None = None,
+        # PKCE configuration
+        forward_pkce: bool = True,
+        # Token endpoint authentication
+        token_endpoint_auth_method: str | None = None,
+        # Extra parameters to forward to authorization endpoint
+        extra_authorize_params: dict[str, str] | None = None,
+        # Extra parameters to forward to token endpoint
+        extra_token_params: dict[str, str] | None = None,
+    ):
+        """Initialize the OAuth proxy provider.
+
+        Args:
+            upstream_authorization_endpoint: URL of upstream authorization endpoint
+            upstream_token_endpoint: URL of upstream token endpoint
+            upstream_client_id: Client ID registered with upstream server
+            upstream_client_secret: Client secret for upstream server
+            upstream_revocation_endpoint: Optional upstream revocation endpoint
+            token_verifier: Token verifier for validating access tokens
+            base_url: Public URL of the server that exposes this FastMCP server; redirect path is
+                relative to this URL
+            redirect_path: Redirect path configured in upstream OAuth app (defaults to "/auth/callback")
+            issuer_url: Issuer URL for OAuth metadata (defaults to base_url)
+            service_documentation_url: Optional service documentation URL
+            allowed_client_redirect_uris: List of allowed redirect URI patterns for MCP clients.
+                Patterns support wildcards (e.g., "http://localhost:*", "https://*.example.com/*").
+                If None (default), only localhost redirect URIs are allowed.
+                If empty list, all redirect URIs are allowed (not recommended for production).
+                These are for MCP clients performing loopback redirects, NOT for the upstream OAuth app.
+            valid_scopes: List of all the possible valid scopes for a client.
+                These are advertised to clients through the `/.well-known` endpoints. Defaults to `required_scopes` if not provided.
+            forward_pkce: Whether to forward PKCE to upstream server (default True).
+                Enable for providers that support/require PKCE (Google, Azure, etc.).
+                Disable only if upstream provider doesn't support PKCE.
+            token_endpoint_auth_method: Token endpoint authentication method for upstream server.
+                Common values: "client_secret_basic", "client_secret_post", "none".
+                If None, authlib will use its default (typically "client_secret_basic").
+            extra_authorize_params: Additional parameters to forward to the upstream authorization endpoint.
+                Useful for provider-specific parameters like Auth0's "audience".
+                Example: {"audience": "https://api.example.com"}
+            extra_token_params: Additional parameters to forward to the upstream token endpoint.
+                Useful for provider-specific parameters during token exchange.
+        """
+        # Always enable DCR since we implement it locally for MCP clients
+        client_registration_options = ClientRegistrationOptions(
+            enabled=True,
+            valid_scopes=valid_scopes or token_verifier.required_scopes,
+        )
+
+        # Enable revocation only if upstream endpoint provided
+        revocation_options = (
+            RevocationOptions(enabled=True) if upstream_revocation_endpoint else None
+        )
+
+        super().__init__(
+            base_url=base_url,
+            issuer_url=issuer_url,
+            service_documentation_url=service_documentation_url,
+            client_registration_options=client_registration_options,
+            revocation_options=revocation_options,
+            required_scopes=token_verifier.required_scopes,
+        )
+
+        # Store upstream configuration
+        self._upstream_authorization_endpoint = upstream_authorization_endpoint
+        self._upstream_token_endpoint = upstream_token_endpoint
+        self._upstream_client_id = upstream_client_id
+        self._upstream_client_secret = SecretStr(upstream_client_secret)
+        self._upstream_revocation_endpoint = upstream_revocation_endpoint
+        self._default_scope_str = " ".join(self.required_scopes or [])
+
+        # Store redirect configuration
+        self._redirect_path = (
+            redirect_path if redirect_path.startswith("/") else f"/{redirect_path}"
+        )
+        self._allowed_client_redirect_uris = allowed_client_redirect_uris
+
+        # PKCE configuration
+        self._forward_pkce = forward_pkce
+
+        # Token endpoint authentication
+        self._token_endpoint_auth_method = token_endpoint_auth_method
+
+        # Extra parameters for authorization and token endpoints
+        self._extra_authorize_params = extra_authorize_params or {}
+        self._extra_token_params = extra_token_params or {}
+
+        # Local state for DCR and token bookkeeping
+        self._clients: dict[str, OAuthClientInformationFull] = {}
+        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] = {}
+
+        # OAuth transaction storage for IdP callback forwarding
+        self._oauth_transactions: dict[
+            str, dict[str, Any]
+        ] = {}  # txn_id -> transaction_data
+        self._client_codes: dict[str, dict[str, Any]] = {}  # client_code -> code_data
+
+        # Use the provided token validator
+        self._token_validator = token_verifier
+
+        logger.debug(
+            "Initialized OAuth proxy provider with upstream server %s",
+            self._upstream_authorization_endpoint,
+        )
+
+    # -------------------------------------------------------------------------
+    # PKCE Helper Methods
+    # -------------------------------------------------------------------------
+
+    def _generate_pkce_pair(self) -> tuple[str, str]:
+        """Generate PKCE code verifier and challenge pair.
+
+        Returns:
+            Tuple of (code_verifier, code_challenge) using S256 method
+        """
+        # Generate code verifier: 43-128 characters from unreserved set
+        code_verifier = generate_token(48)
+
+        # Generate code challenge using S256 (SHA256 + base64url)
+        challenge_bytes = hashlib.sha256(code_verifier.encode()).digest()
+        code_challenge = urlsafe_b64encode(challenge_bytes).decode().rstrip("=")
+
+        return code_verifier, code_challenge
+
+    # -------------------------------------------------------------------------
+    # Client Registration (Local Implementation)
+    # -------------------------------------------------------------------------
+
+    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
+        """Get client information by ID. This is generally the random ID
+        provided to the DCR client during registration, not the upstream client ID.
+
+        For unregistered clients, returns None (which will raise an error in the SDK).
+        """
+        client = self._clients.get(client_id)
+
+        return client
+
+    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
+        """Register a client locally
+
+        When a client registers, we create a ProxyDCRClient that is more
+        forgiving about validating redirect URIs, since the DCR client's
+        redirect URI will likely be localhost or unknown to the proxied IDP. The
+        proxied IDP only knows about this server's fixed redirect URI.
+        """
+
+        # Create a ProxyDCRClient with configured redirect URI validation
+        proxy_client = ProxyDCRClient(
+            client_id=client_info.client_id,
+            client_secret=client_info.client_secret,
+            redirect_uris=client_info.redirect_uris or [AnyUrl("http://localhost")],
+            grant_types=client_info.grant_types
+            or ["authorization_code", "refresh_token"],
+            scope=self._default_scope_str,
+            token_endpoint_auth_method="none",
+            allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
+        )
+
+        # Store the ProxyDCRClient
+        self._clients[client_info.client_id] = proxy_client
+
+        # Log redirect URIs to help users discover what patterns they might need
+        if client_info.redirect_uris:
+            for uri in client_info.redirect_uris:
+                logger.debug(
+                    "Client registered with redirect_uri: %s - if restricting redirect URIs, "
+                    "ensure this pattern is allowed in allowed_client_redirect_uris",
+                    uri,
+                )
+
+        logger.debug(
+            "Registered client %s with %d redirect URIs",
+            client_info.client_id,
+            len(proxy_client.redirect_uris),
+        )
+
+    # -------------------------------------------------------------------------
+    # Authorization Flow (Proxy to Upstream)
+    # -------------------------------------------------------------------------
+
+    async def authorize(
+        self,
+        client: OAuthClientInformationFull,
+        params: AuthorizationParams,
+    ) -> str:
+        """Start OAuth transaction and redirect to upstream IdP.
+
+        This implements the DCR-compliant proxy pattern:
+        1. Store transaction with client details and PKCE challenge
+        2. Generate proxy's own PKCE parameters if forwarding is enabled
+        3. Use transaction ID as state for IdP
+        4. Redirect to IdP with our fixed callback URL and proxy's PKCE
+        """
+        # Generate transaction ID for this authorization request
+        txn_id = secrets.token_urlsafe(32)
+
+        # Generate proxy's own PKCE parameters if forwarding is enabled
+        proxy_code_verifier = None
+        proxy_code_challenge = None
+        if self._forward_pkce and params.code_challenge:
+            proxy_code_verifier, proxy_code_challenge = self._generate_pkce_pair()
+            logger.debug(
+                "Generated proxy PKCE for transaction %s (forwarding client PKCE to upstream)",
+                txn_id,
+            )
+
+        # Store transaction data for IdP callback processing
+        transaction_data = {
+            "client_id": client.client_id,
+            "client_redirect_uri": str(params.redirect_uri),
+            "client_state": params.state,
+            "code_challenge": params.code_challenge,
+            "code_challenge_method": getattr(params, "code_challenge_method", "S256"),
+            "scopes": params.scopes or [],
+            "created_at": time.time(),
+        }
+
+        # Store proxy's PKCE verifier if we're forwarding
+        if proxy_code_verifier:
+            transaction_data["proxy_code_verifier"] = proxy_code_verifier
+
+        self._oauth_transactions[txn_id] = transaction_data
+
+        # Build query parameters for upstream IdP authorization request
+        # Use our fixed IdP callback and transaction ID as state
+        query_params: dict[str, Any] = {
+            "response_type": "code",
+            "client_id": self._upstream_client_id,
+            "redirect_uri": f"{str(self.base_url).rstrip('/')}{self._redirect_path}",
+            "state": txn_id,  # Use txn_id as IdP state
+        }
+
+        # Add scopes - use client scopes or fallback to required scopes
+        scopes_to_use = params.scopes or self.required_scopes or []
+
+        if scopes_to_use:
+            query_params["scope"] = " ".join(scopes_to_use)
+
+        # Forward proxy's PKCE challenge to upstream if enabled
+        if proxy_code_challenge:
+            query_params["code_challenge"] = proxy_code_challenge
+            query_params["code_challenge_method"] = "S256"
+            logger.debug(
+                "Forwarding proxy PKCE challenge to upstream for transaction %s",
+                txn_id,
+            )
+
+        # Forward resource parameter if provided (RFC 8707)
+        if params.resource:
+            query_params["resource"] = params.resource
+            logger.debug(
+                "Forwarding resource indicator '%s' to upstream for transaction %s",
+                params.resource,
+                txn_id,
+            )
+
+        # Add any extra authorization parameters configured for this proxy
+        if self._extra_authorize_params:
+            query_params.update(self._extra_authorize_params)
+            logger.debug(
+                "Adding extra authorization parameters for transaction %s: %s",
+                txn_id,
+                list(self._extra_authorize_params.keys()),
+            )
+
+        # Build the upstream authorization URL
+        separator = "&" if "?" in self._upstream_authorization_endpoint else "?"
+        upstream_url = f"{self._upstream_authorization_endpoint}{separator}{urlencode(query_params)}"
+
+        logger.debug(
+            "Starting OAuth transaction %s for client %s, redirecting to IdP (PKCE forwarding: %s)",
+            txn_id,
+            client.client_id,
+            "enabled" if proxy_code_challenge else "disabled",
+        )
+        return upstream_url
+
+    # -------------------------------------------------------------------------
+    # Authorization Code Handling
+    # -------------------------------------------------------------------------
+
+    async def load_authorization_code(
+        self,
+        client: OAuthClientInformationFull,
+        authorization_code: str,
+    ) -> AuthorizationCode | None:
+        """Load authorization code for validation.
+
+        Look up our client code and return authorization code object
+        with PKCE challenge for validation.
+        """
+        # Look up client code data
+        code_data = self._client_codes.get(authorization_code)
+        if not code_data:
+            logger.debug("Authorization code not found: %s", authorization_code)
+            return None
+
+        # Check if code expired
+        if time.time() > code_data["expires_at"]:
+            logger.debug("Authorization code expired: %s", authorization_code)
+            self._client_codes.pop(authorization_code, None)
+            return None
+
+        # Verify client ID matches
+        if code_data["client_id"] != client.client_id:
+            logger.debug(
+                "Authorization code client ID mismatch: %s vs %s",
+                code_data["client_id"],
+                client.client_id,
+            )
+            return None
+
+        # Create authorization code object with PKCE challenge
+        return AuthorizationCode(
+            code=authorization_code,
+            client_id=client.client_id,
+            redirect_uri=code_data["redirect_uri"],
+            redirect_uri_provided_explicitly=True,
+            scopes=code_data["scopes"],
+            expires_at=code_data["expires_at"],
+            code_challenge=code_data.get("code_challenge", ""),
+        )
+
+    async def exchange_authorization_code(
+        self,
+        client: OAuthClientInformationFull,
+        authorization_code: AuthorizationCode,
+    ) -> OAuthToken:
+        """Exchange authorization code for stored IdP tokens.
+
+        For the DCR-compliant proxy flow, we return the IdP tokens that were obtained
+        during the IdP callback exchange. PKCE validation is handled by the MCP framework.
+        """
+        # Look up stored code data
+        code_data = self._client_codes.get(authorization_code.code)
+        if not code_data:
+            logger.error(
+                "Authorization code not found in client codes: %s",
+                authorization_code.code,
+            )
+            raise TokenError("invalid_grant", "Authorization code not found")
+
+        # Get stored IdP tokens
+        idp_tokens = code_data["idp_tokens"]
+
+        # Clean up client code (one-time use)
+        self._client_codes.pop(authorization_code.code, None)
+
+        # Extract token information for local tracking
+        access_token_value = idp_tokens["access_token"]
+        refresh_token_value = idp_tokens.get("refresh_token")
+        expires_in = int(
+            idp_tokens.get("expires_in", DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS)
+        )
+        expires_at = int(time.time() + expires_in)
+
+        # Store access token locally for tracking
+        access_token = AccessToken(
+            token=access_token_value,
+            client_id=client.client_id,
+            scopes=authorization_code.scopes,
+            expires_at=expires_at,
+        )
+        self._access_tokens[access_token_value] = access_token
+
+        # Store refresh token if provided
+        if refresh_token_value:
+            refresh_token = RefreshToken(
+                token=refresh_token_value,
+                client_id=client.client_id,
+                scopes=authorization_code.scopes,
+                expires_at=None,  # Refresh tokens typically don't expire
+            )
+            self._refresh_tokens[refresh_token_value] = refresh_token
+
+            # Maintain token relationships for cleanup
+            self._access_to_refresh[access_token_value] = refresh_token_value
+            self._refresh_to_access[refresh_token_value] = access_token_value
+
+        logger.debug(
+            "Successfully exchanged client code for stored IdP tokens (client: %s)",
+            client.client_id,
+        )
+
+        return OAuthToken(**idp_tokens)  # type: ignore[arg-type]
+
+    # -------------------------------------------------------------------------
+    # Refresh Token Flow
+    # -------------------------------------------------------------------------
+
+    async def load_refresh_token(
+        self,
+        client: OAuthClientInformationFull,
+        refresh_token: str,
+    ) -> RefreshToken | None:
+        """Load refresh token from local storage."""
+        return self._refresh_tokens.get(refresh_token)
+
+    async def exchange_refresh_token(
+        self,
+        client: OAuthClientInformationFull,
+        refresh_token: RefreshToken,
+        scopes: list[str],
+    ) -> OAuthToken:
+        """Exchange refresh token for new access token using authlib."""
+
+        # Use authlib's AsyncOAuth2Client for refresh token exchange
+        oauth_client = AsyncOAuth2Client(
+            client_id=self._upstream_client_id,
+            client_secret=self._upstream_client_secret.get_secret_value(),
+            token_endpoint_auth_method=self._token_endpoint_auth_method,
+            timeout=HTTP_TIMEOUT_SECONDS,
+        )
+
+        try:
+            logger.debug("Using authlib to refresh token from upstream")
+
+            # Let authlib handle the refresh token exchange
+            token_response: dict[str, Any] = await oauth_client.refresh_token(  # type: ignore[misc]
+                url=self._upstream_token_endpoint,
+                refresh_token=refresh_token.token,
+                scope=" ".join(scopes) if scopes else None,
+            )
+
+            logger.debug(
+                "Successfully refreshed access token via authlib (client: %s)",
+                client.client_id,
+            )
+
+        except Exception as e:
+            logger.error("Authlib refresh token exchange failed: %s", e)
+            raise TokenError(
+                "invalid_grant", f"Upstream refresh token exchange failed: {e}"
+            ) from e
+
+        # Update local token storage
+        new_access_token = token_response["access_token"]
+        expires_in = int(
+            token_response.get("expires_in", DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS)
+        )
+
+        self._access_tokens[new_access_token] = AccessToken(
+            token=new_access_token,
+            client_id=client.client_id,
+            scopes=scopes,
+            expires_at=int(time.time() + expires_in),
+        )
+
+        # Handle refresh token rotation if new one provided
+        if "refresh_token" in token_response:
+            new_refresh_token = token_response["refresh_token"]
+            if new_refresh_token != refresh_token.token:
+                # Remove old refresh token
+                self._refresh_tokens.pop(refresh_token.token, None)
+                old_access = self._refresh_to_access.pop(refresh_token.token, None)
+                if old_access:
+                    self._access_to_refresh.pop(old_access, None)
+
+                # Store new refresh token
+                self._refresh_tokens[new_refresh_token] = RefreshToken(
+                    token=new_refresh_token,
+                    client_id=client.client_id,
+                    scopes=scopes,
+                    expires_at=None,
+                )
+                self._access_to_refresh[new_access_token] = new_refresh_token
+                self._refresh_to_access[new_refresh_token] = new_access_token
+
+        return OAuthToken(**token_response)  # type: ignore[arg-type]
+
+    # -------------------------------------------------------------------------
+    # Token Validation
+    # -------------------------------------------------------------------------
+
+    async def load_access_token(self, token: str) -> AccessToken | None:
+        """Validate access token using upstream JWKS.
+
+        Delegates to the JWT verifier which handles signature validation,
+        expiration checking, and claims validation using the upstream JWKS.
+        """
+        result = await self._token_validator.verify_token(token)
+        if result:
+            logger.debug("Token validated successfully")
+        else:
+            logger.debug("Token validation failed")
+        return result
+
+    # -------------------------------------------------------------------------
+    # Token Revocation
+    # -------------------------------------------------------------------------
+
+    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.
+        """
+        # 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)
+
+        # Attempt upstream revocation if endpoint is configured
+        if self._upstream_revocation_endpoint:
+            try:
+                async with httpx.AsyncClient(
+                    timeout=HTTP_TIMEOUT_SECONDS
+                ) as http_client:
+                    await http_client.post(
+                        self._upstream_revocation_endpoint,
+                        data={"token": token.token},
+                        auth=(
+                            self._upstream_client_id,
+                            self._upstream_client_secret.get_secret_value(),
+                        ),
+                    )
+                    logger.debug("Successfully revoked token with upstream server")
+            except Exception as e:
+                logger.warning("Failed to revoke token with upstream server: %s", e)
+        else:
+            logger.debug("No upstream revocation endpoint configured")
+
+        logger.debug("Token revoked successfully")
+
+    def get_routes(
+        self,
+        mcp_path: str | None = None,
+        mcp_endpoint: Any | None = None,
+    ) -> list[Route]:
+        """Get OAuth routes with custom proxy token handler.
+
+        This method creates standard OAuth routes and replaces the token endpoint
+        with our proxy handler that forwards requests to the upstream OAuth server.
+
+        Args:
+            mcp_path: The path where the MCP endpoint is mounted (e.g., "/mcp")
+            mcp_endpoint: The MCP endpoint handler to protect with auth
+        """
+        # Get standard OAuth routes from parent class
+        routes = super().get_routes(mcp_path, mcp_endpoint)
+        custom_routes = []
+        token_route_found = False
+
+        logger.debug(
+            f"get_routes called - configuring OAuth routes in {len(routes)} routes"
+        )
+
+        for i, route in enumerate(routes):
+            logger.debug(
+                f"Route {i}: {route} - path: {getattr(route, 'path', 'N/A')}, methods: {getattr(route, 'methods', 'N/A')}"
+            )
+
+            # Keep all standard OAuth routes unchanged - our DCR-compliant flow handles everything
+            custom_routes.append(route)
+
+            if (
+                isinstance(route, Route)
+                and route.path == "/token"
+                and route.methods is not None
+                and "POST" in route.methods
+            ):
+                token_route_found = True
+
+        # Add OAuth callback endpoint for forwarding to client callbacks
+        custom_routes.append(
+            Route(
+                path=self._redirect_path,
+                endpoint=self._handle_idp_callback,
+                methods=["GET"],
+            )
+        )
+
+        logger.debug(
+            f"✅ OAuth routes configured: token_endpoint={token_route_found}, total routes={len(custom_routes)} (includes OAuth callback)"
+        )
+        return custom_routes
+
+    # -------------------------------------------------------------------------
+    # IdP Callback Forwarding
+    # -------------------------------------------------------------------------
+
+    async def _handle_idp_callback(self, request: Request) -> RedirectResponse:
+        """Handle callback from upstream IdP and forward to client.
+
+        This implements the DCR-compliant callback forwarding:
+        1. Receive IdP callback with code and txn_id as state
+        2. Exchange IdP code for tokens (server-side)
+        3. Generate our own client code bound to PKCE challenge
+        4. Redirect to client's callback with client code and original state
+        """
+        try:
+            idp_code = request.query_params.get("code")
+            txn_id = request.query_params.get("state")
+            error = request.query_params.get("error")
+
+            if error:
+                logger.error(
+                    "IdP callback error: %s - %s",
+                    error,
+                    request.query_params.get("error_description"),
+                )
+                # TODO: Forward error to client callback
+                return RedirectResponse(
+                    url=f"data:text/html,<h1>OAuth Error</h1><p>{error}: {request.query_params.get('error_description', 'Unknown error')}</p>",
+                    status_code=302,
+                )
+
+            if not idp_code or not txn_id:
+                logger.error("IdP callback missing code or transaction ID")
+                return RedirectResponse(
+                    url="data:text/html,<h1>OAuth Error</h1><p>Missing authorization code or transaction ID</p>",
+                    status_code=302,
+                )
+
+            # Look up transaction data
+            transaction = self._oauth_transactions.get(txn_id)
+            if not transaction:
+                logger.error("IdP callback with invalid transaction ID: %s", txn_id)
+                return RedirectResponse(
+                    url="data:text/html,<h1>OAuth Error</h1><p>Invalid or expired transaction</p>",
+                    status_code=302,
+                )
+
+            # Exchange IdP code for tokens (server-side)
+            oauth_client = AsyncOAuth2Client(
+                client_id=self._upstream_client_id,
+                client_secret=self._upstream_client_secret.get_secret_value(),
+                token_endpoint_auth_method=self._token_endpoint_auth_method,
+                timeout=HTTP_TIMEOUT_SECONDS,
+            )
+
+            try:
+                idp_redirect_uri = (
+                    f"{str(self.base_url).rstrip('/')}{self._redirect_path}"
+                )
+                logger.debug(
+                    f"Exchanging IdP code for tokens with redirect_uri: {idp_redirect_uri}"
+                )
+
+                # Build token exchange parameters
+                token_params = {
+                    "url": self._upstream_token_endpoint,
+                    "code": idp_code,
+                    "redirect_uri": idp_redirect_uri,
+                }
+
+                # Include proxy's code_verifier if we forwarded PKCE
+                proxy_code_verifier = transaction.get("proxy_code_verifier")
+                if proxy_code_verifier:
+                    token_params["code_verifier"] = proxy_code_verifier
+                    logger.debug(
+                        "Including proxy code_verifier in token exchange for transaction %s",
+                        txn_id,
+                    )
+
+                # Add any extra token parameters configured for this proxy
+                if self._extra_token_params:
+                    token_params.update(self._extra_token_params)
+                    logger.debug(
+                        "Adding extra token parameters for transaction %s: %s",
+                        txn_id,
+                        list(self._extra_token_params.keys()),
+                    )
+
+                idp_tokens: dict[str, Any] = await oauth_client.fetch_token(
+                    **token_params
+                )  # type: ignore[misc]
+
+                logger.debug(
+                    f"Successfully exchanged IdP code for tokens (transaction: {txn_id}, PKCE: {bool(proxy_code_verifier)})"
+                )
+
+            except Exception as e:
+                logger.error("IdP token exchange failed: %s", e)
+                # TODO: Forward error to client callback
+                return RedirectResponse(
+                    url=f"data:text/html,<h1>OAuth Error</h1><p>Token exchange failed: {e}</p>",
+                    status_code=302,
+                )
+
+            # Generate our own authorization code for the client
+            client_code = secrets.token_urlsafe(32)
+            code_expires_at = int(time.time() + DEFAULT_AUTH_CODE_EXPIRY_SECONDS)
+
+            # Store client code with PKCE challenge and IdP tokens
+            self._client_codes[client_code] = {
+                "client_id": transaction["client_id"],
+                "redirect_uri": transaction["client_redirect_uri"],
+                "code_challenge": transaction["code_challenge"],
+                "code_challenge_method": transaction["code_challenge_method"],
+                "scopes": transaction["scopes"],
+                "idp_tokens": idp_tokens,
+                "expires_at": code_expires_at,
+                "created_at": time.time(),
+            }
+
+            # Clean up transaction
+            self._oauth_transactions.pop(txn_id, None)
+
+            # Build client callback URL with our code and original state
+            client_redirect_uri = transaction["client_redirect_uri"]
+            client_state = transaction["client_state"]
+
+            callback_params = {
+                "code": client_code,
+                "state": client_state,
+            }
+
+            # Add query parameters to client redirect URI
+            separator = "&" if "?" in client_redirect_uri else "?"
+            client_callback_url = (
+                f"{client_redirect_uri}{separator}{urlencode(callback_params)}"
+            )
+
+            logger.debug(f"Forwarding to client callback for transaction {txn_id}")
+
+            return RedirectResponse(url=client_callback_url, status_code=302)
+
+        except Exception as e:
+            logger.error("Error in IdP callback handler: %s", e, exc_info=True)
+            return RedirectResponse(
+                url="data:text/html,<h1>OAuth Error</h1><p>Internal server error during IdP callback</p>",
+                status_code=302,
+            )