mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(proxy): keep the s3 bucket list for non-admin keys that send a client AWS identity
A client AWS identity field does not prove that the caller can read the bucket. A role or profile name picks credentials the proxy holds, and Bedrock reads the object with them. So the bucket list now applies to every non-admin s3 media part, whatever AWS identity fields the request sends. allow_client_side_credentials: true still turns it off This removes BaseAWSLLM.identity_choosing_params and the router lookup that only the skip needed Ported to lts/1.108.x
This commit is contained in:
mateo-berri
parent
e6531d1b21
commit
c40e4fb2f3
6 changed files with 15 additions and 186 deletions
|
|
@ -232,10 +232,6 @@ class SignsRequestsWithAWS:
|
|||
pass
|
||||
|
||||
|
||||
def _or_env(value: str | None, env_var: str) -> str | None:
|
||||
return value if value is not None else os.getenv(env_var)
|
||||
|
||||
|
||||
class BaseAWSLLM(SignsRequestsWithAWS):
|
||||
# Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request).
|
||||
# Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static
|
||||
|
|
@ -350,19 +346,6 @@ class BaseAWSLLM(SignsRequestsWithAWS):
|
|||
def _is_auth_with_aws_profile(aws_profile_name: str | None) -> bool:
|
||||
return aws_profile_name is not None
|
||||
|
||||
@classmethod
|
||||
def identity_choosing_params(cls, params: AwsAuthParams) -> tuple[str, ...]:
|
||||
token: Final = _or_env(params.aws_web_identity_token, "AWS_WEB_IDENTITY_TOKEN")
|
||||
role: Final = _or_env(params.aws_role_name, "AWS_ROLE_NAME")
|
||||
session: Final = _or_env(params.aws_session_name, "AWS_SESSION_NAME")
|
||||
if cls._is_auth_with_web_identity_token(token, role, session):
|
||||
return ("aws_web_identity_token", "aws_role_name")
|
||||
if cls._is_auth_with_aws_role(role):
|
||||
return ("aws_role_name",)
|
||||
if cls._is_auth_with_aws_profile(_or_env(params.aws_profile_name, "AWS_PROFILE_NAME")):
|
||||
return ("aws_profile_name",)
|
||||
return ()
|
||||
|
||||
@staticmethod
|
||||
def _is_auth_with_aws_session_token_tuple(
|
||||
aws_access_key_id: str | None,
|
||||
|
|
@ -415,7 +398,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
|
|||
aws_sts_endpoint,
|
||||
aws_external_id,
|
||||
) = tuple(
|
||||
_or_env(value, env_var)
|
||||
value if value is not None else os.getenv(env_var)
|
||||
for value, env_var in (
|
||||
(aws_access_key_id, "AWS_ACCESS_KEY_ID"),
|
||||
(aws_secret_access_key, "AWS_SECRET_ACCESS_KEY"),
|
||||
|
|
|
|||
|
|
@ -31,18 +31,15 @@ from litellm.litellm_core_utils.url_utils import (
|
|||
validate_url,
|
||||
)
|
||||
from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata
|
||||
from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
|
||||
)
|
||||
from litellm.types.router import (
|
||||
CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS,
|
||||
Deployment,
|
||||
DeploymentTypedDict,
|
||||
server_owned_wif_fields_present,
|
||||
)
|
||||
from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
|
||||
|
|
@ -241,47 +238,6 @@ def _allow_model_level_clientside_configurable_parameters(
|
|||
)
|
||||
|
||||
|
||||
_CLIENT_AWS_IDENTITY_PARAMS: Final = ("aws_role_name", "aws_web_identity_token", "aws_profile_name")
|
||||
|
||||
|
||||
def _sends_a_value(body: Mapping[str, object], param: str) -> bool:
|
||||
value: Final = body.get(param)
|
||||
return isinstance(value, str) and bool(value)
|
||||
|
||||
|
||||
def _client_aws_fields(body: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return {key: body[key] for key in AWS_AUTH_PARAM_KEYS if _sends_a_value(body, key)}
|
||||
|
||||
|
||||
def _identity_choosing_params(client_fields: Mapping[str, object], deployment: DeploymentTypedDict) -> tuple[str, ...]:
|
||||
return BaseAWSLLM.identity_choosing_params(
|
||||
AwsAuthParams.model_validate({**deployment["litellm_params"], **client_fields})
|
||||
)
|
||||
|
||||
|
||||
def deployment_admits_a_client_aws_identity(body: Mapping[str, object], llm_router: Router | None) -> bool:
|
||||
model: Final = body.get("model")
|
||||
if not isinstance(model, str) or llm_router is None:
|
||||
return False
|
||||
admitted: Final = frozenset(
|
||||
param
|
||||
for param in _CLIENT_AWS_IDENTITY_PARAMS
|
||||
if _sends_a_value(body, param)
|
||||
and _allow_model_level_clientside_configurable_parameters(
|
||||
model=model, param=param, request_body_value=body[param], llm_router=llm_router
|
||||
)
|
||||
)
|
||||
deployments: Final = llm_router.get_model_list(model_name=model) or ()
|
||||
client_fields: Final = _client_aws_fields(body)
|
||||
return (
|
||||
bool(admitted)
|
||||
and bool(deployments)
|
||||
and all(
|
||||
not admitted.isdisjoint(_identity_choosing_params(client_fields, deployment)) for deployment in deployments
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# Config dicts whose entries are spread as ``**dict`` into outbound LLM
|
||||
# API calls. ``litellm_embedding_config`` is consumed by the Milvus
|
||||
# vector store transformer. ``extra_body`` is the OpenAI-SDK passthrough
|
||||
|
|
|
|||
|
|
@ -100,12 +100,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
request_skips_budget_checks,
|
||||
tag_max_budget_check_for_tags,
|
||||
)
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
ClientSideStorageAllowlist,
|
||||
check_response_size_is_safe,
|
||||
deployment_admits_a_client_aws_identity,
|
||||
get_request_route,
|
||||
)
|
||||
from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist, check_response_size_is_safe, get_request_route
|
||||
from litellm.proxy.bug_report_config import build_proxy_bug_report
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
get_logging_caching_headers,
|
||||
|
|
@ -896,7 +891,6 @@ async def refuse_s3_media_the_key_may_not_read(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: Mapping[str, object],
|
||||
managed_files: object,
|
||||
llm_router: Router | None,
|
||||
) -> None:
|
||||
if not s3_media_uris(data) or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return
|
||||
|
|
@ -909,9 +903,7 @@ async def refuse_s3_media_the_key_may_not_read(
|
|||
detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key "
|
||||
f"({', '.join(managed_uris)}). Use a LiteLLM managed file id that this key's user or team created.",
|
||||
)
|
||||
if general_settings.get("allow_client_side_credentials") is True or deployment_admits_a_client_aws_identity(
|
||||
data, llm_router
|
||||
):
|
||||
if general_settings.get("allow_client_side_credentials") is True:
|
||||
return
|
||||
is_listed: Final = ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket
|
||||
unlisted_uris: Final = tuple(uri for uri in s3_uris if names_an_unlisted_s3_bucket(uri, is_listed))
|
||||
|
|
@ -2350,7 +2342,6 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_api_key_dict,
|
||||
general_settings,
|
||||
proxy_logging_obj.get_proxy_hook("managed_files"),
|
||||
llm_router,
|
||||
)
|
||||
self.data = await add_litellm_data_to_request(
|
||||
data=self.data,
|
||||
|
|
|
|||
|
|
@ -17681,7 +17681,7 @@ async def async_queue_request(
|
|||
data = await request.json()
|
||||
data.pop("_litellm_strip_stream_usage", None)
|
||||
await refuse_s3_media_the_key_may_not_read(
|
||||
data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files"), llm_router
|
||||
data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files")
|
||||
)
|
||||
|
||||
# Include original request and headers in the data
|
||||
|
|
|
|||
|
|
@ -1911,12 +1911,11 @@ def _s3_media_frame_check(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: Mapping[str, object],
|
||||
managed_files: object,
|
||||
llm_router: Router | None,
|
||||
) -> ResponsesWsFrameCheck:
|
||||
async def refusal(params: Mapping[str, object], /) -> str | None:
|
||||
try:
|
||||
await refuse_s3_media_the_key_may_not_read(
|
||||
{**params, "model": admitted_model}, user_api_key_dict, general_settings, managed_files, llm_router
|
||||
{**params, "model": admitted_model}, user_api_key_dict, general_settings, managed_files
|
||||
)
|
||||
except HTTPException as refused:
|
||||
return str(refused.detail)
|
||||
|
|
@ -2028,7 +2027,6 @@ async def _responses_websocket_session(
|
|||
user_api_key_dict,
|
||||
general_settings,
|
||||
proxy_logging_obj.get_proxy_hook("managed_files"),
|
||||
llm_router,
|
||||
),
|
||||
**_routing_hints_from_first_ws_frame(first_message),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -419,9 +419,6 @@ _S3_VALUES_ON_NON_MEDIA_ROUTES: Final = (
|
|||
)
|
||||
|
||||
|
||||
_AWS_IDENTITY_ENV_VARS: Final = ("AWS_WEB_IDENTITY_TOKEN", "AWS_ROLE_NAME", "AWS_SESSION_NAME", "AWS_PROFILE_NAME")
|
||||
|
||||
|
||||
async def _send_s3_media(
|
||||
route_type: _S3MediaRoute,
|
||||
path: str,
|
||||
|
|
@ -430,7 +427,7 @@ async def _send_s3_media(
|
|||
general_settings: dict[str, object],
|
||||
managed_files: CustomLogger | None = None,
|
||||
client_auth_params: tuple[str, ...] = (),
|
||||
deployments_aws_params: tuple[Mapping[str, str], ...] = (MappingProxyType({}),),
|
||||
deployment_aws_params: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> object:
|
||||
from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig
|
||||
|
||||
|
|
@ -449,7 +446,6 @@ async def _send_s3_media(
|
|||
**deployment_aws_params,
|
||||
},
|
||||
}
|
||||
for deployment_aws_params in deployments_aws_params
|
||||
]
|
||||
)
|
||||
return await ProxyBaseLLMRequestProcessing(data={"model": "nova", **body}).base_process_llm_request(
|
||||
|
|
@ -848,72 +844,20 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("param", "value", "deployments_aws_params"),
|
||||
[
|
||||
pytest.param("aws_role_name", "arn:aws:iam::111122223333:role/reader", ({},), id="role"),
|
||||
pytest.param(
|
||||
"aws_web_identity_token",
|
||||
"oidc-token",
|
||||
({"aws_role_name": "arn:aws:iam::111122223333:role/web", "aws_session_name": "proxy"},),
|
||||
id="web-identity-token-with-the-deployment-role-and-session",
|
||||
),
|
||||
pytest.param("aws_profile_name", "reader", ({},), id="profile"),
|
||||
],
|
||||
)
|
||||
async def test_non_admin_s3_media_read_with_a_client_aws_identity_the_deployment_admits_needs_no_listed_bucket(
|
||||
self,
|
||||
param: str,
|
||||
value: str,
|
||||
deployments_aws_params: tuple[dict[str, str], ...],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
for env_var in _AWS_IDENTITY_ENV_VARS:
|
||||
monkeypatch.delenv(env_var, raising=False)
|
||||
response: Final = await _send_s3_media(
|
||||
"acompletion",
|
||||
"/v1/chat/completions",
|
||||
{**_chat_image("s3://media-bucket/dir/cat.png"), param: value},
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
{},
|
||||
client_auth_params=(param,),
|
||||
deployments_aws_params=deployments_aws_params,
|
||||
)
|
||||
|
||||
assert "a cat" in str(response)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("identity", "admitted", "deployments_aws_params"),
|
||||
("identity", "deployment_aws_params"),
|
||||
[
|
||||
pytest.param({"aws_role_name": "arn:aws:iam::111122223333:role/reader"}, {}, id="role"),
|
||||
pytest.param(
|
||||
{"aws_web_identity_token": "oidc-token"},
|
||||
("aws_web_identity_token",),
|
||||
({},),
|
||||
id="web-identity-token-without-a-role-and-session",
|
||||
),
|
||||
pytest.param(
|
||||
{"aws_profile_name": "reader"},
|
||||
("aws_profile_name",),
|
||||
({"aws_role_name": "arn:aws:iam::111122223333:role/deployment"},),
|
||||
id="profile-under-a-deployment-role",
|
||||
),
|
||||
pytest.param(
|
||||
{"aws_profile_name": "reader"},
|
||||
("aws_profile_name",),
|
||||
({}, {"aws_role_name": "arn:aws:iam::111122223333:role/deployment"}),
|
||||
id="profile-under-a-role-on-one-deployment-of-the-group",
|
||||
{"aws_role_name": "arn:aws:iam::111122223333:role/web", "aws_session_name": "proxy"},
|
||||
id="web-identity-token-with-the-deployment-role-and-session",
|
||||
),
|
||||
pytest.param({"aws_profile_name": "reader"}, {}, id="profile"),
|
||||
],
|
||||
)
|
||||
async def test_non_admin_s3_media_is_refused_when_the_admitted_client_aws_field_does_not_pick_the_credentials(
|
||||
self,
|
||||
identity: dict[str, str],
|
||||
admitted: tuple[str, ...],
|
||||
deployments_aws_params: tuple[dict[str, str], ...],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
async def test_non_admin_s3_media_needs_a_listed_bucket_even_with_a_client_aws_identity_the_deployment_admits(
|
||||
self, identity: dict[str, str], deployment_aws_params: dict[str, str]
|
||||
) -> None:
|
||||
for env_var in _AWS_IDENTITY_ENV_VARS:
|
||||
monkeypatch.delenv(env_var, raising=False)
|
||||
with pytest.raises(HTTPException) as refused:
|
||||
await _send_s3_media(
|
||||
"acompletion",
|
||||
|
|
@ -921,56 +865,13 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
{**_chat_image("s3://media-bucket/dir/cat.png"), **identity},
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
{},
|
||||
client_auth_params=admitted,
|
||||
deployments_aws_params=deployments_aws_params,
|
||||
client_auth_params=tuple(identity),
|
||||
deployment_aws_params=deployment_aws_params,
|
||||
)
|
||||
|
||||
assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail
|
||||
assert "general_settings.allowed_client_side_s3_buckets" in str(refused.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("uri", "identity", "admitted", "named_in_refusal"),
|
||||
[
|
||||
pytest.param(
|
||||
"s3://media-bucket/dir/cat.png",
|
||||
{"aws_role_name": "arn:aws:iam::111122223333:role/reader"},
|
||||
("aws_region_name",),
|
||||
"general_settings.allowed_client_side_s3_buckets",
|
||||
id="identity-the-deployment-does-not-admit",
|
||||
),
|
||||
pytest.param(
|
||||
"s3://media-bucket/dir/cat.png",
|
||||
{"aws_role_name": ""},
|
||||
("aws_role_name",),
|
||||
"general_settings.allowed_client_side_s3_buckets",
|
||||
id="empty-identity",
|
||||
),
|
||||
pytest.param(
|
||||
"s3://media-bucket/litellm-bedrock-files/other-team.png",
|
||||
{"aws_role_name": "arn:aws:iam::111122223333:role/reader"},
|
||||
("aws_role_name",),
|
||||
"LiteLLM-managed storage",
|
||||
id="managed-key-with-an-admitted-identity",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_non_admin_s3_media_is_refused_unless_the_deployment_admits_the_client_aws_identity(
|
||||
self, uri: str, identity: dict[str, object], admitted: tuple[str, ...], named_in_refusal: str
|
||||
) -> None:
|
||||
with pytest.raises(HTTPException) as refused:
|
||||
await _send_s3_media(
|
||||
"acompletion",
|
||||
"/v1/chat/completions",
|
||||
{**_chat_image(uri), **identity},
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
{},
|
||||
client_auth_params=admitted,
|
||||
)
|
||||
|
||||
assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail
|
||||
assert named_in_refusal in refused.value.detail
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("route_type", "path", "body"), _S3_VALUES_OUTSIDE_MEDIA_PARTS)
|
||||
async def test_non_admin_s3_values_outside_media_parts_pass_with_no_bucket_listed(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue