diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 21abaa3f98..542b6e45d2 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,4 +1,5 @@ import base64 +import binascii import json from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast @@ -498,6 +499,47 @@ async def rotate_mcp_server_credentials_master_key( ) +def _decode_user_credential(stored: str) -> Optional[str]: + """Read back a value persisted in ``LiteLLM_MCPUserCredentials.credential_b64``. + + Tries nacl decryption first (current write format). Falls back to a + plain ``urlsafe_b64decode`` for rows persisted by older code that wrote + the credential without encryption. Returns ``None`` when neither path + yields a valid string. + """ + decrypted = decrypt_value_helper( + value=stored, + key="mcp_user_credential", + exception_type="debug", + return_original_value=False, + ) + if decrypted is not None: + return decrypted + try: + return base64.urlsafe_b64decode(stored).decode() + except (binascii.Error, UnicodeDecodeError, ValueError): + return None + + +def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]: + """Return the OAuth2 payload dict if ``stored`` holds one, else ``None``. + + A row is considered an OAuth2 credential iff its decoded value parses as + a JSON object with ``"type": "oauth2"``. Plain BYOK credentials (which + share the same column) decode to a non-JSON string and return ``None``. + """ + decoded = _decode_user_credential(stored) + if decoded is None: + return None + try: + parsed = json.loads(decoded) + except (ValueError, TypeError): + return None + if isinstance(parsed, dict) and parsed.get("type") == "oauth2": + return parsed + return None + + async def store_user_credential( prisma_client: PrismaClient, user_id: str, @@ -506,7 +548,7 @@ async def store_user_credential( ) -> None: """Store a user credential for a BYOK MCP server.""" - encoded = base64.urlsafe_b64encode(credential.encode()).decode() + encoded = encrypt_value_helper(credential) await prisma_client.db.litellm_mcpusercredentials.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ @@ -532,16 +574,7 @@ async def get_user_credential( ) if row is None: return None - try: - return base64.urlsafe_b64decode(row.credential_b64).decode() - except Exception: - # Fall back to nacl decryption for credentials stored by older code - return decrypt_value_helper( - value=row.credential_b64, - key="byok_credential", - exception_type="debug", - return_original_value=False, - ) + return _decode_user_credential(row.credential_b64) async def has_user_credential( @@ -582,7 +615,7 @@ async def store_user_oauth_credential( ) -> None: """Persist an OAuth2 access token for a user+server pair. - The payload is JSON-serialised and stored base64-encoded in the same + The payload is JSON-serialised and stored encrypted in the same ``credential_b64`` column used by BYOK. A ``"type": "oauth2"`` key differentiates it from plain BYOK API keys. """ @@ -606,29 +639,22 @@ async def store_user_oauth_credential( payload["scopes"] = scopes # Guard against silently overwriting a BYOK credential with an OAuth token. - # BYOK credentials lack a "type" field (or use a non-"oauth2" type). # Skip the guard when the caller knows the row is already an OAuth2 credential # (e.g. during token refresh), saving an extra DB round-trip. if not skip_byok_guard: existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) - if existing is not None: - _byok_error = ValueError( + if ( + existing is not None + and _decode_oauth_payload(existing.credential_b64) is None + ): + raise ValueError( f"A non-OAuth2 credential already exists for user {user_id} " f"and server {server_id}. Refusing to overwrite." ) - try: - raw = json.loads( - base64.urlsafe_b64decode(existing.credential_b64).decode() - ) - except Exception: - # Credential is not base64+JSON — it's a plain-text BYOK key. - raise _byok_error - if raw.get("type") != "oauth2": - raise _byok_error - encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() + encoded = encrypt_value_helper(json.dumps(payload)) await prisma_client.db.litellm_mcpusercredentials.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ @@ -672,15 +698,7 @@ async def get_user_oauth_credential( ) if row is None: return None - try: - decoded = base64.urlsafe_b64decode(row.credential_b64).decode() - parsed = json.loads(decoded) - if isinstance(parsed, dict) and parsed.get("type") == "oauth2": - return parsed - # Row exists but is a BYOK (plain string), not an OAuth token - return None - except Exception: - return None + return _decode_oauth_payload(row.credential_b64) async def list_user_oauth_credentials( @@ -694,14 +712,11 @@ async def list_user_oauth_credentials( ) results: List[Dict[str, Any]] = [] for row in rows: - try: - decoded = base64.urlsafe_b64decode(row.credential_b64).decode() - parsed = json.loads(decoded) - if isinstance(parsed, dict) and parsed.get("type") == "oauth2": - parsed["server_id"] = row.server_id - results.append(parsed) - except Exception: - pass # Skip non-OAuth rows (BYOK plain strings) + payload = _decode_oauth_payload(row.credential_b64) + if payload is None: + continue + payload["server_id"] = row.server_id + results.append(payload) return results diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py new file mode 100644 index 0000000000..ca6992419c --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -0,0 +1,293 @@ +""" +Tests for the encrypted-at-rest persistence of MCP user credentials. + +The ``LiteLLM_MCPUserCredentials.credential_b64`` column previously stored +both BYOK API keys and OAuth2 access tokens as plain ``urlsafe_b64encode`` +of the raw value, leaving credentials readable from any DB read. The fix +runs every write through ``encrypt_value_helper`` (nacl SecretBox) and +keeps a plain-base64 fallback on read so existing rows continue to work. +""" + +import base64 +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.proxy._experimental.mcp_server.db import ( # noqa: E402 + _decode_user_credential, + get_user_credential, + get_user_oauth_credential, + list_user_oauth_credentials, + store_user_credential, + store_user_oauth_credential, +) + + +SALT_KEY = "test-salt-key-for-byok-credential-tests-1234" + + +@pytest.fixture(autouse=True) +def _set_salt_key(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY) + + +def _make_prisma_with_existing(row): + """Build a MagicMock prisma_client whose user-credentials table returns ``row`` + for find_unique and behaves async-correctly for upsert/find_many.""" + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_mcpusercredentials.upsert = AsyncMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[]) + return prisma + + +def _legacy_row(payload: str): + """A row exactly as the pre-fix code would have written it: plain + ``urlsafe_b64encode`` of the raw payload, no encryption.""" + row = MagicMock() + row.credential_b64 = base64.urlsafe_b64encode(payload.encode()).decode() + row.user_id = "alice" + row.server_id = "srv-1" + return row + + +def _stored_value(prisma) -> str: + """Pull the credential_b64 value passed to the most recent upsert call.""" + call = prisma.db.litellm_mcpusercredentials.upsert.call_args + data = call.kwargs["data"] + create_value = data["create"]["credential_b64"] + update_value = data["update"]["credential_b64"] + assert create_value == update_value, "create/update must agree" + return create_value + + +# ── BYOK round-trip ─────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_store_user_credential_does_not_persist_plaintext(): + # Stored bytes must not just be base64 of the secret — that's the regression. + secret = "sk-proj-very-secret-byok-key" + prisma = _make_prisma_with_existing(row=None) + + await store_user_credential(prisma, "alice", "srv-1", secret) + + stored = _stored_value(prisma) + plain_b64 = base64.urlsafe_b64encode(secret.encode()).decode() + assert stored != plain_b64 + # And the secret must not appear anywhere in a plain-b64 decode of the column. + try: + decoded_bytes = base64.urlsafe_b64decode(stored) + except Exception: + decoded_bytes = b"" + assert secret.encode() not in decoded_bytes + + +@pytest.mark.asyncio +async def test_byok_round_trip_returns_plaintext(): + secret = "sk-proj-very-secret-byok-key" + prisma = _make_prisma_with_existing(row=None) + await store_user_credential(prisma, "alice", "srv-1", secret) + + stored = _stored_value(prisma) + row = MagicMock() + row.credential_b64 = stored + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + + result = await get_user_credential(prisma, "alice", "srv-1") + assert result == secret + + +@pytest.mark.asyncio +async def test_byok_get_returns_plaintext_for_legacy_row(): + # Backward-compat: rows persisted by the pre-fix code (plain base64) must + # still decrypt-or-decode cleanly. + legacy_secret = "legacy-byok-key" + prisma = _make_prisma_with_existing(row=_legacy_row(legacy_secret)) + + result = await get_user_credential(prisma, "alice", "srv-1") + assert result == legacy_secret + + +@pytest.mark.asyncio +async def test_byok_get_returns_none_for_missing_row(): + prisma = _make_prisma_with_existing(row=None) + result = await get_user_credential(prisma, "alice", "srv-1") + assert result is None + + +# ── OAuth2 round-trip ───────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_store_user_oauth_credential_does_not_persist_plaintext(): + access_token = "ya29.a0AfH6SMBverysecretaccesstoken" + prisma = _make_prisma_with_existing(row=None) + + await store_user_oauth_credential( + prisma, "alice", "srv-1", access_token, refresh_token="rfr-xyz" + ) + + stored = _stored_value(prisma) + try: + decoded_bytes = base64.urlsafe_b64decode(stored) + except Exception: + decoded_bytes = b"" + assert access_token.encode() not in decoded_bytes + assert b"rfr-xyz" not in decoded_bytes + + +@pytest.mark.asyncio +async def test_oauth_round_trip_returns_payload(): + access_token = "ya29.a0AfH6SMBverysecretaccesstoken" + prisma = _make_prisma_with_existing(row=None) + await store_user_oauth_credential( + prisma, + "alice", + "srv-1", + access_token, + refresh_token="rfr-xyz", + scopes=["a", "b"], + ) + + stored = _stored_value(prisma) + row = MagicMock() + row.credential_b64 = stored + row.server_id = "srv-1" + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is not None + assert result["type"] == "oauth2" + assert result["access_token"] == access_token + assert result["refresh_token"] == "rfr-xyz" + assert result["scopes"] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_oauth_get_returns_payload_for_legacy_row(): + payload = { + "type": "oauth2", + "access_token": "legacy-token", + "connected_at": "2024-01-01T00:00:00Z", + } + legacy_b64 = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() + row = MagicMock() + row.credential_b64 = legacy_b64 + row.server_id = "srv-1" + prisma = _make_prisma_with_existing(row=row) + + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is not None + assert result["access_token"] == "legacy-token" + + +@pytest.mark.asyncio +async def test_oauth_get_returns_none_for_byok_row(): + # A row that holds a BYOK string must not leak as an OAuth payload. + prisma = _make_prisma_with_existing(row=_legacy_row("plain-byok-not-json")) + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is None + + +# ── BYOK guard inside store_user_oauth_credential ───────────────────────────── + + +@pytest.mark.asyncio +async def test_byok_guard_rejects_overwriting_legacy_byok(): + prisma = _make_prisma_with_existing(row=_legacy_row("plain-byok-key")) + with pytest.raises(ValueError, match="non-OAuth2 credential"): + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") + + +@pytest.mark.asyncio +async def test_byok_guard_rejects_overwriting_encrypted_byok(): + # Simulate a row written by the new (encrypted) code path: write a BYOK, + # then attempt to overwrite with an OAuth token. + prisma = _make_prisma_with_existing(row=None) + await store_user_credential(prisma, "alice", "srv-1", "sk-secret-byok") + + encrypted_row = MagicMock() + encrypted_row.credential_b64 = _stored_value(prisma) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock( + return_value=encrypted_row + ) + + with pytest.raises(ValueError, match="non-OAuth2 credential"): + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") + + +@pytest.mark.asyncio +async def test_byok_guard_allows_overwriting_existing_oauth(): + # Refresh path: row already holds an OAuth payload, write must succeed. + prisma = _make_prisma_with_existing(row=None) + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok-1") + + oauth_row = MagicMock() + oauth_row.credential_b64 = _stored_value(prisma) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=oauth_row) + + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok-2") + # Final upsert wrote a new payload (different from the first) + assert _stored_value(prisma) != oauth_row.credential_b64 + + +# ── list_user_oauth_credentials ─────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_list_oauth_credentials_filters_byok_and_returns_payloads(): + # Three rows: one encrypted OAuth, one legacy-plaintext OAuth, one BYOK. + # Only the two OAuth rows should come back. + prisma = _make_prisma_with_existing(row=None) + + await store_user_oauth_credential(prisma, "alice", "srv-encrypted", "tok-enc") + encrypted_b64 = _stored_value(prisma) + encrypted_row = MagicMock() + encrypted_row.credential_b64 = encrypted_b64 + encrypted_row.server_id = "srv-encrypted" + + legacy_payload = { + "type": "oauth2", + "access_token": "tok-legacy", + "connected_at": "2024-01-01T00:00:00Z", + } + legacy_row = MagicMock() + legacy_row.credential_b64 = base64.urlsafe_b64encode( + json.dumps(legacy_payload).encode() + ).decode() + legacy_row.server_id = "srv-legacy" + + byok_row = MagicMock() + byok_row.credential_b64 = base64.urlsafe_b64encode(b"plain-byok-key").decode() + byok_row.server_id = "srv-byok" + + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[encrypted_row, legacy_row, byok_row] + ) + + results = await list_user_oauth_credentials(prisma, "alice") + + server_ids = {r["server_id"] for r in results} + assert server_ids == {"srv-encrypted", "srv-legacy"} + tokens = {r["access_token"] for r in results} + assert tokens == {"tok-enc", "tok-legacy"} + + +# ── _decode_user_credential helper ──────────────────────────────────────────── + + +def test_decode_user_credential_handles_garbage(): + # Malformed input must return None, not raise. + assert _decode_user_credential("not-base64-and-not-encrypted!!!") is None + + +def test_decode_user_credential_legacy_path(): + plain = "legacy-secret" + stored = base64.urlsafe_b64encode(plain.encode()).decode() + assert _decode_user_credential(stored) == plain