mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(mcp): encrypt stored static headers and stdio environment (#40164)
Encrypt secret maps at the shared persistence boundary, preserve plaintext API/runtime views, and extend rotation and migration scanning to legacy rows. Co-authored-by: Claude Code <noreply@anthropic.com>
This commit is contained in:
tin-berri
GitHub
parent
8af624b5f5
commit
cd681a573f
11 changed files with 454 additions and 53 deletions
|
|
@ -6,10 +6,12 @@ Canonical definition for ``litellm_mcpservertable``. Re-exported from
|
|||
"""
|
||||
|
||||
import enum
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic import Field, ValidationInfo, field_validator
|
||||
|
||||
from litellm.types.llms.base import LiteLLMPydanticObjectBase
|
||||
from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType
|
||||
|
|
@ -115,3 +117,12 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
|
|||
submitted_at: datetime | None = None
|
||||
reviewed_at: datetime | None = None
|
||||
review_notes: str | None = None
|
||||
|
||||
@field_validator("static_headers", "env", mode="before")
|
||||
@classmethod
|
||||
def decode_stored_secret_map(cls, value: object, info: ValidationInfo) -> Mapping[str, str] | None:
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decode_secret_map
|
||||
|
||||
if value is None and info.field_name == "env":
|
||||
return MappingProxyType({})
|
||||
return decode_secret_map(value, key=info.field_name or "secret map")
|
||||
|
|
|
|||
|
|
@ -23,8 +23,11 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
SecretMapDecodeError,
|
||||
_get_salt_key,
|
||||
decode_secret_map,
|
||||
decrypt_value_helper,
|
||||
encrypt_secret_map,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
|
@ -360,7 +363,7 @@ def _prepare_mcp_server_data(
|
|||
# exclude_unset filter is respected. Reading back from ``data`` would
|
||||
# reintroduce defaults (e.g. ``env={}``) for fields the caller never set.
|
||||
if data_dict.get("static_headers") is not None:
|
||||
data_dict["static_headers"] = safe_dumps(data_dict["static_headers"])
|
||||
data_dict["static_headers"] = encrypt_secret_map(data_dict["static_headers"])
|
||||
|
||||
# env_vars is read from ``data_dict`` (not ``data``) like every other JSON
|
||||
# column so the exclude_unset filter is respected: a partial update that
|
||||
|
|
@ -376,7 +379,7 @@ def _prepare_mcp_server_data(
|
|||
data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"])
|
||||
|
||||
if data_dict.get("env") is not None:
|
||||
data_dict["env"] = safe_dumps(data_dict["env"])
|
||||
data_dict["env"] = encrypt_secret_map(data_dict["env"])
|
||||
|
||||
if "tool_name_to_display_name" in data_dict:
|
||||
data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {})
|
||||
|
|
@ -589,6 +592,19 @@ def decrypt_credentials(
|
|||
return credentials
|
||||
|
||||
|
||||
def _readable_mcp_servers(
|
||||
rows: Iterable["prisma_db_models.LiteLLM_MCPServerTable"],
|
||||
) -> Iterable[LiteLLM_MCPServerTable]:
|
||||
for row in rows:
|
||||
try:
|
||||
table = LiteLLM_MCPServerTable.model_validate(row.model_dump())
|
||||
except SecretMapDecodeError:
|
||||
verbose_proxy_logger.warning("Skipping MCP server %s: cannot decrypt secret map", row.server_id)
|
||||
continue
|
||||
decrypt_global_env_var_values(table.env_vars)
|
||||
yield table
|
||||
|
||||
|
||||
async def get_all_mcp_servers(
|
||||
prisma_client: PrismaClient,
|
||||
approval_status: str | None = None,
|
||||
|
|
@ -609,10 +625,7 @@ async def get_all_mcp_servers(
|
|||
)
|
||||
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)
|
||||
|
||||
tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
|
||||
for table in tables:
|
||||
decrypt_global_env_var_values(table.env_vars)
|
||||
return tables
|
||||
return list(_readable_mcp_servers(mcp_servers))
|
||||
|
||||
|
||||
async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None:
|
||||
|
|
@ -638,13 +651,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]
|
|||
"server_id": {"in": server_ids},
|
||||
}
|
||||
)
|
||||
final_mcp_servers: Final[list[LiteLLM_MCPServerTable]] = []
|
||||
for _mcp_server in _mcp_servers:
|
||||
table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump())
|
||||
decrypt_global_env_var_values(table.env_vars)
|
||||
final_mcp_servers.append(table)
|
||||
|
||||
return final_mcp_servers
|
||||
return list(_readable_mcp_servers(_mcp_servers))
|
||||
|
||||
|
||||
async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> list[str]:
|
||||
|
|
@ -852,12 +859,10 @@ async def create_mcp_server(
|
|||
data_dict["created_by"] = touched_by
|
||||
data_dict["updated_by"] = touched_by
|
||||
|
||||
new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create(
|
||||
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
|
||||
)
|
||||
new_mcp_server: Final = await MCPServerRepository(prisma_client).table.create(data=data_dict)
|
||||
|
||||
_decrypt_env_vars_on_returned_row(new_mcp_server)
|
||||
return new_mcp_server
|
||||
return LiteLLM_MCPServerTable.model_validate(new_mcp_server.model_dump())
|
||||
|
||||
|
||||
async def create_draft_mcp_server(
|
||||
|
|
@ -1066,13 +1071,13 @@ async def update_mcp_server(
|
|||
|
||||
data_dict["credentials"] = Json(None)
|
||||
|
||||
updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update(
|
||||
updated_mcp_server: Final = await MCPServerRepository(prisma_client).table.update(
|
||||
where={"server_id": data.server_id},
|
||||
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
|
||||
data=data_dict,
|
||||
)
|
||||
|
||||
_decrypt_env_vars_on_returned_row(updated_mcp_server)
|
||||
return updated_mcp_server
|
||||
return LiteLLM_MCPServerTable.model_validate(updated_mcp_server.model_dump()) if updated_mcp_server else None
|
||||
|
||||
|
||||
async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None:
|
||||
|
|
@ -1144,6 +1149,13 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
|
|||
if rotated_env_vars is not None:
|
||||
update_data["env_vars"] = safe_dumps(rotated_env_vars)
|
||||
|
||||
for field in ("static_headers", "env"):
|
||||
try:
|
||||
if secret_map := decode_secret_map(getattr(mcp_server, field, None), key=field):
|
||||
update_data[field] = encrypt_secret_map(secret_map, new_encryption_key=new_master_key)
|
||||
except SecretMapDecodeError:
|
||||
verbose_proxy_logger.warning("Cannot rotate MCP %s for server %s", field, mcp_server.server_id)
|
||||
|
||||
if not update_data:
|
||||
continue
|
||||
|
||||
|
|
@ -1894,9 +1906,7 @@ async def get_mcp_submissions(
|
|||
order={"submitted_at": "desc"},
|
||||
take=500, # safety cap; paginate if needed in a future iteration
|
||||
)
|
||||
items: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows]
|
||||
for item in items:
|
||||
decrypt_global_env_var_values(item.env_vars)
|
||||
items: Final = list(_readable_mcp_servers(rows))
|
||||
|
||||
pending: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review)
|
||||
active: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active)
|
||||
|
|
|
|||
|
|
@ -6272,8 +6272,7 @@ class MCPServerManager:
|
|||
]
|
||||
}
|
||||
)
|
||||
db_mcp_servers: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in raw_rows]
|
||||
verbose_logger.info("Found %s MCP servers in database", len(db_mcp_servers))
|
||||
verbose_logger.info("Found %s MCP servers in database", len(raw_rows))
|
||||
|
||||
previous_registry: Final = self.registry
|
||||
new_registry: Final[dict[str, MCPServer]] = {}
|
||||
|
|
@ -6281,8 +6280,9 @@ class MCPServerManager:
|
|||
# Stage one: build every server. Stage two assigns short prefixes
|
||||
# against the *full* set so dedup is deterministic regardless of
|
||||
# iteration order.
|
||||
for server in db_mcp_servers:
|
||||
for row in raw_rows:
|
||||
try:
|
||||
server = LiteLLM_MCPServerTable.model_validate(row.model_dump())
|
||||
existing_server = previous_registry.get(server.server_id)
|
||||
|
||||
if (
|
||||
|
|
@ -6320,8 +6320,8 @@ class MCPServerManager:
|
|||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"Skipping MCP server %s (%s) during DB reload: %s",
|
||||
server.server_id,
|
||||
getattr(server, "alias", None),
|
||||
getattr(row, "server_id", None),
|
||||
getattr(row, "alias", None),
|
||||
e,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,10 @@
|
|||
import base64
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, Literal, cast
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
# Versioned ciphertext marker for AES-256-GCM values.
|
||||
|
|
@ -203,3 +206,40 @@ def decrypt_value(value: bytes, signing_key: str) -> str:
|
|||
return plaintext
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
||||
class SecretMapDecodeError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
_SECRET_MAP: Final = TypeAdapter(Mapping[str, str])
|
||||
_STORED_SECRET_MAP: Final = TypeAdapter(Mapping[str, str] | str)
|
||||
_SECRET_STRING: Final = TypeAdapter(str)
|
||||
|
||||
|
||||
def encrypt_secret_map(value: Mapping[str, str], new_encryption_key: str | None = None) -> str:
|
||||
if not value:
|
||||
return "{}"
|
||||
ciphertext: Final = _SECRET_STRING.validate_python(
|
||||
encrypt_value_helper(_SECRET_MAP.dump_json(value).decode(), new_encryption_key=new_encryption_key), strict=True
|
||||
)
|
||||
return _SECRET_STRING.dump_json(ciphertext).decode()
|
||||
|
||||
|
||||
def decode_secret_map(value: object, *, key: str) -> Mapping[str, str] | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
stored: Final = (
|
||||
_STORED_SECRET_MAP.validate_json(value, strict=True)
|
||||
if isinstance(value, str) and value.lstrip().startswith(("{", '"'))
|
||||
else _STORED_SECRET_MAP.validate_python(value, strict=True)
|
||||
)
|
||||
if not isinstance(stored, str):
|
||||
return stored
|
||||
decrypted: Final = decrypt_value_helper(
|
||||
value=stored, key=key, exception_type="debug", return_original_value=False
|
||||
)
|
||||
return _SECRET_MAP.validate_json(decrypted, strict=True)
|
||||
except ValidationError:
|
||||
raise SecretMapDecodeError(f"Cannot decode encrypted MCP {key}; check LITELLM_SALT_KEY") from None
|
||||
|
|
|
|||
|
|
@ -43,7 +43,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
_ALGO_AES_GCM,
|
||||
_ENCRYPTION_ALGORITHM_SETTING,
|
||||
_V2_GCM_PREFIX,
|
||||
SecretMapDecodeError,
|
||||
_get_salt_key,
|
||||
decode_secret_map,
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
|
|
@ -65,6 +67,20 @@ class LocationReport:
|
|||
# Used by --check (read-only classification):
|
||||
legacy: int = 0 # nacl ciphertext still awaiting migration
|
||||
|
||||
def count(self, classification: ValueClass | None) -> None:
|
||||
if classification is None:
|
||||
return
|
||||
self.scanned += 1
|
||||
match classification:
|
||||
case "migrated":
|
||||
self.already_v2 += 1
|
||||
case "legacy":
|
||||
self.legacy += 1
|
||||
case "undecryptable":
|
||||
self.undecryptable += 1
|
||||
case _:
|
||||
self.plaintext += 1
|
||||
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
return {
|
||||
"scanned": self.scanned,
|
||||
|
|
@ -441,7 +457,7 @@ def _classify_callback_value(value: object) -> ValueClass:
|
|||
_COVERED_TABLE_SPECS: Final = [
|
||||
("model_table", "litellm_proxymodeltable", ("litellm_params",), ()),
|
||||
("credentials", "litellm_credentialstable", ("credential_values",), ()),
|
||||
("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars"), ()),
|
||||
("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()),
|
||||
("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)),
|
||||
("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)),
|
||||
]
|
||||
|
|
@ -472,14 +488,18 @@ def _classify_into_report(report: LocationReport, value: str) -> None:
|
|||
names, base URLs, …) do not decrypt and fall through to ``plaintext``, so
|
||||
over-scanning a column is harmless to the residual count.
|
||||
"""
|
||||
report.scanned += 1
|
||||
cls: Final = classify_value(value, key="scan")
|
||||
if cls == "migrated":
|
||||
report.already_v2 += 1
|
||||
elif cls == "legacy":
|
||||
report.legacy += 1
|
||||
else: # plaintext / not-a-string
|
||||
report.plaintext += 1
|
||||
report.count(classify_value(value, key="scan"))
|
||||
|
||||
|
||||
def _classify_secret_map(value: object, key: str) -> ValueClass | None:
|
||||
try:
|
||||
decoded: Final = decode_secret_map(value, key=key)
|
||||
except SecretMapDecodeError:
|
||||
return "undecryptable"
|
||||
if not decoded:
|
||||
return None
|
||||
ciphertext: Final = json.loads(value) if isinstance(value, str) and value.lstrip().startswith('"') else value
|
||||
return "migrated" if is_migrated(ciphertext) else "legacy"
|
||||
|
||||
|
||||
async def _scan_one_table(
|
||||
|
|
@ -503,6 +523,9 @@ async def _scan_one_table(
|
|||
raw = getattr(row, col, None)
|
||||
if raw is None:
|
||||
continue
|
||||
if db_attr == "litellm_mcpservertable" and col in ("static_headers", "env"):
|
||||
report.count(_classify_secret_map(raw, col))
|
||||
continue
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
raw = json.loads(raw)
|
||||
|
|
|
|||
|
|
@ -12,28 +12,39 @@ import base64
|
|||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from prisma.models import LiteLLM_MCPServerTable as PrismaMCPServer
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
_decode_user_credential,
|
||||
_prepare_mcp_server_data,
|
||||
create_mcp_server,
|
||||
decrypt_credentials,
|
||||
encrypt_credentials,
|
||||
get_all_mcp_servers,
|
||||
get_mcp_servers,
|
||||
get_mcp_submissions,
|
||||
get_user_credential,
|
||||
get_user_oauth_credential,
|
||||
is_oauth_credential_expired,
|
||||
list_user_oauth_credentials,
|
||||
resolve_valid_user_oauth_token,
|
||||
rotate_mcp_server_credentials_master_key,
|
||||
rotate_mcp_user_credentials_master_key,
|
||||
rotate_mcp_user_env_vars_master_key,
|
||||
store_user_credential,
|
||||
store_user_oauth_credential,
|
||||
update_mcp_server,
|
||||
)
|
||||
from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest
|
||||
from litellm.proxy._types import LiteLLM_MCPServerTable, NewMCPServerRequest, UpdateMCPServerRequest
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
SecretMapDecodeError,
|
||||
decode_secret_map,
|
||||
decrypt_value_helper,
|
||||
encrypt_secret_map,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport
|
||||
|
|
@ -44,6 +55,7 @@ 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)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": "xsalsa20-poly1305"})
|
||||
|
||||
|
||||
def _make_prisma_with_existing(row):
|
||||
|
|
@ -368,6 +380,169 @@ def test_client_private_key_encrypted_at_rest():
|
|||
assert decrypted["client_secret"] == "shh"
|
||||
|
||||
|
||||
@pytest.fixture(params=["xsalsa20-poly1305", "aes-256-gcm"])
|
||||
def map_algorithm(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> str:
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": request.param})
|
||||
return request.param
|
||||
|
||||
|
||||
def _prisma_map_row(data: dict[str, object], quoted: bool = False) -> PrismaMCPServer:
|
||||
return PrismaMCPServer.model_validate({
|
||||
"transport": "http", "mcp_access_groups": [], "allowed_tools": [], "extra_headers": [], "args": [],
|
||||
"allow_all_keys": False, "available_on_public_internet": True, "delegate_auth_to_upstream": False,
|
||||
"oauth_passthrough": False, "per_server_oauth_discovery": False, "is_byok": False, "byok_description": [],
|
||||
**data,
|
||||
**{field: json.dumps(data[field]) for field in ("static_headers", "env") if quoted and data.get(field)},
|
||||
})
|
||||
|
||||
|
||||
class _MapTable:
|
||||
def __init__(self, *rows: dict[str, object], quoted: bool = False) -> None:
|
||||
self.rows = {row["server_id"]: row for row in rows}
|
||||
self.quoted = quoted
|
||||
|
||||
async def create(self, *, data: dict[str, object]) -> PrismaMCPServer:
|
||||
self.rows = {**self.rows, data["server_id"]: dict(data)}
|
||||
return _prisma_map_row(data, self.quoted)
|
||||
|
||||
async def update(self, *, where: dict[str, str], data: dict[str, object]) -> PrismaMCPServer:
|
||||
return await self.create(data={**self.rows[where["server_id"]], **data})
|
||||
|
||||
async def find_many(self, where: object = None) -> list[PrismaMCPServer]:
|
||||
return [_prisma_map_row(row, self.quoted) for row in self.rows.values()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("field", ["static_headers", "env"])
|
||||
@pytest.mark.parametrize("quoted", [False, True])
|
||||
async def test_secret_maps_create_update_round_trip(map_algorithm: str, field: str, quoted: bool) -> None:
|
||||
table: Final = _MapTable(quoted=quoted)
|
||||
prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table))
|
||||
original: Final = {"TOKEN": " sensitive-secret\n", "PREFIX": "v2:gcm:literal", "TEMPLATE": "Bearer ${TOKEN}"}
|
||||
create: Final = NewMCPServerRequest.model_validate({
|
||||
"server_id": "srv-map", "transport": "http", "url": "https://up.example.com/mcp", field: original,
|
||||
})
|
||||
created: Final = await create_mcp_server(prisma, create, touched_by="test")
|
||||
first: Final = table.rows["srv-map"][field]
|
||||
assert isinstance(first, str) and isinstance(json.loads(first), str)
|
||||
assert json.loads(first).startswith("v2:gcm:") is (map_algorithm == "aes-256-gcm")
|
||||
assert "sensitive-secret" not in first and "TEMPLATE" not in first
|
||||
assert getattr(created, field) == original == getattr(create, field)
|
||||
assert decode_secret_map(first, key=field) == original
|
||||
replacement: Final = {**original, "TOKEN": "updated-sensitive-secret"}
|
||||
update: Final = UpdateMCPServerRequest.model_validate({"server_id": "srv-map", field: replacement})
|
||||
updated: Final = await update_mcp_server(prisma, update, touched_by="test")
|
||||
second: Final = table.rows["srv-map"][field]
|
||||
assert second != first and "updated-sensitive-secret" not in second
|
||||
assert decode_secret_map(second, key=field) == replacement
|
||||
assert getattr(updated, field) == replacement == getattr(update, field)
|
||||
assert original["TOKEN"] == " sensitive-secret\n"
|
||||
omitted: Final = await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="srv-map"), touched_by="test")
|
||||
assert table.rows["srv-map"][field] == second and getattr(omitted, field) == replacement
|
||||
cleared: Final = await update_mcp_server(
|
||||
prisma, UpdateMCPServerRequest.model_validate({"server_id": "srv-map", field: {}}), touched_by="test"
|
||||
)
|
||||
assert table.rows["srv-map"][field] == "{}" and getattr(cleared, field) == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["static_headers", "env"])
|
||||
@pytest.mark.parametrize("as_json", [False, True])
|
||||
def test_secret_map_legacy_model_read_preserves_exact_values(field: str, as_json: bool) -> None:
|
||||
original: Final = {"PREFIX": "v2:gcm:literal", "SPACE": " secret\n", "TEMPLATE": "${TOKEN}", "B64": "YWJjZA=="}
|
||||
incoming: Final = {
|
||||
"server_id": "srv-map", "transport": "http", field: json.dumps(original) if as_json else original,
|
||||
}
|
||||
snapshot: Final = json.dumps(incoming)
|
||||
parsed: Final = LiteLLM_MCPServerTable.model_validate(incoming)
|
||||
assert getattr(parsed, field) == original
|
||||
assert json.dumps(incoming) == snapshot
|
||||
assert LiteLLM_MCPServerTable.model_validate(parsed.model_dump()).model_dump() == parsed.model_dump()
|
||||
empty: Final = LiteLLM_MCPServerTable.model_validate({"server_id": "srv-map", "transport": "http", field: None})
|
||||
assert getattr(empty, field) == ({} if field == "env" else None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["static_headers", "env"])
|
||||
@pytest.mark.parametrize("failure", ["wrong-key", "corrupt", "invalid-values", "invalid-shape", "invalid-json"])
|
||||
def test_secret_map_model_read_fails_closed(map_algorithm: str, field: str, failure: str) -> None:
|
||||
plaintext: Final = {"invalid-values": '{"TOKEN": ["sensitive-secret"]}', "invalid-shape": '["sensitive-secret"]',
|
||||
"invalid-json": "sensitive-secret"}.get(failure, '{"TOKEN": "sensitive-secret"}')
|
||||
ciphertext: Final = encrypt_value_helper(
|
||||
plaintext, new_encryption_key="wrong-map-key" if failure == "wrong-key" else None
|
||||
)
|
||||
stored: Final = json.dumps(ciphertext[:-8] if failure == "corrupt" else ciphertext)
|
||||
with pytest.raises(SecretMapDecodeError) as exc:
|
||||
LiteLLM_MCPServerTable.model_validate({"server_id": "srv-map", "transport": "http", field: stored})
|
||||
assert field in str(exc.value) and "LITELLM_SALT_KEY" in str(exc.value)
|
||||
assert all(secret not in str(exc.value) for secret in (plaintext, ciphertext, "sensitive-secret"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("field,other", [("static_headers", "env"), ("env", "static_headers")])
|
||||
async def test_secret_map_rotation_migrates_rekeys_and_preserves_corrupt(
|
||||
map_algorithm: str, field: str, other: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
values: Final = {"TOKEN": "rotation-sensitive-secret", "TEMPLATE": "Bearer ${TOKEN}"}
|
||||
old: Final = encrypt_secret_map(values)
|
||||
corrupt: Final = json.dumps(json.loads(old)[:-8])
|
||||
table: Final = _MapTable(
|
||||
{"server_id": "broken", field: corrupt, other: old},
|
||||
{"server_id": "legacy", field: json.dumps(values), other: "{}"},
|
||||
{"server_id": "encrypted", field: old, other: None},
|
||||
)
|
||||
prisma: Final = SimpleNamespace(db=SimpleNamespace(
|
||||
litellm_mcpservertable=table, litellm_mcpserveroauthclient=SimpleNamespace(find_many=AsyncMock(return_value=[]))
|
||||
))
|
||||
await rotate_mcp_server_credentials_master_key(prisma, touched_by="test", new_master_key="rotated-map-key")
|
||||
assert table.rows["broken"][field] == corrupt
|
||||
assert table.rows["legacy"][other] == "{}" and table.rows["encrypted"][other] is None
|
||||
for server_id, map_field in (("broken", other), ("legacy", field), ("encrypted", field)):
|
||||
stored: Final = table.rows[server_id][map_field]
|
||||
assert isinstance(json.loads(stored), str) and stored != old and "rotation-sensitive-secret" not in stored
|
||||
with pytest.raises(SecretMapDecodeError):
|
||||
decode_secret_map(stored, key=map_field)
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "rotated-map-key")
|
||||
for server_id, map_field in (("broken", other), ("legacy", field), ("encrypted", field)):
|
||||
assert decode_secret_map(table.rows[server_id][map_field], key=map_field) == values
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("reader", [get_all_mcp_servers, get_mcp_servers, get_mcp_submissions])
|
||||
@pytest.mark.parametrize("field", ["static_headers", "env"])
|
||||
async def test_bulk_reads_isolate_corrupt_secret_maps(reader, field, map_algorithm, caplog):
|
||||
secret = {"TOKEN": "bulk-sensitive-secret"}
|
||||
encrypted = encrypt_secret_map(secret)
|
||||
corrupt = encrypt_secret_map(secret, new_encryption_key="wrong-bulk-key")
|
||||
rows = [
|
||||
_prisma_map_row({"server_id": "broken", field: corrupt, "approval_status": "pending_review"}),
|
||||
_prisma_map_row({"server_id": "healthy", field: encrypted, "approval_status": "active"}),
|
||||
]
|
||||
snapshot = [row.model_dump() for row in rows]
|
||||
table = SimpleNamespace(find_many=AsyncMock(return_value=rows))
|
||||
prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table))
|
||||
result = await reader(prisma, ["broken", "healthy"]) if reader is get_mcp_servers else await reader(prisma)
|
||||
items = result.items if reader is get_mcp_submissions else result
|
||||
assert [row.server_id for row in items] == ["healthy"]
|
||||
assert getattr(items[0], field) == secret
|
||||
assert [row.model_dump() for row in rows] == snapshot
|
||||
assert "broken" in caplog.text
|
||||
assert all(value not in caplog.text for value in ("bulk-sensitive-secret", corrupt, encrypted))
|
||||
if reader is get_mcp_submissions:
|
||||
assert (result.total, result.pending_review, result.active, result.rejected) == (1, 0, 1, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("reader", [get_all_mcp_servers, get_mcp_servers, get_mcp_submissions])
|
||||
async def test_bulk_reads_do_not_swallow_unrelated_validation_errors(reader):
|
||||
from pydantic import ValidationError
|
||||
|
||||
row = _prisma_map_row({"server_id": "invalid", "transport": "unsupported"})
|
||||
table = SimpleNamespace(find_many=AsyncMock(return_value=[row]))
|
||||
prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table))
|
||||
request = reader(prisma, ["invalid"]) if reader is get_mcp_servers else reader(prisma)
|
||||
with pytest.raises(ValidationError, match="transport"):
|
||||
await request
|
||||
|
||||
|
||||
# ── BYOK round-trip ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -861,6 +861,7 @@ _SALT_KEY = "test-salt-key-for-env-vars-tests-1234"
|
|||
@pytest.fixture
|
||||
def env_vars_salt_key(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", _SALT_KEY)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": "xsalsa20-poly1305"})
|
||||
|
||||
|
||||
def _mock_env_vars_prisma(row=None):
|
||||
|
|
@ -1518,9 +1519,16 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri
|
|||
assert "s3cr3t-p@ss" not in encrypted_env_vars_str
|
||||
|
||||
def _prisma_row_with_json_string_env_vars():
|
||||
row = MagicMock()
|
||||
row.env_vars = encrypted_env_vars_str
|
||||
return row
|
||||
import json
|
||||
|
||||
from prisma.models import LiteLLM_MCPServerTable
|
||||
|
||||
return LiteLLM_MCPServerTable.model_validate({
|
||||
"server_id": "srv-returned", "transport": "http", "mcp_access_groups": [], "allowed_tools": [],
|
||||
"extra_headers": [], "args": [], "allow_all_keys": False, "available_on_public_internet": True,
|
||||
"delegate_auth_to_upstream": False, "oauth_passthrough": False, "per_server_oauth_discovery": False,
|
||||
"is_byok": False, "byok_description": [], "env_vars": json.dumps(encrypted_env_vars_str),
|
||||
})
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.create = AsyncMock(
|
||||
|
|
@ -1537,7 +1545,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri
|
|||
touched_by="test-user",
|
||||
)
|
||||
assert isinstance(created.env_vars, list)
|
||||
assert created.env_vars[0]["value"] == "s3cr3t-p@ss"
|
||||
assert created.env_vars[0].value == "s3cr3t-p@ss"
|
||||
assert created.env_vars[0].name == "DB_PASSWORD"
|
||||
assert created.env == {}
|
||||
|
||||
mock_prisma_upd = MagicMock()
|
||||
mock_prisma_upd.db.litellm_mcpservertable.update = AsyncMock(
|
||||
|
|
@ -1549,7 +1559,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri
|
|||
touched_by="test-user",
|
||||
)
|
||||
assert isinstance(updated.env_vars, list)
|
||||
assert updated.env_vars[0]["value"] == "s3cr3t-p@ss"
|
||||
assert updated.env_vars[0].value == "s3cr3t-p@ss"
|
||||
assert updated.env_vars[0].name == "DB_PASSWORD"
|
||||
assert updated.env == {}
|
||||
|
||||
|
||||
def test_reencrypt_global_env_var_values_handles_json_string(env_vars_salt_key):
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import json
|
|||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from prisma import Json
|
||||
from prisma import Json, models
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
create_mcp_server,
|
||||
|
|
@ -28,8 +28,11 @@ def _credentials_cleared(value) -> bool:
|
|||
def _mock_prisma():
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable = AsyncMock()
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
|
||||
mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=MagicMock())
|
||||
row = models.LiteLLM_MCPServerTable.model_construct(
|
||||
server_id="test-server", transport="http", env={}, env_vars=[]
|
||||
)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=row)
|
||||
mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=row)
|
||||
return mock_prisma
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -905,6 +905,68 @@ class TestMCPServerManager:
|
|||
assert retry_slot is not None
|
||||
assert retry_slot.generation > old_generation
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("corrupt_column", ("static_headers", "env"))
|
||||
async def test_database_reload_drops_cached_server_whose_secret_map_stops_decoding(
|
||||
self, monkeypatch, caplog, corrupt_column
|
||||
):
|
||||
from types import SimpleNamespace
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_secret_map
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-reload-secret-map-salt")
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": "aes-256-gcm"})
|
||||
headers = {"Authorization": "Bearer dummy-header-secret-4f1c"}
|
||||
env = {"UPSTREAM_TOKEN": "dummy-env-secret-9a2b"}
|
||||
stamp = datetime.now()
|
||||
cached = MCPServer(
|
||||
server_id="cached-server",
|
||||
name="cached_server",
|
||||
url="https://up.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
static_headers=dict(headers),
|
||||
env=dict(env),
|
||||
updated_at=stamp,
|
||||
)
|
||||
manager = MCPServerManager()
|
||||
manager.registry[cached.server_id] = cached
|
||||
stored = {"static_headers": encrypt_secret_map(headers), "env": encrypt_secret_map(env)}
|
||||
corrupted = {**stored, corrupt_column: stored[corrupt_column][:-6] + 'AAAAA"'}
|
||||
|
||||
def _row(server_id, maps):
|
||||
row = MagicMock()
|
||||
row.server_id = server_id
|
||||
row.alias = server_id
|
||||
row.model_dump.return_value = {
|
||||
"server_id": server_id,
|
||||
"alias": server_id,
|
||||
"server_name": server_id,
|
||||
"url": "https://up.example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
"updated_at": stamp,
|
||||
**maps,
|
||||
}
|
||||
return row
|
||||
|
||||
table = SimpleNamespace(
|
||||
find_many=AsyncMock(return_value=[_row(cached.server_id, corrupted), _row("healthy-sibling", stored)])
|
||||
)
|
||||
prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table))
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
|
||||
with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
|
||||
await manager.reload_servers_from_database()
|
||||
|
||||
assert set(manager.registry) == {"healthy-sibling"}
|
||||
sibling = manager.registry["healthy-sibling"]
|
||||
assert dict(sibling.static_headers) == headers
|
||||
assert dict(sibling.env) == env
|
||||
logged = "\n".join(caplog.messages)
|
||||
assert cached.server_id in logged
|
||||
for secret in (*headers.values(), *env.values(), *stored.values(), corrupted[corrupt_column]):
|
||||
assert secret not in logged
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_oauth_discovery_preserves_manual_authorization_url_gate(self):
|
||||
with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "false"}):
|
||||
|
|
|
|||
|
|
@ -15,6 +15,11 @@ import httpx
|
|||
|
||||
from litellm.experimental_mcp_client.client import MCPSigV4Auth, MCPClient
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport
|
||||
from prisma import models
|
||||
|
||||
|
||||
def _updated_row() -> models.LiteLLM_MCPServerTable:
|
||||
return models.LiteLLM_MCPServerTable.model_construct(server_id="test-server", transport="http", env={}, env_vars=[])
|
||||
|
||||
|
||||
class TestMCPSigV4Auth:
|
||||
|
|
@ -600,7 +605,7 @@ class TestCredentialMergeOnUpdate:
|
|||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
|
|
@ -639,7 +644,7 @@ class TestCredentialMergeOnUpdate:
|
|||
from litellm.proxy._types import UpdateMCPServerRequest
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
|
|
@ -667,7 +672,7 @@ class TestCredentialMergeOnUpdate:
|
|||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
|
|
@ -709,7 +714,7 @@ class TestCredentialMergeOnUpdate:
|
|||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
|
|
@ -752,7 +757,7 @@ class TestCredentialMergeOnUpdate:
|
|||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
|
|
@ -1083,7 +1088,7 @@ class TestAuthTypeSwitchClearsCredentials:
|
|||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row())
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ proof-of-fix (real proxy + DB) is performed separately on the repro server.
|
|||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -457,6 +458,65 @@ async def test_scan_covered_tables_classifies_legacy_and_v2(salt_key, monkeypatc
|
|||
assert by_loc["credentials"].legacy == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("column", ("static_headers", "env"))
|
||||
@pytest.mark.parametrize("algorithm", ("xsalsa20-poly1305", "aes-256-gcm"))
|
||||
@pytest.mark.parametrize("as_json", (False, True))
|
||||
@pytest.mark.parametrize(
|
||||
"case", ("legacy", "encrypted", "wrong-key", "corrupt", "invalid-shape", "invalid-scalar", "empty", "null")
|
||||
)
|
||||
async def test_check_classifies_mcp_secret_maps(
|
||||
salt_key: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
column: str,
|
||||
algorithm: str,
|
||||
as_json: bool,
|
||||
case: str,
|
||||
) -> None:
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": algorithm})
|
||||
plaintext: Final = {"Authorization": "v2:gcm:operator-text", "CUSTOM": "litellm_enc::literal\n café "}
|
||||
ciphertext: Final = encrypt_value_helper(json.dumps(plaintext))
|
||||
cases: Final[dict[str, object]] = {
|
||||
"legacy": plaintext,
|
||||
"encrypted": ciphertext,
|
||||
"wrong-key": encrypt_value_helper(json.dumps(plaintext), new_encryption_key="different-map-salt"),
|
||||
"corrupt": ciphertext[:-4] + "AAAA",
|
||||
"invalid-shape": encrypt_value_helper(json.dumps({"Authorization": 42})),
|
||||
"invalid-scalar": "null",
|
||||
"empty": {},
|
||||
"null": None,
|
||||
}
|
||||
value: Final = json.dumps(cases[case]) if as_json and case != "null" else cases[case]
|
||||
row: Final = SimpleNamespace(**{column: value})
|
||||
client: Final = MagicMock()
|
||||
_empty_covered_tables(client)
|
||||
client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row])
|
||||
client.db.litellm_mcpservertable.update = AsyncMock()
|
||||
client.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
|
||||
client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None)
|
||||
client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
|
||||
report: Final = await cm.check_encryption(client)
|
||||
expected_legacy: Final = int(case == "legacy" or (case == "encrypted" and algorithm == "xsalsa20-poly1305"))
|
||||
expected_v2: Final = int(case == "encrypted" and algorithm == "aes-256-gcm")
|
||||
expected_invalid: Final = int(case in ("wrong-key", "corrupt", "invalid-shape", "invalid-scalar"))
|
||||
|
||||
assert report.as_dict()["locations"]["mcp_server"] == {
|
||||
"scanned": int(case not in ("empty", "null")),
|
||||
"migrated": 0,
|
||||
"already_v2": expected_v2,
|
||||
"plaintext": 0,
|
||||
"undecryptable": expected_invalid,
|
||||
"legacy": expected_legacy,
|
||||
}
|
||||
assert report.residual_legacy == expected_legacy
|
||||
assert report.total_undecryptable == expected_invalid
|
||||
assert getattr(row, column) == value
|
||||
client.db.litellm_mcpservertable.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_counts_covered_table_residual(salt_key, monkeypatch):
|
||||
"""check_encryption now scans the rotation-covered tables (model table here),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue