refactor(oauth-config): use shared authorization attempts

Author: wenxi-onyx

backend/onyx/auth/oauth_token_manager.py

@@ -276,8 +276,14 @@ def is_token_expired(cls, token_data: dict[str, Any]) -> bool:
         # Add 60 second buffer to avoid race conditions
         return int(time.time()) + 60 >= expires_at
 
-    def exchange_code_for_token(self, code: str, redirect_uri: str) -> dict[str, Any]:
-        """Exchange authorization code for access token"""
+    def exchange_code_for_token(
+        self,
+        code: str,
+        redirect_uri: str,
+        *,
+        code_verifier: str | None = None,
+    ) -> dict[str, Any]:
+        """Exchange an authorization code, including a PKCE verifier when supplied."""
         if (
             self.oauth_config.client_id is None
             or self.oauth_config.client_secret is None
@@ -287,22 +293,33 @@ def exchange_code_for_token(self, code: str, redirect_uri: str) -> dict[str, Any
             )
 
         return exchange_oauth_code_for_token(
-            self._flow_params(self.oauth_config), code, redirect_uri
+            self.flow_params(self.oauth_config),
+            code,
+            redirect_uri,
+            code_verifier=code_verifier,
         )
 
     @staticmethod
     def build_authorization_url(
-        oauth_config: OAuthConfig, redirect_uri: str, state: str
+        oauth_config: OAuthConfig,
+        redirect_uri: str,
+        state: str,
+        *,
+        code_challenge: str | None = None,
     ) -> str:
-        """Build OAuth authorization URL"""
+        """Build an authorization URL, including a PKCE challenge when supplied."""
         if oauth_config.client_id is None:
             raise ValueError("OAuth client_id is required to build authorization URL")
         return build_oauth_authorization_url(
-            OAuthTokenManager._flow_params(oauth_config), redirect_uri, state
+            OAuthTokenManager.flow_params(oauth_config),
+            redirect_uri,
+            state,
+            code_challenge=code_challenge,
         )
 
     @staticmethod
-    def _flow_params(oauth_config: OAuthConfig) -> OAuthFlowParams:
+    def flow_params(oauth_config: OAuthConfig) -> OAuthFlowParams:
+        """Return the protocol inputs represented by an OAuthConfig."""
         if oauth_config.client_id is None:
             raise ValueError("OAuth client_id is required")
         client_secret = (

backend/onyx/server/features/oauth_config/api.py

@@ -1,10 +1,18 @@
 """API endpoints for OAuth configuration management."""
 
+import secrets
+
 from fastapi import APIRouter, Depends, HTTPException
+from pydantic import BaseModel, ConfigDict
 from sqlalchemy.orm import Session
 
-from onyx.auth.oauth_token_manager import OAuthTokenManager
+from onyx.auth.oauth_token_manager import (
+    OAuthTokenManager,
+    conflicting_authorization_params,
+)
 from onyx.auth.permissions import has_global_permission, require_permission
+from onyx.auth.pkce import generate_pkce_pair
+from onyx.cache.factory import get_cache_backend
 from onyx.configs.app_configs import WEB_DOMAIN
 from onyx.db.engine.sql_engine import get_session
 from onyx.db.enums import Permission
@@ -21,9 +29,15 @@
 )
 from onyx.error_handling.error_codes import OnyxErrorCode
 from onyx.error_handling.exceptions import OnyxError
-from onyx.federated_connectors.oauth_utils import (
-    generate_oauth_state,
-    verify_oauth_state,
+from onyx.oauth.authorization_attempt import (
+    AuthorizationAttemptStore,
+    canonical_json_fingerprint,
+    generate_authorization_state,
+)
+from onyx.oauth.models import (
+    OAuthConfigurationFingerprint,
+    PKCECodeVerifier,
+    SafeOAuthReturnPath,
 )
 from onyx.server.features.oauth_config.models import (
     OAuthCallbackResponse,
@@ -40,6 +54,47 @@
 admin_router = APIRouter(prefix="/admin/oauth-config")
 router = APIRouter(prefix="/oauth-config")
 
+_OAUTH_CALLBACK_PATH = "/oauth-config/callback"
+
+
+class _OAuthConfigAttemptPayload(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    oauth_config_id: int
+    return_path: SafeOAuthReturnPath
+    configuration_fingerprint: OAuthConfigurationFingerprint
+    code_verifier: PKCECodeVerifier
+
+
+_AUTHORIZATION_ATTEMPTS = AuthorizationAttemptStore(
+    cache_backend_provider=lambda: get_cache_backend(),
+    namespace="oauth-config",
+    payload_type=_OAuthConfigAttemptPayload,
+)
+
+
+def _oauth_callback_url() -> str:
+    return f"{WEB_DOMAIN}{_OAUTH_CALLBACK_PATH}"
+
+
+def _oauth_config_fingerprint(oauth_config: OAuthConfig) -> str:
+    return canonical_json_fingerprint(
+        {
+            "redirect_uri": _oauth_callback_url(),
+            "flow": OAuthTokenManager.flow_params(oauth_config).model_dump(mode="json"),
+        },
+    )
+
+
+def _validate_additional_authorization_params(oauth_config: OAuthConfig) -> None:
+    reserved = conflicting_authorization_params(oauth_config.additional_params)
+    if reserved:
+        raise OnyxError(
+            OnyxErrorCode.INVALID_INPUT,
+            "OAuth additional parameters cannot override: "
+            f"{', '.join(sorted(reserved))}",
+        )
+
 
 def _oauth_config_to_snapshot(
     oauth_config: OAuthConfig, db_session: Session
@@ -215,18 +270,25 @@ def initiate_oauth_flow(
             detail=f"OAuth config with id {request.oauth_config_id} not found",
         )
 
-    # Generate state parameter and store in Redis
-    state = generate_oauth_state(
-        federated_connector_id=request.oauth_config_id,
-        user_id=str(user.id),
-        redirect_uri=request.return_path,
-        additional_data={"oauth_config_id": request.oauth_config_id},
-    )
+    _validate_additional_authorization_params(oauth_config)
+    code_verifier, code_challenge = generate_pkce_pair()
+    state = generate_authorization_state()
 
-    # Build authorization URL
-    redirect_uri = f"{WEB_DOMAIN}/oauth-config/callback"
     authorization_url = OAuthTokenManager.build_authorization_url(
-        oauth_config, redirect_uri, state
+        oauth_config,
+        _oauth_callback_url(),
+        state,
+        code_challenge=code_challenge,
+    )
+    _AUTHORIZATION_ATTEMPTS.store(
+        owner_id=str(user.id),
+        state=state,
+        payload=_OAuthConfigAttemptPayload(
+            oauth_config_id=oauth_config.id,
+            return_path=request.return_path,
+            configuration_fingerprint=_oauth_config_fingerprint(oauth_config),
+            code_verifier=code_verifier,
+        ),
     )
 
     return OAuthInitiateResponse(authorization_url=authorization_url, state=state)
@@ -245,39 +307,39 @@ def handle_oauth_callback(
     Exchanges the authorization code for an access token and stores it.
     Accepts code and state as query parameters (standard OAuth flow).
     """
-    try:
-        # Verify state and retrieve session data
-        session = verify_oauth_state(state)
-
-        # Verify the user_id matches
-        if str(user.id) != session.user_id:
-            raise HTTPException(
-                status_code=403, detail="User mismatch in OAuth callback"
-            )
-
-        # Extract oauth_config_id from session (stored during initiate)
-        oauth_config_id = session.federated_connector_id
-
-        # Get OAuth config
-        oauth_config = get_oauth_config(oauth_config_id, db_session)
-        if not oauth_config:
-            raise HTTPException(
-                status_code=404,
-                detail=f"OAuth config with id {oauth_config_id} not found",
-            )
+    attempt = _AUTHORIZATION_ATTEMPTS.consume(owner_id=str(user.id), state=state)
+    payload = attempt.payload
+
+    oauth_config = get_oauth_config(payload.oauth_config_id, db_session)
+    if not oauth_config:
+        raise OnyxError(
+            OnyxErrorCode.NOT_FOUND,
+            f"OAuth config with id {payload.oauth_config_id} not found",
+        )
+    if not secrets.compare_digest(
+        payload.configuration_fingerprint,
+        _oauth_config_fingerprint(oauth_config),
+    ):
+        raise OnyxError(
+            OnyxErrorCode.INVALID_INPUT,
+            "OAuth configuration changed while authorization was pending.",
+        )
 
+    try:
         # Exchange code for token
-        redirect_uri = f"{WEB_DOMAIN}/oauth-config/callback"
         token_manager = OAuthTokenManager(oauth_config, user.id, db_session)
-        token_data = token_manager.exchange_code_for_token(code, redirect_uri)
+        token_data = token_manager.exchange_code_for_token(
+            code,
+            _oauth_callback_url(),
+            code_verifier=payload.code_verifier,
+        )
 
         # Store token
         upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)
 
         # Return success with redirect
-        return_path = session.redirect_uri or "/chat"
         return OAuthCallbackResponse(
-            redirect_url=return_path,
+            redirect_url=payload.return_path,
         )
 
     except ValueError as e:

backend/onyx/server/features/oauth_config/models.py

@@ -3,6 +3,8 @@
 
 from pydantic import BaseModel
 
+from onyx.oauth.models import SafeOAuthReturnPath
+
 
 class OAuthConfigCreate(BaseModel):
     name: str
@@ -40,7 +42,7 @@ class OAuthConfigSnapshot(BaseModel):
 
 class OAuthInitiateRequest(BaseModel):
     oauth_config_id: int
-    return_path: str = "/chat"  # Where to redirect after OAuth flow
+    return_path: SafeOAuthReturnPath = "/chat"
 
 
 class OAuthInitiateResponse(BaseModel):