mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
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:
parent
3726bceb53
commit
dfb7424b4b
2 changed files with 49 additions and 19 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue