mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge pull request #31435 from kingdoooo/litellm_bedrock_output_bucket
fix(bedrock): validate file-content retrieval against the configured output bucket (#26335)
This commit is contained in:
Mateo Wang
GitHub
commit
96d2ceef3b
3 changed files with 233 additions and 16 deletions
|
|
@ -3,6 +3,7 @@ import json
|
|||
import os
|
||||
import time
|
||||
from collections.abc import Iterable, Mapping, MutableMapping, Sequence
|
||||
from contextlib import suppress
|
||||
from functools import cache
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
|
|
@ -149,11 +150,12 @@ class _BedrockS3RequestParams(BaseModel):
|
|||
|
||||
|
||||
class _TrustedS3ModelCredentials(BaseModel):
|
||||
"""The S3 bucket the server trusts file ids against, from the deployment snapshot."""
|
||||
"""The S3 buckets the server trusts file ids against, from the deployment snapshot."""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
s3_bucket_name: str | None = None
|
||||
s3_output_bucket_name: str | None = None
|
||||
|
||||
|
||||
def extract_s3_uri_from_file_id(file_id: str) -> str:
|
||||
|
|
@ -179,6 +181,18 @@ def extract_s3_uri_from_file_id(file_id: str) -> str:
|
|||
raise ValueError("file_id must be a managed LiteLLM S3 file id")
|
||||
|
||||
|
||||
_S3_BUCKET_REQUIRED_ERROR: Final = "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval."
|
||||
|
||||
|
||||
def _trusted_s3_model_credentials(litellm_params: Mapping[str, object]) -> _TrustedS3ModelCredentials:
|
||||
trusted_model_credentials: Final = litellm_params.get("_litellm_internal_model_credentials")
|
||||
if not isinstance(trusted_model_credentials, MappingProxyType):
|
||||
return _TrustedS3ModelCredentials()
|
||||
snapshot: Final[dict[str, object]] = {}
|
||||
snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot
|
||||
return _TrustedS3ModelCredentials.model_validate(snapshot)
|
||||
|
||||
|
||||
def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str:
|
||||
"""
|
||||
Resolve the server-configured S3 bucket for Bedrock file operations.
|
||||
|
|
@ -187,20 +201,50 @@ def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str:
|
|||
environment; never a request-supplied param, since the bucket is what
|
||||
`validate_managed_cloud_file_id` checks file ids against.
|
||||
"""
|
||||
trusted_model_credentials: Final = litellm_params.get("_litellm_internal_model_credentials")
|
||||
bucket_name: str | None = None
|
||||
if isinstance(trusted_model_credentials, MappingProxyType):
|
||||
snapshot: Final[dict[str, object]] = {}
|
||||
snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot
|
||||
bucket_name = _TrustedS3ModelCredentials.model_validate(snapshot).s3_bucket_name
|
||||
bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
|
||||
bucket_name: Final = _trusted_s3_model_credentials(litellm_params).s3_bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
|
||||
if not bucket_name:
|
||||
raise ValueError(
|
||||
"S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval."
|
||||
)
|
||||
raise ValueError(_S3_BUCKET_REQUIRED_ERROR)
|
||||
return bucket_name
|
||||
|
||||
|
||||
def get_configured_s3_bucket_names(litellm_params: Mapping[str, object]) -> tuple[str, ...]:
|
||||
"""
|
||||
Resolve the server-configured S3 buckets a Bedrock file id may live in.
|
||||
|
||||
Bedrock batch outputs land in ``s3_output_bucket_name`` when it differs from
|
||||
the input bucket, so retrieval validates against both. Same trust rules as
|
||||
``get_configured_s3_bucket_name``: only the immutable credential snapshot or
|
||||
the environment, never a request param.
|
||||
"""
|
||||
trusted: Final = _trusted_s3_model_credentials(litellm_params)
|
||||
input_bucket: Final = trusted.s3_bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
|
||||
output_bucket: Final = trusted.s3_output_bucket_name or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME")
|
||||
buckets: Final = tuple(dict.fromkeys(bucket for bucket in (input_bucket, output_bucket) if bucket))
|
||||
if not buckets:
|
||||
raise ValueError(_S3_BUCKET_REQUIRED_ERROR)
|
||||
return buckets
|
||||
|
||||
|
||||
def _validate_file_id_against_configured_buckets(
|
||||
s3_uri: str,
|
||||
configured_bucket_names: tuple[str, ...],
|
||||
allow_legacy_cloud_file_ids: bool,
|
||||
) -> tuple[str, str]:
|
||||
def validate_against(configured_bucket_name: str) -> tuple[str, str]:
|
||||
return validate_managed_cloud_file_id(
|
||||
file_id=s3_uri,
|
||||
scheme="s3://",
|
||||
configured_bucket_name=configured_bucket_name,
|
||||
allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES,
|
||||
allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids,
|
||||
)
|
||||
|
||||
for candidate_bucket_name in configured_bucket_names[:-1]:
|
||||
with suppress(ValueError):
|
||||
return validate_against(candidate_bucket_name)
|
||||
return validate_against(configured_bucket_names[-1])
|
||||
|
||||
|
||||
def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Response) -> int:
|
||||
"""
|
||||
S3 answers PutObject with an empty body, so the stored object size comes from the
|
||||
|
|
@ -1186,11 +1230,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
raise ValueError("file_id is required for Bedrock file content retrieval")
|
||||
|
||||
s3_uri: Final = extract_s3_uri_from_file_id(file_id)
|
||||
bucket_name, object_key = validate_managed_cloud_file_id(
|
||||
file_id=s3_uri,
|
||||
scheme="s3://",
|
||||
configured_bucket_name=get_configured_s3_bucket_name(litellm_params),
|
||||
allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES,
|
||||
bucket_name, object_key = _validate_file_id_against_configured_buckets(
|
||||
s3_uri=s3_uri,
|
||||
configured_bucket_names=get_configured_s3_bucket_names(litellm_params),
|
||||
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1992,6 +1992,113 @@ class TestBedrockFileContentTransformation:
|
|||
litellm_params=self._litellm_params(),
|
||||
)
|
||||
|
||||
def _trusted(self, **deployment_litellm_params) -> dict:
|
||||
"""Build the trusted snapshot the way the proxy does: deployment
|
||||
litellm_params funneled through ``CredentialLiteLLMParams`` (the strict
|
||||
allowlist ``get_deployment_credentials_with_provider`` applies) before
|
||||
retrieval ever sees them. Injecting a raw ``MappingProxyType`` would
|
||||
bypass that filter and hide whether a bucket field actually survives
|
||||
into the snapshot in production."""
|
||||
from types import MappingProxyType
|
||||
|
||||
from litellm.types.router import CredentialLiteLLMParams
|
||||
|
||||
snapshot = CredentialLiteLLMParams(**deployment_litellm_params).model_dump(
|
||||
exclude_none=True
|
||||
)
|
||||
params = self._litellm_params()
|
||||
params["_litellm_internal_model_credentials"] = MappingProxyType(snapshot)
|
||||
return params
|
||||
|
||||
def test_retrieves_from_distinct_output_bucket(self, monkeypatch):
|
||||
"""Batch outputs can land in a separate s3_output_bucket_name. Retrieval
|
||||
must validate the file id against the output bucket too, not just the
|
||||
input bucket, or the very outputs the feature serves are unreachable.
|
||||
The snapshot is built through the production credential filter, so this
|
||||
fails if s3_output_bucket_name is dropped from that allowlist."""
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
||||
monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False)
|
||||
|
||||
url, _ = BedrockFilesConfig().transform_file_content_request(
|
||||
file_content_request={
|
||||
"file_id": "s3://out-bucket/litellm-batch-outputs/job/in.jsonl.out"
|
||||
},
|
||||
optional_params={},
|
||||
litellm_params=self._trusted(
|
||||
s3_bucket_name="in-bucket", s3_output_bucket_name="out-bucket"
|
||||
),
|
||||
)
|
||||
|
||||
assert (
|
||||
url
|
||||
== "https://s3.us-west-2.amazonaws.com/out-bucket/litellm-batch-outputs/job/in.jsonl.out"
|
||||
)
|
||||
|
||||
def test_output_bucket_falls_back_to_env(self, monkeypatch):
|
||||
"""The output bucket resolves from AWS_S3_OUTPUT_BUCKET_NAME when not in
|
||||
the trusted snapshot, mirroring the input-bucket env fallback."""
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "in-bucket")
|
||||
monkeypatch.setenv("AWS_S3_OUTPUT_BUCKET_NAME", "env-out-bucket")
|
||||
|
||||
url, _ = BedrockFilesConfig().transform_file_content_request(
|
||||
file_content_request={
|
||||
"file_id": "s3://env-out-bucket/litellm-batch-outputs/job/in.jsonl.out"
|
||||
},
|
||||
optional_params={},
|
||||
litellm_params=self._litellm_params(),
|
||||
)
|
||||
|
||||
assert (
|
||||
url
|
||||
== "https://s3.us-west-2.amazonaws.com/env-out-bucket/litellm-batch-outputs/job/in.jsonl.out"
|
||||
)
|
||||
|
||||
def test_input_bucket_still_validates_when_output_bucket_set(self, monkeypatch):
|
||||
"""Adding output-bucket support must not break retrieval of input-bucket
|
||||
objects when both buckets are configured."""
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
||||
monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False)
|
||||
|
||||
url, _ = BedrockFilesConfig().transform_file_content_request(
|
||||
file_content_request={
|
||||
"file_id": "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out"
|
||||
},
|
||||
optional_params={},
|
||||
litellm_params=self._trusted(
|
||||
s3_bucket_name="in-bucket", s3_output_bucket_name="out-bucket"
|
||||
),
|
||||
)
|
||||
|
||||
assert (
|
||||
url
|
||||
== "https://s3.us-west-2.amazonaws.com/in-bucket/litellm-batch-outputs/job/in.jsonl.out"
|
||||
)
|
||||
|
||||
def test_rejects_bucket_outside_input_and_output(self, monkeypatch):
|
||||
"""A file id whose bucket is neither the input nor the output bucket is
|
||||
still rejected (SSRF / bucket-confusion guard)."""
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
||||
monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False)
|
||||
|
||||
with pytest.raises(ValueError, match="configured storage bucket"):
|
||||
BedrockFilesConfig().transform_file_content_request(
|
||||
file_content_request={
|
||||
"file_id": "s3://other-bucket/litellm-batch-outputs/job/x.jsonl.out"
|
||||
},
|
||||
optional_params={},
|
||||
litellm_params=self._trusted(
|
||||
s3_bucket_name="in-bucket", s3_output_bucket_name="out-bucket"
|
||||
),
|
||||
)
|
||||
|
||||
def test_sign_request_without_botocore_raises_helpful_error(self, monkeypatch):
|
||||
"""A missing botocore must surface an actionable 'install boto3' error
|
||||
rather than a raw import failure."""
|
||||
|
|
|
|||
|
|
@ -143,3 +143,71 @@ class TestRouterCredentialResolution:
|
|||
assert credentials is not None
|
||||
assert credentials.get("api_key") == "sk-static-key"
|
||||
assert "azure_ad_token" not in credentials
|
||||
|
||||
|
||||
class TestRouterCredentialResolutionS3OutputBucket:
|
||||
"""Same strict-dump trap as azure_ad_token (#30235), for Bedrock batch
|
||||
file retrieval (#26335). Bedrock batch outputs land in a per-model
|
||||
``s3_output_bucket_name`` when it differs from the input bucket. The
|
||||
file-content retrieval path validates a file id against the buckets in the
|
||||
trusted credential snapshot, and that snapshot is built by round-tripping
|
||||
the deployment's ``litellm_params`` through ``CredentialLiteLLMParams``. If
|
||||
the field is undeclared it is dropped, so the output bucket never reaches
|
||||
retrieval and output-bucket file ids are rejected as foreign."""
|
||||
|
||||
def test_credentials_preserve_s3_output_bucket_name(self):
|
||||
from litellm import Router
|
||||
|
||||
deployment_id = "bedrock-batch-output-bucket-fixed-uuid"
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-batch",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"s3_bucket_name": "in-bucket",
|
||||
"s3_output_bucket_name": "out-bucket",
|
||||
"aws_region_name": "us-west-2",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id=deployment_id
|
||||
)
|
||||
assert credentials is not None
|
||||
assert credentials.get("s3_output_bucket_name") == "out-bucket", (
|
||||
"Router credential resolution dropped s3_output_bucket_name; "
|
||||
"Bedrock batch file-content retrieval will reject output-bucket "
|
||||
"file ids as foreign for model-routed deployments (#26335)"
|
||||
)
|
||||
assert credentials.get("s3_bucket_name") == "in-bucket"
|
||||
|
||||
def test_credentials_without_output_bucket_unaffected(self):
|
||||
"""A deployment that configures only the input bucket keeps it and does
|
||||
not gain a phantom output bucket in the resolved credentials."""
|
||||
from litellm import Router
|
||||
|
||||
deployment_id = "bedrock-batch-input-only-fixed-uuid"
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-batch-input-only",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"s3_bucket_name": "in-bucket",
|
||||
"aws_region_name": "us-west-2",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id=deployment_id
|
||||
)
|
||||
assert credentials is not None
|
||||
assert credentials.get("s3_bucket_name") == "in-bucket"
|
||||
assert "s3_output_bucket_name" not in credentials
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue