From 1d407c2f26d7587bf184f7cf19dfcaede7e860d7 Mon Sep 17 00:00:00 2001 From: Kent Date: Fri, 26 Jun 2026 17:24:21 +0800 Subject: [PATCH] fix(bedrock): validate file-content retrieval against the configured output bucket Bedrock batch jobs write their results to s3_output_bucket_name when it differs from the input bucket, but the file-content retrieval path validated the file id only against the input bucket (s3_bucket_name). A deployment that configures a separate output bucket therefore could not retrieve its own batch outputs: the id validated against the input bucket and was rejected as a foreign bucket. Resolve the trusted output bucket alongside the input bucket from the immutable credential snapshot (or AWS_S3_OUTPUT_BUCKET_NAME), and try the file id against each configured bucket, returning the first that validates. The SSRF guard is preserved: only server-configured buckets are tried, never a request param, and an id outside both is still rejected. --- litellm/llms/bedrock/files/transformation.py | 67 +++++++++++-- .../test_bedrock_files_transformation.py | 94 +++++++++++++++++++ 2 files changed, 151 insertions(+), 10 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 6cfaa88275d..df06c333d1f 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -81,11 +81,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: @@ -135,6 +136,41 @@ def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str: 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_model_credentials = litellm_params.get( + "_litellm_internal_model_credentials" + ) + input_bucket: str | None = None + output_bucket: str | None = None + if isinstance(trusted_model_credentials, MappingProxyType): + snapshot: dict[str, object] = {} + snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot + trusted = _TrustedS3ModelCredentials.model_validate(snapshot) + input_bucket = trusted.s3_bucket_name + output_bucket = trusted.s3_output_bucket_name + input_bucket = input_bucket or os.getenv("AWS_S3_BUCKET_NAME") + output_bucket = output_bucket or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") + + buckets = tuple( + dict.fromkeys(bucket for bucket in (input_bucket, output_bucket) if bucket) + ) + if not buckets: + raise ValueError( + "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval." + ) + return buckets + + class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Config for Bedrock Files - handles S3 uploads for Bedrock batch processing @@ -1042,15 +1078,26 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): raise ValueError("file_id is required for Bedrock file content retrieval") s3_uri = 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, - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( - litellm_params - ), - ) + allow_legacy = should_allow_legacy_cloud_file_ids(litellm_params) + last_error: ValueError | None = None + bucket_name: str | None = None + object_key: str | None = None + for configured_bucket in get_configured_s3_bucket_names(litellm_params): + try: + bucket_name, object_key = validate_managed_cloud_file_id( + file_id=s3_uri, + scheme="s3://", + configured_bucket_name=configured_bucket, + allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, + allow_legacy_cloud_file_ids=allow_legacy, + ) + break + except ValueError as e: + last_error = e + if bucket_name is None or object_key is None: + raise last_error or ValueError( + "file_id must reference a LiteLLM-managed storage object" + ) # The shared file-content handler passes optional_params={}, so AWS # credentials/region arrive via litellm_params here (unlike the upload diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e15..dd111969555 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1308,6 +1308,100 @@ class TestBedrockFileContentTransformation: litellm_params=self._litellm_params(), ) + def _trusted(self, **creds) -> dict: + from types import MappingProxyType + + params = self._litellm_params() + params["_litellm_internal_model_credentials"] = MappingProxyType(dict(creds)) + 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.""" + 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."""