refactor(oauth): consolidate authorization flow primitives

Author: wenxi-onyx

backend/onyx/auth/oauth_token_manager.py

@@ -1,4 +1,5 @@
 import time
+from collections.abc import Mapping
 from typing import Any
 from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
 from uuid import UUID
@@ -52,6 +53,26 @@ def validate_oauth_endpoint_url(url: str, *, resolve_dns: bool = True) -> None:
     "accounts.google.com": {"access_type": "offline", "prompt": "consent"},
 }
 
+_PROTOCOL_AUTHORIZATION_PARAMS = frozenset(
+    {
+        "client_id",
+        "code_challenge",
+        "code_challenge_method",
+        "redirect_uri",
+        "resource",
+        "response_type",
+        "scope",
+        "state",
+    }
+)
+
+
+def conflicting_authorization_params(
+    additional_params: Mapping[str, object] | None,
+) -> frozenset[str]:
+    """Return configured parameters whose values the OAuth flow must own."""
+    return _PROTOCOL_AUTHORIZATION_PARAMS.intersection(additional_params or {})
+
 
 def ensure_offline_access_auth_params(authorization_url: str) -> str:
     """Merge missing offline-access params into a built authorize URL.

backend/onyx/oauth/authorization_attempt.py

@@ -1,10 +1,11 @@
 import hashlib
+import json
 import re
 import secrets
 from datetime import datetime, timedelta, timezone
 from typing import Generic, cast
 
-from pydantic import ValidationError
+from pydantic import JsonValue, ValidationError
 
 from onyx.cache.interface import CacheBackend
 from onyx.error_handling.error_codes import OnyxErrorCode
@@ -20,6 +21,18 @@
 MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS = 10 * 60
 
 
+def canonical_json_fingerprint(value: JsonValue) -> str:
+    """Fingerprint configuration captured by a pending authorization attempt."""
+    serialized = json.dumps(
+        value,
+        ensure_ascii=True,
+        allow_nan=False,
+        sort_keys=True,
+        separators=(",", ":"),
+    )
+    return hashlib.sha256(serialized.encode()).hexdigest()
+
+
 class AuthorizationAttemptStore(Generic[PayloadT]):
     """Stores typed OAuth attempts in a tenant-scoped cache."""
 

backend/onyx/oauth/models.py

@@ -1,10 +1,27 @@
-from typing import Generic, TypeVar
+from typing import Annotated, Generic, TypeVar
 
-from pydantic import AwareDatetime, BaseModel, ConfigDict, Field
+from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field
+
+from onyx.utils.url import sanitize_next_url
 
 PayloadT = TypeVar("PayloadT", bound=BaseModel)
 
 
+def _validate_safe_oauth_return_path(value: str) -> str:
+    if sanitize_next_url(value) != value or any(
+        not character.isprintable() for character in value
+    ):
+        raise ValueError("OAuth return path must be a safe internal path")
+    return value
+
+
+SafeOAuthReturnPath = Annotated[
+    str,
+    Field(max_length=2048),
+    AfterValidator(_validate_safe_oauth_return_path),
+]
+
+
 class AuthorizationAttempt(BaseModel, Generic[PayloadT]):
     """One pending authorization request for an authenticated user.
 

backend/onyx/server/documents/standard_oauth.py

@@ -2,7 +2,7 @@
 from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
 from fastapi import APIRouter, Depends, Query, Request
-from pydantic import BaseModel, ConfigDict, ValidationError, field_validator
+from pydantic import BaseModel, ConfigDict, ValidationError
 from sqlalchemy.orm import Session
 
 from onyx.auth.permissions import require_permission
@@ -21,10 +21,10 @@
     AuthorizationAttemptStore,
     generate_authorization_state,
 )
+from onyx.oauth.models import SafeOAuthReturnPath
 from onyx.server.documents.models import CredentialBase
 from onyx.utils.logger import setup_logger
 from onyx.utils.subclasses import find_all_subclasses_in_package
-from onyx.utils.url import sanitize_next_url
 
 logger = setup_logger()
 
@@ -83,19 +83,10 @@ def _validate_additional_kwargs(
 class _ConnectorOAuthAttemptPayload(BaseModel):
     model_config = ConfigDict(extra="forbid", frozen=True)
 
-    desired_return_path: str
+    desired_return_path: SafeOAuthReturnPath
     additional_kwargs: dict[str, str]
     code_verifier: str | None = None
 
-    @field_validator("desired_return_path")
-    @classmethod
-    def validate_return_path(cls, value: str) -> str:
-        if sanitize_next_url(value) != value or any(
-            not character.isprintable() for character in value
-        ):
-            raise ValueError("OAuth return path must be a local application path")
-        return value
-
 
 def _authorization_attempt_store(
     source: DocumentSource,

backend/onyx/server/features/mcp/credentials.py

@@ -19,8 +19,6 @@
 in ``onyx.db.mcp``. It must not import ``onyx.server.features.mcp.oauth``.
 """
 
-import hashlib
-import json
 import time
 from collections.abc import Mapping
 from typing import cast
@@ -32,6 +30,7 @@
 from onyx.db.enums import MCPAuthenticationPerformer, MCPAuthenticationType
 from onyx.db.mcp import get_user_connection_config
 from onyx.db.models import MCPConnectionConfig, MCPServer, User
+from onyx.oauth.authorization_attempt import canonical_json_fingerprint
 from onyx.server.features.mcp.models import (
     DENYLISTED_MCP_HEADERS,
     MCPAuthTemplate,
@@ -76,26 +75,19 @@ def mcp_oauth_connection_headers_fingerprint(headers: dict[str, str]) -> str:
         for key, value in headers.items()
         if key.lower() != "authorization"
     )
-    serialized_headers = json.dumps(
-        routing_headers, ensure_ascii=True, separators=(",", ":")
-    )
-    return hashlib.sha256(serialized_headers.encode()).hexdigest()
+    return canonical_json_fingerprint(routing_headers)
 
 
 def mcp_oauth_client_information_fingerprint(
     client_information: OAuthClientInformationFull,
 ) -> str:
-    serialized_client = json.dumps(
+    return canonical_json_fingerprint(
         client_information.model_dump(
             mode="json",
             exclude_none=True,
             by_alias=True,
         ),
-        ensure_ascii=True,
-        sort_keys=True,
-        separators=(",", ":"),
     )
-    return hashlib.sha256(serialized_client.encode()).hexdigest()
 
 
 def requires_user_authentication(

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

@@ -11,6 +11,7 @@
 from mcp.types import Tool as MCPLibTool
 from pydantic import AnyUrl, BaseModel, Field, model_validator
 
+from onyx.auth.oauth_token_manager import conflicting_authorization_params
 from onyx.db.enums import (
     EndpointPolicy,
     MCPAuthenticationPerformer,
@@ -19,21 +20,12 @@
     MCPServerStatus,
     MCPTransport,
 )
+from onyx.oauth.models import SafeOAuthReturnPath
 
 # Matches `{placeholder_name}` inside header value templates.
 _PLACEHOLDER_RE = re.compile(r"\{([^}]+)\}")
 # RFC 9110 field-name syntax: a non-empty sequence of HTTP token characters.
 _HTTP_FIELD_NAME_RE = re.compile(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]+")
-RESERVED_MCP_OAUTH_AUTHORIZATION_PARAMS = {
-    "client_id",
-    "code_challenge",
-    "code_challenge_method",
-    "redirect_uri",
-    "resource",
-    "response_type",
-    "scope",
-    "state",
-}
 
 
 def _build_auto_substitution_map(*, user_email: str) -> dict[str, str]:
@@ -387,8 +379,8 @@ def validate_auth_configuration(self) -> "MCPToolCreateRequest":
                 raise ValueError(
                     "oauth_token_endpoint is required for known-provider OAuth mode"
                 )
-            reserved_params = RESERVED_MCP_OAUTH_AUTHORIZATION_PARAMS.intersection(
-                self.oauth_additional_auth_params or {}
+            reserved_params = conflicting_authorization_params(
+                self.oauth_additional_auth_params
             )
             if reserved_params:
                 raise ValueError(
@@ -511,7 +503,9 @@ class MCPOAuthConnectResponse(BaseModel):
 
 class MCPUserOAuthConnectRequest(BaseModel):
     server_id: int = Field(..., description="ID of the MCP server")
-    return_path: str = Field(..., description="Path to redirect to after callback")
+    return_path: SafeOAuthReturnPath = Field(
+        ..., description="Path to redirect to after callback"
+    )
     include_resource_param: bool = Field(..., description="Include resource parameter")
     force_reauthentication: bool = Field(
         default=False,
@@ -539,17 +533,6 @@ class MCPUserOAuthConnectRequest(BaseModel):
         ),
     )
 
-    @model_validator(mode="after")
-    def validate_return_path(self) -> "MCPUserOAuthConnectRequest":
-        if (
-            not self.return_path.startswith("/")
-            or self.return_path.startswith("//")
-            or "\\" in self.return_path
-            or any(not character.isprintable() for character in self.return_path)
-        ):
-            raise ValueError("return_path must be a safe internal path")
-        return self
-
 
 class MCPUserOAuthConnectResponse(BaseModel):
     server_id: int
@@ -592,7 +575,7 @@ class MCPOAuthServerSnapshot(BaseModel):
 class MCPOAuthFlowState(BaseModel):
     server_id: int
     connection_config_id: int
-    return_path: str
+    return_path: SafeOAuthReturnPath
     code_verifier: str
     redirect_uri: AnyUrl
     server_snapshot: MCPOAuthServerSnapshot

backend/onyx/server/features/mcp/oauth_flow.py

@@ -18,6 +18,7 @@
 from onyx.auth.oauth_token_manager import (
     OAuthFlowParams,
     build_oauth_authorization_url,
+    conflicting_authorization_params,
 )
 from onyx.cache.factory import get_cache_backend
 from onyx.db.enums import MCPTransport
@@ -36,7 +37,6 @@
 )
 from onyx.server.features.mcp.models import (
     DENYLISTED_MCP_HEADERS,
-    RESERVED_MCP_OAUTH_AUTHORIZATION_PARAMS,
     MCPOAuthFlowState,
     MCPOAuthServerSnapshot,
     MCPPendingOAuthAuthorization,
@@ -85,8 +85,8 @@ def _known_provider_flow_params(
             "Known-provider OAuth requires oauth_authorization_endpoint, "
             "oauth_token_endpoint, and a non-empty client_id.",
         )
-    reserved_params = RESERVED_MCP_OAUTH_AUTHORIZATION_PARAMS.intersection(
-        mcp_server.oauth_additional_auth_params or {}
+    reserved_params = conflicting_authorization_params(
+        mcp_server.oauth_additional_auth_params
     )
     if reserved_params:
         raise OnyxError(

backend/tests/unit/onyx/oauth/test_authorization_attempt.py

@@ -5,7 +5,10 @@
 
 from onyx.error_handling.error_codes import OnyxErrorCode
 from onyx.error_handling.exceptions import OnyxError
-from onyx.oauth.authorization_attempt import AuthorizationAttemptStore
+from onyx.oauth.authorization_attempt import (
+    AuthorizationAttemptStore,
+    canonical_json_fingerprint,
+)
 from onyx.oauth.models import AuthorizationAttempt
 from tests.unit.fakes import FakeCache
 
@@ -34,6 +37,18 @@ def _store(
     )
 
 
+def test_configuration_fingerprint_is_canonical_and_value_sensitive() -> None:
+    first = {"client": {"id": "client-id", "secret": "secret"}, "scopes": ["a"]}
+    reordered = {"scopes": ["a"], "client": {"secret": "secret", "id": "client-id"}}
+    changed = {
+        "client": {"id": "client-id", "secret": "rotated"},
+        "scopes": ["a"],
+    }
+
+    assert canonical_json_fingerprint(first) == canonical_json_fingerprint(reordered)
+    assert canonical_json_fingerprint(first) != canonical_json_fingerprint(changed)
+
+
 def test_store_generates_isolated_one_time_attempts() -> None:
     cache = FakeCache()
     store = _store(cache)