mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
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:
@@ -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``.
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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"
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
@@ -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"]),
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
@@ -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
@@ -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),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
@@ -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}",
|
||||
|
||||
Reference in New Issue
Block a user