From dfb7424b4b3176903476816adb797cb3e0fbbcdf Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:45:08 -0500 Subject: [PATCH 1/3] fix(bedrock): sign rerank requests with the shared, header-filtered SigV4 helper BedrockRerankHandler._prepare_request duplicated ad-hoc SigV4 signing instead of using BaseAWSLLM.get_request_headers, the helper every other Bedrock handler (embeddings, converse, invoke, image) already uses. The duplicate skipped header filtering before signing, so any forwarded header (e.g. x-forwarded-for) got included in the signed set and could invalidate the signature if rewritten downstream between signing and delivery, the same class of bug fixed for the invoke path in #19111. --- litellm/llms/bedrock/rerank/handler.py | 29 +++++--------- .../test_bedrock_rerank_header_forwarding.py | 39 +++++++++++++++++++ 2 files changed, 49 insertions(+), 19 deletions(-) diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 1cc72f265eb..79b70c47a9a 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -135,11 +135,6 @@ class BedrockRerankHandler(BaseAWSLLM): data: dict, optional_params: dict, ) -> BedrockPreparedRequest: - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model) ### SET RUNTIME ENDPOINT ### @@ -150,24 +145,20 @@ class BedrockRerankHandler(BaseAWSLLM): ) proxy_endpoint_url = proxy_endpoint_url.replace("bedrock-runtime", "bedrock-agent-runtime") proxy_endpoint_url = f"{proxy_endpoint_url}/rerank" - sigv4: Final = SigV4Auth( - boto3_credentials_info.credentials, - "bedrock", - boto3_credentials_info.aws_region_name, - ) - # Make POST Request - body: Final = json.dumps(data).encode("utf-8") + body: Final = json.dumps(data).encode("utf-8") headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - request: Final = AWSRequest(method="POST", url=proxy_endpoint_url, data=body, headers=headers) - sigv4.add_auth(request) - if ( - extra_headers is not None and "Authorization" in extra_headers - ): # prevent sigv4 from overwriting the auth header - request.headers["Authorization"] = extra_headers["Authorization"] - prepped: Final = request.prepare() + + prepped: Final = self.get_request_headers( + credentials=boto3_credentials_info.credentials, + aws_region_name=boto3_credentials_info.aws_region_name, + extra_headers=extra_headers, + endpoint_url=proxy_endpoint_url, + data=body, + headers=headers, + ) return BedrockPreparedRequest( endpoint_url=proxy_endpoint_url, diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index 17443ca899e..748d46af895 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -17,6 +17,7 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo +from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler # Mock response for Bedrock rerank @@ -408,3 +409,41 @@ def test_bedrock_rerank_extra_headers_and_headers_merge(): except Exception as e: pytest.fail(f"Failed to merge and forward headers: {str(e)}") + + +def test_bedrock_rerank_forwarded_headers_excluded_from_sigv4_signature(): + """ + A forwarded header like x-forwarded-for can be rewritten between LiteLLM + signing the request and AWS receiving it (e.g. by an intermediate load + balancer), which invalidates the signature if that header was part of + the signed set. It must still reach Bedrock, just unsigned. + """ + from botocore.credentials import Credentials + + handler = BedrockRerankHandler() + mock_credentials_info = Boto3CredentialsInfo( + credentials=Credentials("test-access-key", "test-secret-key", "test-token"), + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=None, + ) + + with patch.object( + BedrockRerankHandler, + "_get_boto_credentials_from_optional_params", + return_value=mock_credentials_info, + ): + prepared_request = handler._prepare_request( + model="cohere.rerank-v3-5:0", + api_base=None, + extra_headers={"x-forwarded-for": "203.0.113.5"}, + data={"query": test_query, "documents": test_documents}, + optional_params={}, + ) + + headers = prepared_request["prepped"].headers + signed_headers = headers["Authorization"].split("SignedHeaders=")[1].split(",")[0].split(";") + + assert "x-forwarded-for" not in signed_headers, ( + f"x-forwarded-for must not be part of the SigV4 signature, got SignedHeaders={signed_headers}" + ) + assert headers["x-forwarded-for"] == "203.0.113.5", "forwarded header must still reach Bedrock, unsigned" From d80608eca6e9a1a98c3b0c2f7620c6d6496712e6 Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:58:12 -0500 Subject: [PATCH 2/3] test(bedrock): drop class-level monkeypatch in rerank signature test Pass static AWS credentials through optional_params so the real credential-resolution path runs locally instead of patching BedrockRerankHandler._get_boto_credentials_from_optional_params. --- .../test_bedrock_rerank_header_forwarding.py | 30 +++++++------------ 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index 748d46af895..ebe0df2a1c7 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -418,27 +418,19 @@ def test_bedrock_rerank_forwarded_headers_excluded_from_sigv4_signature(): balancer), which invalidates the signature if that header was part of the signed set. It must still reach Bedrock, just unsigned. """ - from botocore.credentials import Credentials - handler = BedrockRerankHandler() - mock_credentials_info = Boto3CredentialsInfo( - credentials=Credentials("test-access-key", "test-secret-key", "test-token"), - aws_region_name="us-east-1", - aws_bedrock_runtime_endpoint=None, - ) - with patch.object( - BedrockRerankHandler, - "_get_boto_credentials_from_optional_params", - return_value=mock_credentials_info, - ): - prepared_request = handler._prepare_request( - model="cohere.rerank-v3-5:0", - api_base=None, - extra_headers={"x-forwarded-for": "203.0.113.5"}, - data={"query": test_query, "documents": test_documents}, - optional_params={}, - ) + prepared_request = handler._prepare_request( + model="cohere.rerank-v3-5:0", + api_base=None, + extra_headers={"x-forwarded-for": "203.0.113.5"}, + data={"query": test_query, "documents": test_documents}, + optional_params={ + "aws_access_key_id": "test-access-key", + "aws_secret_access_key": "test-secret-key", + "aws_region_name": "us-east-1", + }, + ) headers = prepared_request["prepped"].headers signed_headers = headers["Authorization"].split("SignedHeaders=")[1].split(",")[0].split(";") From 729a95232204599e550f46c3de8aec1af9455673 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:17:24 -0700 Subject: [PATCH 3/3] fix(bedrock): keep rerank on SigV4 when a Bedrock API key is set Routing rerank through get_request_headers also picked up its AWS_BEARER_TOKEN_BEDROCK branch. Bedrock API keys are only valid for Bedrock and Bedrock Runtime actions, not for Agents for Amazon Bedrock Runtime ones, and rerank is served by bedrock-agent-runtime, so AWS rejects a bearer-signed rerank call. Opt the rerank handler out of the bearer path so it keeps signing with SigV4. --- litellm/llms/bedrock/base_aws_llm.py | 7 +++-- litellm/llms/bedrock/rerank/handler.py | 1 + .../test_bedrock_rerank_header_forwarding.py | 30 +++++++++++++++++++ 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index db6f2c0d491..4332848e545 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1434,9 +1434,12 @@ class BaseAWSLLM: data: str | bytes, headers: dict, api_key: str | None = None, + supports_bearer_token: bool = True, ) -> AWSPreparedRequest: - if api_key is not None: - aws_bearer_token: str | None = api_key + if not supports_bearer_token: + aws_bearer_token: str | None = None + elif api_key is not None: + aws_bearer_token = api_key else: aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK") diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 79b70c47a9a..cb0473887ea 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -158,6 +158,7 @@ class BedrockRerankHandler(BaseAWSLLM): endpoint_url=proxy_endpoint_url, data=body, headers=headers, + supports_bearer_token=False, ) return BedrockPreparedRequest( diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index ebe0df2a1c7..dd14b38f07a 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -439,3 +439,33 @@ def test_bedrock_rerank_forwarded_headers_excluded_from_sigv4_signature(): f"x-forwarded-for must not be part of the SigV4 signature, got SignedHeaders={signed_headers}" ) assert headers["x-forwarded-for"] == "203.0.113.5", "forwarded header must still reach Bedrock, unsigned" + + +def test_bedrock_rerank_signs_with_sigv4_even_when_bedrock_api_key_is_set(monkeypatch): + """ + Bedrock API keys are only valid for Bedrock and Bedrock Runtime actions, not for + Agents for Amazon Bedrock Runtime ones. Rerank is served by bedrock-agent-runtime, + so it has to keep signing with SigV4 even when AWS_BEARER_TOKEN_BEDROCK is set. + """ + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "test-bedrock-api-key") + + handler = BedrockRerankHandler() + + prepared_request = handler._prepare_request( + model="cohere.rerank-v3-5:0", + api_base=None, + extra_headers=None, + data={"query": test_query, "documents": test_documents}, + optional_params={ + "aws_access_key_id": "test-access-key", + "aws_secret_access_key": "test-secret-key", + "aws_region_name": "us-east-1", + }, + ) + + assert prepared_request["endpoint_url"].startswith("https://bedrock-agent-runtime.") + + authorization = prepared_request["prepped"].headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256"), ( + f"rerank must sign with SigV4, got Authorization={authorization[:30]}" + )