PR Feedback - Make requestAlgorithmMember context more specific.

jonathan343 · Nov 22, 2024 · 503ead6 · 503ead6
1 parent a2b19bd
commit 503ead6
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 13 deletions.
diff --git a/botocore/httpchecksum.py b/botocore/httpchecksum.py
@@ -263,7 +263,7 @@ def resolve_request_checksum_algorithm(
     if has_checksum_header(request):
         return
 
-    extra_headers = {}
+    checksum_context = request["context"].get("checksum", {})
     request_checksum_calculation = request["context"][
         "client_config"
     ].request_checksum_calculation
@@ -299,8 +299,10 @@ def resolve_request_checksum_algorithm(
         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
+        checksum_context["request_algorithm_header"] = {
+            "name": algorithm_member_header,
+            "value": DEFAULT_CHECKSUM_ALGORITHM,
+        }
     else:
         return
 
@@ -313,17 +315,13 @@ 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
 
@@ -362,8 +360,11 @@ 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"])
+    if "request_algorithm_header" in checksum_context:
+        request_algorithm_header = checksum_context["request_algorithm_header"]
+        request["headers"][request_algorithm_header["name"]] = (
+            request_algorithm_header["value"]
+        )
 
 
 def _apply_request_header_checksum(request):

diff --git a/tests/unit/test_httpchecksum.py b/tests/unit/test_httpchecksum.py
@@ -125,10 +125,19 @@ def test_request_checksum_algorithm_model_default(self):
             "algorithm": "crc32",
             "in": "header",
             "name": "x-amz-checksum-crc32",
-            "extra_headers": {'x-amz-request-algorithm': 'CRC32'},
+        }
+        expected_request_algorithm_header = {
+            "name": "x-amz-request-algorithm",
+            "value": "CRC32",
         }
         actual_algorithm = request["context"]["checksum"]["request_algorithm"]
+        actual_request_algorithm_header = request["context"]["checksum"][
+            "request_algorithm_header"
+        ]
         self.assertEqual(actual_algorithm, expected_algorithm)
+        self.assertEqual(
+            actual_request_algorithm_header, expected_request_algorithm_header
+        )
 
         # Param is present, sha256 checksum will be set
         params = {"Algorithm": "sha256"}
@@ -163,10 +172,19 @@ 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'},
+        }
+        expected_request_algorithm_header = {
+            "name": "x-amz-request-algorithm",
+            "value": "CRC32",
         }
         actual_algorithm = request["context"]["checksum"]["request_algorithm"]
+        actual_request_algorithm_header = request["context"]["checksum"][
+            "request_algorithm_header"
+        ]
         self.assertEqual(actual_algorithm, expected_algorithm)
+        self.assertEqual(
+            actual_request_algorithm_header, expected_request_algorithm_header
+        )
 
         # Param is present, sha256 checksum will be set in the trailer
         params = {"Algorithm": "sha256"}
@@ -403,8 +421,11 @@ def test_apply_request_checksum_extra_headers(self):
                 "in": "trailer",
                 "algorithm": "crc32",
                 "name": "x-amz-checksum-crc32",
-                "extra_headers": {"foo": "bar"},
-            }
+            },
+            "request_algorithm_header": {
+                "name": "foo",
+                "value": "bar",
+            },
         }
         apply_request_checksum(request)
         self.assertEqual(request["headers"]["foo"], "bar")