From 5f37b2800d5d09457f20f171260f9f9a6b63f421 Mon Sep 17 00:00:00 2001 From: Micah Mills <13383869+micahmills88@users.noreply.github.com> Date: Fri, 9 Oct 2026 14:13:10 -0400 Subject: [PATCH 1/2] fix(bedrock): discard stale payload hashes when resigning --- litellm/llms/bedrock/base_aws_llm.py | 6 +- tests/unit/llms/bedrock/test_base_aws_llm.py | 95 +++++++++++++- .../custom_httpx/test_llm_http_handler.py | 122 ++++++++++-------- 3 files changed, 164 insertions(+), 59 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 91d18c651da..d0e2ee08000 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -55,7 +55,9 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile( r"(?:^|\.)sts(?:-fips)?\.([a-z0-9-]+)\.(?:amazonaws\.com(?:\.cn)?|vpce\.amazonaws\.com)" ) -SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"}) +SIGV4_COMPUTED_HEADERS: Final = frozenset( + {"authorization", "x-amz-date", "x-amz-security-token", "date", "x-amz-content-sha256"} +) _AWS_SESSION_TAGS_ADAPTER: Final[TypeAdapter[tuple[AwsSessionTag, ...]]] = TypeAdapter(tuple[AwsSessionTag, ...]) @@ -1675,6 +1677,8 @@ class BaseAWSLLM(SignsRequestsWithAWS): if header_value is None: continue header_lower = header_name.lower() + if header_lower in SIGV4_COMPUTED_HEADERS: + continue if header_lower in aws_headers or header_lower.startswith("x-amz-") or header_lower.startswith("x-amzn-"): aws_signature_headers[header_name] = header_value diff --git a/tests/unit/llms/bedrock/test_base_aws_llm.py b/tests/unit/llms/bedrock/test_base_aws_llm.py index cc625dcf3f3..90f16ad6014 100644 --- a/tests/unit/llms/bedrock/test_base_aws_llm.py +++ b/tests/unit/llms/bedrock/test_base_aws_llm.py @@ -12,7 +12,7 @@ from fastapi.testclient import TestClient from collections.abc import Callable from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Literal, Optional from unittest.mock import MagicMock, Mock, patch from botocore.awsrequest import AWSPreparedRequest, AWSRequest @@ -3073,7 +3073,7 @@ def test_filter_headers_skips_none_values(): filtered = llm._filter_headers_for_aws_signature(headers) assert filtered["Content-Type"] == "application/json" - assert filtered["x-amz-date"] == "20240101T000000Z" + assert "x-amz-date" not in filtered assert "x-amz-security-token" not in filtered assert "x-amzn-bedrock-kb-session-id" not in filtered assert "host" not in filtered @@ -4215,3 +4215,94 @@ def test_shared_json_signer_preserves_body_and_signs_for_the_requested_service() ) assert request.body == body assert "/us-west-2/s3/aws4_request" in request.headers["Authorization"] + + +def _request_with_payload_hash( + entrypoint: Literal["invoke", "prepared", "mantle"], + headers: dict[str, str], + request_data: dict[str, str], + url: str, +) -> tuple[dict[str, str], str | bytes]: + from litellm.llms.bedrock_mantle.messages.transformation import BedrockMantleAnthropicMessagesConfig + + optional_params: Final = { + "aws_access_key_id": "payload-hash-access", + "aws_secret_access_key": "payload-hash-secret", + "aws_region_name": "us-east-1", + } + if entrypoint == "invoke": + invoke_headers, invoke_body = BaseAWSLLM()._sign_request( + service_name="bedrock", + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=url, + ) + assert invoke_body is not None + return invoke_headers, invoke_body + if entrypoint == "mantle": + mantle_headers, mantle_body = BedrockMantleAnthropicMessagesConfig().sign_request( + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=url, + ) + assert mantle_body is not None + return mantle_headers, mantle_body + prepared: Final = BaseAWSLLM().get_request_headers( + credentials=Credentials("payload-hash-access", "payload-hash-secret"), + aws_region_name="us-east-1", + extra_headers=None, + endpoint_url=url, + data=json.dumps(request_data), + headers=headers, + supports_bearer_token=False, + ) + assert isinstance(prepared.body, (str, bytes)) + return dict(prepared.headers), prepared.body + + +@pytest.mark.parametrize("hash_header", (None, "x-amz-content-sha256", "X-Amz-Content-SHA256")) +@pytest.mark.parametrize( + ("entrypoint", "url"), + ( + ("invoke", "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/invoke"), + ("prepared", "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/converse"), + ("mantle", "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages"), + ), +) +def test_sigv4_resign_uses_current_body_hash_and_preserves_forwarded_headers( + entrypoint: Literal["invoke", "prepared", "mantle"], url: str, hash_header: str | None +) -> None: + import hashlib + + original: Final = {"input": "before retry"} + original_hash: Final = hashlib.sha256(json.dumps(original).encode()).hexdigest() + input_headers: Final = { + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + **({hash_header: original_hash} if hash_header is not None else {}), + } + first_headers, first_body = _request_with_payload_hash(entrypoint, input_headers, original, url) + assert first_headers["Authorization"].split("Signature=")[1] == _recomputed_sigv4_signature( + url=url, + secret_key="payload-hash-secret", + authorization=first_headers["Authorization"], + headers=first_headers, + body=first_body, + ) + previous_headers: Final = dict(first_headers) + retry_headers, retry_body = _request_with_payload_hash(entrypoint, first_headers, {"input": "after retry: é"}, url) + retry_bytes: Final = retry_body if isinstance(retry_body, bytes) else retry_body.encode() + expected_hash: Final = hashlib.sha256(retry_bytes).hexdigest() + normalized: Final = {name.lower(): value for name, value in retry_headers.items()} + assert normalized.get("x-amz-content-sha256", expected_hash) == expected_hash + assert retry_headers["Authorization"].split("Signature=")[1] == _recomputed_sigv4_signature( + url=url, + secret_key="payload-hash-secret", + authorization=retry_headers["Authorization"], + headers=retry_headers, + body=retry_body, + ) + assert retry_headers["anthropic-version"] == "2023-06-01" + assert first_headers == previous_headers diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 25d2d11b5ca..243cb26a76f 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -2772,26 +2772,25 @@ async def test_alist_input_items_surfaces_upstream_error_status(): @pytest.mark.asyncio -async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_request(monkeypatch): - """Regression: after Bedrock rejects a replayed thinking block (400 invalid signature), - the strip-and-retry re-sign must not inherit attempt 1's SigV4 Authorization/X-Amz-Date; - reusing them over the new stripped body makes AWS return 403 SignatureDoesNotMatch.""" - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( - AmazonAnthropicClaudeMessagesConfig, - ) +@pytest.mark.parametrize("hash_header", (None, "x-amz-content-sha256", "X-Amz-Content-SHA256")) +@pytest.mark.parametrize("stream", (False, True)) +async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_request( + monkeypatch: pytest.MonkeyPatch, hash_header: str | None, stream: bool +) -> None: + import hashlib for env_var in ("AWS_BEARER_TOKEN_BEDROCK", "AWS_SESSION_TOKEN", "AWS_PROFILE"): monkeypatch.delenv(env_var, raising=False) - handler = BaseLLMHTTPHandler() - provider_config = AmazonAnthropicClaudeMessagesConfig() - litellm_params = GenericLiteLLMParams( - aws_access_key_id="AKIAIOSFODNN7EXAMPLE", - aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + handler: Final = BaseLLMHTTPHandler() + provider_config: Final = AmazonAnthropicClaudeMessagesConfig() + litellm_params: Final = GenericLiteLLMParams( + aws_access_key_id="payload-hash-access", + aws_secret_access_key="payload-hash-secret", aws_region_name="us-east-1", ) - request_url = "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/invoke" - request_body = { + request_url: Final = "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/invoke" + request_body: Final = { "anthropic_version": "bedrock-2023-05-31", "max_tokens": 100, "messages": [ @@ -2806,60 +2805,71 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques {"role": "user", "content": "continue"}, ], } - first_attempt_headers, signed_json_body = provider_config.sign_request( - headers={"Content-Type": "application/json"}, + original_hash: Final = hashlib.sha256(json.dumps(request_body).encode()).hexdigest() + first_headers, signed_body = provider_config.sign_request( + headers={ + "Content-Type": "application/json", + **({hash_header: original_hash} if hash_header is not None else {}), + }, optional_params=dict(litellm_params), request_data=request_body, api_base=request_url, api_key=None, - stream=False, + stream=stream, fake_stream=False, model="test-model", ) + posts: Final[list[httpx.Request]] = [] # mutable-ok: the injected transport records request order - posts: list = [] - invalid_signature_response = httpx.Response( - 400, - text='{"message": "messages.1.content.0: Invalid `signature` in `thinking` block"}', - request=httpx.Request("POST", request_url), - ) - ok_response = httpx.Response(200, json={"id": "msg_1"}, request=httpx.Request("POST", request_url)) + def bedrock_peer(request: httpx.Request) -> httpx.Response: + posts.append(request) + body_hash: Final = hashlib.sha256(request.content).hexdigest() + if request.headers.get("x-amz-content-sha256", body_hash) != body_hash: + return httpx.Response( + 403, + headers={"x-amzn-errortype": "SignatureDoesNotMatch"}, + json={"message": "SignatureDoesNotMatch: payload hash does not match request bytes"}, + ) + if len(posts) == 1: + return httpx.Response( + 400, + json={"message": "messages.1.content.0: Invalid `signature` in `thinking` block"}, + ) + return httpx.Response(200, json={"id": "msg_1"}) - class FakeAsyncClient: - async def post(self, url, headers, data, stream=False, logging_obj=None, timeout=None): - posts.append({"headers": dict(headers), "data": data}) - return invalid_signature_response if len(posts) == 1 else ok_response - - logging_obj: Final = Mock(baseline_cache_context=None) - logging_obj.model_call_details = {} - - response = await handler._async_post_anthropic_messages_with_http_error_retry( - async_httpx_client=FakeAsyncClient(), - request_url=request_url, - headers=dict(first_attempt_headers), - signed_json_body=signed_json_body, - request_body=request_body, - stream=False, - logging_obj=logging_obj, - provider_config=provider_config, - litellm_params=litellm_params, - api_key=None, - model="test-model", - ) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(bedrock_peer)) + logging_obj: Final = Mock(baseline_cache_context=None, model_call_details={}) + try: + response: Final = await handler._async_post_anthropic_messages_with_http_error_retry( + async_httpx_client=client, + request_url=request_url, + headers=first_headers, + signed_json_body=signed_body, + request_body=request_body, + stream=stream, + logging_obj=logging_obj, + provider_config=provider_config, + litellm_params=litellm_params, + api_key=None, + model="test-model", + ) + await response.aclose() + finally: + await client.close() assert response.status_code == 200 assert len(posts) == 2 - retry_payload = json.loads(posts[1]["data"]) - retry_blocks = [ - block - for message in retry_payload["messages"] - if isinstance(message.get("content"), list) - for block in message["content"] - ] - assert retry_blocks and all(block["type"] != "thinking" for block in retry_blocks) - retry_authorization = posts[1]["headers"]["Authorization"] - assert retry_authorization.startswith("AWS4-HMAC-SHA256") - assert retry_authorization != first_attempt_headers["Authorization"] + assert posts[0].content != posts[1].content + assert json.loads(posts[1].content) == { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": "ok"}]}, + {"role": "user", "content": "continue"}, + ], + } + assert posts[1].headers["authorization"] != posts[0].headers["authorization"] def test_aws_signing_overrides_only_fills_missing_credentials(): From 294362f85ba0e116bc47c6df18d0c9524f49bead Mon Sep 17 00:00:00 2001 From: Micah Mills <13383869+micahmills88@users.noreply.github.com> Date: Fri, 9 Oct 2026 14:48:04 -0400 Subject: [PATCH 2/2] test(bedrock): expect computed headers to be filtered before signing --- tests/unit/llms/bedrock/chat/test_invoke_handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index d5fe5d8cf65..2f7942ae6e9 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -1407,11 +1407,11 @@ def test_filter_headers_for_aws_signature(): expected_aws_headers = { "Content-Type": "application/json", "Host": "bedrock-runtime.us-east-1.amazonaws.com", - "x-amz-date": "20240101T120000Z", - "x-amz-security-token": "test-token", } assert filtered_headers == expected_aws_headers, f"Expected {expected_aws_headers}, got {filtered_headers}" excluded_headers = [ + "x-amz-date", + "x-amz-security-token", "x-custom-header", "x-litellm-user-id", "x-forwarded-for",