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.
This commit is contained in:
Noah Nistler 2026-08-10 15:45:08 -05:00
parent 3726bceb53
commit dfb7424b4b
2 changed files with 49 additions and 19 deletions

View file

@ -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,

View file

@ -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"