Files
turnstone/tests/test_mcp_token_store.py
T
Patrick Buckley 29c42c1427 feat(mcp): per-(user, server) OAuth 2.1 + PKCE flow
Lands the OAuth flow that uses the token-at-rest store from the prior
commit: discovery (RFC 9728 PRM + RFC 8414 AS metadata with operator-
override precedence), PKCE S256 (mandatory — refuse AS without it),
RFC 8707 resource indicator on every authorize and token request,
RFC 7591 minimal one-shot dynamic client registration, authorization-
code exchange, refresh-token grant with re-read-after-acquire single-
flight lock, and the /v1/api/mcp/oauth/{start,callback} endpoints
mounted on both server and console.

Refactored:
- validate_url_no_ssrf, validate_discovered_endpoint, is_localhost,
  effective_port, sanitize_log_text moved out of oidc.py into a shared
  oauth_ssrf module; oidc.py re-exports for compatibility. The shared
  helpers also expose async wrappers (validate_url_no_ssrf_async,
  validate_discovered_endpoint_async) so OAuth-MCP discovery — invoked
  from async handlers — does not block the event loop on the
  synchronous socket.getaddrinfo call.
- MCPTokenStore.get_oauth_client_secret reader path added (the prior
  commit was write-only)
- Storage protocol gains create/pop/cleanup_*_mcp_oauth_pending_state
  and get_mcp_oauth_client_secret_ct (mirror OIDC pending-state
  pattern: SQLite BEGIN IMMEDIATE select-then-delete, Postgres atomic
  DELETE...RETURNING)

Refresh-grant correctness:
- When the AS omits refresh_token (RFC 6749 §6 — MAY rotate), the
  existing refresh value is preserved at the OAuth-flow layer rather
  than cleared, so production ASes (Google, Auth0 default, Okta) don't
  force re-consent every hour
- expires_in accepts int, float, str-with-decimal — earlier int-coerce
  through str() failed on float and silently dropped expiry tracking
- The refresh-grant `resource=` parameter (RFC 8707) is the canonical
  MCP server URL, not the audience. Audience and resource are distinct
  concepts; using audience as resource would mismatch the AS RS
  allowlist.

Audience handling:
- _validate_token_audience accepts str or tuple; the callback resolves
  accepted_audiences = {server_url, oauth_audience} and validates
  against the set, so Auth0-style ASes that honor `audience=` (not
  RFC 8707 `resource=`) issue tokens that pass audience-bound
  validation
- build_authorize_url emits both `resource=` (RFC 8707) and
  `audience=` (Auth0-style) per server config; comment documents which
  AS implementations need which form

Security hardening:
- redirect_uri pinned to oidc_config.redirect_base instead of the
  request Host header — closes the same Host-header injection PR #476
  fixed for OIDC. Both /start and /callback return 503 with operator-
  actionable hint when redirect_base is unset
- DCR registration runs under per-server asyncio.Lock with re-fetch
  inside the lock, so concurrent /start callers don't both register
  and overwrite each other's client_id (the second user's code is no
  longer rejected on callback)
- /callback error branch pops the pending state row before redirecting
  so a leaked state can't be replayed against a separately-obtained
  code in the 60s cleanup window
- WWW-Authenticate Bearer parser handles RFC 7235 quoted-string
  escapes (\" and \\) instead of the naive [^"]+ regex
- AS-controlled response bodies and error_description query params go
  through sanitize_log_text before reaching exception messages or
  audit details. AS error responses are parsed for the standard
  RFC 6749 fields (error, error_description, error_uri), each
  capped at 80 chars and run through redact_credentials to defend
  against ASes that echo the request body back into their error
  payload.
- oauth_as_issuer_cached is re-validated against the SSRF guard on
  read; on rejection the column is cleared and PRM rediscovery runs
- DCR / token-endpoint / refresh-endpoint response bodies cap at 64
  KiB (PRM/AS metadata cap stays at 256 KiB) so a hostile or
  malfunctioning AS can't exhaust client memory.
- oauth_client_secret operator input capped at 1024 chars at the
  admin-form boundary; longer plaintext rejected with 400.
- /start and /callback responses stamp `X-Frame-Options: DENY` so the
  redirected pages can't be framed by attacker sites.
- delete_user cascades to mcp_user_tokens and mcp_oauth_pending so
  user deletion no longer leaves dangling per-user OAuth state.
- Renaming or deleting an oauth_user MCP server purges per-user
  tokens and pending OAuth state for the previous server name
  (delete_mcp_oauth_rows_by_server_name). The OAuth tables key on the
  mutable server_name; without this purge, a future server with the
  same name (and an attacker-controlled URL) would silently rebind
  prior user tokens. A future schema migration will replace the
  server_name key with a server_id FK + ON DELETE CASCADE.
- get_user_access_token catches MCPTokenDecryptError (raised when no
  installed key can decrypt the row, e.g. after key rotation) and
  falls through to None so dispatch surfaces a re-consent rather than
  crashing.
- oauth_user MCP server rows are skipped in the static auto-connect
  path. Auto-connecting them at startup with empty headers fails the
  AS check and trips the circuit breaker; per-user tokens come online
  lazily once the user has consented.

Audit (mcp_server.oauth.* prefix):
- consent_started, consent_completed, consent_failed, token_refreshed,
  token_revoked, dcr_registered. _audit_event is async and wraps
  record_audit in asyncio.to_thread so the audit write doesn't block
  the event loop. resource_id on the audit row is the immutable
  server_id (PK UUID) so admin-driven server renames don't break
  event correlation; server_name is exposed in detail for cross-
  reference. dcr_registered detail.has_secret reflects whether the
  DCR-issued secret was actually persisted (the prior code reported
  has_secret=true even on persistence failure).
- _admin_mcp_action audits the immutable server_id, not the mutable
  server_name (which is what the column is — the table's PK was
  always server_id).
- All OAuth-flow log keys use the mcp_server.oauth.* prefix to match
  the audit-action taxonomy.

Lifespan close-order in turnstone.server and turnstone.console.server
is reversed (LIFO) — mcp_oauth → mcp_crypto → oidc — to match init
order.

Deferred until the upcoming per-user pool integration:
- Multi-node refresh-lock contention via pg_advisory_lock
- DCR re-register on token-endpoint 401 (the dispatch path surfaces
  those 401s)
- TTL-LRU caching of decrypted plaintext access tokens
- DNS-rebinding hardening (httpx Transport pin) — documented as
  limitation in oauth_ssrf module docstring

Tests: 7 new test files / ~85 new tests covering discovery precedence
+ PRM quoted-string parsing, PKCE round-trip, SSRF helper extraction,
authorize/callback handlers including 503-on-no-redirect-base + DCR
concurrency + JWT audience polymorphism + callback-error-pops-pending,
refresh single-flight lock, refresh resource-vs-audience regression,
decrypt-error fallthrough, _db_servers_to_config skipping oauth_user,
pending-state CRUD round-trip.
2026-05-04 22:00:23 -07:00

287 lines
10 KiB
Python

"""Tests for ``MCPTokenStore`` ciphertext-aware CRUD.
Phase 3 of the OAuth-MCP RFC: validates the encrypt/decrypt boundary
between :class:`MCPTokenStore` and the storage protocol's ciphertext-only
columns. Exercises the row-not-deleted-on-decrypt-failure invariant.
"""
from __future__ import annotations
import base64
import pytest
from cryptography.fernet import Fernet
from turnstone.core.mcp_crypto import (
MCPTokenCipher,
MCPTokenCipherConfig,
MCPTokenDecryptError,
MCPTokenStore,
)
def _make_cipher() -> MCPTokenCipher:
raw = base64.urlsafe_b64decode(Fernet.generate_key())
return MCPTokenCipher(MCPTokenCipherConfig(keys=(raw,)))
def _make_store(backend, *, audit: bool = False) -> tuple[MCPTokenStore, MCPTokenCipher]:
cipher = _make_cipher()
store = MCPTokenStore(
backend,
cipher,
node_id="test-node",
audit_storage=backend if audit else None,
)
return store, cipher
def _seed_server(backend, *, server_id: str = "srv-id-1", name: str = "srv-a") -> str:
backend.create_mcp_server(
server_id=server_id,
name=name,
transport="streamable-http",
url="https://mcp.example.com/sse",
auth_type="oauth_user",
)
return server_id
# ---------------------------------------------------------------------------
# User-token CRUD
# ---------------------------------------------------------------------------
class TestUserTokenCRUD:
def test_create_and_get_round_trip(self, backend) -> None:
store, _ = _make_store(backend)
store.create_user_token(
"u1",
"srv-a",
access_token="access-aaa",
refresh_token="refresh-bbb",
expires_at="2026-05-04T12:00:00",
scopes="openid profile",
as_issuer="https://auth.example.com",
audience="https://mcp.example.com",
)
plain = store.get_user_token("u1", "srv-a")
assert plain is not None
assert plain["user_id"] == "u1"
assert plain["server_name"] == "srv-a"
assert plain["access_token"] == "access-aaa"
assert plain["refresh_token"] == "refresh-bbb"
assert plain["scopes"] == "openid profile"
assert plain["audience"] == "https://mcp.example.com"
def test_create_with_no_refresh_token(self, backend) -> None:
store, _ = _make_store(backend)
store.create_user_token(
"u1",
"srv-a",
access_token="access-only",
refresh_token=None,
expires_at=None,
scopes=None,
as_issuer="https://auth.example.com",
audience="https://mcp.example.com",
)
plain = store.get_user_token("u1", "srv-a")
assert plain is not None
assert plain["access_token"] == "access-only"
assert plain["refresh_token"] is None
def test_get_missing_returns_none(self, backend) -> None:
store, _ = _make_store(backend)
assert store.get_user_token("nobody", "srv-a") is None
def test_update_after_refresh(self, backend) -> None:
store, _ = _make_store(backend)
store.create_user_token(
"u1",
"srv-a",
access_token="old-access",
refresh_token="old-refresh",
expires_at="2026-05-04T12:00:00",
scopes="openid",
as_issuer="https://auth.example.com",
audience="https://mcp.example.com",
)
ok = store.update_user_token_after_refresh(
"u1",
"srv-a",
access_token="new-access",
refresh_token="new-refresh",
expires_at="2026-05-04T13:00:00",
)
assert ok is True
plain = store.get_user_token("u1", "srv-a")
assert plain is not None
assert plain["access_token"] == "new-access"
assert plain["refresh_token"] == "new-refresh"
assert plain["expires_at"] == "2026-05-04T13:00:00"
# Preserved columns:
assert plain["scopes"] == "openid"
assert plain["as_issuer"] == "https://auth.example.com"
# last_refreshed got stamped:
assert plain["last_refreshed"] is not None
def test_update_after_refresh_missing_row_returns_false(self, backend) -> None:
store, _ = _make_store(backend)
ok = store.update_user_token_after_refresh(
"u1",
"srv-a",
access_token="x",
refresh_token=None,
expires_at=None,
)
assert ok is False
def test_delete(self, backend) -> None:
store, _ = _make_store(backend)
store.create_user_token(
"u1",
"srv-a",
access_token="a",
refresh_token=None,
expires_at=None,
scopes=None,
as_issuer="https://auth.example.com",
audience="https://mcp.example.com",
)
assert store.delete_user_token("u1", "srv-a") is True
assert store.get_user_token("u1", "srv-a") is None
# Idempotent: deleting again returns False.
assert store.delete_user_token("u1", "srv-a") is False
# ---------------------------------------------------------------------------
# Client-secret writer
# ---------------------------------------------------------------------------
class TestClientSecretWriter:
def test_set_oauth_client_secret_round_trip(self, backend) -> None:
store, cipher = _make_store(backend)
server_id = _seed_server(backend)
ok = store.set_oauth_client_secret(server_id, "plaintext-secret")
assert ok is True
# Read raw via get_mcp_server: ciphertext != plaintext, decrypts back.
raw = backend.get_mcp_server(server_id)
assert raw is not None
ct = raw["oauth_client_secret_ct"]
assert isinstance(ct, (bytes, bytearray, memoryview))
ct_bytes = bytes(ct)
assert ct_bytes != b"plaintext-secret"
assert cipher.decrypt(ct_bytes) == b"plaintext-secret"
def test_set_oauth_client_secret_clear_with_none(self, backend) -> None:
store, _ = _make_store(backend)
server_id = _seed_server(backend)
store.set_oauth_client_secret(server_id, "x")
assert store.set_oauth_client_secret(server_id, None) is True
raw = backend.get_mcp_server(server_id)
assert raw is not None
assert raw["oauth_client_secret_ct"] is None
def test_set_oauth_client_secret_missing_server_returns_false(self, backend) -> None:
store, _ = _make_store(backend)
ok = store.set_oauth_client_secret("does-not-exist", "x")
assert ok is False
# ---------------------------------------------------------------------------
# Decrypt failure: row preservation invariant
# ---------------------------------------------------------------------------
class TestDecryptFailureInvariant:
def test_get_user_token_with_wrong_key_raises_decrypt_error(self, backend) -> None:
"""CRITICAL: when no installed key can decrypt a stored row,
``get_user_token`` MUST NOT auto-delete the row. The row is
still valid; this node just doesn't have the right key.
"""
# Write under cipher A.
store_a, _cipher_a = _make_store(backend)
store_a.create_user_token(
"u1",
"srv-a",
access_token="secret-access",
refresh_token="secret-refresh",
expires_at="2026-05-04T12:00:00",
scopes="openid",
as_issuer="https://auth.example.com",
audience="https://mcp.example.com",
)
raw_before = backend.get_mcp_user_token("u1", "srv-a")
assert raw_before is not None
ct_before = bytes(raw_before["access_token_ct"])
# Read under cipher B (different key).
store_b, cipher_b = _make_store(backend)
with pytest.raises(MCPTokenDecryptError) as exc_info:
store_b.get_user_token("u1", "srv-a")
# The exception carries the keys we tried — useful for audit.
assert exc_info.value.key_fingerprints_attempted == cipher_b.key_fingerprints
# Row MUST still exist with ciphertext intact.
raw_after = backend.get_mcp_user_token("u1", "srv-a")
assert raw_after is not None
assert bytes(raw_after["access_token_ct"]) == ct_before
def test_decrypt_failure_emits_audit_when_configured(self, backend) -> None:
"""When ``audit_storage`` is set, decrypt failures emit a
``mcp_server.oauth.token_decrypt_failure`` audit event."""
store_a, _ = _make_store(backend)
store_a.create_user_token(
"u1",
"srv-a",
access_token="x",
refresh_token=None,
expires_at=None,
scopes=None,
as_issuer="https://a",
audience="https://m",
)
store_b, cipher_b = _make_store(backend, audit=True)
with pytest.raises(MCPTokenDecryptError):
store_b.get_user_token("u1", "srv-a")
events = backend.list_audit_events(limit=10)
actions = {ev.get("action") for ev in events}
assert "mcp_server.oauth.token_decrypt_failure" in actions
# ---------------------------------------------------------------------------
# Client-secret reader — q-9
# ---------------------------------------------------------------------------
class TestClientSecretReader:
def test_get_oauth_client_secret_returns_none_when_row_absent(self, backend) -> None:
store, _ = _make_store(backend)
assert store.get_oauth_client_secret("does-not-exist") is None
def test_get_oauth_client_secret_returns_none_when_column_null(self, backend) -> None:
store, _ = _make_store(backend)
server_id = _seed_server(backend)
# No set_oauth_client_secret call — column stays NULL.
assert store.get_oauth_client_secret(server_id) is None
def test_get_oauth_client_secret_round_trip(self, backend) -> None:
store, _ = _make_store(backend)
server_id = _seed_server(backend)
store.set_oauth_client_secret(server_id, "shhh-its-secret")
assert store.get_oauth_client_secret(server_id) == "shhh-its-secret"
def test_get_oauth_client_secret_raises_on_key_mismatch(self, backend) -> None:
store_a, _ = _make_store(backend)
server_id = _seed_server(backend)
store_a.set_oauth_client_secret(server_id, "secret-under-key-a")
# Cipher B has a different key — decrypt fails loudly.
store_b, _ = _make_store(backend)
with pytest.raises(MCPTokenDecryptError):
store_b.get_oauth_client_secret(server_id)