feat(auth): store PKCE verifiers per flow to survive overlapping flows

sdk-compliance.yaml

@@ -68,6 +68,12 @@ features:
     symbols:
       - AsyncGoTrueClient.exchange_code_for_session
       - SyncGoTrueClient.exchange_code_for_session
+  auth.sign_in.concurrent_pkce_flows:
+    status: implemented
+    symbols:
+      - OAuthResponse.flow_id
+      - AsyncGoTrueClient.exchange_code_for_session
+      - SyncGoTrueClient.exchange_code_for_session
 
   auth.session.get_session:
     status: implemented

src/auth/scripts/run-unasync.py

@@ -9,7 +9,12 @@
     unasync.Rule(
         fromdir="/_async/",
         todir="/_sync/",
-        additional_replacements={"AsyncClient": "Client", "aclose": "close"},
+        additional_replacements={
+            "AsyncClient": "Client",
+            "aclose": "close",
+            # asyncio is only used for the Lock in the PKCE verifier store
+            "asyncio": "threading",
+        },
     ),
     unasync._DEFAULT_RULE,
 )

src/auth/src/supabase_auth/_async/gotrue_client.py

@@ -24,13 +24,15 @@
     AuthImplicitGrantRedirectError,
     AuthInvalidCredentialsError,
     AuthInvalidJwtError,
+    AuthPKCECodeVerifierMissingError,
     AuthRetryableError,
     AuthSessionMissingError,
     UserDoesntExist,
 )
 from ..helpers import (
     decode_jwt,
     generate_pkce_challenge,
+    generate_pkce_flow_id,
     generate_pkce_verifier,
     model_dump_json,
     model_validate,
@@ -41,6 +43,7 @@
     parse_sso_response,
     parse_user_response,
     validate_exp,
+    validate_pkce_flow_id,
 )
 from ..timer import Timer
 from ..types import (
@@ -94,6 +97,7 @@
 from .gotrue_admin_api import AsyncGoTrueAdminAPI
 from .gotrue_base_api import AsyncGoTrueBaseAPI
 from .gotrue_mfa_api import AsyncGoTrueMFAAPI
+from .pkce_verifier_store import AsyncPKCEVerifierStore
 from .storage import AsyncMemoryStorage, AsyncSupportedStorage
 
 
@@ -153,6 +157,9 @@ def __init__(
         self._network_retries = 0
         self._state_change_emitters: Dict[str, Subscription] = {}
         self._flow_type = flow_type
+        self._pkce_verifier_store = AsyncPKCEVerifierStore(
+            self._storage, self._storage_key
+        )
 
         self.admin = AsyncGoTrueAdminAPI(
             url=self._url,
@@ -443,6 +450,11 @@ async def sign_in_with_oauth(
     ) -> OAuthResponse:
         """
         Log in an existing user via a third-party provider.
+
+        On the PKCE flow the returned ``flow_id`` identifies the code verifier
+        stored for this call. Pass it to ``exchange_code_for_session`` so that
+        the exchange keeps working even when other PKCE flows are started before
+        this one completes.
         """
         await self._remove_session()
 
@@ -455,10 +467,10 @@ async def sign_in_with_oauth(
             params["redirect_to"] = redirect_to
         if scopes:
             params["scopes"] = scopes
-        url_with_qs, _ = await self._get_url_for_provider(
+        url_with_qs, _, flow_id = await self._get_url_for_provider(
             f"{self._url}/authorize", provider, params
         )
-        return OAuthResponse(provider=provider, url=url_with_qs)
+        return OAuthResponse(provider=provider, url=url_with_qs, flow_id=flow_id)
 
     async def link_identity(
         self, credentials: SignInWithOAuthCredentials
@@ -474,7 +486,7 @@ async def link_identity(
             params["scopes"] = scopes
         params["skip_http_redirect"] = "true"
         url = "user/identities/authorize"
-        _, query = await self._get_url_for_provider(url, provider, params)
+        _, query, flow_id = await self._get_url_for_provider(url, provider, params)
 
         session = await self.get_session()
         if not session:
@@ -487,7 +499,7 @@ async def link_identity(
             jwt=session.access_token,
         )
         link_identity = parse_link_identity_response(response)
-        return OAuthResponse(provider=provider, url=link_identity.url)
+        return OAuthResponse(provider=provider, url=link_identity.url, flow_id=flow_id)
 
     async def get_user_identities(self) -> IdentitiesResponse:
         response = await self.get_user()
@@ -801,6 +813,7 @@ async def sign_out(self, options: Optional[SignOutOptions] = None) -> None:
 
         if signout_options["scope"] != "others":
             await self._remove_session()
+            await self._pkce_verifier_store.remove_all()
             self._notify_all_subscribers("SIGNED_OUT", None)
 
     def on_auth_state_change(
@@ -1173,41 +1186,71 @@ async def _get_url_for_provider(
         url: str,
         provider: Provider,
         params: Dict[str, str],
-    ) -> Tuple[str, QueryParams]:
+    ) -> Tuple[str, QueryParams, Optional[str]]:
+        """
+        Build the authorize URL for ``provider``.
+
+        On the PKCE flow a fresh code verifier is generated and stored in a
+        slot of its own, identified by the returned flow id, so that several
+        flows can be pending at the same time. Returns ``None`` as the flow id
+        on the implicit flow.
+        """
         query = QueryParams(params)
+        flow_id: Optional[str] = None
         if self._flow_type == "pkce":
             code_verifier = generate_pkce_verifier()
             code_challenge = generate_pkce_challenge(code_verifier)
-            await self._storage.set_item(
-                f"{self._storage_key}-code-verifier", code_verifier
-            )
+            flow_id = generate_pkce_flow_id()
+            await self._pkce_verifier_store.store(flow_id, code_verifier)
             code_challenge_method = (
                 "plain" if code_verifier == code_challenge else "s256"
             )
             query = query.set("code_challenge", code_challenge).set(
                 "code_challenge_method", code_challenge_method
             )
         query = query.set("provider", provider)
-        return f"{url}?{query}", query
+        return f"{url}?{query}", query, flow_id
 
     async def exchange_code_for_session(
         self, params: CodeExchangeParams
     ) -> AuthResponse:
-        code_verifier = params.get("code_verifier") or await self._storage.get_item(
-            f"{self._storage_key}-code-verifier"
-        )
-        response = await self._request(
-            "POST",
-            "token",
-            query=QueryParams(grant_type="pkce"),
-            body={
-                "auth_code": params.get("auth_code"),
-                "code_verifier": code_verifier,
-            },
-            redirect_to=params.get("redirect_to"),
-        )
+        """
+        Exchange an authorization code obtained through the PKCE flow for a session.
+
+        Pass the ``flow_id`` returned by ``sign_in_with_oauth`` or
+        ``link_identity`` to use the code verifier of that specific flow. With a
+        flow id, only that flow's verifier is used; if it is missing the
+        exchange fails with ``AuthPKCECodeVerifierMissingError`` rather than
+        spending the single-use code on another flow's verifier. Without a
+        flow id the verifier of the most recently started flow is used.
+        """
+        flow_id: Optional[str] = None
+        requested_flow_id = params.get("flow_id")
+        if requested_flow_id is not None:
+            flow_id = validate_pkce_flow_id(requested_flow_id)
+            if flow_id is None:
+                raise AuthPKCECodeVerifierMissingError()
+
+        code_verifier = params.get("code_verifier")
+        if not code_verifier:
+            code_verifier = await self._pkce_verifier_store.retrieve(flow_id)
+            if flow_id is not None and code_verifier is None:
+                raise AuthPKCECodeVerifierMissingError()
+
+        try:
+            response = await self._request(
+                "POST",
+                "token",
+                query=QueryParams(grant_type="pkce"),
+                body={
+                    "auth_code": params.get("auth_code"),
+                    "code_verifier": code_verifier,
+                },
+                redirect_to=params.get("redirect_to"),
+            )
+        finally:
+            await self._pkce_verifier_store.remove(flow_id)
         auth_response = parse_auth_response(response)
-        await self._storage.remove_item(f"{self._storage_key}-code-verifier")
         if auth_response.session:
             await self._save_session(auth_response.session)
             self._notify_all_subscribers("SIGNED_IN", auth_response.session)

src/auth/src/supabase_auth/_async/pkce_verifier_store.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import json
+from asyncio import Lock
+from typing import List, Optional
+
+from ..constants import PKCE_MAX_CONCURRENT_FLOWS
+from ..helpers import (
+    pkce_flow_index_key,
+    pkce_legacy_verifier_key,
+    pkce_verifier_slot_key,
+    validate_pkce_flow_id,
+)
+from .storage import AsyncSupportedStorage
+
+
+class AsyncPKCEVerifierStore:
+    """Keeps one PKCE code verifier per pending flow.
+
+    Each flow gets its own storage slot keyed by a flow id, so starting a
+    second PKCE flow no longer overwrites the verifier of the first one. The
+    store is bounded: an ordered index of pending flow ids is kept in storage
+    (storage backends cannot enumerate keys) and the oldest flow is evicted
+    once more than ``PKCE_MAX_CONCURRENT_FLOWS`` are pending.
+
+    The single legacy key (``{storage_key}-code-verifier``) is still written
+    with the verifier of the most recently started flow and is used whenever a
+    caller does not know its flow id, which keeps older call sites working.
+    """
+
+    def __init__(self, storage: AsyncSupportedStorage, storage_key: str) -> None:
+        self._storage = storage
+        self._storage_key = storage_key
+        # The index is read, modified and written back across several awaits;
+        # without serialization two concurrent flow starts could each write a
+        # different index and one slot would become unreachable for eviction
+        # and clean-up.
+        self._lock = Lock()
+
+    @property
+    def legacy_key(self) -> str:
+        return pkce_legacy_verifier_key(self._storage_key)
+
+    @property
+    def index_key(self) -> str:
+        return pkce_flow_index_key(self._storage_key)
+
+    def slot_key(self, flow_id: str) -> str:
+        return pkce_verifier_slot_key(self._storage_key, flow_id)
+
+    async def _read_index(self) -> List[str]:
+        raw = await self._storage.get_item(self.index_key)
+        if not raw:
+            return []
+        try:
+            parsed = json.loads(raw)
+        except ValueError:
+            return []
+        if not isinstance(parsed, list):
+            return []
+        # Ids read back from storage are as untrusted as caller input.
+        return [
+            flow_id
+            for flow_id in (validate_pkce_flow_id(item) for item in parsed)
+            if flow_id is not None
+        ]
+
+    async def _write_index(self, index: List[str]) -> None:
+        if index:
+            await self._storage.set_item(self.index_key, json.dumps(index))
+        else:
+            await self._storage.remove_item(self.index_key)
+
+    async def store(self, flow_id: str, code_verifier: str) -> List[str]:
+        """Store ``code_verifier`` for ``flow_id`` and return any evicted flow ids.
+
+        Raises ``ValueError`` for a malformed flow id: a slot the index could
+        not validate on the way back out could never be evicted or removed.
+        """
+        if validate_pkce_flow_id(flow_id) is None:
+            raise ValueError("Invalid PKCE flow id")
+
+        async with self._lock:
+            await self._storage.set_item(self.slot_key(flow_id), code_verifier)
+
+            index = [item for item in await self._read_index() if item != flow_id]
+            index.append(flow_id)
+            evicted: List[str] = []
+            while len(index) > PKCE_MAX_CONCURRENT_FLOWS:
+                oldest = index.pop(0)
+                await self._storage.remove_item(self.slot_key(oldest))
+                evicted.append(oldest)
+            await self._write_index(index)
+
+            # Mirror the newest verifier into the legacy key for callers that
+            # never learned their flow id.
+            await self._storage.set_item(self.legacy_key, code_verifier)
+            return evicted
+
+    async def retrieve(self, flow_id: Optional[str]) -> Optional[str]:
+        """Return the verifier for ``flow_id``, or the legacy verifier if it is ``None``.
+
+        With a flow id only that flow's slot is consulted. Falling back to
+        another flow's verifier would spend the single-use auth code on a
+        request that cannot succeed, so a miss is reported as ``None`` instead.
+        """
+        if flow_id is None:
+            return await self._storage.get_item(self.legacy_key)
+        if validate_pkce_flow_id(flow_id) is None:
+            return None
+        return await self._storage.get_item(self.slot_key(flow_id))
+
+    async def remove(self, flow_id: Optional[str]) -> None:
+        """Forget the verifier of ``flow_id``, or only the legacy verifier if it is ``None``.
+
+        The legacy key is cleared alongside a slot only when it still holds
+        that slot's verifier, so removing a finished flow never destroys the
+        fallback of a newer pending flow.
+        """
+        async with self._lock:
+            if flow_id is None:
+                await self._storage.remove_item(self.legacy_key)
+                return
+            if validate_pkce_flow_id(flow_id) is None:
+                return
+
+            slot_key = self.slot_key(flow_id)
+            removed_verifier = await self._storage.get_item(slot_key)
+            await self._storage.remove_item(slot_key)
+
+            index = await self._read_index()
+            remaining = [item for item in index if item != flow_id]
+            if len(remaining) != len(index):
+                await self._write_index(remaining)
+
+            if removed_verifier is not None:
+                legacy_verifier = await self._storage.get_item(self.legacy_key)
+                if legacy_verifier == removed_verifier:
+                    await self._storage.remove_item(self.legacy_key)
+
+    async def remove_all(self) -> None:
+        """Forget every pending verifier, the index and the legacy verifier."""
+        async with self._lock:
+            for flow_id in await self._read_index():
+                await self._storage.remove_item(self.slot_key(flow_id))
+            await self._storage.remove_item(self.index_key)
+            await self._storage.remove_item(self.legacy_key)