S3 - UploadFile with ChecksumAlgorithm changes body content (#5498)

moto/s3/responses.py

@@ -1104,6 +1104,15 @@ class S3Response(BaseResponse):
             line = body_io.readline()
         return bytes(new_body)
 
+    def _handle_encoded_body(self, body, content_length):
+        body_io = io.BytesIO(body)
+        # first line should equal '{content_length}\r\n
+        body_io.readline()
+        # Body contains actual data next
+        return body_io.read(content_length)
+        # last line should equal
+        # amz-checksum-sha256:<..>\r\n
+
     @amzn_request_id
     def key_response(self, request, full_url, headers):
         # Key and Control are lumped in because splitting out the regex is too much of a pain :/
@@ -1384,12 +1393,22 @@ class S3Response(BaseResponse):
         checksum_header = f"x-amz-checksum-{checksum_algorithm.lower()}"
         checksum_value = request.headers.get(checksum_header)
         if not checksum_value and checksum_algorithm:
-            search = re.search(r"x-amz-checksum-\w+:(\w+={1,2})", body.decode())
+            # Extract the checksum-value from the body first
+            search = re.search(b"x-amz-checksum-\w+:(\w+={1,2})", body)
             checksum_value = search.group(1) if search else None
 
         if checksum_value:
             response_headers.update({checksum_header: checksum_value})
 
+        # Extract the actual data from the body second
+        if (
+            request.headers.get("x-amz-content-sha256", None)
+            == "STREAMING-UNSIGNED-PAYLOAD-TRAILER"
+        ):
+            body = self._handle_encoded_body(
+                body, int(request.headers["x-amz-decoded-content-length"])
+            )
+
         bucket_key_enabled = request.headers.get(
             "x-amz-server-side-encryption-bucket-key-enabled", None
         )

tests/test_s3/test_s3.py

@@ -808,6 +808,24 @@ def test_upload_from_file_to_presigned_url():
     os.remove("text.txt")
 
 
+@mock_s3
+def test_upload_file_with_checksum_algorithm():
+    random_bytes = b"\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00\xff\n\xdb\x00C\x00\x08\x06\x06\x07\x06\x05\x08\x07\x07\n"
+    with open("rb.tmp", mode="wb") as f:
+        f.write(random_bytes)
+    s3 = boto3.resource("s3", region_name=DEFAULT_REGION_NAME)
+    s3_client = boto3.client("s3", region_name=DEFAULT_REGION_NAME)
+    bucket = "mybucket"
+    s3_client.create_bucket(Bucket=bucket)
+    s3_client.upload_file(
+        "rb.tmp", bucket, "my_key.csv", ExtraArgs={"ChecksumAlgorithm": "SHA256"}
+    )
+    os.remove("rb.tmp")
+
+    actual_content = s3.Object(bucket, "my_key.csv").get()["Body"].read()
+    assert random_bytes == actual_content
+
+
 @mock_s3
 def test_put_chunked_with_v4_signature_in_body():
     bucket_name = "mybucket"