mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-09 14:22:12 +00:00
chore(mcp): encrypt user-scoped credentials at rest
LiteLLM_MCPUserCredentials.credential_b64 stored both BYOK API keys and OAuth2 access tokens as plain urlsafe-base64 of the raw value. Any DB read could recover the upstream-provider key. Run all writes through encrypt_value_helper (nacl SecretBox, the same helper used for the server-level credentials column) and read back via a small dual-path helper that tries decryption first, then falls back to plain base64 so existing rows keep working until they get rewritten. Folds the three near-identical "decode -> json.loads -> check type == oauth2" sites into _decode_oauth_payload, which simplifies the BYOK guard inside store_user_oauth_credential.
This commit is contained in:
user
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user