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,12 @@
 import hashlib
+import json
 import re
 import secrets
+from collections.abc import Callable
 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,16 +22,28 @@
 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."""
 
     def __init__(
         self,
-        cache: CacheBackend,
         *,
+        cache_backend_provider: Callable[[], CacheBackend],
         namespace: str,
         payload_type: type[PayloadT],
-        ttl_seconds: int,
+        ttl_seconds: int = MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS,
     ) -> None:
         if not _NAMESPACE_PATTERN.fullmatch(namespace) or len(namespace) > 64:
             raise ValueError("OAuth attempt namespace is invalid")
@@ -39,7 +53,7 @@ def __init__(
                 f"{MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS} seconds"
             )
 
-        self._cache = cache
+        self._cache_backend_provider = cache_backend_provider
         self._namespace = namespace
         self._attempt_type = cast(
             type[AuthorizationAttempt[PayloadT]],
@@ -67,7 +81,7 @@ def store(
             expires_at=expires_at,
             payload=payload,
         )
-        stored = self._cache.set_if_absent(
+        stored = self._cache_backend_provider().set_if_absent(
             self._key(attempt.owner_id, attempt.state),
             attempt.model_dump_json(),
             ex=self._ttl_seconds,
@@ -85,7 +99,7 @@ def consume(
         owner_id: str,
         state: str,
     ) -> AuthorizationAttempt[PayloadT]:
-        stored = self._cache.getdel(self._key(owner_id, state))
+        stored = self._cache_backend_provider().getdel(self._key(owner_id, state))
         if stored is None:
             raise _invalid_attempt_error()
 

backend/onyx/oauth/models.py

@@ -1,10 +1,29 @@
-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),
+]
+OAuthConfigurationFingerprint = Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")]
+PKCECodeVerifier = Annotated[str, Field(min_length=43, max_length=128)]
+
+
 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 PKCECodeVerifier, 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,25 +83,16 @@ 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
+    code_verifier: PKCECodeVerifier | None = None
 
 
 def _authorization_attempt_store(
     source: DocumentSource,
 ) -> AuthorizationAttemptStore[_ConnectorOAuthAttemptPayload]:
     return AuthorizationAttemptStore(
-        get_cache_backend(),
+        cache_backend_provider=get_cache_backend,
         namespace=f"{_OAUTH_ATTEMPT_NAMESPACE_PREFIX}-{source.value}",
         payload_type=_ConnectorOAuthAttemptPayload,
         ttl_seconds=_OAUTH_STATE_EXPIRATION_SECONDS,

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

@@ -19,19 +19,18 @@
 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
 
 from mcp.shared.auth import OAuthClientInformationFull
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, JsonValue
 from sqlalchemy.orm import Session
 
 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,
@@ -71,31 +70,27 @@ def mcp_token_expired(config_data: MCPConnectionData) -> bool:
 
 
 def mcp_oauth_connection_headers_fingerprint(headers: dict[str, str]) -> str:
-    routing_headers = sorted(
-        (key.lower(), value)
-        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()
+    routing_headers: JsonValue = [
+        [key, value]
+        for key, value in sorted(
+            (key.lower(), value)
+            for key, value in headers.items()
+            if key.lower() != "authorization"
+        )
+    ]
+    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,16 @@
     MCPServerStatus,
     MCPTransport,
 )
+from onyx.oauth.models import (
+    OAuthConfigurationFingerprint,
+    PKCECodeVerifier,
+    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 +383,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 +507,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 +537,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
@@ -574,7 +561,7 @@ def validate_outcome(self) -> "MCPUserOAuthConnectResponse":
 class MCPPendingOAuthAuthorization(BaseModel):
     authorization_url: str
     state: str
-    code_verifier: str
+    code_verifier: PKCECodeVerifier
 
 
 class MCPOAuthServerSnapshot(BaseModel):
@@ -592,12 +579,12 @@ class MCPOAuthServerSnapshot(BaseModel):
 class MCPOAuthFlowState(BaseModel):
     server_id: int
     connection_config_id: int
-    return_path: str
-    code_verifier: str
+    return_path: SafeOAuthReturnPath
+    code_verifier: PKCECodeVerifier
     redirect_uri: AnyUrl
     server_snapshot: MCPOAuthServerSnapshot
-    connection_headers_fingerprint: str
-    client_information_fingerprint: str
+    connection_headers_fingerprint: OAuthConfigurationFingerprint
+    client_information_fingerprint: OAuthConfigurationFingerprint
     protected_resource_metadata: ProtectedResourceMetadata | None = None
     oauth_metadata: OAuthMetadata | None = None
     authorization_server_url: str | None = None

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(
@@ -180,7 +180,7 @@ def validate_mcp_oauth_flow_configuration(
 
 def mcp_oauth_attempt_store() -> AuthorizationAttemptStore[MCPOAuthFlowState]:
     return AuthorizationAttemptStore(
-        get_cache_backend(),
+        cache_backend_provider=get_cache_backend,
         namespace="mcp",
         payload_type=MCPOAuthFlowState,
         ttl_seconds=MCP_OAUTH_FLOW_TTL_SECONDS,