refactor(oauth): share authorization attempts across OAuth flows

backend/onyx/cache/interface.py

@@ -80,6 +80,11 @@ class CacheBackend(abc.ABC):
     def get(self, key: str) -> bytes | None:
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def getdel(self, key: str) -> bytes | None:
+        """Atomically return and remove an unexpired value."""
+        raise NotImplementedError
+
     @abc.abstractmethod
     def set(
         self,
@@ -89,6 +94,16 @@ def set(
     ) -> None:
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def set_if_absent(
+        self,
+        key: str,
+        value: str | bytes | int | float,
+        ex: int | None = None,
+    ) -> bool:
+        """Set a value only if no unexpired value uses the key."""
+        raise NotImplementedError
+
     @abc.abstractmethod
     def delete(self, key: str) -> None:
         raise NotImplementedError

backend/onyx/cache/postgres_backend.py

@@ -163,6 +163,27 @@ def get(self, key: str) -> bytes | None:
             return None
         return bytes(value)
 
+    def getdel(self, key: str) -> bytes | None:
+        from onyx.db.engine.sql_engine import get_session_with_tenant
+
+        stmt = (
+            delete(CacheStore)
+            .where(
+                CacheStore.key == key,
+                or_(
+                    CacheStore.expires_at.is_(None),
+                    CacheStore.expires_at > func.now(),
+                ),
+            )
+            .returning(CacheStore.value)
+        )
+        with get_session_with_tenant(tenant_id=self._tenant_id) as session:
+            value = session.execute(stmt).scalar_one_or_none()
+            session.commit()
+        if value is None:
+            return None
+        return bytes(value)
+
     def set(
         self,
         key: str,
@@ -189,6 +210,38 @@ def set(
             session.execute(stmt)
             session.commit()
 
+    def set_if_absent(
+        self,
+        key: str,
+        value: str | bytes | int | float,
+        ex: int | None = None,
+    ) -> bool:
+        from onyx.db.engine.sql_engine import get_session_with_tenant
+
+        value_bytes = _to_bytes(value)
+        expires_at = (
+            datetime.now(timezone.utc) + timedelta(seconds=ex)
+            if ex is not None
+            else None
+        )
+        stmt = (
+            pg_insert(CacheStore)
+            .values(key=key, value=value_bytes, expires_at=expires_at)
+            .on_conflict_do_update(
+                index_elements=[CacheStore.key],
+                set_={"value": value_bytes, "expires_at": expires_at},
+                where=(
+                    CacheStore.expires_at.is_not(None)
+                    & (CacheStore.expires_at <= func.now())
+                ),
+            )
+            .returning(CacheStore.key)
+        )
+        with get_session_with_tenant(tenant_id=self._tenant_id) as session:
+            stored_key = session.execute(stmt).scalar_one_or_none()
+            session.commit()
+        return stored_key is not None
+
     def delete(self, key: str) -> None:
         from onyx.db.engine.sql_engine import get_session_with_tenant
 

backend/onyx/cache/redis_backend.py

@@ -55,6 +55,9 @@ def __init__(self, redis_client: TenantRedisClient) -> None:
     def get(self, key: str) -> bytes | None:
         return self._r.get(key)
 
+    def getdel(self, key: str) -> bytes | None:
+        return self._r.getdel(key)
+
     def set(
         self,
         key: str,
@@ -63,6 +66,14 @@ def set(
     ) -> None:
         self._r.set(key, value, ex=ex)
 
+    def set_if_absent(
+        self,
+        key: str,
+        value: str | bytes | int | float,
+        ex: int | None = None,
+    ) -> bool:
+        return bool(self._r.set(key, value, ex=ex, nx=True))
+
     def delete(self, key: str) -> None:
         self._r.delete(key)
 

backend/onyx/oauth/authorization_attempt.py

@@ -0,0 +1,122 @@
+import hashlib
+import re
+import secrets
+from datetime import datetime, timedelta, timezone
+from typing import Generic, cast
+
+from pydantic import ValidationError
+
+from onyx.cache.interface import CacheBackend
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
+from onyx.oauth.models import AuthorizationAttempt, PayloadT
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+_NAMESPACE_PATTERN = re.compile(r"^[a-z][a-z0-9_-]*$")
+_KEY_PREFIX = "oauth:authorization_attempt:v1"
+_INVALID_ATTEMPT_MESSAGE = "Invalid or expired OAuth authorization attempt"
+MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS = 10 * 60
+
+
+class AuthorizationAttemptStore(Generic[PayloadT]):
+    """Stores typed OAuth attempts in a tenant-scoped cache."""
+
+    def __init__(
+        self,
+        cache: CacheBackend,
+        *,
+        namespace: str,
+        payload_type: type[PayloadT],
+        ttl_seconds: int,
+    ) -> None:
+        if not _NAMESPACE_PATTERN.fullmatch(namespace) or len(namespace) > 64:
+            raise ValueError("OAuth attempt namespace is invalid")
+        if not 0 < ttl_seconds <= MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS:
+            raise ValueError(
+                "OAuth attempt TTL must be between 1 and "
+                f"{MAX_AUTHORIZATION_ATTEMPT_TTL_SECONDS} seconds"
+            )
+
+        self._cache = cache
+        self._namespace = namespace
+        self._attempt_type = cast(
+            type[AuthorizationAttempt[PayloadT]],
+            AuthorizationAttempt.__class_getitem__(payload_type),
+        )
+        self._ttl_seconds = ttl_seconds
+
+    def store(
+        self,
+        *,
+        owner_id: str,
+        payload: PayloadT,
+        state: str | None = None,
+    ) -> AuthorizationAttempt[PayloadT]:
+        """Persist a new attempt.
+
+        When supplied, ``state`` must come from a cryptographically secure
+        generator. Its entropy cannot be established through format validation.
+        """
+        expires_at = datetime.now(timezone.utc) + timedelta(seconds=self._ttl_seconds)
+        attempt = self._attempt_type(
+            namespace=self._namespace,
+            owner_id=owner_id,
+            state=state if state is not None else generate_authorization_state(),
+            expires_at=expires_at,
+            payload=payload,
+        )
+        stored = self._cache.set_if_absent(
+            self._key(attempt.owner_id, attempt.state),
+            attempt.model_dump_json(),
+            ex=self._ttl_seconds,
+        )
+        if not stored:
+            raise OnyxError(
+                OnyxErrorCode.CONFLICT,
+                "An OAuth authorization attempt already uses this state",
+            )
+        return attempt
+
+    def consume(
+        self,
+        *,
+        owner_id: str,
+        state: str,
+    ) -> AuthorizationAttempt[PayloadT]:
+        stored = self._cache.getdel(self._key(owner_id, state))
+        if stored is None:
+            raise _invalid_attempt_error()
+
+        try:
+            attempt = self._attempt_type.model_validate_json(stored)
+        except (ValidationError, ValueError, TypeError) as error:
+            logger.warning(
+                "Rejected malformed OAuth authorization attempt (%s)",
+                type(error).__name__,
+            )
+            raise _invalid_attempt_error() from error
+
+        if (
+            attempt.namespace != self._namespace
+            or attempt.owner_id != owner_id
+            or not secrets.compare_digest(attempt.state, state)
+            or attempt.expires_at <= datetime.now(timezone.utc)
+        ):
+            raise _invalid_attempt_error()
+        return attempt
+
+    def _key(self, owner_id: str, state: str) -> str:
+        owner_hash = hashlib.sha256(owner_id.encode()).hexdigest()
+        state_hash = hashlib.sha256(state.encode()).hexdigest()
+        return f"{_KEY_PREFIX}:{self._namespace}:{owner_hash}:{state_hash}"
+
+
+def generate_authorization_state() -> str:
+    """Generate a 256-bit, URL-safe OAuth state value."""
+    return secrets.token_urlsafe(32)
+
+
+def _invalid_attempt_error() -> OnyxError:
+    return OnyxError(OnyxErrorCode.INVALID_INPUT, _INVALID_ATTEMPT_MESSAGE)

backend/onyx/oauth/models.py

@@ -0,0 +1,22 @@
+from typing import Generic, TypeVar
+
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field
+
+PayloadT = TypeVar("PayloadT", bound=BaseModel)
+
+
+class AuthorizationAttempt(BaseModel, Generic[PayloadT]):
+    """One pending authorization request for an authenticated user.
+
+    Payloads can contain identifiers, protocol context, and short-lived
+    transaction secrets such as PKCE verifiers. They must not contain long-lived
+    credentials or provider client secrets.
+    """
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    namespace: str = Field(min_length=1, max_length=64)
+    owner_id: str = Field(min_length=1, max_length=256)
+    state: str = Field(pattern=r"^[A-Za-z0-9_-]{22,1024}$")
+    expires_at: AwareDatetime
+    payload: PayloadT