Skip to content

Commit

Permalink
Resolve the "requestAlgorithmMember" header in "apply_request_checksum".
Browse files Browse the repository at this point in the history
  • Loading branch information
jonathan343 committed Nov 21, 2024
1 parent ff4f122 commit a2b19bd
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 16 deletions.
26 changes: 16 additions & 10 deletions botocore/httpchecksum.py
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,7 @@ def resolve_request_checksum_algorithm(
if has_checksum_header(request):
return

extra_headers = {}
request_checksum_calculation = request["context"][
"client_config"
].request_checksum_calculation
Expand Down Expand Up @@ -295,9 +296,11 @@ def resolve_request_checksum_algorithm(
algorithm_member and request_checksum_calculation == "when_supported"
):
algorithm_name = DEFAULT_CHECKSUM_ALGORITHM.lower()
_set_request_algorithm_member_header(
algorithm_member_header = _get_request_algorithm_member_header(
operation_model, request, algorithm_member
)
if algorithm_member_header is not None:
extra_headers[algorithm_member_header] = DEFAULT_CHECKSUM_ALGORITHM
else:
return

Expand All @@ -310,37 +313,38 @@ def resolve_request_checksum_algorithm(
# We only support unsigned trailer checksums currently. As this
# disables payload signing we'll only use trailers over TLS.
location_type = "trailer"
pass

algorithm = {
"algorithm": algorithm_name,
"in": location_type,
"name": f"x-amz-checksum-{algorithm_name}",
}
if extra_headers:
algorithm["extra_headers"] = extra_headers

checksum_context = request["context"].get("checksum", {})
checksum_context["request_algorithm"] = algorithm
request["context"]["checksum"] = checksum_context


def _set_request_algorithm_member_header(
def _get_request_algorithm_member_header(
operation_model, request, algorithm_member
):
"""Set the header targeted by the "requestAlgorithmMember" to the default checksum algorithm."""
"""Get the name of the header targeted by the "requestAlgorithmMember"."""
operation_input_shape = operation_model.input_shape
if not isinstance(operation_input_shape, StructureShape):
return

algorithm_member_shape = operation_input_shape.members.get(
algorithm_member
)
if not algorithm_member_shape:
return

field_name = algorithm_member_shape.serialization.get("name")
if not field_name:
return

request["headers"].setdefault(field_name, DEFAULT_CHECKSUM_ALGORITHM)
return (
algorithm_member_shape.serialization.get("name")
if algorithm_member_shape
else None
)


def apply_request_checksum(request):
Expand All @@ -358,6 +362,8 @@ def apply_request_checksum(request):
raise FlexibleChecksumError(
error_msg="Unknown checksum variant: {}".format(algorithm["in"])
)
if "extra_headers" in algorithm:
request["headers"].update(algorithm["extra_headers"])


def _apply_request_header_checksum(request):
Expand Down
21 changes: 15 additions & 6 deletions tests/unit/test_httpchecksum.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,12 +125,10 @@ def test_request_checksum_algorithm_model_default(self):
"algorithm": "crc32",
"in": "header",
"name": "x-amz-checksum-crc32",
"extra_headers": {'x-amz-request-algorithm': 'CRC32'},
}
actual_algorithm = request["context"]["checksum"]["request_algorithm"]
self.assertEqual(actual_algorithm, expected_algorithm)
self.assertEqual(
request["headers"]["x-amz-request-algorithm"], "CRC32"
)

# Param is present, sha256 checksum will be set
params = {"Algorithm": "sha256"}
Expand Down Expand Up @@ -165,12 +163,10 @@ def test_request_checksum_algorithm_model_default_streaming(self):
"algorithm": "crc32",
"in": "trailer",
"name": "x-amz-checksum-crc32",
"extra_headers": {'x-amz-request-algorithm': 'CRC32'},
}
actual_algorithm = request["context"]["checksum"]["request_algorithm"]
self.assertEqual(actual_algorithm, expected_algorithm)
self.assertEqual(
request["headers"]["x-amz-request-algorithm"], "CRC32"
)

# Param is present, sha256 checksum will be set in the trailer
params = {"Algorithm": "sha256"}
Expand Down Expand Up @@ -400,6 +396,19 @@ def test_apply_request_checksum_content_encoding_default(self):
apply_request_checksum(request)
self.assertEqual(request["headers"]["Content-Encoding"], "aws-chunked")

def test_apply_request_checksum_extra_headers(self):
request = self._build_request(b"")
request["context"]["checksum"] = {
"request_algorithm": {
"in": "trailer",
"algorithm": "crc32",
"name": "x-amz-checksum-crc32",
"extra_headers": {"foo": "bar"},
}
}
apply_request_checksum(request)
self.assertEqual(request["headers"]["foo"], "bar")

def test_response_checksum_algorithm_no_model(self):
request = self._build_request(b"")
operation_model = self._make_operation_model()
Expand Down

0 comments on commit a2b19bd

Please sign in to comment.