mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(vector-stores): restore redacted placeholders on update and drop forged admin marker on REST stores
Echoing a GET /vector_store/info body back into /vector_store/update stored the literal REDACTED placeholder as the credential. The update now restores every placeholder from the existing row, nested fields included, and drops placeholders for keys the row never had. The non-gRPC persistence branch also strips the internal admin marker so a REST store cannot carry a forged approval into a later transport switch.
This commit is contained in:
parent
b87186a717
commit
3dfd4420c6
4 changed files with 97 additions and 4 deletions
|
|
@ -94,9 +94,11 @@ def prepare_connection_for_persistence(
|
|||
previous_is_grpc: Final = _is_grpc_connection(existing_custom_llm_provider, existing)
|
||||
effective_is_grpc: Final = _is_grpc_connection(custom_llm_provider, effective)
|
||||
if not previous_is_grpc and not effective_is_grpc:
|
||||
return ( # mutable-ok: persistence requires an isolated JSON-serializable dict
|
||||
dict(supplied) if isinstance(litellm_params, dict) else dict(existing)
|
||||
)
|
||||
return { # mutable-ok: persistence requires an isolated JSON-serializable dict
|
||||
key: value
|
||||
for key, value in (supplied if isinstance(litellm_params, dict) else existing).items()
|
||||
if key != MILVUS_ADMIN_CONFIGURED_CONNECTION
|
||||
}
|
||||
|
||||
is_create: Final = existing_custom_llm_provider is None
|
||||
provider_changed: Final = not is_create and custom_llm_provider != existing_custom_llm_provider
|
||||
|
|
|
|||
|
|
@ -112,6 +112,20 @@ def _redact_sensitive_litellm_params(litellm_params: object, _depth: int = 0) ->
|
|||
return out
|
||||
|
||||
|
||||
def _restore_redacted_litellm_params(supplied: object, existing: object, _depth: int = 0) -> object:
|
||||
if supplied == REDACTED_BY_LITELM_STRING:
|
||||
return existing
|
||||
if _depth >= _REDACT_LITELLM_PARAMS_MAX_DEPTH or not isinstance(supplied, dict) or not isinstance(existing, dict):
|
||||
return supplied
|
||||
supplied_params: Final = deserialize_litellm_params(supplied)
|
||||
existing_params: Final = deserialize_litellm_params(existing)
|
||||
return {
|
||||
key: _restore_redacted_litellm_params(value, existing_params.get(key), _depth + 1)
|
||||
for key, value in supplied_params.items()
|
||||
if value != REDACTED_BY_LITELM_STRING or key in existing_params
|
||||
}
|
||||
|
||||
|
||||
def _validated_litellm_params(
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> Mapping[str, object]:
|
||||
|
|
@ -625,12 +639,15 @@ async def update_vector_store(
|
|||
)
|
||||
|
||||
existing_litellm_params: Final = deserialize_litellm_params(existing_vector_store.get("litellm_params"))
|
||||
supplied_litellm_params: Final = _restore_redacted_litellm_params(
|
||||
update_data.get("litellm_params"), existing_litellm_params
|
||||
)
|
||||
effective_provider: Final = update_data.get("custom_llm_provider") or existing_vector_store.get(
|
||||
"custom_llm_provider"
|
||||
)
|
||||
effective_litellm_params: Final = prepare_vector_store_connection_for_persistence(
|
||||
custom_llm_provider=effective_provider,
|
||||
litellm_params=update_data.get("litellm_params"),
|
||||
litellm_params=supplied_litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
existing_custom_llm_provider=existing_vector_store.get("custom_llm_provider"),
|
||||
existing_litellm_params=existing_litellm_params,
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ IGNORE_FUNCTIONS = [
|
|||
"_read_image_bytes", # max depth set.
|
||||
"_get_masked_values", # max depth set (default 20) to prevent infinite recursion while masking nested sensitive config dicts.
|
||||
"_redact_sensitive_litellm_params", # max depth set (default 10).
|
||||
"_restore_redacted_litellm_params", # max depth set (default 10), mirrors _redact_sensitive_litellm_params on the write side.
|
||||
"_redact_secret_values_in_obj", # max depth set (default 10, _REDACT_SECRET_MAX_DEPTH); fails closed by returning "REDACTED" at the cap.
|
||||
"_resolve", # OCI: $ref resolver bounded by `resolving_stack` cycle guard.
|
||||
"resolve_oci_schema_anyof", # OCI: bounded by JSON-schema tree depth (no cycles possible in well-formed input).
|
||||
|
|
|
|||
|
|
@ -1090,6 +1090,18 @@ def test_non_grpc_connection_update_keeps_replacement_semantics():
|
|||
assert params == {"api_key": "new-key"}
|
||||
|
||||
|
||||
def test_non_grpc_connection_update_drops_forged_admin_marker():
|
||||
params = prepare_vector_store_connection_for_persistence(
|
||||
custom_llm_provider="openai",
|
||||
litellm_params={"api_key": "new-key", MILVUS_ADMIN_CONFIGURED_CONNECTION: True},
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
existing_custom_llm_provider="openai",
|
||||
existing_litellm_params={"api_key": "old-key"},
|
||||
)
|
||||
|
||||
assert params == {"api_key": "new-key"}
|
||||
|
||||
|
||||
class TestCheckVectorStorePermission:
|
||||
"""Test suite for check_vector_store_permission function."""
|
||||
|
||||
|
|
@ -3334,6 +3346,67 @@ class TestUpdateVectorStoreAccessControlAndRedaction:
|
|||
assert params["api_key"] == REDACTED_BY_LITELM_STRING
|
||||
assert params["api_base"] == "https://api.openai.com/v1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_restores_redacted_placeholder_from_existing_row(self):
|
||||
from litellm.constants import REDACTED_BY_LITELM_STRING
|
||||
from litellm.proxy.vector_store_endpoints.management_endpoints import (
|
||||
update_vector_store,
|
||||
)
|
||||
from litellm.types.vector_stores import VectorStoreUpdateRequest
|
||||
|
||||
existing_row = MagicMock()
|
||||
existing_row.model_dump = MagicMock(
|
||||
return_value={
|
||||
"vector_store_id": "vs_owned",
|
||||
"team_id": "team-A",
|
||||
"custom_llm_provider": "openai",
|
||||
"litellm_params": {
|
||||
"api_key": "sk-real-openai-key-123",
|
||||
"api_base": "https://old.example",
|
||||
"litellm_embedding_config": {"api_key": "sk-real-embedding-key", "model": "text-embedding-3-small"},
|
||||
},
|
||||
}
|
||||
)
|
||||
updated_row = MagicMock()
|
||||
updated_row.model_dump = MagicMock(return_value={"vector_store_id": "vs_owned", "team_id": "team-A"})
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_row)
|
||||
mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=updated_row)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||
patch("litellm.vector_store_registry", None),
|
||||
):
|
||||
await update_vector_store(
|
||||
data=VectorStoreUpdateRequest(
|
||||
vector_store_id="vs_owned",
|
||||
litellm_params={
|
||||
"api_key": REDACTED_BY_LITELM_STRING,
|
||||
"api_base": "https://new.example",
|
||||
"litellm_embedding_config": {
|
||||
"api_key": REDACTED_BY_LITELM_STRING,
|
||||
"model": "text-embedding-3-small",
|
||||
},
|
||||
"never_stored_secret": REDACTED_BY_LITELM_STRING,
|
||||
},
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin", team_id="team-A"
|
||||
),
|
||||
)
|
||||
|
||||
update_call = mock_prisma_client.db.litellm_managedvectorstorestable.update.call_args
|
||||
persisted = json.loads(update_call.kwargs["data"]["litellm_params"])
|
||||
assert persisted["api_key"] == "sk-real-openai-key-123"
|
||||
assert persisted["api_base"] == "https://new.example"
|
||||
assert persisted["litellm_embedding_config"] == {
|
||||
"api_key": "sk-real-embedding-key",
|
||||
"model": "text-embedding-3-small",
|
||||
}
|
||||
assert "never_stored_secret" not in persisted
|
||||
assert REDACTED_BY_LITELM_STRING not in update_call.kwargs["data"]["litellm_params"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_row_deleted_mid_update_returns_404(self):
|
||||
"""A concurrent delete between the authorization read and the write makes Prisma's
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue