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.

(cherry picked from commit 29c42c1427)
This commit is contained in:
Patrick Buckley
2026-05-04 20:05:01 -07:00
parent a4c335d7bf
commit b0f7029ff1
25 changed files with 5185 additions and 164 deletions
+18
View File
@@ -8,9 +8,27 @@ import pytest
if TYPE_CHECKING:
from turnstone.core.mcp_client import MCPClientManager, StaticServerState
from turnstone.core.mcp_crypto import MCPTokenCipher
from turnstone.core.oidc import OIDCConfig
def make_mcp_token_cipher() -> MCPTokenCipher:
"""Build a single-key MCP token cipher for tests.
Used by test files that need to exercise ``MCPTokenStore`` round-
trips without the lifespan-side configuration loader; centralised
here so the key/material defaults stay aligned across files.
"""
import base64
from cryptography.fernet import Fernet
from turnstone.core.mcp_crypto import MCPTokenCipher, MCPTokenCipherConfig
raw = base64.urlsafe_b64decode(Fernet.generate_key())
return MCPTokenCipher(MCPTokenCipherConfig(keys=(raw,)))
def _seed_static_state(mgr: MCPClientManager, name: str, **overrides: Any) -> StaticServerState:
"""Get-or-create a ``StaticServerState`` on ``mgr`` and apply ``overrides``.
+6 -6
View File
@@ -355,9 +355,9 @@ class TestCreateMcpServer:
assert "name" in r.json()["error"].lower()
def test_admin_create_oauth_server(self, client):
"""Phase 3: admin can POST a server with auth_type=oauth_user;
the seven OAuth text fields round-trip via GET and the plaintext
client secret is encrypted-at-rest via the dedicated writer."""
"""Admin can POST a server with auth_type=oauth_user; the seven
OAuth text fields round-trip via GET and the plaintext client
secret is encrypted-at-rest via the dedicated writer."""
r = client.post(
"/v1/api/admin/mcp-servers",
json={
@@ -381,7 +381,7 @@ class TestCreateMcpServer:
assert data["oauth_audience"] == "https://mcp.example.com"
assert data["oauth_registration_mode"] == "preregistered"
assert data["oauth_authorization_server_url"] == "https://auth.example.com"
# Phase 3: ciphertext is persisted; the response masks it to "***".
# Ciphertext is persisted; the response masks it to "***".
assert data["oauth_client_secret_ct"] == "***"
def test_admin_create_oauth_server_without_token_store_returns_503(self, client_no_token_store):
@@ -537,8 +537,8 @@ class TestUpdateMcpServer:
assert "transport" in r.json()["error"].lower()
def test_admin_update_auth_type_static_to_oauth(self, client):
"""Phase 2: an existing static row can be flipped to oauth_user
with OAuth fields supplied alongside."""
"""An existing static row can be flipped to oauth_user with
OAuth fields supplied alongside."""
created = _create_server(
client,
name="flip-to-oauth",
+56
View File
@@ -15,6 +15,7 @@ import pytest
from tests.conftest import _seed_static_state
from turnstone.core.mcp_client import (
MCPClientManager,
_db_servers_to_config,
_mcp_to_openai,
load_mcp_config,
)
@@ -261,6 +262,61 @@ class TestLoadMcpConfig:
assert result == {}
class TestDBServersToConfig:
"""``_db_servers_to_config`` shapes DB rows for the MCP client."""
def test_static_streamable_http_row_passes_through(self) -> None:
rows = [
{
"name": "static-srv",
"transport": "streamable-http",
"url": "https://mcp.example.com",
"headers": '{"Authorization": "Bearer token"}',
"auth_type": "static",
}
]
result = _db_servers_to_config(rows)
assert "static-srv" in result
assert result["static-srv"]["url"] == "https://mcp.example.com"
assert result["static-srv"]["headers"] == {"Authorization": "Bearer token"}
def test_db_servers_to_config_skips_oauth_user_rows(self) -> None:
"""Rows with auth_type=oauth_user must be invisible to the static
auto-connect path.
Auto-connecting these with empty headers fails the AS check and
trips the circuit breaker on startup. Per-user OAuth servers
come online lazily once the user has consented.
"""
rows = [
{
"name": "static-srv",
"transport": "streamable-http",
"url": "https://static.example.com",
"headers": "{}",
"auth_type": "static",
},
{
"name": "oauth-srv",
"transport": "streamable-http",
"url": "https://oauth.example.com",
"headers": "{}",
"auth_type": "oauth_user",
},
{
"name": "stdio-srv",
"transport": "stdio",
"command": "echo",
"args": "[]",
"env": "{}",
"auth_type": "none",
},
]
result = _db_servers_to_config(rows)
assert set(result) == {"static-srv", "stdio-srv"}
assert "oauth-srv" not in result
# ---------------------------------------------------------------------------
# merge_mcp_tools
# ---------------------------------------------------------------------------
+1 -2
View File
@@ -1,7 +1,6 @@
"""Tests for ``turnstone.core.mcp_crypto`` cipher + config loading.
Covers Phase 3 of the OAuth-MCP RFC: token-at-rest encryption.
See ``docs/design/oauth-mcp.md`` §5.3.
Covers token-at-rest encryption for OAuth-MCP.
"""
from __future__ import annotations
+546
View File
@@ -0,0 +1,546 @@
"""Discovery tests for the per-(user, server) MCP OAuth flow.
Covers PRM (RFC 9728) and AS metadata (RFC 8414) discovery, including:
- override URL takes precedence
- PRM happy path: server URL -> .well-known/oauth-protected-resource
-> ``authorization_servers[0]``
- PRM 401 + ``WWW-Authenticate: Bearer resource_metadata="..."`` follows
the URL.
- AS metadata without S256 -> :class:`MCPOAuthDiscoveryError`.
- SSRF rejection on AS issuer URL.
- In-memory cache hit/miss + persistent cache write to
``mcp_servers.oauth_as_issuer_cached``.
"""
from __future__ import annotations
import asyncio
import time
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from turnstone.core.mcp_oauth import (
ASMetadata,
MCPOAuthDiscoveryError,
_parse_prm_url_from_www_authenticate,
discover_authorization_server,
)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _mk_response(
status_code: int = 200,
json_body: Any = None,
headers: dict[str, str] | None = None,
) -> MagicMock:
"""Build a MagicMock that quacks like ``httpx.Response``."""
resp = MagicMock(spec=httpx.Response)
resp.status_code = status_code
resp.headers = headers or {}
resp.content = (str(json_body) if json_body is not None else "").encode("utf-8")
if json_body is not None:
resp.json.return_value = json_body
else:
resp.json.side_effect = ValueError("no body")
resp.text = str(json_body) if json_body is not None else ""
return resp
def _good_as_metadata_doc() -> dict[str, Any]:
return {
"issuer": "https://as.example.com",
"authorization_endpoint": "https://as.example.com/authorize",
"token_endpoint": "https://as.example.com/token",
"jwks_uri": "https://as.example.com/jwks",
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["none", "client_secret_basic"],
"registration_endpoint": "https://as.example.com/register",
}
def _public_addr_patch():
return patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("93.184.216.34", 0))])
def _mk_storage_mock(server_id: str = "srv-id") -> MagicMock:
storage = MagicMock()
storage.update_mcp_server.return_value = True
return storage
# ---------------------------------------------------------------------------
# PRM parser
# ---------------------------------------------------------------------------
class TestParsePRMUrl:
def test_extracts_resource_metadata_url(self) -> None:
header = (
'Bearer error="invalid_token", '
'resource_metadata="https://srv.example.com/.well-known/oauth-protected-resource"'
)
url = _parse_prm_url_from_www_authenticate(header)
assert url == "https://srv.example.com/.well-known/oauth-protected-resource"
def test_returns_none_when_absent(self) -> None:
assert _parse_prm_url_from_www_authenticate('Bearer realm="x"') is None
def test_handles_empty_header(self) -> None:
assert _parse_prm_url_from_www_authenticate("") is None
def test_handles_escaped_quote_in_value(self) -> None:
"""RFC 7230 quoted-string allows ``\\"`` — naive ``[^"]+`` truncates.
A malicious or buggy resource server could send an embedded
escaped quote; the parser must yield the unescaped value, not
the prefix up to the escaped quote.
"""
header = 'Bearer resource_metadata="https://srv.example.com/with\\"quote"'
url = _parse_prm_url_from_www_authenticate(header)
assert url == 'https://srv.example.com/with"quote'
def test_handles_escaped_backslash(self) -> None:
header = 'Bearer resource_metadata="https://srv.example.com/back\\\\slash"'
url = _parse_prm_url_from_www_authenticate(header)
assert url == "https://srv.example.com/back\\slash"
def test_unterminated_quoted_string_returns_none(self) -> None:
# Closing quote missing — naive regex would still match, but
# the proper parser should reject malformed input.
header = 'Bearer resource_metadata="https://srv.example.com/no-close'
assert _parse_prm_url_from_www_authenticate(header) is None
# ---------------------------------------------------------------------------
# discover_authorization_server happy paths
# ---------------------------------------------------------------------------
class TestDiscoveryOverride:
def test_override_url_skips_prm(self) -> None:
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
return await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url="https://as.example.com",
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
meta = asyncio.run(_run())
assert isinstance(meta, ASMetadata)
assert meta.token_endpoint == "https://as.example.com/token"
# Only the AS metadata URL was hit, not PRM.
called_urls = [c.args[0] for c in client.get.call_args_list]
assert all("oauth-authorization-server" in u for u in called_urls)
class TestDiscoveryPRM:
def test_prm_happy_path(self) -> None:
async def _get(url, *args, **kwargs):
if url.endswith("/oauth-protected-resource"):
return _mk_response(
200,
{
"resource": "https://mcp.example.com",
"authorization_servers": ["https://as.example.com"],
},
)
if url.endswith("/oauth-authorization-server"):
return _mk_response(200, _good_as_metadata_doc())
raise AssertionError(f"unexpected URL: {url}")
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(side_effect=_get)
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
return await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url=None,
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
meta = asyncio.run(_run())
assert meta.issuer == "https://as.example.com"
def test_prm_401_follows_www_authenticate(self) -> None:
async def _get(url, *args, **kwargs):
if url == "https://mcp.example.com/.well-known/oauth-protected-resource":
return _mk_response(
401,
headers={
"www-authenticate": (
'Bearer error="invalid_token", '
"resource_metadata="
'"https://meta.example.com/prm"'
)
},
json_body=None,
)
if url == "https://meta.example.com/prm":
return _mk_response(
200,
{
"authorization_servers": ["https://as.example.com"],
},
)
if url.endswith("/oauth-authorization-server"):
return _mk_response(200, _good_as_metadata_doc())
raise AssertionError(f"unexpected URL: {url}")
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(side_effect=_get)
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
return await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url=None,
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
meta = asyncio.run(_run())
assert meta.token_endpoint == "https://as.example.com/token"
def test_prm_401_without_resource_metadata_raises(self) -> None:
async def _get(url, *args, **kwargs):
return _mk_response(401, headers={"www-authenticate": "Basic realm=x"})
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(side_effect=_get)
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url=None,
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
with pytest.raises(MCPOAuthDiscoveryError, match="resource_metadata"):
asyncio.run(_run())
# ---------------------------------------------------------------------------
# AS metadata validation
# ---------------------------------------------------------------------------
class TestASMetadataValidation:
def test_no_s256_raises(self) -> None:
doc = _good_as_metadata_doc()
doc["code_challenge_methods_supported"] = ["plain"]
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, doc))
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url="https://as.example.com",
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
with pytest.raises(MCPOAuthDiscoveryError, match="S256"):
asyncio.run(_run())
def test_missing_endpoints_raises(self) -> None:
doc = _good_as_metadata_doc()
del doc["token_endpoint"]
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, doc))
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url="https://as.example.com",
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
with pytest.raises(MCPOAuthDiscoveryError, match="missing required"):
asyncio.run(_run())
def test_third_party_endpoint_rejected(self) -> None:
doc = _good_as_metadata_doc()
doc["token_endpoint"] = "https://attacker.example.com/token"
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, doc))
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url="https://as.example.com",
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
with pytest.raises(MCPOAuthDiscoveryError, match="token_endpoint"):
asyncio.run(_run())
def test_ssrf_on_override_rejected(self) -> None:
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock()
storage = _mk_storage_mock()
async def _run():
# Resolve to private 10.x — SSRF guard fires before any HTTP call.
with patch(
"socket.getaddrinfo",
return_value=[(2, 1, 6, "", ("10.0.0.1", 0))],
):
await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url="https://internal.corp.example.com",
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
with pytest.raises(MCPOAuthDiscoveryError):
asyncio.run(_run())
client.get.assert_not_called()
# ---------------------------------------------------------------------------
# Caching
# ---------------------------------------------------------------------------
class TestMetadataCache:
def test_cache_miss_then_hit(self) -> None:
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
storage = _mk_storage_mock()
cache: dict[str, tuple[ASMetadata, float]] = {}
async def _run():
with _public_addr_patch():
first = await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url="https://as.example.com",
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
metadata_cache=cache,
)
second = await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url="https://as.example.com",
cached_issuer="https://as.example.com",
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
metadata_cache=cache,
)
return first, second
first, second = asyncio.run(_run())
assert first.token_endpoint == second.token_endpoint
# First call hit AS metadata; second call hit the cache.
assert client.get.call_count == 1
def test_cache_expiry_refetches(self) -> None:
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
storage = _mk_storage_mock()
# Pre-populate cache with a very stale entry.
stale_meta = ASMetadata(
issuer="https://as.example.com",
authorization_endpoint="https://as.example.com/authorize",
token_endpoint="https://as.example.com/token",
registration_endpoint=None,
jwks_uri=None,
code_challenge_methods_supported=("S256",),
token_endpoint_auth_methods_supported=(),
)
cache = {"https://as.example.com": (stale_meta, time.monotonic() - 10**6)}
async def _run():
with _public_addr_patch():
return await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url="https://as.example.com",
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
metadata_cache=cache,
)
meta = asyncio.run(_run())
# Stale entry was bypassed -> we hit the network.
assert client.get.call_count == 1
assert meta.token_endpoint == "https://as.example.com/token"
def test_persistent_cache_write_on_first_resolution(self) -> None:
async def _get(url, *args, **kwargs):
if url.endswith("/oauth-protected-resource"):
return _mk_response(200, {"authorization_servers": ["https://as.example.com"]})
return _mk_response(200, _good_as_metadata_doc())
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(side_effect=_get)
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
return await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url=None,
cached_issuer=None,
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
asyncio.run(_run())
# update_mcp_server was called once with the cached issuer.
storage.update_mcp_server.assert_called_once_with(
"srv-id", oauth_as_issuer_cached="https://as.example.com"
)
def test_persistent_cache_skip_when_already_cached(self) -> None:
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
storage = _mk_storage_mock()
async def _run():
with _public_addr_patch():
return await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url=None,
cached_issuer="https://as.example.com",
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
asyncio.run(_run())
storage.update_mcp_server.assert_not_called()
# ---------------------------------------------------------------------------
# sec-3 — cached_issuer re-validated on read
# ---------------------------------------------------------------------------
class TestCachedIssuerSSRFRevalidation:
"""A cached issuer URL must still pass SSRF validation on every read.
Defense-in-depth: an admin who points ``oauth_as_issuer_cached`` at a
private address (or a hostname that has rebound to one) should not
bypass the guard just because the value was already in the row.
"""
def test_cached_issuer_rejected_clears_row_and_falls_through_to_prm(self) -> None:
async def _get(url: str, *args: Any, **kwargs: Any) -> MagicMock:
if url.endswith("/oauth-protected-resource"):
return _mk_response(200, {"authorization_servers": ["https://as.example.com"]})
if url.endswith("/oauth-authorization-server"):
return _mk_response(200, _good_as_metadata_doc())
raise AssertionError(f"unexpected URL: {url}")
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(side_effect=_get)
storage = _mk_storage_mock()
# cached_issuer points at a private host. SSRF guard fires on
# the cached value first, the row is cleared, and PRM
# discovery runs as a fallback.
async def _run() -> Any:
with patch(
"socket.getaddrinfo",
# Private resolution for "internal.corp", public for everything else.
side_effect=lambda host, *a, **kw: [
(2, 1, 6, "", ("10.0.0.1" if "internal" in host else "93.184.216.34", 0))
],
):
return await discover_authorization_server(
server_name="srv-x",
server_url="https://mcp.example.com/sse",
override_url=None,
cached_issuer="https://internal.corp.example.com",
http_client=client,
storage=storage,
server_id="srv-id",
trusted_hosts=frozenset(),
)
meta = asyncio.run(_run())
assert meta.token_endpoint == "https://as.example.com/token"
# The bad cached_issuer was cleared from the row.
clear_calls = [
c
for c in storage.update_mcp_server.call_args_list
if c.kwargs.get("oauth_as_issuer_cached") is None
]
assert clear_calls, "cached_issuer should have been cleared"
+870
View File
@@ -0,0 +1,870 @@
"""Integration tests for the MCP OAuth HTTP handlers.
Mirrors the structure of ``tests/test_oidc_handlers.py``: builds a small
Starlette app with the ``/api/mcp/oauth/start`` and ``/api/mcp/oauth/callback``
routes wired in, mocks the ``httpx.AsyncClient`` calls into the AS, and
exercises both happy and failure paths.
"""
from __future__ import annotations
import asyncio
import base64
import json
import urllib.parse
from typing import TYPE_CHECKING, Any
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.routing import Mount, Route
from starlette.testclient import TestClient
from tests.conftest import make_mcp_token_cipher
from turnstone.core.auth import AuthResult
from turnstone.core.mcp_crypto import MCPTokenStore
from turnstone.core.mcp_oauth import (
handle_mcp_oauth_authorize,
handle_mcp_oauth_callback,
)
from turnstone.core.oidc import OIDCConfig
from turnstone.core.storage._sqlite import SQLiteBackend
if TYPE_CHECKING:
from starlette.requests import Request
from starlette.responses import Response
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class _InjectAuthMiddleware(BaseHTTPMiddleware):
"""Stamp a fixed authenticated user on every request."""
async def dispatch(self, request: Request, call_next: Any) -> Response:
request.state.auth_result = AuthResult(
user_id="user-1",
scopes=frozenset({"write"}),
token_source="config",
permissions=frozenset({"read", "write"}),
)
return await call_next(request)
async def _mcp_authorize(request: Request) -> Response:
return await handle_mcp_oauth_authorize(request)
async def _mcp_callback(request: Request) -> Response:
return await handle_mcp_oauth_callback(request)
def _build_app(
*,
storage: SQLiteBackend,
http_client: httpx.AsyncClient,
token_store: MCPTokenStore | None,
redirect_base: str = "https://testserver",
) -> Starlette:
app = Starlette(
routes=[
Mount(
"/v1",
routes=[
Route("/api/mcp/oauth/start", _mcp_authorize),
Route("/api/mcp/oauth/callback", _mcp_callback),
],
),
],
middleware=[Middleware(_InjectAuthMiddleware)],
)
app.state.auth_storage = storage
app.state.mcp_token_store = token_store
app.state.mcp_oauth_http_client = http_client
app.state.mcp_oauth_refresh_locks = {}
app.state.mcp_oauth_dcr_locks = {}
app.state.mcp_oauth_metadata_cache = {}
app.state.mcp_oauth_last_cleanup_monotonic = 0.0
# Mirror the OIDC redirect_base contract — the MCP OAuth handlers
# reuse it to pin the callback URL against Host-header injection.
app.state.oidc_config = OIDCConfig(
enabled=False,
redirect_base=redirect_base,
)
return app
def _make_token_store(backend: SQLiteBackend) -> MCPTokenStore:
return MCPTokenStore(backend, make_mcp_token_cipher(), node_id="test")
def _seed_oauth_user_server(
backend: SQLiteBackend,
*,
name: str = "srv-oauth",
server_id: str = "srv-id-1",
client_id: str | None = "client-abc",
cached_issuer: str | None = "https://as.example.com",
registration_mode: str | None = None,
) -> str:
backend.create_mcp_server(
server_id=server_id,
name=name,
transport="streamable-http",
url="https://mcp.example.com/sse",
auth_type="oauth_user",
oauth_client_id=client_id,
oauth_scopes="openid profile",
oauth_audience="https://mcp.example.com",
oauth_authorization_server_url=None,
oauth_registration_mode=registration_mode,
)
if cached_issuer is not None:
backend.update_mcp_server(server_id, oauth_as_issuer_cached=cached_issuer)
return server_id
def _good_as_metadata_doc() -> dict[str, Any]:
return {
"issuer": "https://as.example.com",
"authorization_endpoint": "https://as.example.com/authorize",
"token_endpoint": "https://as.example.com/token",
"registration_endpoint": "https://as.example.com/register",
"jwks_uri": "https://as.example.com/jwks",
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["none", "client_secret_basic"],
}
def _mk_response(
status_code: int = 200,
json_body: Any = None,
headers: dict[str, str] | None = None,
) -> MagicMock:
resp = MagicMock(spec=httpx.Response)
resp.status_code = status_code
resp.headers = headers or {}
body_str = (
json.dumps(json_body) if json_body is not None and not isinstance(json_body, str) else ""
)
resp.content = body_str.encode("utf-8")
if json_body is not None:
resp.json.return_value = json_body
else:
resp.json.side_effect = ValueError("no body")
resp.text = body_str
return resp
def _public_addr_patch():
return patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("93.184.216.34", 0))])
@pytest.fixture
def storage(tmp_path: Any) -> SQLiteBackend:
backend = SQLiteBackend(str(tmp_path / "test.db"))
backend.create_user("user-1", "user1", "User One", "hash")
return backend
@pytest.fixture
def http_client_mock() -> MagicMock:
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock()
client.post = AsyncMock()
return client
# ---------------------------------------------------------------------------
# /start
# ---------------------------------------------------------------------------
class TestAuthorize:
def test_happy_path_redirects_to_as(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
token_store = _make_token_store(storage)
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/start?server=srv-oauth&return_url=/admin/mcp-servers",
follow_redirects=False,
)
assert resp.status_code == 302
location = resp.headers["location"]
parsed = urllib.parse.urlparse(location)
assert parsed.scheme == "https"
assert parsed.netloc == "as.example.com"
params = urllib.parse.parse_qs(parsed.query)
assert params["client_id"] == ["client-abc"]
assert params["response_type"] == ["code"]
assert params["code_challenge_method"] == ["S256"]
assert "code_challenge" in params
assert "state" in params
assert params["resource"] == ["https://mcp.example.com/sse"]
def test_unknown_server_404(self, storage: SQLiteBackend, http_client_mock: MagicMock) -> None:
token_store = _make_token_store(storage)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get("/v1/api/mcp/oauth/start?server=nope")
assert resp.status_code == 404
def test_wrong_auth_type_400(self, storage: SQLiteBackend, http_client_mock: MagicMock) -> None:
storage.create_mcp_server(
server_id="static-1",
name="srv-static",
transport="streamable-http",
url="https://x.example.com",
auth_type="static",
)
token_store = _make_token_store(storage)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get("/v1/api/mcp/oauth/start?server=srv-static")
assert resp.status_code == 400
assert "per-user OAuth" in resp.json()["error"]
def test_missing_server_param_400(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
token_store = _make_token_store(storage)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get("/v1/api/mcp/oauth/start")
assert resp.status_code == 400
def test_as_without_s256_redirects_with_error(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
token_store = _make_token_store(storage)
bad_doc = _good_as_metadata_doc()
bad_doc["code_challenge_methods_supported"] = ["plain"]
http_client_mock.get.return_value = _mk_response(200, bad_doc)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get("/v1/api/mcp/oauth/start?server=srv-oauth", follow_redirects=False)
assert resp.status_code == 502
assert "S256" in resp.json()["error"]
def test_pending_state_persisted(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
token_store = _make_token_store(storage)
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get("/v1/api/mcp/oauth/start?server=srv-oauth", follow_redirects=False)
location = resp.headers["location"]
params = urllib.parse.parse_qs(urllib.parse.urlparse(location).query)
state = params["state"][0]
pending = storage.pop_mcp_oauth_pending_state(state)
assert pending is not None
assert pending["user_id"] == "user-1"
assert pending["server_name"] == "srv-oauth"
def test_return_url_cross_origin_falls_back(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
token_store = _make_token_store(storage)
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/start"
"?server=srv-oauth&return_url=https://attacker.example.com/x",
follow_redirects=False,
)
assert resp.status_code == 302
location = resp.headers["location"]
params = urllib.parse.parse_qs(urllib.parse.urlparse(location).query)
# Pull pending and verify the return_url was sanitised to "/".
pending = storage.pop_mcp_oauth_pending_state(params["state"][0])
assert pending is not None
assert pending["return_url"] == "/"
# ---------------------------------------------------------------------------
# /callback
# ---------------------------------------------------------------------------
class TestCallback:
def _seed_pending(
self,
storage: SQLiteBackend,
*,
state: str = "valid-state",
user_id: str = "user-1",
server_name: str = "srv-oauth",
verifier: str = "verifier-blob",
return_url: str = "/admin/mcp-servers",
) -> None:
storage.create_mcp_oauth_pending_state(state, user_id, server_name, verifier, return_url)
def test_happy_path_persists_token_and_redirects(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
self._seed_pending(storage)
token_store = _make_token_store(storage)
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.return_value = _mk_response(
200,
{
"access_token": "opaque-access-aaa",
"refresh_token": "refresh-bbb",
"token_type": "Bearer",
"expires_in": 3600,
"scope": "openid profile",
},
)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/callback?code=auth-code&state=valid-state",
follow_redirects=False,
)
assert resp.status_code == 302
assert resp.headers["location"] == "/admin/mcp-servers"
# Token row was persisted.
plain = token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
assert plain["access_token"] == "opaque-access-aaa"
assert plain["refresh_token"] == "refresh-bbb"
def test_state_mismatch_redirects(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
token_store = _make_token_store(storage)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get(
"/v1/api/mcp/oauth/callback?code=auth-code&state=does-not-exist",
follow_redirects=False,
)
assert resp.status_code == 302
assert "session+expired" in resp.headers["location"]
def test_user_id_mismatch_redirects(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
# Pending row attributes the flow to a different user.
self._seed_pending(storage, user_id="other-user")
token_store = _make_token_store(storage)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get(
"/v1/api/mcp/oauth/callback?code=c&state=valid-state",
follow_redirects=False,
)
assert resp.status_code == 302
assert "user+mismatch" in resp.headers["location"]
def test_as_error_redirects_with_message(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
token_store = _make_token_store(storage)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get(
"/v1/api/mcp/oauth/callback?error=access_denied&error_description=user+cancelled",
follow_redirects=False,
)
assert resp.status_code == 302
assert "mcp_oauth_error" in resp.headers["location"]
def test_jwt_audience_mismatch_redirects(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
self._seed_pending(storage)
token_store = _make_token_store(storage)
# Build a JWT with a wrong audience.
header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode()
payload = (
base64.urlsafe_b64encode(json.dumps({"aud": "https://wrong.example.com"}).encode())
.rstrip(b"=")
.decode()
)
bad_jwt = f"{header}.{payload}.sig"
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.return_value = _mk_response(
200,
{"access_token": bad_jwt, "expires_in": 3600},
)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/callback?code=c&state=valid-state",
follow_redirects=False,
)
assert resp.status_code == 302
assert "audience+mismatch" in resp.headers["location"]
# And no token row written.
assert token_store.get_user_token("user-1", "srv-oauth") is None
def test_opaque_token_logs_and_trusts(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
self._seed_pending(storage)
token_store = _make_token_store(storage)
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.return_value = _mk_response(
200, {"access_token": "opaque-no-dots", "expires_in": 3600}
)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/callback?code=c&state=valid-state",
follow_redirects=False,
)
# Opaque tokens are trusted per the documented contract.
assert resp.status_code == 302
assert resp.headers["location"] == "/admin/mcp-servers"
plain = token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
assert plain["access_token"] == "opaque-no-dots"
def test_refresh_token_omitted_creates_row_without_refresh(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
self._seed_pending(storage)
token_store = _make_token_store(storage)
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.return_value = _mk_response(
200, {"access_token": "opaque-access", "expires_in": 3600}
)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
client.get(
"/v1/api/mcp/oauth/callback?code=c&state=valid-state",
follow_redirects=False,
)
plain = token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
assert plain["refresh_token"] is None
# ---------------------------------------------------------------------------
# 503 paths when mcp_token_store is None
# ---------------------------------------------------------------------------
class TestNoTokenStore:
def test_start_503(self, storage: SQLiteBackend) -> None:
client_obj = MagicMock(spec=httpx.AsyncClient)
app = _build_app(storage=storage, http_client=client_obj, token_store=None)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get("/v1/api/mcp/oauth/start?server=anything")
assert resp.status_code == 503
body = resp.json()
assert "mcp_token_encryption_key" in body["hint"]
def test_callback_503(self, storage: SQLiteBackend) -> None:
client_obj = MagicMock(spec=httpx.AsyncClient)
app = _build_app(storage=storage, http_client=client_obj, token_store=None)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get("/v1/api/mcp/oauth/callback?code=x&state=y")
assert resp.status_code == 503
# ---------------------------------------------------------------------------
# DCR
# ---------------------------------------------------------------------------
class TestDCR:
def test_registration_endpoint_hit_when_dcr_and_no_client_id(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(
storage,
client_id=None,
registration_mode="dcr",
)
token_store = _make_token_store(storage)
async def _post(url, *args, **kwargs):
if url.endswith("/register"):
return _mk_response(
201, {"client_id": "dcr-client-xyz", "client_secret": "dcr-secret"}
)
raise AssertionError(f"unexpected POST {url}")
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.side_effect = _post
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get("/v1/api/mcp/oauth/start?server=srv-oauth", follow_redirects=False)
assert resp.status_code == 302
# client_id was persisted on the row.
row = storage.get_mcp_server_by_name("srv-oauth")
assert row is not None
assert row["oauth_client_id"] == "dcr-client-xyz"
# Client secret was encrypted + stored.
secret = token_store.get_oauth_client_secret(row["server_id"])
assert secret == "dcr-secret"
def test_concurrent_dcr_callers_register_once(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
"""Two concurrent /start calls against a DCR server must register exactly once.
Without the per-server lock + re-fetch, both callers race past
the NULL ``oauth_client_id`` check and both POST to /register.
The second registration's client_id then overwrites the first
in the row, leaving the first user's authorize URL pointing at
a client_id that never makes it back to /callback (the row has
a different one), so the AS rejects the code exchange.
"""
_seed_oauth_user_server(
storage,
client_id=None,
registration_mode="dcr",
)
token_store = _make_token_store(storage)
register_calls = 0
register_started = asyncio.Event()
register_release = asyncio.Event()
async def _post(url, *args, **kwargs):
nonlocal register_calls
if url.endswith("/register"):
register_calls += 1
register_started.set()
# Block first caller inside the AS POST so the second
# caller is forced to take the DCR lock contended.
await register_release.wait()
return _mk_response(
201,
{"client_id": f"dcr-client-{register_calls}", "client_secret": "secret"},
)
raise AssertionError(f"unexpected POST {url}")
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.side_effect = _post
# Drive both /start handlers directly through ASGI to share the
# same app.state.mcp_oauth_dcr_locks dict — TestClient spawns its
# own thread loop per call, so direct invocation is the
# cleanest way to exercise the lock.
from starlette.requests import Request
async def _make_request(app: Starlette) -> Request:
scope = {
"type": "http",
"method": "GET",
"path": "/v1/api/mcp/oauth/start",
"raw_path": b"/v1/api/mcp/oauth/start",
"query_string": b"server=srv-oauth",
"headers": [(b"host", b"app.example.com")],
"scheme": "https",
"server": ("app.example.com", 443),
"app": app,
}
async def _receive() -> dict[str, Any]:
return {"type": "http.request", "body": b"", "more_body": False}
req = Request(scope, _receive)
req.state.auth_result = AuthResult(
user_id="user-1",
scopes=frozenset({"write"}),
token_source="config",
permissions=frozenset({"read", "write"}),
)
return req
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
async def _two_concurrent() -> tuple[Any, Any]:
with _public_addr_patch():
req1 = await _make_request(app)
req2 = await _make_request(app)
t1 = asyncio.create_task(handle_mcp_oauth_authorize(req1))
# Wait until the first caller is inside the /register POST
# so the second caller has to queue on the DCR lock.
await register_started.wait()
t2 = asyncio.create_task(handle_mcp_oauth_authorize(req2))
# Give t2 a chance to enter _register_dynamic_client_if_needed
# and block on the lock.
await asyncio.sleep(0.05)
register_release.set()
return await asyncio.gather(t1, t2)
resp1, resp2 = asyncio.run(_two_concurrent())
# Exactly one /register POST hit the AS.
assert register_calls == 1
# Both /start handlers redirected (302) using the SAME client_id.
assert resp1.status_code == 302
assert resp2.status_code == 302
loc1 = urllib.parse.parse_qs(urllib.parse.urlparse(resp1.headers["location"]).query)
loc2 = urllib.parse.parse_qs(urllib.parse.urlparse(resp2.headers["location"]).query)
assert loc1["client_id"] == loc2["client_id"]
# The persisted row shows the single client_id from the first
# (only) registration.
row = storage.get_mcp_server_by_name("srv-oauth")
assert row is not None
assert row["oauth_client_id"] == "dcr-client-1"
# ---------------------------------------------------------------------------
# JWT audience handling — bug-3 + sec-1 + bug-5
# ---------------------------------------------------------------------------
class TestJWTAudienceHandling:
"""Operator-set ``oauth_audience`` must be honored by JWT aud validation.
Auth0 (and similar non-RFC-8707 ASes) read the ``audience=`` URL
parameter and mint tokens with ``aud=<oauth_audience>`` rather than
the canonical resource URL. Phase 4 used the canonical URL only,
so legitimate Auth0 tokens were rejected.
"""
def _seed_pending(
self,
storage: SQLiteBackend,
*,
state: str = "valid-state",
verifier: str = "verifier-blob",
) -> None:
storage.create_mcp_oauth_pending_state(
state, "user-1", "srv-oauth", verifier, "/admin/mcp-servers"
)
@staticmethod
def _make_jwt(aud: Any) -> str:
header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode()
payload = base64.urlsafe_b64encode(json.dumps({"aud": aud}).encode()).rstrip(b"=").decode()
return f"{header}.{payload}.sig"
def test_jwt_aud_matches_oauth_audience_when_set(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
"""JWT ``aud=oauth_audience`` is accepted (Auth0 form)."""
_seed_oauth_user_server(storage)
# ``oauth_audience`` was seeded as ``https://mcp.example.com``;
# ``server_url`` is ``https://mcp.example.com/sse``. A JWT whose
# aud matches the configured oauth_audience must be accepted
# even though it differs from the canonical resource URL.
self._seed_pending(storage)
token_store = _make_token_store(storage)
good_jwt = self._make_jwt("https://mcp.example.com")
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.return_value = _mk_response(
200, {"access_token": good_jwt, "expires_in": 3600}
)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/callback?code=c&state=valid-state",
follow_redirects=False,
)
# Accepted — redirect to return_url, token row written.
assert resp.status_code == 302
assert resp.headers["location"] == "/admin/mcp-servers"
plain = token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
def test_jwt_aud_matches_canonical_resource_url(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
"""JWT ``aud=server_url`` is also accepted (RFC 8707 form)."""
_seed_oauth_user_server(storage)
self._seed_pending(storage)
token_store = _make_token_store(storage)
good_jwt = self._make_jwt("https://mcp.example.com/sse")
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.return_value = _mk_response(
200, {"access_token": good_jwt, "expires_in": 3600}
)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/callback?code=c&state=valid-state",
follow_redirects=False,
)
assert resp.status_code == 302
assert resp.headers["location"] == "/admin/mcp-servers"
def test_jwt_aud_list_with_one_match_accepted(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
"""JWT ``aud`` may be a list — accepted when ANY entry matches."""
_seed_oauth_user_server(storage)
self._seed_pending(storage)
token_store = _make_token_store(storage)
good_jwt = self._make_jwt(["https://other.example.com", "https://mcp.example.com"])
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
http_client_mock.post.return_value = _mk_response(
200, {"access_token": good_jwt, "expires_in": 3600}
)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/callback?code=c&state=valid-state",
follow_redirects=False,
)
assert resp.status_code == 302
assert resp.headers["location"] == "/admin/mcp-servers"
# ---------------------------------------------------------------------------
# sec-1 — redirect_uri pinned to oidc_config.redirect_base
# ---------------------------------------------------------------------------
class TestRedirectBasePinning:
def test_start_503_when_redirect_base_unset(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
"""``oidc_config.redirect_base`` empty → ``/start`` returns 503."""
_seed_oauth_user_server(storage)
token_store = _make_token_store(storage)
app = _build_app(
storage=storage,
http_client=http_client_mock,
token_store=token_store,
redirect_base="",
)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get("/v1/api/mcp/oauth/start?server=srv-oauth")
assert resp.status_code == 503
body = resp.json()
assert "redirect" in body["error"].lower()
def test_redirect_uri_uses_pinned_base_not_host_header(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
"""Spoofed Host header MUST NOT influence redirect_uri."""
_seed_oauth_user_server(storage)
token_store = _make_token_store(storage)
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
app = _build_app(
storage=storage,
http_client=http_client_mock,
token_store=token_store,
redirect_base="https://app.example.com",
)
client = TestClient(app, raise_server_exceptions=False)
with _public_addr_patch():
resp = client.get(
"/v1/api/mcp/oauth/start?server=srv-oauth",
headers={"Host": "attacker.example.com"},
follow_redirects=False,
)
assert resp.status_code == 302
location = resp.headers["location"]
params = urllib.parse.parse_qs(urllib.parse.urlparse(location).query)
assert params["redirect_uri"] == ["https://app.example.com/v1/api/mcp/oauth/callback"]
# ---------------------------------------------------------------------------
# bug-5 — error /callback pops pending state
# ---------------------------------------------------------------------------
class TestCallbackErrorPopsPendingState:
def test_callback_error_pops_pending_state(
self, storage: SQLiteBackend, http_client_mock: MagicMock
) -> None:
_seed_oauth_user_server(storage)
storage.create_mcp_oauth_pending_state(
"valid-state", "user-1", "srv-oauth", "verifier-blob", "/admin/mcp-servers"
)
token_store = _make_token_store(storage)
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
resp = client.get(
"/v1/api/mcp/oauth/callback"
"?error=access_denied&error_description=user+cancelled&state=valid-state",
follow_redirects=False,
)
assert resp.status_code == 302
# Pending state was popped — a replay with the same state should
# now miss and redirect to "session expired".
replay = client.get(
"/v1/api/mcp/oauth/callback?code=any&state=valid-state",
follow_redirects=False,
)
assert replay.status_code == 302
assert "session+expired" in replay.headers["location"]
+53
View File
@@ -0,0 +1,53 @@
"""PKCE pair-generation tests for the MCP OAuth flow.
Verifies the contract documented in RFC 7636 §4.1 and §4.2:
- ``code_verifier`` is a high-entropy 43..128 character urlsafe-base64 string.
- ``code_challenge`` is the BASE64URL-NO-PADDING encoding of
``SHA256(verifier)``.
"""
from __future__ import annotations
import base64
import hashlib
import string
from turnstone.core.mcp_oauth import generate_pkce_pair
_URLSAFE_CHARS = set(string.ascii_letters + string.digits + "-_")
class TestGeneratePkcePair:
def test_returns_tuple_of_strings(self) -> None:
verifier, challenge = generate_pkce_pair()
assert isinstance(verifier, str)
assert isinstance(challenge, str)
def test_verifier_length_in_rfc_range(self) -> None:
for _ in range(20):
verifier, _ = generate_pkce_pair()
assert 43 <= len(verifier) <= 128
def test_verifier_is_urlsafe(self) -> None:
for _ in range(20):
verifier, _ = generate_pkce_pair()
assert all(ch in _URLSAFE_CHARS for ch in verifier)
def test_challenge_matches_sha256_of_verifier(self) -> None:
for _ in range(20):
verifier, challenge = generate_pkce_pair()
digest = hashlib.sha256(verifier.encode("ascii")).digest()
expected = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
assert challenge == expected
def test_challenge_has_no_padding(self) -> None:
for _ in range(20):
_, challenge = generate_pkce_pair()
assert "=" not in challenge
def test_pairs_are_unique(self) -> None:
pairs = {generate_pkce_pair() for _ in range(50)}
# 50 random draws shouldn't collide; if they do we have a much
# bigger problem than this assertion.
assert len(pairs) == 50
+580
View File
@@ -0,0 +1,580 @@
"""Refresh-grant tests for ``get_user_access_token``.
The refresh path is the hottest hot-path in OAuth-MCP: every dispatch
call funnels through it, and any bug — double-refresh, swallowed
``revoke``, lost ``refresh_token`` — manifests as either a thundering
herd against the AS or a stuck "consent required" loop.
The concurrency test is the protocol-correctness highlight: TWO
``asyncio.create_task(get_user_access_token(...))`` against an expired
token, and we assert that the AS sees exactly ONE refresh POST and
both coroutines return the same access_token.
"""
from __future__ import annotations
import asyncio
from datetime import UTC, datetime, timedelta
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from tests.conftest import make_mcp_token_cipher
from turnstone.core.mcp_crypto import MCPTokenStore
from turnstone.core.mcp_oauth import get_user_access_token
from turnstone.core.storage._sqlite import SQLiteBackend
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
def _make_app_state(storage: SQLiteBackend, *, http_client: httpx.AsyncClient) -> SimpleNamespace:
cipher = make_mcp_token_cipher()
state = SimpleNamespace(
auth_storage=storage,
mcp_token_store=MCPTokenStore(storage, cipher, node_id="test"),
mcp_oauth_http_client=http_client,
mcp_oauth_refresh_locks={},
mcp_oauth_metadata_cache={},
)
return state
def _seed_server(backend: SQLiteBackend, *, server_id: str = "srv-id") -> None:
backend.create_mcp_server(
server_id=server_id,
name="srv-oauth",
transport="streamable-http",
url="https://mcp.example.com/sse",
auth_type="oauth_user",
oauth_client_id="client-abc",
oauth_scopes="openid profile",
oauth_audience="https://mcp.example.com",
)
backend.update_mcp_server(server_id, oauth_as_issuer_cached="https://as.example.com")
def _seed_token(
state: SimpleNamespace,
*,
user_id: str = "user-1",
server_name: str = "srv-oauth",
expires_in_seconds: int = 3600,
refresh: str | None = "refresh-rrr",
) -> None:
expires_at = (datetime.now(UTC) + timedelta(seconds=expires_in_seconds)).strftime(
"%Y-%m-%dT%H:%M:%S"
)
state.mcp_token_store.create_user_token(
user_id,
server_name,
access_token="access-aaa",
refresh_token=refresh,
expires_at=expires_at,
scopes="openid profile",
as_issuer="https://as.example.com",
audience="https://mcp.example.com",
)
@pytest.fixture
def storage(tmp_path: Any) -> SQLiteBackend:
return SQLiteBackend(str(tmp_path / "test.db"))
def _good_as_metadata_doc() -> dict[str, Any]:
return {
"issuer": "https://as.example.com",
"authorization_endpoint": "https://as.example.com/authorize",
"token_endpoint": "https://as.example.com/token",
"jwks_uri": "https://as.example.com/jwks",
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["none"],
}
def _mk_response(status_code: int = 200, json_body: Any = None) -> MagicMock:
resp = MagicMock(spec=httpx.Response)
resp.status_code = status_code
resp.headers = {}
body = "" if json_body is None else str(json_body)
resp.content = body.encode("utf-8")
if json_body is not None:
resp.json.return_value = json_body
else:
resp.json.side_effect = ValueError("no body")
resp.text = body
return resp
def _public_addr_patch():
return patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("93.184.216.34", 0))])
# ---------------------------------------------------------------------------
# Happy paths
# ---------------------------------------------------------------------------
class TestUnchangedToken:
def test_returns_existing_token_when_not_expired(self, storage: SQLiteBackend) -> None:
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock()
client.post = AsyncMock()
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=3600)
async def _run():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
token = asyncio.run(_run())
assert token == "access-aaa"
# No AS calls were made.
client.get.assert_not_called()
client.post.assert_not_called()
def test_no_token_returns_none(self, storage: SQLiteBackend) -> None:
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
state = _make_app_state(storage, http_client=client)
async def _run():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
assert asyncio.run(_run()) is None
def test_get_user_access_token_handles_decrypt_error(self, storage: SQLiteBackend) -> None:
"""A decrypt failure on the stored token must not crash dispatch.
When the operator rotates ``mcp_token_encryption_key`` and drops
the prior key, every existing user-token row decrypts to
:class:`MCPTokenDecryptError`. ``get_user_access_token`` MUST
catch that and return ``None`` (forcing the user back through
the consent flow) rather than propagating the exception up to
the dispatch caller.
"""
from turnstone.core.mcp_crypto import MCPTokenDecryptError
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=3600)
# Replace the get_user_token method on the store to raise the
# canonical key-mismatch error.
original_get = state.mcp_token_store.get_user_token
def _raise_decrypt(*args, **kwargs):
raise MCPTokenDecryptError(
"no installed key can decrypt",
key_fingerprints_attempted=("aabbccdd",),
)
state.mcp_token_store.get_user_token = _raise_decrypt
try:
async def _run():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
result = asyncio.run(_run())
assert result is None
finally:
state.mcp_token_store.get_user_token = original_get
# ---------------------------------------------------------------------------
# Refresh path
# ---------------------------------------------------------------------------
class TestRefresh:
def test_refreshes_when_expired(self, storage: SQLiteBackend) -> None:
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(
200,
{
"access_token": "access-NEW",
"refresh_token": "refresh-NEW",
"expires_in": 3600,
},
)
)
state = _make_app_state(storage, http_client=client)
# Seed an expired token.
_seed_token(state, expires_in_seconds=-1000)
async def _run():
with _public_addr_patch():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
token = asyncio.run(_run())
assert token == "access-NEW"
# The refresh endpoint was hit exactly once.
assert client.post.call_count == 1
# Verify the new tokens were persisted.
plain = state.mcp_token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
assert plain["access_token"] == "access-NEW"
assert plain["refresh_token"] == "refresh-NEW"
def test_refresh_failure_deletes_row_and_returns_none(self, storage: SQLiteBackend) -> None:
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
# AS rejects the refresh — token should be revoked.
client.post = AsyncMock(
return_value=_mk_response(
400,
{"error": "invalid_grant"},
)
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
emitted: list[str] = []
async def _run():
from turnstone.core import mcp_oauth as mod
real_record_audit = mod.record_audit
def _capture(*args, **kwargs):
# signature: (storage, user_id, action, resource_type, resource_id, detail)
emitted.append(args[2] if len(args) >= 3 else kwargs.get("action", ""))
return real_record_audit(*args, **kwargs)
with (
patch.object(mod, "record_audit", side_effect=_capture),
_public_addr_patch(),
):
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
result = asyncio.run(_run())
assert result is None
# Row was deleted.
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is None
# Audit emitted the revoke event.
assert "mcp_server.oauth.token_revoked" in emitted
def test_concurrent_callers_via_lock(self, storage: SQLiteBackend) -> None:
"""Two concurrent refresh calls must produce exactly one AS POST.
Both callers see the same access_token.
"""
_seed_server(storage)
# Coordinate the AS POST so both callers race the lock.
post_started = asyncio.Event()
post_release = asyncio.Event()
async def _post(url, *args, **kwargs):
post_started.set()
await post_release.wait()
return _mk_response(
200,
{
"access_token": "access-ONE",
"refresh_token": "refresh-ONE",
"expires_in": 3600,
},
)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(side_effect=_post)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
async def _both():
with _public_addr_patch():
t1 = asyncio.create_task(
get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
)
# Wait for the first task to enter the AS POST so the
# second task is forced to take the lock contended.
await post_started.wait()
t2 = asyncio.create_task(
get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
)
# Give t2 a chance to queue on the lock.
await asyncio.sleep(0.05)
post_release.set()
return await asyncio.gather(t1, t2)
a, b = asyncio.run(_both())
assert a == "access-ONE"
assert b == "access-ONE"
# Exactly one POST.
assert client.post.call_count == 1
def test_refresh_omitted_refresh_token_preserves_existing(self, storage: SQLiteBackend) -> None:
"""RFC 6749 §6 — AS MAY omit refresh_token; we PRESERVE the existing one.
Production ASes (Google, default Auth0, default Okta) do NOT
rotate refresh tokens. Clearing the column on every refresh
would force the user to re-consent every hour. The contract:
replace the persisted refresh token only when the AS issues a
new one; otherwise pass the prior refresh token through to
``update_user_token_after_refresh``. The ``refresh_token=None``
sentinel still means "clear" at the storage layer (Phase 3
contract preserved); the *_refresh_and_persist_* layer
translates "omitted" into "pass through existing".
"""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(
200,
# No refresh_token in response.
{"access_token": "access-NEW", "expires_in": 3600},
)
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000, refresh="refresh-original")
async def _run():
with _public_addr_patch():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
token = asyncio.run(_run())
assert token == "access-NEW"
plain = state.mcp_token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
# Refresh column was PRESERVED — the original refresh token is
# still usable for the next refresh cycle.
assert plain["refresh_token"] == "refresh-original"
def test_refresh_rotated_refresh_token_replaces_existing(self, storage: SQLiteBackend) -> None:
"""When AS issues a fresh refresh_token, the new value REPLACES the prior."""
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(
200,
{
"access_token": "access-NEW",
"refresh_token": "refresh-NEW",
"expires_in": 3600,
},
)
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000, refresh="refresh-original")
async def _run():
with _public_addr_patch():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
token = asyncio.run(_run())
assert token == "access-NEW"
plain = state.mcp_token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
assert plain["refresh_token"] == "refresh-NEW"
def test_expired_no_refresh_token_revokes(self, storage: SQLiteBackend) -> None:
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000, refresh=None)
async def _run():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
result = asyncio.run(_run())
assert result is None
# Row was deleted (re-consent path).
assert state.mcp_token_store.get_user_token("user-1", "srv-oauth") is None
def test_refresh_grant_sends_server_url_as_resource_not_audience(
self, storage: SQLiteBackend
) -> None:
"""RFC 8707 ``resource=`` is the canonical MCP server URL.
Earlier code passed ``oauth_audience`` as the resource value;
Auth0-style ASes that honor a separate ``audience=`` parameter
would then receive the wrong URL in ``resource=``, and ASes that
validate ``resource`` against their RS allowlist would reject
the refresh. The refresh-grant MUST send the canonical server
URL on ``resource=``.
"""
_seed_server(storage)
# Override the audience on the seed server so it diverges from
# the canonical server URL.
backend_row = storage.get_mcp_server_by_name("srv-oauth")
assert backend_row is not None
storage.update_mcp_server(
backend_row["server_id"],
oauth_audience="https://different-audience.example.com/api",
)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(
200,
{
"access_token": "access-NEW",
"refresh_token": "refresh-NEW",
"expires_in": 3600,
},
)
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
async def _run():
with _public_addr_patch():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
asyncio.run(_run())
# The refresh POST was made; assert the form payload sent
# ``resource=server_url`` not ``resource=audience``.
assert client.post.call_count == 1
post_kwargs = client.post.call_args.kwargs
sent_resource = post_kwargs["data"]["resource"]
assert sent_resource == "https://mcp.example.com/sse"
assert sent_resource != "https://different-audience.example.com/api"
# ---------------------------------------------------------------------------
# expires_in parsing — bug-2
# ---------------------------------------------------------------------------
class TestExpiresInParsing:
"""``_expires_at_from_response`` must accept int, float, and string.
Real ASes have been seen returning ``3600.0`` (float) and ``"3600"``
(string) — the prior ``int(str(3600.0))`` raised ValueError, leaving
``expires_at=None``. ``None`` then made ``_token_needs_refresh``
return False, so the token was never refreshed and effectively never
expired (it accumulated until the AS revoked it server-side).
"""
def test_expires_in_float_parsed_correctly(self, storage: SQLiteBackend) -> None:
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(
200,
{
"access_token": "access-NEW",
"refresh_token": "refresh-NEW",
"expires_in": 3600.0,
},
)
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
async def _run():
with _public_addr_patch():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
token = asyncio.run(_run())
assert token == "access-NEW"
plain = state.mcp_token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
# expires_at must be populated — a None value here means the
# float was rejected and the next refresh cycle would skip it.
assert plain["expires_at"] is not None
def test_expires_in_str_with_decimal_parsed_correctly(self, storage: SQLiteBackend) -> None:
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(
200,
{
"access_token": "access-NEW",
"refresh_token": "refresh-NEW",
"expires_in": "3600.0",
},
)
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
async def _run():
with _public_addr_patch():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
token = asyncio.run(_run())
assert token == "access-NEW"
plain = state.mcp_token_store.get_user_token("user-1", "srv-oauth")
assert plain is not None
assert plain["expires_at"] is not None
def test_expires_in_int_string_parsed_correctly(self, storage: SQLiteBackend) -> None:
_seed_server(storage)
client = MagicMock(spec=httpx.AsyncClient)
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
client.post = AsyncMock(
return_value=_mk_response(
200,
{
"access_token": "access-NEW",
"refresh_token": "refresh-NEW",
"expires_in": "3600",
},
)
)
state = _make_app_state(storage, http_client=client)
_seed_token(state, expires_in_seconds=-1000)
async def _run():
with _public_addr_patch():
return await get_user_access_token(
app_state=state, user_id="user-1", server_name="srv-oauth"
)
token = asyncio.run(_run())
assert token == "access-NEW"
def test_expires_in_garbage_returns_none(self) -> None:
from turnstone.core.mcp_oauth import _expires_at_from_response
assert _expires_at_from_response({}) is None
assert _expires_at_from_response({"expires_in": "abc"}) is None
assert _expires_at_from_response({"expires_in": None}) is None
assert _expires_at_from_response({"expires_in": True}) is None
assert _expires_at_from_response({"expires_in": 0}) is None
assert _expires_at_from_response({"expires_in": -5}) is None
+125
View File
@@ -0,0 +1,125 @@
"""Storage CRUD tests for the per-(user, server) MCP OAuth pending-state table.
Validates the storage-protocol additions for the per-(user, server)
OAuth flow:
- ``create_mcp_oauth_pending_state``
- ``pop_mcp_oauth_pending_state`` (atomic, with TTL)
- ``cleanup_expired_mcp_oauth_pending_states``
- ``get_mcp_oauth_client_secret_ct``
"""
from __future__ import annotations
import sqlalchemy as sa
class TestCreateAndPop:
def test_round_trip(self, backend) -> None:
backend.create_mcp_oauth_pending_state(
"state-1",
"user-a",
"srv-x",
"verifier-blob",
"/admin/mcp-servers",
)
row = backend.pop_mcp_oauth_pending_state("state-1", max_age_seconds=600)
assert row is not None
assert row["state"] == "state-1"
assert row["user_id"] == "user-a"
assert row["server_name"] == "srv-x"
assert row["code_verifier"] == "verifier-blob"
assert row["return_url"] == "/admin/mcp-servers"
def test_pop_consumes_row(self, backend) -> None:
backend.create_mcp_oauth_pending_state("s2", "u", "s", "v", "/r")
first = backend.pop_mcp_oauth_pending_state("s2")
assert first is not None
# Second pop must miss — row was consumed.
second = backend.pop_mcp_oauth_pending_state("s2")
assert second is None
def test_pop_missing_returns_none(self, backend) -> None:
assert backend.pop_mcp_oauth_pending_state("never-existed") is None
class TestTTL:
def test_pop_rejects_expired_row(self, backend) -> None:
backend.create_mcp_oauth_pending_state("old-state", "u", "s", "v", "/r")
# Backdate it so it's older than the TTL window.
with backend._engine.connect() as conn:
conn.execute(
sa.text(
"UPDATE mcp_oauth_pending SET created_at = '2020-01-01T00:00:00' "
"WHERE state = 'old-state'"
)
)
conn.commit()
# Default TTL is 600s — the row is decades old.
row = backend.pop_mcp_oauth_pending_state("old-state")
assert row is None
# Even though pop returned None, the row must have been wiped — a
# second pop with a giant TTL must still see nothing.
again = backend.pop_mcp_oauth_pending_state("old-state", max_age_seconds=10**9)
assert again is None
def test_pop_accepts_fresh_row(self, backend) -> None:
backend.create_mcp_oauth_pending_state("fresh", "u", "s", "v", "/r")
row = backend.pop_mcp_oauth_pending_state("fresh", max_age_seconds=600)
assert row is not None
assert row["state"] == "fresh"
class TestCleanup:
def test_cleanup_deletes_only_expired(self, backend) -> None:
backend.create_mcp_oauth_pending_state("old", "u", "s", "v", "/r")
backend.create_mcp_oauth_pending_state("new", "u", "s", "v", "/r")
with backend._engine.connect() as conn:
conn.execute(
sa.text(
"UPDATE mcp_oauth_pending SET created_at = '2020-01-01T00:00:00' "
"WHERE state = 'old'"
)
)
conn.commit()
deleted = backend.cleanup_expired_mcp_oauth_pending_states(max_age_seconds=600)
assert deleted == 1
# Old gone, new still around.
assert backend.pop_mcp_oauth_pending_state("old") is None
survivor = backend.pop_mcp_oauth_pending_state("new")
assert survivor is not None
def test_cleanup_no_rows(self, backend) -> None:
assert backend.cleanup_expired_mcp_oauth_pending_states() == 0
class TestGetOAuthClientSecretCt:
def test_returns_none_when_unset(self, backend) -> None:
backend.create_mcp_server(
server_id="srv-id",
name="srv-x",
transport="streamable-http",
url="https://mcp.example.com/sse",
auth_type="oauth_user",
)
assert backend.get_mcp_oauth_client_secret_ct("srv-id") is None
def test_returns_ciphertext_after_set(self, backend) -> None:
backend.create_mcp_server(
server_id="srv-id",
name="srv-x",
transport="streamable-http",
url="https://mcp.example.com/sse",
auth_type="oauth_user",
)
ct = b"\x00\xff\x42encrypted-blob"
ok = backend.set_mcp_oauth_client_secret_ct("srv-id", ct)
assert ok is True
out = backend.get_mcp_oauth_client_secret_ct("srv-id")
assert out == ct
def test_returns_none_for_missing_server(self, backend) -> None:
assert backend.get_mcp_oauth_client_secret_ct("does-not-exist") is None
+33
View File
@@ -251,3 +251,36 @@ class TestDecryptFailureInvariant:
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)
+172
View File
@@ -0,0 +1,172 @@
"""Direct tests for the shared SSRF helpers in :mod:`turnstone.core.oauth_ssrf`.
The OIDC test suite already exercises these via the OIDC adapter
(``OIDCError`` re-raises). This file pins the canonical
:class:`OAuthSSRFError` exception so callers that don't go through OIDC
(notably ``mcp_oauth``) can rely on a stable contract.
"""
from __future__ import annotations
import urllib.parse
from unittest.mock import patch
import pytest
from turnstone.core.oauth_ssrf import (
OAuthSSRFError,
effective_port,
is_localhost,
validate_discovered_endpoint,
validate_url_no_ssrf,
)
class TestIsLocalhost:
def test_loopback_names(self) -> None:
assert is_localhost("localhost")
assert is_localhost("127.0.0.1")
assert is_localhost("::1")
assert is_localhost("foo.localhost")
def test_non_loopback(self) -> None:
assert not is_localhost("example.com")
assert not is_localhost("internal.corp")
class TestEffectivePort:
def test_explicit_port(self) -> None:
p = urllib.parse.urlparse("https://idp.example.com:9443/foo")
assert effective_port(p) == 9443
def test_default_https(self) -> None:
p = urllib.parse.urlparse("https://idp.example.com/foo")
assert effective_port(p) == 443
def test_default_http(self) -> None:
p = urllib.parse.urlparse("http://idp.example.com/foo")
assert effective_port(p) == 80
def test_unknown_scheme(self) -> None:
p = urllib.parse.urlparse("ftp://idp.example.com/foo")
assert effective_port(p) is None
class TestValidateUrlNoSSRF:
_PUBLIC_ADDR = [(2, 1, 6, "", ("93.184.216.34", 0))]
_PRIVATE_ADDR = [(2, 1, 6, "", ("10.0.0.1", 0))]
_LOOPBACK_ADDR = [(2, 1, 6, "", ("127.0.0.1", 0))]
def test_valid_https(self) -> None:
with patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR):
parsed = validate_url_no_ssrf("https://idp.example.com/foo", allow_http=False)
assert parsed.scheme == "https"
assert parsed.hostname == "idp.example.com"
def test_rejects_http_when_not_allowed(self) -> None:
with pytest.raises(OAuthSSRFError, match="must use HTTPS"):
validate_url_no_ssrf("http://idp.example.com", allow_http=False)
def test_allows_http_localhost_with_flag(self) -> None:
with patch("socket.getaddrinfo", return_value=self._LOOPBACK_ADDR):
validate_url_no_ssrf("http://localhost:8080", allow_http=True)
def test_rejects_http_non_localhost_even_with_flag(self) -> None:
with pytest.raises(OAuthSSRFError, match="must use HTTPS"):
validate_url_no_ssrf("http://idp.example.com", allow_http=True)
def test_rejects_userinfo(self) -> None:
with pytest.raises(OAuthSSRFError, match="embedded credentials"):
validate_url_no_ssrf("https://user:pass@idp.example.com", allow_http=False)
def test_rejects_private_address(self) -> None:
with (
patch("socket.getaddrinfo", return_value=self._PRIVATE_ADDR),
pytest.raises(OAuthSSRFError, match="non-public address"),
):
validate_url_no_ssrf("https://corp.example.com", allow_http=False)
def test_rejects_unresolvable(self) -> None:
import socket
with (
patch("socket.getaddrinfo", side_effect=socket.gaierror("fail")),
pytest.raises(OAuthSSRFError, match="cannot be resolved"),
):
validate_url_no_ssrf("https://no.such.host.invalid", allow_http=False)
class TestValidateDiscoveredEndpoint:
_PUBLIC_ADDR = [(2, 1, 6, "", ("93.184.216.34", 0))]
def test_same_origin_passes(self) -> None:
issuer = urllib.parse.urlparse("https://idp.example.com")
with patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR):
validate_discovered_endpoint(
"https://idp.example.com/token",
issuer,
allow_http=False,
trusted_endpoint_hosts=frozenset(),
)
def test_third_party_host_rejected(self) -> None:
issuer = urllib.parse.urlparse("https://idp.example.com")
with (
patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR),
pytest.raises(OAuthSSRFError, match="not trusted"),
):
validate_discovered_endpoint(
"https://attacker.example.com/token",
issuer,
allow_http=False,
trusted_endpoint_hosts=frozenset(),
)
def test_trusted_endpoint_host_passes(self) -> None:
issuer = urllib.parse.urlparse("https://idp.example.com")
with patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR):
validate_discovered_endpoint(
"https://shard.example.com/token",
issuer,
allow_http=False,
trusted_endpoint_hosts=frozenset({"shard.example.com"}),
)
def test_known_google_alias_passes(self) -> None:
"""The hard-coded Google alias map covers oauth2.googleapis.com."""
issuer = urllib.parse.urlparse("https://accounts.google.com")
with patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR):
validate_discovered_endpoint(
"https://oauth2.googleapis.com/token",
issuer,
allow_http=False,
trusted_endpoint_hosts=frozenset(),
)
def test_scheme_mismatch_rejected(self) -> None:
# When the issuer is http://localhost (allow_http=True), an
# https:// endpoint must still be rejected as a scheme mismatch.
issuer = urllib.parse.urlparse("http://localhost:8080")
with (
patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("127.0.0.1", 0))]),
pytest.raises(OAuthSSRFError, match="scheme"),
):
validate_discovered_endpoint(
"https://localhost:8080/token",
issuer,
allow_http=True,
trusted_endpoint_hosts=frozenset(),
)
def test_port_mismatch_rejected(self) -> None:
issuer = urllib.parse.urlparse("https://idp.example.com")
with (
patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR),
pytest.raises(OAuthSSRFError, match="port"),
):
validate_discovered_endpoint(
"https://idp.example.com:9443/token",
issuer,
allow_http=False,
trusted_endpoint_hosts=frozenset(),
)
+82 -8
View File
@@ -1621,6 +1621,20 @@ async def oidc_callback(request: Request) -> Response:
return await handle_oidc_callback(request, JWT_AUD_CONSOLE)
async def mcp_oauth_authorize(request: Request) -> Response:
"""GET /v1/api/mcp/oauth/start — begin per-(user, server) OAuth flow."""
from turnstone.core.mcp_oauth import handle_mcp_oauth_authorize
return await handle_mcp_oauth_authorize(request)
async def mcp_oauth_callback(request: Request) -> Response:
"""GET /v1/api/mcp/oauth/callback — AS-redirected OAuth callback."""
from turnstone.core.mcp_oauth import handle_mcp_oauth_callback
return await handle_mcp_oauth_callback(request)
# ---------------------------------------------------------------------------
# Route handlers — available models (lightweight, no admin permission)
# ---------------------------------------------------------------------------
@@ -4214,14 +4228,20 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
await initialize_oidc_state(app.state)
# MCP-OAuth token-at-rest encryption — fail-loud on misconfiguration
# when any mcp_servers row has auth_type='oauth_user'. See
# docs/design/oauth-mcp.md §5.3. The console acts as the cluster
# admin surface and is the canonical writer for OAuth client secrets,
# so it needs the cipher even when no node currently dispatches.
# when any mcp_servers row has auth_type='oauth_user'. The console
# acts as the cluster admin surface and is the canonical writer for
# OAuth client secrets, so it needs the cipher even when no node
# currently dispatches.
from turnstone.core.mcp_crypto import initialize_mcp_crypto_state
initialize_mcp_crypto_state(app.state, node_id="console")
# Per-(user, server) OAuth flow state — long-lived HTTP client +
# in-process refresh lock + metadata cache.
from turnstone.core.mcp_oauth import initialize_mcp_oauth_state
await initialize_mcp_oauth_state(app.state)
# Register console in service registry so other services can discover it
console_url = getattr(app.state, "console_url", "")
_console_heartbeat_task: Any = None
@@ -4482,12 +4502,17 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
log.debug("console.coord_ui_refs_reset_failed", exc_info=True)
await app.state.proxy_sse_client.aclose()
await app.state.proxy_client.aclose()
from turnstone.core.oidc import close_oidc_state
# Close in reverse order of initialization (mcp_oauth → mcp_crypto →
# oidc) per LIFO teardown discipline.
from turnstone.core.mcp_oauth import close_mcp_oauth_state
await close_oidc_state(app.state)
await close_mcp_oauth_state(app.state)
from turnstone.core.mcp_crypto import close_mcp_crypto_state
close_mcp_crypto_state(app.state)
from turnstone.core.oidc import close_oidc_state
await close_oidc_state(app.state)
app.state.collector.stop()
audit_exec_shutdown = getattr(app.state, "audit_executor", None)
if audit_exec_shutdown is not None:
@@ -8123,6 +8148,14 @@ def _apply_oauth_client_secret(
{"error": "oauth_client_secret must be a string or null"},
status_code=400,
)
# Cap plaintext at 1024 chars — defends against pathological input
# blowing up the Fernet ciphertext column. Real OAuth client secrets
# from production AS implementations are well under 256 chars.
if isinstance(secret_input, str) and len(secret_input) > 1024:
return None, JSONResponse(
{"error": "oauth_client_secret must be 1024 characters or fewer"},
status_code=400,
)
token_store = getattr(request.app.state, "mcp_token_store", None)
if token_store is None:
return None, JSONResponse({"error": _OAUTH_TOKEN_STORE_503_MSG}, status_code=503)
@@ -8547,6 +8580,11 @@ async def admin_update_mcp_server(request: Request) -> JSONResponse:
transitioning_away_from_oauth_user = (
updates.get("auth_type") in {"static", "none"} and existing.get("auth_type") == "oauth_user"
)
# Renaming an oauth_user row needs the same per-user-token purge as
# delete: the tokens are keyed on the OLD ``server_name`` and a row
# later created with that old name (with attacker-controlled URL)
# would otherwise silently rebind them.
name_changing = "name" in updates and existing.get("name", "") != updates["name"]
if updates.get("auth_type") and updates["auth_type"] != "oauth_user":
updates.update(
{
@@ -8572,6 +8610,22 @@ async def admin_update_mcp_server(request: Request) -> JSONResponse:
audit_uid, ip = _audit_context(request)
# Purge per-user tokens + pending OAuth states keyed on the OLD
# name on rename, or on ``oauth_user → static/none`` transition.
# Both cases would otherwise leave tokens orphaned-and-rebindable
# because the OAuth tables key on the mutable ``server_name``.
if name_changing or transitioning_away_from_oauth_user:
purge_target = existing.get("name", "")
if purge_target:
try:
storage.delete_mcp_oauth_rows_by_server_name(purge_target)
except Exception:
log.warning(
"admin.mcp.purge_oauth_rows_failed server_id=%s",
server_id,
exc_info=True,
)
# sec-2: when the row is leaving ``oauth_user``, also clear the encrypted
# client secret column. Operators expect "disable OAuth" to revoke
# credentials; leaving stale ciphertext that resurfaces if the row is
@@ -8656,6 +8710,24 @@ async def admin_delete_mcp_server(request: Request) -> JSONResponse:
if existing is None:
return JSONResponse({"error": "MCP server not found"}, status_code=404)
# Purge per-user tokens + pending OAuth states keyed on this name
# before the row goes away. The OAuth tables are keyed on the
# mutable ``server_name``; without this purge, a future server with
# the same name (and an attacker-controlled URL) would silently
# rebind those tokens. A future schema migration will replace this
# with a server_id FK + ON DELETE CASCADE; until then, explicit
# purge is the only safe path.
existing_name = existing.get("name", "")
if existing_name:
try:
storage.delete_mcp_oauth_rows_by_server_name(existing_name)
except Exception:
log.warning(
"admin.mcp.purge_oauth_rows_failed server_id=%s",
server_id,
exc_info=True,
)
storage.delete_mcp_server(server_id)
audit_uid, ip = _audit_context(request)
@@ -8665,7 +8737,7 @@ async def admin_delete_mcp_server(request: Request) -> JSONResponse:
"mcp_server.delete",
"mcp_server",
server_id,
{"name": existing.get("name", "")},
{"name": existing_name},
ip,
)
@@ -8792,7 +8864,7 @@ async def _admin_mcp_action(request: Request, action: str) -> JSONResponse:
return JSONResponse({"error": "invalid server name"}, status_code=400)
existing = storage.get_mcp_server_by_name(name)
target_id = existing.get("id", name) if existing else name
target_id = existing.get("server_id", name) if existing else name
audit_uid, ip = _audit_context(request)
record_audit(
@@ -11438,6 +11510,8 @@ def create_app(
Route("/api/auth/refresh", auth_refresh, methods=["POST"]),
Route("/api/auth/oidc/authorize", oidc_authorize),
Route("/api/auth/oidc/callback", oidc_callback),
Route("/api/mcp/oauth/start", mcp_oauth_authorize),
Route("/api/mcp/oauth/callback", mcp_oauth_callback),
Route("/api/admin/users", admin_list_users),
Route("/api/admin/users", admin_create_user, methods=["POST"]),
Route("/api/admin/users/{user_id}", admin_delete_user, methods=["DELETE"]),
+19
View File
@@ -3323,6 +3323,12 @@ function _renderMcpServers(items) {
'<button class="admin-btn-action" data-mcp-reconnect="' +
escapeHtml(s.name) +
'">reconnect</button>';
if (s.auth_type === "oauth_user") {
actionBtns +=
'<button class="admin-btn-action" data-mcp-oauth-connect="' +
escapeHtml(s.name) +
'">connect</button>';
}
var actions = isConfig
? actionBtns
: actionBtns +
@@ -3431,6 +3437,19 @@ function _renderMcpServers(items) {
});
});
});
el.querySelectorAll("[data-mcp-oauth-connect]").forEach(function (btn) {
btn.addEventListener("click", function () {
var name = this.getAttribute("data-mcp-oauth-connect");
// Open the OAuth /start endpoint in a new window so the redirect
// chain (AS → callback → return_url) doesn't displace the admin UI.
var url =
"/v1/api/mcp/oauth/start?server=" +
encodeURIComponent(name) +
"&return_url=" +
encodeURIComponent(window.location.href);
window.open(url, "_blank", "noopener");
});
});
el.querySelectorAll("[data-mcp-delete]").forEach(function (btn) {
btn.addEventListener("click", function () {
var sid = this.getAttribute("data-mcp-delete");
+11 -2
View File
@@ -40,8 +40,17 @@ Action-name conventions (non-exhaustive — grep
clears a per-server OAuth client secret;
``mcp_server.oauth.token_decrypt_failure`` from
``MCPTokenStore.get_user_token`` when no
installed key can decrypt a stored token).
See ``docs/design/oauth-mcp.md`` §5.3.
installed key can decrypt a stored token;
``mcp_server.oauth.consent_started`` /
``.consent_completed`` / ``.consent_failed`` for
the per-user authorization-flow handlers;
``mcp_server.oauth.token_refreshed`` /
``.token_revoked`` from
``get_user_access_token`` when the
refresh-grant exchange runs;
``mcp_server.oauth.dcr_registered`` when a
client_id was provisioned dynamically against
an AS that exposes ``registration_endpoint``).
When adding a new namespace, prefer extending an existing prefix over
inventing a synonym (e.g. ``mcp_server.refresh`` rather than
+28 -16
View File
@@ -78,11 +78,11 @@ def _mcp_to_openai(server_name: str, tool: Any) -> dict[str, Any]:
class StaticServerState:
"""Per-server state for auth_type ∈ {none, static}. Name-keyed only.
Phase 5 introduces PoolEntryState as the (user, server)-keyed sibling
for auth_type=oauth_user. Together with the typed map declarations
(dict[str, StaticServerState] vs dict[tuple[str, str], PoolEntryState]),
this makes accidental cross-keying lookups easier to catch in review
and rejected by mypy.
The upcoming per-user pool integration introduces PoolEntryState as
the (user, server)-keyed sibling for auth_type=oauth_user. Together
with the typed map declarations (dict[str, StaticServerState] vs
dict[tuple[str, str], PoolEntryState]), this makes accidental
cross-keying lookups easier to catch in review and rejected by mypy.
"""
name: str
@@ -103,10 +103,10 @@ class StaticServerState:
class PoolEntryState:
"""Per-(user, server) state for auth_type = oauth_user.
Defined for Phase 5 use; not instantiated anywhere in Phase 0.
open_lock has no default RFC §2.0 invariant 2 forbids allocating an
asyncio.Lock outside the mcp-loop. Phase 5 allocates lazily inside
connect coroutines.
Defined for use by the upcoming per-user pool integration; currently
no production caller. open_lock has no default the mcp-loop
invariant forbids allocating an asyncio.Lock outside the mcp-loop;
the pool integration allocates lazily inside connect coroutines.
"""
key: tuple[str, str] # (user_id, server_name)
@@ -143,8 +143,9 @@ class MCPClientManager:
# Per-server state for auth_type ∈ {none, static}. Each entry holds
# session/stack/streams/catalog/capability flags for one name-keyed
# connection. Phase 5 introduces a sibling pool-entry map for
# auth_type=oauth_user; static entries always live here.
# connection. The upcoming per-user pool integration introduces a
# sibling pool-entry map for auth_type=oauth_user; static entries
# always live here.
self._static_servers: dict[str, StaticServerState] = {}
self._tools: list[dict[str, Any]] = []
@@ -318,8 +319,9 @@ class MCPClientManager:
Pre-closing unblocks anyio transport tasks stuck on zero-buffer
``send()`` calls, preventing the CPU busy-loop from SDK #2147.
Parameter is ``str`` today; Phase 5 widens to ``str | tuple[str, str]``
once ``PoolEntryState`` is wired.
Parameter is ``str`` today; the upcoming per-user pool
integration widens it to ``str | tuple[str, str]`` once
``PoolEntryState`` is wired.
"""
state = self._static_servers.get(key)
if state is None or state.streams is None:
@@ -336,8 +338,9 @@ class MCPClientManager:
Fails fast when the server is unreachable, avoiding the anyio
cancel-scope orphan bug that causes 100% CPU spin.
Parameter is ``str`` today; Phase 5 widens to ``str | tuple[str, str]``
once ``PoolEntryState`` is wired.
Parameter is ``str`` today; the upcoming per-user pool
integration widens it to ``str | tuple[str, str]`` once
``PoolEntryState`` is wired.
"""
from urllib.parse import urlparse
@@ -1764,9 +1767,18 @@ class MCPClientManager:
def _db_servers_to_config(rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
"""Convert mcp_servers DB rows to the config dict format."""
"""Convert mcp_servers DB rows to the config dict format.
Skips ``auth_type='oauth_user'`` rows: they need per-user bearer
tokens fetched at dispatch time via the OAuth flow, so auto-connecting
them at startup with empty headers fails handshake and trips the
circuit breaker. The upcoming per-user pool integration brings them
online lazily once a user has consented.
"""
result: dict[str, dict[str, Any]] = {}
for row in rows:
if row.get("auth_type") == "oauth_user":
continue
name = row["name"]
cfg: dict[str, Any] = {"type": row["transport"]}
if row["transport"] == "stdio":
+35 -11
View File
@@ -1,4 +1,4 @@
"""Token-at-rest encryption for OAuth-MCP. See docs/design/oauth-mcp.md §5.3.
"""Token-at-rest encryption for OAuth-MCP.
Uses cryptography.fernet (AES-128-CBC + HMAC-SHA256, 256-bit total key
material, encrypt-then-MAC). Single-key chosen for v1; rotation supported
@@ -342,10 +342,9 @@ class MCPTokenStore:
``refresh_token=None`` CLEARS the column it does NOT preserve
the existing value. Per RFC 6749 §6, an authorization server MAY
omit ``refresh_token`` from the refresh response; in that case
the caller (the OAuth flow that lands in a future phase) MUST
pre-resolve whether to keep the existing refresh token or drop
it before invoking this method. This API has no
"leave unchanged" sentinel.
the OAuth-flow caller MUST pre-resolve whether to keep the
existing refresh token or drop it before invoking this method.
This API has no "leave unchanged" sentinel.
"""
access_ct = self._cipher.encrypt(access_token.encode("utf-8"))
refresh_ct = self._cipher.encrypt(refresh_token.encode("utf-8")) if refresh_token else None
@@ -376,21 +375,46 @@ class MCPTokenStore:
secret_ct = self._cipher.encrypt(plaintext_secret.encode("utf-8"))
return self._storage.set_mcp_oauth_client_secret_ct(server_id, secret_ct)
def get_oauth_client_secret(self, server_id: str) -> str | None:
"""Decrypt and return the per-server OAuth client secret, or None.
Returns ``None`` when the row is missing or the column is NULL.
Raises :class:`MCPTokenDecryptError` on key mismatch the caller
decides whether to treat that as a missing-secret case (e.g. log +
prompt re-consent) or surface as a configuration failure.
"""
secret_ct = self._storage.get_mcp_oauth_client_secret_ct(server_id)
if secret_ct is None:
return None
return self._cipher.decrypt(secret_ct).decode("utf-8")
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _audit_decrypt_failure(self, server_name: str, fingerprints: tuple[str, ...]) -> None:
"""Best-effort audit emit on decrypt failure (no-op when unconfigured)."""
"""Best-effort audit emit on decrypt failure (no-op when unconfigured).
Uses ``server_id`` (PK UUID) as ``resource_id`` so admin-driven
server renames don't break event correlation. Falls back to
``server_name`` when the lookup misses.
"""
if self._audit_storage is None:
return
resource_id = server_name
try:
row = self._audit_storage.get_mcp_server_by_name(server_name)
except Exception:
row = None
if row is not None:
resource_id = str(row.get("server_id") or server_name)
try:
record_audit(
self._audit_storage,
user_id="",
action="mcp_server.oauth.token_decrypt_failure",
resource_type="mcp_server",
resource_id=server_name,
resource_id=resource_id,
detail={
"server_name": server_name,
"key_fingerprints_attempted": list(fingerprints),
@@ -399,7 +423,7 @@ class MCPTokenStore:
)
except Exception:
log.warning(
"mcp.oauth.audit_emit_failed",
"mcp_server.oauth.audit_emit_failed",
action="token_decrypt_failure",
server_name=server_name,
exc_info=True,
@@ -438,7 +462,7 @@ def initialize_mcp_crypto_state(app_state: object, *, node_id: str = "") -> None
try:
cipher_cfg = load_mcp_token_cipher_config()
except MCPTokenKeyConfigError as exc:
log.error("mcp.oauth.key_config_invalid: %s", exc)
log.error("mcp_server.oauth.key_config_invalid: %s", exc)
raise SystemExit(1) from exc
storage = get_storage()
@@ -462,7 +486,7 @@ def initialize_mcp_crypto_state(app_state: object, *, node_id: str = "") -> None
# exercised; install None sentinels so callers can fast-path.
app_state.mcp_token_cipher = None # type: ignore[attr-defined]
app_state.mcp_token_store = None # type: ignore[attr-defined]
log.debug("mcp.oauth.disabled (no key configured, no oauth_user rows)")
log.debug("mcp_server.oauth.disabled (no key configured, no oauth_user rows)")
return
cipher = MCPTokenCipher(cipher_cfg)
@@ -474,7 +498,7 @@ def initialize_mcp_crypto_state(app_state: object, *, node_id: str = "") -> None
audit_storage=storage,
)
log.info(
"mcp.oauth.cipher_installed",
"mcp_server.oauth.cipher_installed",
keys=len(cipher.key_fingerprints),
active_fp=cipher.key_fingerprints[0],
)
File diff suppressed because it is too large Load Diff
+224
View File
@@ -0,0 +1,224 @@
"""Shared SSRF and same-origin validation for OAuth/OIDC endpoint URLs.
Extracted from :mod:`turnstone.core.oidc` so the per-(user, server) MCP
OAuth flow (see :mod:`turnstone.core.mcp_oauth`) can reuse the exact same
guards without depending on the OIDC module.
The canonical exception is :class:`OAuthSSRFError`. The OIDC module wraps
calls to these helpers and re-raises ``OIDCError`` so its public API is
unchanged. The MCP OAuth module catches :class:`OAuthSSRFError` directly.
DNS-rebinding limitation: this module resolves the hostname during
validation, but the subsequent ``httpx`` call resolves again. A hostname
the operator points at could in principle rebind between the two resolves
to expose an internal address. Callers must ensure the AS / IdP hostname
is operator-controlled the SSRF guard prevents private-IP responses for
hostnames the operator points at, but does not prevent rebinding by a
hostile DNS authority. Pinning a single resolution into the ``httpx``
transport is a future hardening step.
"""
from __future__ import annotations
import asyncio
import ipaddress
import socket
import urllib.parse
# ---------------------------------------------------------------------------
# Trusted-host allowlist for well-known multi-origin IdPs / authorization
# servers whose discovery documents legitimately reference endpoints on
# hostnames distinct from the issuer hostname. eTLD+1 matching does not
# work here (e.g. google.com vs googleapis.com), so an explicit allow-map
# is the only safe option.
# ---------------------------------------------------------------------------
KNOWN_TRUSTED_OAUTH_ENDPOINT_HOSTS: dict[str, frozenset[str]] = {
"accounts.google.com": frozenset(
{
"accounts.google.com",
"oauth2.googleapis.com",
"www.googleapis.com",
"openidconnect.googleapis.com",
}
),
}
# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------
class OAuthSSRFError(Exception):
"""Raised when an SSRF/same-origin validation fails.
OIDC callers wrap this and re-raise as ``OIDCError`` to preserve the
existing public API.
"""
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def is_localhost(hostname: str) -> bool:
"""Return True if *hostname* refers to the loopback interface."""
return hostname in ("localhost", "127.0.0.1", "::1") or hostname.endswith(".localhost")
def sanitize_log_text(text: str, limit: int = 200) -> str:
"""Escape control characters and truncate untrusted text for log/audit inclusion.
Untrusted bytes (e.g. an AS error body, ``error_description`` from a
callback redirect) embedded in log lines or exception messages must not
be able to forge fake log records via CR/LF or hide content via NULs /
other control characters. ``unicode_escape`` renders these as visible
``\\r``, ``\\n``, ``\\x00`` etc., and *limit* caps the *rendered* length.
Shared with the OIDC module its private ``_sanitize_log_text`` is a
legacy alias that forwards here.
"""
if not text:
return ""
return text.encode("unicode_escape").decode("ascii")[:limit]
def effective_port(parsed: urllib.parse.ParseResult) -> int | None:
"""Return the explicit port if set, else the scheme default."""
if parsed.port is not None:
return parsed.port
return {"http": 80, "https": 443}.get(parsed.scheme)
def validate_url_no_ssrf(url: str, *, allow_http: bool) -> urllib.parse.ParseResult:
"""Run the scheme/userinfo/SSRF checks shared by issuer and discovered URLs.
Returns the parsed URL on success. Raises :class:`OAuthSSRFError` on
failure. The ``allow_http`` flag is the only knob: when ``True``,
``http://`` is accepted *if* the hostname is also a localhost form;
when ``False``, only ``https://`` is accepted.
"""
parsed = urllib.parse.urlparse(url)
hostname = parsed.hostname
if not hostname:
raise OAuthSSRFError(f"endpoint URL has no hostname: {url}")
if parsed.username or parsed.password:
raise OAuthSSRFError("endpoint URL must not contain embedded credentials (userinfo)")
if parsed.scheme != "https":
if allow_http and parsed.scheme == "http" and is_localhost(hostname):
pass
else:
raise OAuthSSRFError(f"endpoint URL must use HTTPS (got {parsed.scheme}://): {url}")
try:
addr_infos = socket.getaddrinfo(hostname, None, proto=socket.IPPROTO_TCP)
except socket.gaierror as exc:
raise OAuthSSRFError(f"endpoint hostname cannot be resolved: {hostname}") from exc
for _family, _type, _proto, _canonname, sockaddr in addr_infos:
try:
addr = ipaddress.ip_address(sockaddr[0])
except ValueError as exc:
raise OAuthSSRFError(
f"endpoint hostname resolved to invalid IP {sockaddr[0]!r}: {hostname}"
) from exc
if not addr.is_global and not is_localhost(hostname):
raise OAuthSSRFError(f"endpoint URL resolves to non-public address ({addr}): {url}")
return parsed
def validate_discovered_endpoint(
url: str,
issuer_parsed: urllib.parse.ParseResult,
*,
allow_http: bool,
trusted_endpoint_hosts: frozenset[str],
) -> None:
"""Validate an endpoint URL pulled from an OIDC/OAuth discovery document.
Applies :func:`validate_url_no_ssrf` plus the same-origin / trusted-host
constraint: the endpoint host must equal the issuer host, be in the
well-known trust map, or be in the operator-supplied
``trusted_endpoint_hosts``. Effective port (with scheme defaults
applied) and scheme must match the issuer.
Raises :class:`OAuthSSRFError` on validation failure.
"""
parsed = validate_url_no_ssrf(url, allow_http=allow_http)
issuer_hostname = (issuer_parsed.hostname or "").lower()
endpoint_hostname = (parsed.hostname or "").lower()
if parsed.scheme != issuer_parsed.scheme:
raise OAuthSSRFError(
f"discovered endpoint scheme ({parsed.scheme}) "
f"does not match issuer ({issuer_parsed.scheme}): {url}"
)
known_trusted = KNOWN_TRUSTED_OAUTH_ENDPOINT_HOSTS.get(issuer_hostname, frozenset())
host_allowed = (
endpoint_hostname == issuer_hostname
or endpoint_hostname in known_trusted
or endpoint_hostname in trusted_endpoint_hosts
)
if not host_allowed:
raise OAuthSSRFError(
f"discovered endpoint host ({endpoint_hostname}) "
f"does not match issuer ({issuer_hostname}) and is not trusted: {url}"
)
endpoint_port = effective_port(parsed)
issuer_port = effective_port(issuer_parsed)
if endpoint_port != issuer_port:
raise OAuthSSRFError(
f"discovered endpoint port ({endpoint_port}) "
f"does not match issuer ({issuer_port}): {url}"
)
async def validate_url_no_ssrf_async(url: str, *, allow_http: bool) -> urllib.parse.ParseResult:
"""Async variant of :func:`validate_url_no_ssrf` for hot-path callers.
The synchronous variant calls ``socket.getaddrinfo``, which blocks
the event loop. Async OAuth flows (notably
:mod:`turnstone.core.mcp_oauth`) wrap their validation calls in
:func:`asyncio.to_thread` to keep the loop responsive. This wrapper
centralises that wrapping so callers don't repeat the idiom.
"""
return await asyncio.to_thread(validate_url_no_ssrf, url, allow_http=allow_http)
async def validate_discovered_endpoint_async(
url: str,
issuer_parsed: urllib.parse.ParseResult,
*,
allow_http: bool,
trusted_endpoint_hosts: frozenset[str],
) -> None:
"""Async variant of :func:`validate_discovered_endpoint`."""
await asyncio.to_thread(
validate_discovered_endpoint,
url,
issuer_parsed,
allow_http=allow_http,
trusted_endpoint_hosts=trusted_endpoint_hosts,
)
__all__ = [
"KNOWN_TRUSTED_OAUTH_ENDPOINT_HOSTS",
"OAuthSSRFError",
"effective_port",
"is_localhost",
"sanitize_log_text",
"validate_discovered_endpoint",
"validate_discovered_endpoint_async",
"validate_url_no_ssrf",
"validate_url_no_ssrf_async",
]
+34 -103
View File
@@ -11,11 +11,9 @@ import asyncio
import base64
import dataclasses
import hashlib
import ipaddress
import os
import re
import secrets
import socket
import urllib.parse
import uuid
from dataclasses import dataclass, field
@@ -27,6 +25,17 @@ if TYPE_CHECKING:
from collections.abc import Mapping
from turnstone.core.log import get_logger
from turnstone.core.oauth_ssrf import (
KNOWN_TRUSTED_OAUTH_ENDPOINT_HOSTS,
OAuthSSRFError,
is_localhost,
)
from turnstone.core.oauth_ssrf import (
validate_discovered_endpoint as _ssrf_validate_discovered_endpoint,
)
from turnstone.core.oauth_ssrf import (
validate_url_no_ssrf as _ssrf_validate_url_no_ssrf,
)
log = get_logger(__name__)
@@ -58,21 +67,9 @@ _ALLOWED_ID_TOKEN_ALGS = [
"PS512",
]
# Well-known IdPs whose discovery documents legitimately reference endpoints on
# hostnames distinct from the issuer hostname. Keys are issuer hostnames; values
# are the set of additional endpoint hostnames the issuer is allowed to delegate
# to. eTLD+1 matching does not work here (e.g. google.com vs googleapis.com),
# so an explicit allow-map is the only safe option.
_KNOWN_TRUSTED_ENDPOINT_HOSTS: dict[str, frozenset[str]] = {
"accounts.google.com": frozenset(
{
"accounts.google.com",
"oauth2.googleapis.com",
"www.googleapis.com",
"openidconnect.googleapis.com",
}
),
}
# Re-export the shared trusted-host map under the legacy OIDC name so existing
# callers / tests continue to work without churn.
_KNOWN_TRUSTED_ENDPOINT_HOSTS = KNOWN_TRUSTED_OAUTH_ENDPOINT_HOSTS
# ---------------------------------------------------------------------------
@@ -94,14 +91,10 @@ class OIDCKeyNotFoundError(OIDCError):
def _sanitize_log_text(s: str, limit: int) -> str:
"""Escape control characters and truncate untrusted text for log inclusion.
"""Legacy alias for the shared :func:`oauth_ssrf.sanitize_log_text`."""
from turnstone.core.oauth_ssrf import sanitize_log_text
Untrusted bytes (e.g. an IdP error body) embedded in log lines must not be
able to forge fake log records via CR/LF or hide content via NULs / other
control characters. ``unicode_escape`` renders these as visible ``\\r``,
``\\n``, ``\\x00`` etc., and the limit caps the *rendered* length.
"""
return s.encode("unicode_escape").decode("ascii")[:limit]
return sanitize_log_text(s, limit)
# ---------------------------------------------------------------------------
@@ -293,60 +286,19 @@ def load_oidc_config() -> OIDCConfig:
# ---------------------------------------------------------------------------
# SSRF validation
#
# The actual checks live in :mod:`turnstone.core.oauth_ssrf` so the MCP OAuth
# flow can reuse them. The wrappers below preserve OIDC's public API by
# converting :class:`OAuthSSRFError` to :class:`OIDCError`.
# ---------------------------------------------------------------------------
def _is_localhost(hostname: str) -> bool:
"""Return True if *hostname* refers to the loopback interface."""
return hostname in ("localhost", "127.0.0.1", "::1") or hostname.endswith(".localhost")
def _effective_port(parsed: urllib.parse.ParseResult) -> int | None:
"""Return the explicit port if set, else the scheme default."""
if parsed.port is not None:
return parsed.port
return {"http": 80, "https": 443}.get(parsed.scheme)
def _validate_url_no_ssrf(url: str, *, allow_http: bool) -> urllib.parse.ParseResult:
"""Run the scheme/userinfo/SSRF checks shared by issuer and discovered URLs.
Returns the parsed URL on success. Raises :class:`OIDCError` on failure.
The ``allow_http`` flag is the only knob: when ``True``, ``http://`` is
accepted *if* the hostname is also a localhost form; when ``False``,
only ``https://`` is accepted.
"""
parsed = urllib.parse.urlparse(url)
hostname = parsed.hostname
if not hostname:
raise OIDCError(f"OIDC URL has no hostname: {url}")
if parsed.username or parsed.password:
raise OIDCError("OIDC URL must not contain embedded credentials (userinfo)")
if parsed.scheme != "https":
if allow_http and parsed.scheme == "http" and _is_localhost(hostname):
pass
else:
raise OIDCError(f"OIDC URL must use HTTPS (got {parsed.scheme}://): {url}")
"""OIDC-flavoured wrapper around :func:`oauth_ssrf.validate_url_no_ssrf`."""
try:
addr_infos = socket.getaddrinfo(hostname, None, proto=socket.IPPROTO_TCP)
except socket.gaierror as exc:
raise OIDCError(f"OIDC hostname cannot be resolved: {hostname}") from exc
for _family, _type, _proto, _canonname, sockaddr in addr_infos:
try:
addr = ipaddress.ip_address(sockaddr[0])
except ValueError as exc:
raise OIDCError(
f"OIDC hostname resolved to invalid IP {sockaddr[0]!r}: {hostname}"
) from exc
if not addr.is_global and not _is_localhost(hostname):
raise OIDCError(f"OIDC URL resolves to non-public address ({addr}): {url}")
return parsed
return _ssrf_validate_url_no_ssrf(url, allow_http=allow_http)
except OAuthSSRFError as exc:
raise OIDCError(str(exc)) from exc
def validate_issuer_url(url: str) -> None:
@@ -388,36 +340,15 @@ def validate_discovered_endpoint(
Raises :class:`OIDCError` on validation failure.
"""
parsed = _validate_url_no_ssrf(url, allow_http=allow_http)
issuer_hostname = (issuer_parsed.hostname or "").lower()
endpoint_hostname = (parsed.hostname or "").lower()
if parsed.scheme != issuer_parsed.scheme:
raise OIDCError(
f"OIDC discovered endpoint scheme ({parsed.scheme}) "
f"does not match issuer ({issuer_parsed.scheme}): {url}"
)
known_trusted = _KNOWN_TRUSTED_ENDPOINT_HOSTS.get(issuer_hostname, frozenset())
host_allowed = (
endpoint_hostname == issuer_hostname
or endpoint_hostname in known_trusted
or endpoint_hostname in trusted_endpoint_hosts
)
if not host_allowed:
raise OIDCError(
f"OIDC discovered endpoint host ({endpoint_hostname}) "
f"does not match issuer ({issuer_hostname}) and is not trusted: {url}"
)
endpoint_port = _effective_port(parsed)
issuer_port = _effective_port(issuer_parsed)
if endpoint_port != issuer_port:
raise OIDCError(
f"OIDC discovered endpoint port ({endpoint_port}) "
f"does not match issuer ({issuer_port}): {url}"
try:
_ssrf_validate_discovered_endpoint(
url,
issuer_parsed,
allow_http=allow_http,
trusted_endpoint_hosts=trusted_endpoint_hosts,
)
except OAuthSSRFError as exc:
raise OIDCError(str(exc)) from exc
# ---------------------------------------------------------------------------
@@ -483,7 +414,7 @@ async def discover_oidc(
)
return dataclasses.replace(config, enabled=False)
allow_http = _is_localhost(issuer_parsed.hostname or "")
allow_http = is_localhost(issuer_parsed.hostname or "")
trusted_hosts = frozenset(h.lower() for h in config.trusted_endpoint_hosts)
required = (
("authorization_endpoint", authorization_endpoint),
+104 -1
View File
@@ -13,7 +13,12 @@ if TYPE_CHECKING:
import sqlalchemy as sa
from turnstone.core.log import get_logger
from turnstone.core.storage._protocol import MCPUserToken, OIDCIdentity, OIDCPendingState
from turnstone.core.storage._protocol import (
MCPOAuthPendingState,
MCPUserToken,
OIDCIdentity,
OIDCPendingState,
)
from turnstone.core.storage._schema import (
api_tokens,
audit_events,
@@ -22,6 +27,7 @@ from turnstone.core.storage._schema import (
conversations,
heuristic_rules,
intent_verdicts,
mcp_oauth_pending,
mcp_servers,
mcp_user_tokens,
metadata,
@@ -1177,6 +1183,8 @@ class PostgreSQLBackend:
conn.execute(sa.delete(channel_users).where(channel_users.c.user_id == user_id))
conn.execute(sa.delete(api_tokens).where(api_tokens.c.user_id == user_id))
conn.execute(sa.delete(oidc_identities).where(oidc_identities.c.user_id == user_id))
conn.execute(sa.delete(mcp_user_tokens).where(mcp_user_tokens.c.user_id == user_id))
conn.execute(sa.delete(mcp_oauth_pending).where(mcp_oauth_pending.c.user_id == user_id))
result = conn.execute(sa.delete(users).where(users.c.user_id == user_id))
conn.commit()
return result.rowcount > 0
@@ -3937,6 +3945,101 @@ class PostgreSQLBackend:
conn.commit()
return result.rowcount > 0
def delete_mcp_oauth_rows_by_server_name(self, server_name: str) -> int:
"""Purge user tokens + pending OAuth state for *server_name*."""
with self._conn() as conn:
tokens_result = conn.execute(
sa.delete(mcp_user_tokens).where(mcp_user_tokens.c.server_name == server_name)
)
pending_result = conn.execute(
sa.delete(mcp_oauth_pending).where(mcp_oauth_pending.c.server_name == server_name)
)
conn.commit()
return int(tokens_result.rowcount or 0) + int(pending_result.rowcount or 0)
def get_mcp_oauth_client_secret_ct(self, server_id: str) -> bytes | None:
"""Return the encrypted OAuth client secret column or None."""
with self._conn() as conn:
row = conn.execute(
sa.select(mcp_servers.c.oauth_client_secret_ct).where(
mcp_servers.c.server_id == server_id
)
).fetchone()
if row is None or row[0] is None:
return None
return bytes(row[0])
# -- MCP OAuth pending state (per-(user, server) flow) ---------------------
def create_mcp_oauth_pending_state(
self,
state: str,
user_id: str,
server_name: str,
code_verifier: str,
return_url: str,
) -> None:
now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
with self._conn() as conn:
conn.execute(
sa.insert(mcp_oauth_pending),
{
"state": state,
"user_id": user_id,
"server_name": server_name,
"code_verifier": code_verifier,
"return_url": return_url,
"created_at": now,
},
)
conn.commit()
def pop_mcp_oauth_pending_state(
self, state: str, max_age_seconds: int = 600
) -> MCPOAuthPendingState | None:
cutoff = (datetime.now(UTC) - timedelta(seconds=max_age_seconds)).strftime(
"%Y-%m-%dT%H:%M:%S"
)
with self._conn() as conn:
# Atomic DELETE...RETURNING for true one-time consumption
row = conn.execute(
sa.text(
"DELETE FROM mcp_oauth_pending "
"WHERE state = :state AND created_at > :cutoff "
"RETURNING state, user_id, server_name, code_verifier, "
"return_url, created_at"
),
{"state": state, "cutoff": cutoff},
).fetchone()
# Also clean up the row if it existed but was expired
if not row:
conn.execute(sa.delete(mcp_oauth_pending).where(mcp_oauth_pending.c.state == state))
conn.commit()
if not row:
return None
return MCPOAuthPendingState(
state=row[0],
user_id=row[1],
server_name=row[2],
code_verifier=row[3],
return_url=row[4],
created_at=row[5],
)
def cleanup_expired_mcp_oauth_pending_states(self, max_age_seconds: int = 600) -> int:
cutoff = (datetime.now(UTC) - timedelta(seconds=max_age_seconds)).strftime(
"%Y-%m-%dT%H:%M:%S"
)
with self._conn() as conn:
result = conn.execute(
sa.delete(mcp_oauth_pending).where(mcp_oauth_pending.c.created_at < cutoff)
)
conn.commit()
return result.rowcount
# -- Model definitions -----------------------------------------------------
def create_model_definition(
+54 -4
View File
@@ -42,10 +42,9 @@ class OIDCPendingState(TypedDict):
class MCPUserToken(TypedDict):
"""Row shape returned by per-(user, MCP server) OAuth token lookups.
See ``docs/design/oauth-mcp.md`` §5.1. ``access_token_ct`` and
``refresh_token_ct`` are Fernet ciphertext blobs; the storage layer
returns them verbatim and ``MCPTokenStore`` (Phase 3) handles
encrypt/decrypt.
``access_token_ct`` and ``refresh_token_ct`` are Fernet ciphertext
blobs; the storage layer returns them verbatim and ``MCPTokenStore``
handles encrypt/decrypt.
"""
user_id: str
@@ -1708,6 +1707,57 @@ class StorageBackend(Protocol):
"""Delete the per-(user, server) token row. Returns True if existed."""
...
def delete_mcp_oauth_rows_by_server_name(self, server_name: str) -> int:
"""Purge per-(user, server) tokens and pending OAuth states for *server_name*.
Used when the operator renames or deletes an MCP server row to
prevent old user tokens from rebinding to a freshly-created
server with the same ``name``. Returns the total number of rows
deleted across both tables.
Both ``mcp_user_tokens`` and ``mcp_oauth_pending`` are keyed on
the mutable ``server_name`` rather than the immutable
``server_id``; until those tables migrate to a server_id FK with
ON DELETE CASCADE (a future schema migration), explicit purge on
rename/delete is the only safe path.
"""
...
def get_mcp_oauth_client_secret_ct(self, server_id: str) -> bytes | None:
"""Return the encrypted OAuth client secret column or None.
Mirror of :meth:`set_mcp_oauth_client_secret_ct` for the read path.
Returns ``None`` when the row does not exist or the column is NULL.
"""
...
# -- MCP OAuth pending state (per-(user, server) flow) ---------------------
def create_mcp_oauth_pending_state(
self,
state: str,
user_id: str,
server_name: str,
code_verifier: str,
return_url: str,
) -> None:
"""Insert a pending MCP OAuth flow row for callback validation."""
...
def pop_mcp_oauth_pending_state(
self, state: str, max_age_seconds: int = 600
) -> MCPOAuthPendingState | None:
"""Atomically fetch+delete a pending MCP OAuth row.
Returns ``None`` when the row is missing or older than
``max_age_seconds``.
"""
...
def cleanup_expired_mcp_oauth_pending_states(self, max_age_seconds: int = 600) -> int:
"""Bulk-delete expired pending MCP OAuth rows. Returns count deleted."""
...
# -- Model definitions -----------------------------------------------------
def create_model_definition(
+3 -3
View File
@@ -607,7 +607,7 @@ mcp_servers = sa.Table(
sa.Column("registry_name", sa.Text, nullable=True),
sa.Column("registry_version", sa.Text, nullable=False, server_default=""),
sa.Column("registry_meta", sa.Text, nullable=False, server_default="{}"),
# Per-(user, server) OAuth 2.1 — see docs/design/oauth-mcp.md §5.2.
# Per-(user, server) OAuth 2.1 columns.
# `auth_type` is one of: 'none', 'static', 'oauth_user'. The other
# `oauth_*` columns are NULL when auth_type != 'oauth_user'.
# `oauth_client_secret_ct` is Fernet ciphertext; never decrypted on
@@ -708,8 +708,8 @@ oidc_pending_states = sa.Table(
# ---------------------------------------------------------------------------
# MCP per-(user, server) OAuth tokens and pending authorization-flow state.
# See docs/design/oauth-mcp.md §5.1. No FKs at the schema level (matches
# `oidc_*` tables; tests avoid orphan rows via fixtures).
# No FKs at the schema level (matches `oidc_*` tables; tests avoid orphan
# rows via fixtures).
# ---------------------------------------------------------------------------
mcp_user_tokens = sa.Table(
+107 -1
View File
@@ -13,7 +13,12 @@ if TYPE_CHECKING:
from collections.abc import Iterator
from turnstone.core.log import get_logger
from turnstone.core.storage._protocol import MCPUserToken, OIDCIdentity, OIDCPendingState
from turnstone.core.storage._protocol import (
MCPOAuthPendingState,
MCPUserToken,
OIDCIdentity,
OIDCPendingState,
)
from turnstone.core.storage._schema import (
api_tokens,
audit_events,
@@ -22,6 +27,7 @@ from turnstone.core.storage._schema import (
conversations,
heuristic_rules,
intent_verdicts,
mcp_oauth_pending,
mcp_servers,
mcp_user_tokens,
metadata,
@@ -1332,6 +1338,8 @@ class SQLiteBackend:
conn.execute(sa.delete(channel_users).where(channel_users.c.user_id == user_id))
conn.execute(sa.delete(api_tokens).where(api_tokens.c.user_id == user_id))
conn.execute(sa.delete(oidc_identities).where(oidc_identities.c.user_id == user_id))
conn.execute(sa.delete(mcp_user_tokens).where(mcp_user_tokens.c.user_id == user_id))
conn.execute(sa.delete(mcp_oauth_pending).where(mcp_oauth_pending.c.user_id == user_id))
result = conn.execute(sa.delete(users).where(users.c.user_id == user_id))
conn.commit()
return result.rowcount > 0
@@ -4083,6 +4091,104 @@ class SQLiteBackend:
conn.commit()
return result.rowcount > 0
def delete_mcp_oauth_rows_by_server_name(self, server_name: str) -> int:
"""Purge user tokens + pending OAuth state for *server_name*."""
with self._conn() as conn:
tokens_result = conn.execute(
sa.delete(mcp_user_tokens).where(mcp_user_tokens.c.server_name == server_name)
)
pending_result = conn.execute(
sa.delete(mcp_oauth_pending).where(mcp_oauth_pending.c.server_name == server_name)
)
conn.commit()
return int(tokens_result.rowcount or 0) + int(pending_result.rowcount or 0)
def get_mcp_oauth_client_secret_ct(self, server_id: str) -> bytes | None:
"""Return the encrypted OAuth client secret column or None."""
with self._conn() as conn:
row = conn.execute(
sa.select(mcp_servers.c.oauth_client_secret_ct).where(
mcp_servers.c.server_id == server_id
)
).fetchone()
if row is None or row[0] is None:
return None
return bytes(row[0])
# -- MCP OAuth pending state (per-(user, server) flow) ---------------------
def create_mcp_oauth_pending_state(
self,
state: str,
user_id: str,
server_name: str,
code_verifier: str,
return_url: str,
) -> None:
now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
with self._conn() as conn:
conn.execute(
sa.insert(mcp_oauth_pending),
{
"state": state,
"user_id": user_id,
"server_name": server_name,
"code_verifier": code_verifier,
"return_url": return_url,
"created_at": now,
},
)
conn.commit()
def pop_mcp_oauth_pending_state(
self, state: str, max_age_seconds: int = 600
) -> MCPOAuthPendingState | None:
cutoff = (datetime.now(UTC) - timedelta(seconds=max_age_seconds)).strftime(
"%Y-%m-%dT%H:%M:%S"
)
with self._conn() as conn:
# Acquire write lock before SELECT to prevent TOCTOU race
conn.execute(sa.text("BEGIN IMMEDIATE"))
row = conn.execute(
sa.select(
mcp_oauth_pending.c.state,
mcp_oauth_pending.c.user_id,
mcp_oauth_pending.c.server_name,
mcp_oauth_pending.c.code_verifier,
mcp_oauth_pending.c.return_url,
mcp_oauth_pending.c.created_at,
).where(
(mcp_oauth_pending.c.state == state) & (mcp_oauth_pending.c.created_at > cutoff)
)
).fetchone()
# Always delete the row (whether valid, expired, or missing is fine)
conn.execute(sa.delete(mcp_oauth_pending).where(mcp_oauth_pending.c.state == state))
conn.commit()
if not row:
return None
return MCPOAuthPendingState(
state=row[0],
user_id=row[1],
server_name=row[2],
code_verifier=row[3],
return_url=row[4],
created_at=row[5],
)
def cleanup_expired_mcp_oauth_pending_states(self, max_age_seconds: int = 600) -> int:
cutoff = (datetime.now(UTC) - timedelta(seconds=max_age_seconds)).strftime(
"%Y-%m-%dT%H:%M:%S"
)
with self._conn() as conn:
result = conn.execute(
sa.delete(mcp_oauth_pending).where(mcp_oauth_pending.c.created_at < cutoff)
)
conn.commit()
return result.rowcount
# -- Model definitions -----------------------------------------------------
def create_model_definition(
@@ -1,9 +1,8 @@
"""Add OAuth-MCP schema (Phase 2).
"""Add OAuth-MCP schema.
Adds two new tables (``mcp_user_tokens``, ``mcp_oauth_pending``) and
eight new columns on ``mcp_servers`` to support per-(user, server)
OAuth 2.1 + PKCE authorization for MCP servers. See
``docs/design/oauth-mcp.md`` §5.1 / §5.2.
OAuth 2.1 + PKCE authorization for MCP servers.
Existing rows continue working: every new column is either nullable or
defaults to ``'static'`` (the existing behavior). After the new
+32 -4
View File
@@ -2702,6 +2702,20 @@ async def oidc_callback(request: Request) -> Response:
return await handle_oidc_callback(request, JWT_AUD_SERVER)
async def mcp_oauth_authorize(request: Request) -> Response:
"""GET /v1/api/mcp/oauth/start — begin per-(user, server) OAuth flow."""
from turnstone.core.mcp_oauth import handle_mcp_oauth_authorize
return await handle_mcp_oauth_authorize(request)
async def mcp_oauth_callback(request: Request) -> Response:
"""GET /v1/api/mcp/oauth/callback — AS-redirected OAuth callback."""
from turnstone.core.mcp_oauth import handle_mcp_oauth_callback
return await handle_mcp_oauth_callback(request)
def list_interface_settings(request: Request) -> JSONResponse:
"""GET /v1/api/admin/settings — return interface settings from ConfigStore.
@@ -3482,12 +3496,17 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
await initialize_oidc_state(app.state)
# MCP-OAuth token-at-rest encryption — fail-loud on misconfiguration
# when any mcp_servers row has auth_type='oauth_user'. See
# docs/design/oauth-mcp.md §5.3.
# when any mcp_servers row has auth_type='oauth_user'.
from turnstone.core.mcp_crypto import initialize_mcp_crypto_state
initialize_mcp_crypto_state(app.state, node_id=getattr(app.state, "node_id", ""))
# Per-(user, server) OAuth flow state — long-lived HTTP client +
# in-process refresh lock + metadata cache.
from turnstone.core.mcp_oauth import initialize_mcp_oauth_state
await initialize_mcp_oauth_state(app.state)
# TLS: start auto-renewal if client was initialized
tls_client = getattr(app.state, "tls_client", None)
if tls_client is not None:
@@ -3635,12 +3654,19 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
app.state.mcp_client.shutdown()
if app.state.registry:
app.state.registry.shutdown()
from turnstone.core.oidc import close_oidc_state
# Close in reverse order of initialization (mcp_oauth → mcp_crypto →
# oidc). The OAuth flow holds a long-lived httpx.AsyncClient that
# depends on no later-initialised state, but reversing init order
# is the conventional LIFO discipline.
from turnstone.core.mcp_oauth import close_mcp_oauth_state
await close_oidc_state(app.state)
await close_mcp_oauth_state(app.state)
from turnstone.core.mcp_crypto import close_mcp_crypto_state
close_mcp_crypto_state(app.state)
from turnstone.core.oidc import close_oidc_state
await close_oidc_state(app.state)
app.state.sse_executor.shutdown(wait=True, cancel_futures=True)
@@ -3876,6 +3902,8 @@ def create_app(
Route("/api/auth/refresh", auth_refresh, methods=["POST"]),
Route("/api/auth/oidc/authorize", oidc_authorize),
Route("/api/auth/oidc/callback", oidc_callback),
Route("/api/mcp/oauth/start", mcp_oauth_authorize),
Route("/api/mcp/oauth/callback", mcp_oauth_callback),
Route("/api/admin/settings", list_interface_settings),
Route(
"/api/admin/settings/{key:path}",