scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

fix(proxy): harden resource ownership fallbacks

Browse source
This commit is contained in:
user 2026-04-30 19:29:08 -07:00
parent 6aac4552f9
commit f3cff9338e
5 changed files with 294 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/common_utils/resource_ownership.py
View file

@ -2,6 +2,8 @@ from typing import List, Optional
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
UNSCOPED_RESOURCE_OWNER_SCOPE = "__litellm_unscoped_proxy__"
def is_proxy_admin(user_api_key_dict: Optional[UserAPIKeyAuth]) -> bool:
if user_api_key_dict is None:
@ -41,6 +43,10 @@ def get_resource_owner_scopes(
_add(f"org:{user_api_key_dict.org_id}")
if user_api_key_dict.api_key:
_add(f"key:{user_api_key_dict.api_key}")
if user_api_key_dict.token:
_add(f"key:{user_api_key_dict.token}")
if not scopes:
_add(UNSCOPED_RESOURCE_OWNER_SCOPE)
return scopes
@ -59,7 +65,9 @@ def get_primary_resource_owner_scope(
return f"org:{user_api_key_dict.org_id}"
if user_api_key_dict.api_key:
return f"key:{user_api_key_dict.api_key}"
return None
if user_api_key_dict.token:
return f"key:{user_api_key_dict.token}"
return UNSCOPED_RESOURCE_OWNER_SCOPE
def user_can_access_resource_owner(

20
litellm/proxy/container_endpoints/endpoints.py
View file

@ -122,6 +122,11 @@ async def create_container(
user_api_base=user_api_base,
version=version,
)
return await record_container_owner(
response=response,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
@ -129,11 +134,6 @@ async def create_container(
proxy_logging_obj=proxy_logging_obj,
version=version,
)
return await record_container_owner(
response=response,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
@router.get(
@ -220,6 +220,11 @@ async def list_containers(
user_api_base=user_api_base,
version=version,
)
return await filter_container_list_response(
response=response,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
@ -227,11 +232,6 @@ async def list_containers(
proxy_logging_obj=proxy_logging_obj,
version=version,
)
return await filter_container_list_response(
response=response,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
@router.get(

114
litellm/proxy/container_endpoints/ownership.py
View file

@ -94,10 +94,14 @@ async def record_container_owner(
) -> Any:
container_id = _get_response_id(response)
owner = get_primary_resource_owner_scope(user_api_key_dict)
prisma_client = await _get_prisma_client()
if is_proxy_admin(user_api_key_dict) and (container_id is None or owner is None):
return response
if container_id is None or owner is None:
if container_id is None:
verbose_proxy_logger.warning(
"Skipping container ownership tracking because provider response has no id"
)
return response
if owner is None:
raise HTTPException(status_code=500, detail="Unable to track container")
original_container_id, resolved_provider = decode_container_id_for_ownership(
@ -112,16 +116,17 @@ async def record_container_owner(
file_object["custom_llm_provider"] = resolved_provider
file_object["provider_container_id"] = original_container_id
if prisma_client is None:
existing_owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id)
if existing_owner is not None and not user_can_access_resource_owner(
existing_owner, user_api_key_dict
):
raise HTTPException(status_code=403, detail="Forbidden")
_IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner
return response
try:
prisma_client = await _get_prisma_client()
if prisma_client is None:
existing_owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id)
if existing_owner is not None and not user_can_access_resource_owner(
existing_owner, user_api_key_dict
):
raise HTTPException(status_code=403, detail="Forbidden")
_IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner
return response
existing = await prisma_client.db.litellm_managedobjecttable.find_unique(
where={"model_object_id": model_object_id}
)
@ -169,25 +174,30 @@ async def _get_container_owner(
original_container_id: str,
custom_llm_provider: str,
) -> Optional[str]:
prisma_client = await _get_prisma_client()
if prisma_client is None:
return _IN_MEMORY_CONTAINER_OWNERS.get(
_container_model_object_id(
original_container_id,
custom_llm_provider,
)
)
row = await prisma_client.db.litellm_managedobjecttable.find_first(
where={
"model_object_id": _container_model_object_id(
original_container_id,
custom_llm_provider,
),
"file_purpose": CONTAINER_OBJECT_PURPOSE,
}
model_object_id = _container_model_object_id(
original_container_id,
custom_llm_provider,
)
return getattr(row, "created_by", None) if row is not None else None
try:
prisma_client = await _get_prisma_client()
if prisma_client is None:
return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id)
row = await prisma_client.db.litellm_managedobjecttable.find_first(
where={
"model_object_id": model_object_id,
"file_purpose": CONTAINER_OBJECT_PURPOSE,
}
)
return getattr(row, "created_by", None) if row is not None else None
except Exception as e:
verbose_proxy_logger.warning(
"Failed to load container ownership for container_id=%s; "
"falling back to in-process tracking: %s",
model_object_id,
e,
)
return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id)
async def assert_user_can_access_container(
@ -257,30 +267,38 @@ async def _get_allowed_container_ids(
user_api_key_dict: UserAPIKeyAuth,
custom_llm_provider: str,
) -> Set[str]:
prisma_client = await _get_prisma_client()
if prisma_client is None:
owner_scopes = get_resource_owner_scopes(user_api_key_dict)
return {
model_object_id
for model_object_id, owner in _IN_MEMORY_CONTAINER_OWNERS.items()
if owner in owner_scopes
}
owner_scopes = get_resource_owner_scopes(user_api_key_dict)
if not owner_scopes:
return set()
rows = await prisma_client.db.litellm_managedobjecttable.find_many(
where={
"file_purpose": CONTAINER_OBJECT_PURPOSE,
"created_by": {"in": owner_scopes},
}
)
return {
row.model_object_id
for row in rows
if getattr(row, "model_object_id", None) is not None
in_memory_allowed_ids = {
model_object_id
for model_object_id, owner in _IN_MEMORY_CONTAINER_OWNERS.items()
if owner in owner_scopes
}
try:
prisma_client = await _get_prisma_client()
if prisma_client is None:
return in_memory_allowed_ids
rows = await prisma_client.db.litellm_managedobjecttable.find_many(
where={
"file_purpose": CONTAINER_OBJECT_PURPOSE,
"created_by": {"in": owner_scopes},
}
)
return {
row.model_object_id
for row in rows
if getattr(row, "model_object_id", None) is not None
}
except Exception as e:
verbose_proxy_logger.warning(
"Failed to load allowed container ids; falling back to in-process "
"tracking: %s",
e,
)
return in_memory_allowed_ids
async def filter_container_list_response(

158
tests/test_litellm/containers/test_container_proxy_ownership.py
View file

@ -82,6 +82,83 @@ async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch):
assert data["updated_by"] == "team:team-1"
@pytest.mark.asyncio
async def test_should_record_token_owner_for_keys_without_user_team_or_org(monkeypatch):
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=prisma_client),
)
auth = UserAPIKeyAuth(token="hashed-token")
await ownership.record_container_owner(
response=_container("cntr_provider"),
user_api_key_dict=auth,
custom_llm_provider="openai",
)
data = table.create.await_args.kwargs["data"]
assert data["created_by"] == "key:hashed-token"
assert data["updated_by"] == "key:hashed-token"
@pytest.mark.asyncio
async def test_should_record_unscoped_owner_for_identityless_proxy_auth(monkeypatch):
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=None),
)
auth = UserAPIKeyAuth()
await ownership.record_container_owner(
response=_container("cntr_provider"),
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert (
ownership._IN_MEMORY_CONTAINER_OWNERS["container:openai:cntr_provider"]
== "__litellm_unscoped_proxy__"
)
original_id, provider = await ownership.assert_user_can_access_container(
container_id="cntr_provider",
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert original_id == "cntr_provider"
assert provider == "openai"
@pytest.mark.asyncio
async def test_should_skip_owner_record_when_provider_response_has_no_id(monkeypatch):
table = AsyncMock()
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=prisma_client),
)
response = {"object": "container"}
returned = await ownership.record_container_owner(
response=response,
user_api_key_dict=UserAPIKeyAuth(user_id="user-1"),
custom_llm_provider="openai",
)
assert returned == response
table.find_unique.assert_not_awaited()
table.create.assert_not_awaited()
@pytest.mark.asyncio
async def test_should_fallback_to_memory_when_persistent_owner_record_fails(
monkeypatch,
@ -178,6 +255,55 @@ async def test_should_deny_untracked_container_access_by_default(monkeypatch):
assert exc.value.status_code == 403
@pytest.mark.asyncio
async def test_should_fallback_to_memory_when_owner_lookup_fails(monkeypatch):
table = AsyncMock()
table.find_first.side_effect = Exception("db unavailable")
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=prisma_client),
)
ownership._IN_MEMORY_CONTAINER_OWNERS["container:openai:cntr_owned"] = "user-1"
auth = UserAPIKeyAuth(user_id="user-1")
original_id, provider = await ownership.assert_user_can_access_container(
container_id="cntr_owned",
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert original_id == "cntr_owned"
assert provider == "openai"
@pytest.mark.asyncio
async def test_should_fail_closed_when_owner_lookup_fails_without_memory(monkeypatch):
table = AsyncMock()
table.find_first.side_effect = Exception("db unavailable")
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=prisma_client),
)
auth = UserAPIKeyAuth(user_id="user-1")
with pytest.raises(HTTPException) as exc:
await ownership.assert_user_can_access_container(
container_id="cntr_owned",
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert exc.value.status_code == 403
@pytest.mark.asyncio
async def test_should_allow_untracked_container_access_when_enabled(monkeypatch):
monkeypatch.setattr(
@ -362,6 +488,38 @@ async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch
assert [item.id for item in filtered.data] == ["cntr_owned"]
@pytest.mark.asyncio
async def test_should_filter_container_list_with_memory_when_db_lookup_fails(
monkeypatch,
):
table = AsyncMock()
table.find_many.side_effect = Exception("db unavailable")
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=prisma_client),
)
ownership._IN_MEMORY_CONTAINER_OWNERS["container:openai:cntr_owned"] = "user-1"
auth = UserAPIKeyAuth(user_id="user-1")
response = ContainerListResponse(
object="list",
data=[_container("cntr_owned"), _container("cntr_other")],
has_more=True,
)
filtered = await ownership.filter_container_list_response(
response=response,
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert [item.id for item in filtered.data] == ["cntr_owned"]
assert filtered.has_more is False
@pytest.mark.asyncio
async def test_should_forward_decoded_container_id_for_proxy_forwarding(monkeypatch):
from litellm.proxy.container_endpoints import handler_factory

51
tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py
View file

@ -44,6 +44,57 @@ async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch):
assert table.create.await_args.kwargs["data"]["updated_by"] == "team:team-1"
@pytest.mark.asyncio
async def test_should_store_token_owner_for_keys_without_user_team_or_org(monkeypatch):
table = AsyncMock()
table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"])
prisma_client = type(
"Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()}
)()
monkeypatch.setattr(
LiteLLMSkillsHandler,
"_get_prisma_client",
AsyncMock(return_value=prisma_client),
)
auth = UserAPIKeyAuth(token="hashed-token")
skill = await LiteLLMSkillsHandler.create_skill(
data=NewSkillRequest(display_title="skill"),
user_api_key_dict=auth,
)
assert skill.created_by == "key:hashed-token"
assert table.create.await_args.kwargs["data"]["updated_by"] == "key:hashed-token"
@pytest.mark.asyncio
async def test_should_store_unscoped_owner_for_identityless_proxy_auth(monkeypatch):
table = AsyncMock()
table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"])
prisma_client = type(
"Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()}
)()
monkeypatch.setattr(
LiteLLMSkillsHandler,
"_get_prisma_client",
AsyncMock(return_value=prisma_client),
)
auth = UserAPIKeyAuth()
skill = await LiteLLMSkillsHandler.create_skill(
data=NewSkillRequest(display_title="skill"),
user_api_key_dict=auth,
)
assert skill.created_by == "__litellm_unscoped_proxy__"
assert (
table.create.await_args.kwargs["data"]["updated_by"]
== "__litellm_unscoped_proxy__"
)
@pytest.mark.asyncio
async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypatch):
table = AsyncMock()