fix(auth): Preserve browser consent bindings across concurrent OAuth approvals

fastmcp_slim/fastmcp/server/auth/oauth_proxy/consent.py

@@ -50,6 +50,8 @@
 # rewriting them — the read-modify-write that concurrent renders race on.
 _MAX_CSRF_TOKENS = 10
 
+_CONSENT_BINDING_COOKIE_BASE = "MCP_CONSENT_BINDING"
+
 # Base name of the consent-state cookie. One cookie is set per issued CSRF
 # token (`MCP_CONSENT_STATE_<digest>`) rather than one list shared by all of
 # them: two renders in flight at once both build their Set-Cookie from the same
@@ -323,11 +325,11 @@ def _expire_cookie(
         )
 
     def _read_consent_bindings(self: OAuthProxy, request: Request) -> dict[str, str]:
-        """Read the consent binding map from the signed cookie.
+        """Read a pre-upgrade consent binding map from its signed cookie.
 
-        Returns a dict of {txn_id: consent_token} for all pending flows.
+        Kept temporarily so OAuth flows approved before an upgrade can finish.
         """
-        cookie_name = self._cookie_name("MCP_CONSENT_BINDING")
+        cookie_name = self._cookie_name(_CONSENT_BINDING_COOKIE_BASE)
         raw = request.cookies.get(cookie_name)
         # Only fall back to the non-__Host- name over plain HTTP. On HTTPS,
         # __Host- enforces host-only scope; accepting the weaker name would
@@ -347,64 +349,46 @@ def _read_consent_bindings(self: OAuthProxy, request: Request) -> dict[str, str]
             logger.debug("Failed to decode consent binding cookie")
         return {}
 
-    def _write_consent_bindings(
+    def _consent_binding_cookie_name(self: OAuthProxy, txn_id: str) -> str:
+        """Return an independent cookie name for one consent transaction."""
+        digest = hashlib.sha256(txn_id.encode()).hexdigest()[:32]
+        return self._cookie_name(f"{_CONSENT_BINDING_COOKIE_BASE}_{digest}")
+
+    def _set_consent_binding_cookie(
         self: OAuthProxy,
         response: HTMLResponse | RedirectResponse,
-        bindings: dict[str, str],
+        txn_id: str,
+        consent_token: str,
     ) -> None:
-        """Write the consent binding map to a signed cookie."""
-        name = self._cookie_name("MCP_CONSENT_BINDING")
-        if not bindings:
-            response.set_cookie(
-                name,
-                "",
-                max_age=0,
-                secure=self._is_https,
-                httponly=True,
-                samesite="lax",
-                path="/",
-            )
-            return
-        payload_bytes = json.dumps(bindings, separators=(",", ":")).encode()
-        payload_b64 = base64.b64encode(payload_bytes).decode()
-        signed_value = self._sign_cookie(payload_b64)
+        """Set an independently signed cookie for one consent transaction.
+
+        This cookie binds the browser that approved consent to the IdP callback,
+        ensuring a different browser cannot complete the OAuth flow. Its name
+        is derived from the transaction ID, so concurrent approvals from the
+        same cookie snapshot cannot overwrite one another.
+        """
+        name = self._consent_binding_cookie_name(txn_id)
         response.set_cookie(
             name,
-            signed_value,
+            self._sign_cookie(consent_token),
             max_age=15 * 60,
             secure=self._is_https,
             httponly=True,
             samesite="lax",
             path="/",
         )
 
-    def _set_consent_binding_cookie(
+    def _clear_consent_binding_cookie(
         self: OAuthProxy,
-        request: Request,
         response: HTMLResponse | RedirectResponse,
         txn_id: str,
-        consent_token: str,
     ) -> None:
-        """Add a consent binding entry for a transaction.
+        """Expire this transaction's binding after its successful callback.
 
-        This cookie binds the browser that approved consent to the IdP callback,
-        ensuring a different browser cannot complete the OAuth flow. Multiple
-        concurrent flows are supported by storing a map of txn_id → consent_token.
+        The legacy shared cookie is left to its 15-minute expiry: rewriting it
+        from a callback's cookie snapshot could erase another concurrent flow.
         """
-        bindings = self._read_consent_bindings(request)
-        bindings[txn_id] = consent_token
-        self._write_consent_bindings(response, bindings)
-
-    def _clear_consent_binding_cookie(
-        self: OAuthProxy,
-        request: Request,
-        response: HTMLResponse | RedirectResponse,
-        txn_id: str,
-    ) -> None:
-        """Remove a specific consent binding entry after successful callback."""
-        bindings = self._read_consent_bindings(request)
-        bindings.pop(txn_id, None)
-        self._write_consent_bindings(response, bindings)
+        self._expire_cookie(response, self._consent_binding_cookie_name(txn_id))
 
     def _verify_consent_binding_cookie(
         self: OAuthProxy,
@@ -413,11 +397,16 @@ def _verify_consent_binding_cookie(
         expected_token: str,
     ) -> bool:
         """Verify the consent binding for a specific transaction."""
-        bindings = self._read_consent_bindings(request)
-        actual = bindings.get(txn_id)
-        if not actual:
-            return False
-        return hmac.compare_digest(actual, expected_token)
+        name = self._consent_binding_cookie_name(txn_id)
+        raw = request.cookies.get(name)
+        if raw:
+            actual = self._verify_cookie(raw)
+            if actual and hmac.compare_digest(actual, expected_token):
+                return True
+
+        # Accept signed shared-map cookies for flows approved before an upgrade.
+        actual = self._read_consent_bindings(request).get(txn_id)
+        return bool(actual) and hmac.compare_digest(actual, expected_token)
 
     async def _handle_consent(
         self: OAuthProxy, request: Request
@@ -474,9 +463,7 @@ async def _show_consent_page(
                     )
                     upstream_url = self._build_upstream_authorize_url(txn_id, txn)
                     response = RedirectResponse(url=upstream_url, status_code=302)
-                    self._set_consent_binding_cookie(
-                        request, response, txn_id, consent_token
-                    )
+                    self._set_consent_binding_cookie(response, txn_id, consent_token)
                     return response
 
                 if client_key in denied:
@@ -706,7 +693,7 @@ async def _submit_consent(
             self._clear_consent_state_for_transaction(
                 request, response, txn_id, include_legacy=legacy_csrf
             )
-            self._set_consent_binding_cookie(request, response, txn_id, consent_token)
+            self._set_consent_binding_cookie(response, txn_id, consent_token)
             return response
 
         elif action == "deny":

fastmcp_slim/fastmcp/server/auth/oauth_proxy/proxy.py

@@ -2881,7 +2881,7 @@ async def _handle_idp_callback(
             logger.debug(f"Forwarding to client callback for transaction {txn_id}")
 
             response = RedirectResponse(url=client_callback_url, status_code=302)
-            self._clear_consent_binding_cookie(request, response, txn_id)
+            self._clear_consent_binding_cookie(response, txn_id)
             return response
 
         except Exception as e:

tests/server/auth/test_oauth_consent_page.py

@@ -3,6 +3,7 @@
 import re
 import secrets
 import time
+from http.cookies import SimpleCookie
 from unittest.mock import Mock
 from urllib.parse import parse_qs, urlparse
 
@@ -13,6 +14,8 @@
 from mcp_types import Icon
 from pydantic import AnyUrl
 from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import RedirectResponse
 from starlette.testclient import TestClient
 
 from fastmcp import FastMCP
@@ -100,6 +103,155 @@ def _extract_csrf(html: str) -> str | None:
     return m.group(1) if m else None
 
 
+def _request_with_cookies(cookies: dict[str, str]) -> Request:
+    cookie_header = "; ".join(f"{name}={value}" for name, value in cookies.items())
+    return Request(
+        {
+            "type": "http",
+            "method": "GET",
+            "path": "/",
+            "headers": [(b"cookie", cookie_header.encode())] if cookie_header else [],
+        }
+    )
+
+
+def _apply_set_cookies(cookies: dict[str, str], response) -> None:
+    for header in response.headers.getlist("set-cookie"):
+        parsed = SimpleCookie()
+        parsed.load(header)
+        for name, morsel in parsed.items():
+            if morsel["max-age"] == "0":
+                cookies.pop(name, None)
+            else:
+                cookies[name] = morsel.value
+
+
+def _set_cookie_value(response, name: str) -> str:
+    for header in response.headers.getlist("set-cookie"):
+        parsed = SimpleCookie()
+        parsed.load(header)
+        if name in parsed:
+            return parsed[name].value
+    raise AssertionError(f"Set-Cookie header for {name!r} was not found")
+
+
+def _cookie_max_age(response, name: str) -> str | None:
+    for header in response.headers.getlist("set-cookie"):
+        parsed = SimpleCookie()
+        parsed.load(header)
+        if name in parsed:
+            return parsed[name]["max-age"]
+    return None
+
+
+class TestConsentBindingCookies:
+    def test_concurrent_bindings_are_independent_and_order_independent(
+        self, oauth_proxy_https
+    ):
+        """Two responses from one cookie snapshot preserve both browser bindings."""
+        first = RedirectResponse("https://idp.example/authorize")
+        second = RedirectResponse("https://idp.example/authorize")
+        oauth_proxy_https._set_consent_binding_cookie(
+            first, "transaction-one", "token-one"
+        )
+        oauth_proxy_https._set_consent_binding_cookie(
+            second, "transaction-two", "token-two"
+        )
+
+        first_name = oauth_proxy_https._consent_binding_cookie_name("transaction-one")
+        second_name = oauth_proxy_https._consent_binding_cookie_name("transaction-two")
+        first_value = _set_cookie_value(first, first_name)
+        assert first_name != second_name
+
+        for responses in ((first, second), (second, first)):
+            browser_cookies: dict[str, str] = {}
+            for response in responses:
+                _apply_set_cookies(browser_cookies, response)
+
+            request = _request_with_cookies(browser_cookies)
+            assert oauth_proxy_https._verify_consent_binding_cookie(
+                request, "transaction-one", "token-one"
+            )
+            assert oauth_proxy_https._verify_consent_binding_cookie(
+                request, "transaction-two", "token-two"
+            )
+            assert not oauth_proxy_https._verify_consent_binding_cookie(
+                _request_with_cookies({}), "transaction-one", "token-one"
+            )
+
+            tampered = dict(browser_cookies)
+            tampered[first_name] = (
+                "A" if first_value[0] != "A" else "B"
+            ) + first_value[1:]
+            assert not oauth_proxy_https._verify_consent_binding_cookie(
+                _request_with_cookies(tampered), "transaction-one", "token-one"
+            )
+
+            cleared = RedirectResponse("https://client.example/callback")
+            oauth_proxy_https._clear_consent_binding_cookie(cleared, "transaction-one")
+            assert _cookie_max_age(cleared, first_name) == "0"
+            assert _cookie_max_age(cleared, second_name) is None
+
+    def test_concurrent_bindings_from_one_snapshot_fit_cookie_header(
+        self, oauth_proxy_https
+    ):
+        """Concurrent Set-Cookie responses keep every compact binding."""
+        bindings = {
+            f"{index:032x}-0000-0000-0000-000000000000": secrets.token_urlsafe(32)
+            for index in range(20)
+        }
+        responses = []
+        for txn_id, token in bindings.items():
+            response = RedirectResponse("https://idp.example/authorize")
+            oauth_proxy_https._set_consent_binding_cookie(
+                response,
+                txn_id,
+                token,
+            )
+            responses.append(response)
+
+        browser_cookies: dict[str, str] = {}
+        for response in responses:
+            _apply_set_cookies(browser_cookies, response)
+
+        assert len(browser_cookies) == len(bindings)
+        cookie_header = "; ".join(
+            f"{name}={value}" for name, value in browser_cookies.items()
+        )
+        assert len(cookie_header.encode()) < 4096
+
+        request = _request_with_cookies(browser_cookies)
+        for txn_id, token in bindings.items():
+            assert oauth_proxy_https._verify_consent_binding_cookie(
+                request, txn_id, token
+            )
+
+    def test_legacy_binding_cookie_remains_valid_during_upgrade(
+        self, oauth_proxy_https
+    ):
+        """An in-flight flow using the old shared cookie can still finish."""
+        import base64
+        import json
+
+        bindings = {"legacy-one": "token-one", "legacy-two": "token-two"}
+        encoded = base64.b64encode(
+            json.dumps(bindings, separators=(",", ":")).encode()
+        ).decode()
+        legacy_value = oauth_proxy_https._sign_cookie(encoded)
+        request = _request_with_cookies({"__Host-MCP_CONSENT_BINDING": legacy_value})
+
+        assert oauth_proxy_https._verify_consent_binding_cookie(
+            request, "legacy-one", "token-one"
+        )
+        assert oauth_proxy_https._verify_consent_binding_cookie(
+            request, "legacy-two", "token-two"
+        )
+
+        cleared = RedirectResponse("https://client.example/callback")
+        oauth_proxy_https._clear_consent_binding_cookie(cleared, "legacy-one")
+        assert _cookie_max_age(cleared, "__Host-MCP_CONSENT_BINDING") is None
+
+
 class TestConsentPageServerIcon:
     """Tests for server icon display in OAuth consent screen."""