diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 0885775932c..455197acc67 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1508,14 +1508,34 @@ class BaseAWSLLM: sigv4.add_auth(request) request_headers_dict = dict(request.headers) - # Add back original headers after signing. Only headers in SignedHeaders - # are integrity-protected; forwarded headers (x-forwarded-*) must remain unsigned. + # Replay caller headers only when they are not bound to the fresh SigV4 signature. + authorization_header = next( + ( + str(header_value) + for header_name, header_value in request_headers_dict.items() + if header_name.lower() == "authorization" and header_value is not None + ), + "", + ) + signed_header_names = {"authorization", "host"} + for auth_part in authorization_header.split(","): + auth_part = auth_part.strip() + if auth_part.startswith("SignedHeaders="): + # SigV4 binds these header names to the current body; replaying stale + # retry headers would make AWS verify a different canonical request. + signed_header_names.update( + signed_header.strip().lower() + for signed_header in auth_part.removeprefix("SignedHeaders=").split( + ";" + ) + if signed_header.strip() + ) + break for header_name, header_value in headers.items(): - if header_value is not None: - request_headers_dict[header_name] = header_value - if ( - headers is not None and "Authorization" in headers - ): # prevent sigv4 from overwriting the auth header - request_headers_dict["Authorization"] = headers["Authorization"] + if header_value is None: + continue + if header_name.lower() in signed_header_names: + continue + request_headers_dict[header_name] = header_value return request_headers_dict, request.body diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index a4969e5dacc..a8be3854641 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -74,6 +74,15 @@ def _os_environ_without_aws_keys() -> Dict[str, str]: return {k: v for k, v in os.environ.items() if not k.startswith("AWS_")} +def _get_header_value_case_insensitive( + headers: Dict[str, Any], header_name: str +) -> Any: + for key, value in headers.items(): + if key.lower() == header_name.lower(): + return value + raise AssertionError(f"Missing header {header_name}; found {list(headers.keys())}") + + def test_ambient_env_credentials_use_iam_cache_across_instances(): """Else-branch env path uses shared iam_cache; second call on another instance does not refetch.""" base_a = BaseAWSLLM() @@ -539,6 +548,73 @@ def test_sign_request_with_sigv4(): assert result_body == mock_request.body +def test_sign_request_replays_only_unsigned_original_headers(): + llm = BaseAWSLLM() + mock_credentials = Credentials("test_key", "test_secret", "test_token") + stale_authorization = ( + "AWS4-HMAC-SHA256 Credential=old, " + "SignedHeaders=content-type;host;x-amz-date, Signature=old" + ) + stale_amz_date = "19990101T000000Z" + stale_security_token = "old_token" + headers = { + "Authorization": stale_authorization, + "x-amz-date": stale_amz_date, + "x-amz-security-token": stale_security_token, + "X-Forwarded-For": "203.0.113.10", + "X-Custom-Header": "caller-value", + } + + with ( + patch("litellm.llms.bedrock.base_aws_llm.get_secret_str", return_value=None), + patch.object(llm, "get_credentials", return_value=mock_credentials), + patch.object(llm, "_get_aws_region_name", return_value="us-west-2"), + ): + result_headers, result_body = llm._sign_request( + service_name="bedrock", + headers=headers, + optional_params={ + "aws_access_key_id": "test_key", + "aws_secret_access_key": "test_secret", + "aws_region_name": "us-west-2", + }, + request_data={"prompt": "test"}, + api_base="https://bedrock-runtime.us-west-2.amazonaws.com/model/test/invoke", + ) + + authorization = _get_header_value_case_insensitive(result_headers, "authorization") + amz_date = _get_header_value_case_insensitive(result_headers, "x-amz-date") + security_token = _get_header_value_case_insensitive( + result_headers, "x-amz-security-token" + ) + + assert authorization != stale_authorization + assert authorization.startswith("AWS4-HMAC-SHA256") + assert "SignedHeaders=" in authorization + assert "x-amz-date" in authorization + assert "x-amz-security-token" in authorization + assert amz_date != stale_amz_date + assert security_token != stale_security_token + assert all( + value != stale_authorization + for key, value in result_headers.items() + if key.lower() == "authorization" + ) + assert all( + value != stale_amz_date + for key, value in result_headers.items() + if key.lower() == "x-amz-date" + ) + assert all( + value != stale_security_token + for key, value in result_headers.items() + if key.lower() == "x-amz-security-token" + ) + assert result_headers["X-Forwarded-For"] == "203.0.113.10" + assert result_headers["X-Custom-Header"] == "caller-value" + assert result_body is not None + + def test_sign_request_with_api_key_bearer_token(): """ Test that _sign_request uses the api_key parameter as a bearer token when provided