scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): own terminal-batch managed files by creator and fail closed on unconverted ids

Browse source 
Address security review on #35487: register lazily-created managed-file rows
under the batch creator (from the stored batch object) rather than whoever
retrieves the terminal batch first, and redact any raw provider output/error
file id that could not be rewritten to a managed unified id so it can never
reach the client and bypass the /v1/files ownership check

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-01 21:12:50 +00:00
parent 6924f90133
commit f809db81ce
2 changed files with 149 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

84
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -926,6 +926,41 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None:
pass
def _resolve_batch_file_owner_auth(db_batch_object, requester_api_key_dict):
"""Return the auth whose identity should own lazily-registered managed-file rows.
These rows must belong to the batch creator (as the async poller / live path
would have written them), never to whoever happens to retrieve the terminal
batch first, so ownership is taken from the stored batch object when available
and only falls back to the requester when there is no stored object.
"""
if db_batch_object is not None:
from litellm.proxy._types import UserAPIKeyAuth
return UserAPIKeyAuth(
user_id=getattr(db_batch_object, "created_by", None) or "default-user-id",
team_id=getattr(db_batch_object, "team_id", None),
)
return requester_api_key_dict
def _redact_unconverted_batch_output_file_ids(response, verbose_proxy_logger) -> None:
"""Fail closed: never return a raw provider output/error file ID.
Any ID that could not be rewritten to a managed unified ID (hook missing,
model metadata unresolved, or registration raised) is dropped so it can't be
replayed against /v1/files to bypass the managed-file ownership check.
"""
for file_attr in ("output_file_id", "error_file_id"):
file_id = getattr(response, file_attr, None)
if file_id and not _is_base64_encoded_unified_file_id(file_id):
setattr(response, file_attr, None)
verbose_proxy_logger.warning(
f"Redacted unconverted raw batch {file_attr}={file_id!r} to avoid leaking a provider file ID "
"that bypasses the managed-file ownership check"
)
async def ensure_batch_response_managed_file_ids(
response,
managed_files_obj,
@ -935,10 +970,36 @@ async def ensure_batch_response_managed_file_ids(
db_batch_object=None,
unified_batch_id: Optional[str] = None,
) -> None:
"""Normalize batch file IDs to managed unified IDs before DB persistence."""
"""Normalize batch file IDs to managed unified IDs before returning/persisting.
Registers missing managed-file rows for raw provider output/error IDs (owned
by the batch creator) and rewrites them to unified IDs. Any raw ID that still
could not be converted is redacted so a raw provider file ID never reaches the
client, where it would bypass the /v1/files ownership check.
"""
await resolve_input_file_id_to_unified(response, prisma_client)
await resolve_output_file_ids_to_unified(response, prisma_client)
await _register_raw_batch_output_file_ids(
response=response,
managed_files_obj=managed_files_obj,
requester_api_key_dict=user_api_key_dict,
db_batch_object=db_batch_object,
unified_batch_id=unified_batch_id,
verbose_proxy_logger=verbose_proxy_logger,
)
_redact_unconverted_batch_output_file_ids(response, verbose_proxy_logger)
async def _register_raw_batch_output_file_ids(
response,
managed_files_obj,
requester_api_key_dict,
db_batch_object,
unified_batch_id: Optional[str],
verbose_proxy_logger,
) -> None:
if managed_files_obj is None:
return
@ -957,16 +1018,17 @@ async def ensure_batch_response_managed_file_ids(
if target_model_names:
model_name = ",".join(target_model_names)
if user_api_key_dict is None and db_batch_object is not None:
from litellm.proxy._types import UserAPIKeyAuth
user_api_key_dict = UserAPIKeyAuth(
user_id=getattr(db_batch_object, "created_by", None) or "default-user-id",
team_id=getattr(db_batch_object, "team_id", None),
)
if user_api_key_dict is None:
owner_api_key_dict = _resolve_batch_file_owner_auth(
db_batch_object=db_batch_object,
requester_api_key_dict=requester_api_key_dict,
)
if owner_api_key_dict is None:
return
parent_otel_span = getattr(requester_api_key_dict, "parent_otel_span", None) or getattr(
owner_api_key_dict, "parent_otel_span", None
)
for file_attr in ("output_file_id", "error_file_id"):
raw_file_id = getattr(response, file_attr, None)
if not raw_file_id or _is_base64_encoded_unified_file_id(raw_file_id):
@ -980,9 +1042,9 @@ async def ensure_batch_response_managed_file_ids(
await managed_files_obj.store_unified_file_id(
file_id=new_unified_file_id,
file_object=None,
litellm_parent_otel_span=getattr(user_api_key_dict, "parent_otel_span", None),
litellm_parent_otel_span=parent_otel_span,
model_mappings={model_id: raw_file_id},
user_api_key_dict=user_api_key_dict,
user_api_key_dict=owner_api_key_dict,
)
setattr(response, file_attr, new_unified_file_id)
verbose_proxy_logger.debug(f"Converted batch {file_attr} {raw_file_id!r} to managed ID before DB write")

95
tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py
View file

@ -1,5 +1,6 @@
"""Regression: update_batch_in_database must not persist raw provider output_file_id."""
import base64
import json
from types import SimpleNamespace
from typing import Optional
@ -11,7 +12,28 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
ensure_batch_response_managed_file_ids,
update_batch_in_database,
)
from litellm.types.utils import LiteLLMBatch
from litellm.types.utils import LiteLLMBatch, SpecialEnums
def _make_unified_file_id(
raw_output_file_id: str = "file-managed-output",
model_id: str = "my-model",
model_name: str = "openai/gpt-4o",
) -> str:
"""Build a base64 unified file id the way get_unified_output_file_id does, so
_is_base64_encoded_unified_file_id recognizes it (a fake id that fails that
check would be redacted as fail-closed)."""
plain = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
"application/json",
"00000000-0000-0000-0000-000000000abc",
model_name,
raw_output_file_id,
model_id,
)
return base64.urlsafe_b64encode(plain.encode()).decode().rstrip("=")
UNIFIED_FILE_ID = _make_unified_file_id()
def _build_batch_response(
@ -38,7 +60,7 @@ def _build_batch_response(
return batch
def _build_managed_files_mock(unified_id: str = "file-bWFuYWdlZF9vdXRwdXRfaWQ="):
def _build_managed_files_mock(unified_id: str = UNIFIED_FILE_ID):
mock = MagicMock()
mock.get_unified_output_file_id = MagicMock(return_value=unified_id)
mock.store_unified_file_id = AsyncMock()
@ -55,7 +77,7 @@ def _build_prisma_mock():
@pytest.mark.asyncio
async def test_update_batch_in_database_stores_unified_output_file_id():
raw_output_file_id = "file-rawoutput789"
unified_output_file_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ="
unified_output_file_id = UNIFIED_FILE_ID
batch_id = "batch_managed_ids_test"
unified_batch_id = (
"litellm_proxy;model_id:my-model;llm_batch_id:batch_managed_ids_test"
@ -92,7 +114,7 @@ async def test_update_batch_in_database_stores_unified_output_file_id():
@pytest.mark.asyncio
async def test_ensure_batch_response_normalizes_error_file_id():
"""Both output_file_id and error_file_id must be normalized to managed IDs."""
unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ="
unified_id = UNIFIED_FILE_ID
response = _build_batch_response(
output_file_id="file-raw-output",
error_file_id="file-raw-error",
@ -116,8 +138,9 @@ async def test_ensure_batch_response_normalizes_error_file_id():
@pytest.mark.asyncio
async def test_ensure_batch_response_swallows_conversion_errors():
"""When the managed-files conversion raises, the failure is logged, not propagated."""
async def test_ensure_batch_response_redacts_id_when_conversion_errors():
"""Fail closed: when conversion raises, the raw provider id must be redacted,
not returned, so it can't be replayed against /v1/files to bypass ownership."""
raw_output_file_id = "file-raw-output"
response = _build_batch_response(
output_file_id=raw_output_file_id,
@ -139,14 +162,14 @@ async def test_ensure_batch_response_swallows_conversion_errors():
user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"),
)
assert response.output_file_id == raw_output_file_id
assert response.output_file_id is None
mock_logger.warning.assert_called()
@pytest.mark.asyncio
async def test_ensure_batch_response_builds_auth_from_db_batch_object():
"""If user_api_key_dict is omitted, fall back to created_by/team_id on db_batch_object."""
unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ="
unified_id = UNIFIED_FILE_ID
response = _build_batch_response(
output_file_id="file-raw-output",
hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"},
@ -175,7 +198,7 @@ async def test_ensure_batch_response_builds_auth_from_db_batch_object():
@pytest.mark.asyncio
async def test_ensure_batch_response_resolves_model_name_from_unified_file_id():
"""When hidden_params lacks model_name, derive it from unified_file_id."""
unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ="
unified_id = UNIFIED_FILE_ID
response = _build_batch_response(
output_file_id="file-raw-output",
hidden_params={
@ -208,7 +231,7 @@ async def test_ensure_batch_response_derives_model_id_from_unified_batch_id():
ids as managed files by deriving model_id from the unified batch id, otherwise
the terminal-retrieve short-circuit leaks raw provider file ids.
"""
unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ="
unified_id = UNIFIED_FILE_ID
response = _build_batch_response(
output_file_id="file-raw-output",
error_file_id="file-raw-error",
@ -235,8 +258,9 @@ async def test_ensure_batch_response_derives_model_id_from_unified_batch_id():
@pytest.mark.asyncio
async def test_ensure_batch_response_returns_early_without_managed_files_obj():
"""Without managed_files_obj, the helper is a no-op (no conversion attempted)."""
async def test_ensure_batch_response_redacts_raw_id_without_managed_files_obj():
"""Fail closed: without the managed-files hook, conversion is impossible, so the
raw provider id must be redacted rather than returned."""
response = _build_batch_response(
output_file_id="file-raw-output",
hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"},
@ -250,12 +274,13 @@ async def test_ensure_batch_response_returns_early_without_managed_files_obj():
user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"),
)
assert response.output_file_id == "file-raw-output"
assert response.output_file_id is None
@pytest.mark.asyncio
async def test_ensure_batch_response_returns_early_without_model_id():
"""Without model_id in hidden_params, the helper cannot create managed IDs."""
async def test_ensure_batch_response_redacts_raw_id_without_model_id():
"""Fail closed: without a resolvable model_id no managed id can be built, so the
raw provider id must be redacted rather than returned."""
response = _build_batch_response(
output_file_id="file-raw-output",
hidden_params={"model_name": "openai/gpt-4o"},
@ -270,13 +295,14 @@ async def test_ensure_batch_response_returns_early_without_model_id():
user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"),
)
assert response.output_file_id == "file-raw-output"
assert response.output_file_id is None
mock_managed_files.get_unified_output_file_id.assert_not_called()
@pytest.mark.asyncio
async def test_ensure_batch_response_returns_early_without_auth():
"""Without user_api_key_dict or db_batch_object, no conversion is attempted."""
async def test_ensure_batch_response_redacts_raw_id_without_auth():
"""Fail closed: without any owner identity no row can be registered, so the raw
provider id must be redacted rather than returned."""
response = _build_batch_response(
output_file_id="file-raw-output",
hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"},
@ -290,5 +316,36 @@ async def test_ensure_batch_response_returns_early_without_auth():
verbose_proxy_logger=MagicMock(),
)
assert response.output_file_id == "file-raw-output"
assert response.output_file_id is None
mock_managed_files.get_unified_output_file_id.assert_not_called()
@pytest.mark.asyncio
async def test_ensure_batch_response_owns_registered_files_by_batch_creator_not_requester():
"""Regression for GH #33989 review: when a non-creator retrieves a terminal batch
before its files are registered, the new managed-file rows must be owned by the
batch creator (from db_batch_object), not the requester, otherwise the requester
could grant themselves access and the real creator couldn't manage the file."""
response = _build_batch_response(
output_file_id="file-raw-output",
hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"},
)
mock_managed_files = _build_managed_files_mock()
await ensure_batch_response_managed_file_ids(
response=response,
managed_files_obj=mock_managed_files,
prisma_client=_build_prisma_mock(),
verbose_proxy_logger=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(user_id="attacker-user", team_id="attacker-team"),
db_batch_object=SimpleNamespace(
created_by="victim-user", team_id="victim-team", status="failed"
),
)
forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[
"user_api_key_dict"
]
assert forwarded_auth.user_id == "victim-user"
assert forwarded_auth.team_id == "victim-team"
assert response.output_file_id == UNIFIED_FILE_ID