From 6ad549d04ec33d38f79ff0aa103c7a74f3e93cfa Mon Sep 17 00:00:00 2001
From: SamRemis <sjremis94@gmail.com>
Date: Wed, 9 Aug 2023 14:58:03 -0400
Subject: [PATCH] Bugfix multipart uploader with flexible checksums (#2754)

Enable use of flexible checksums on the multipart uploader and multipart copy
---
 .../bugfix-flexible-checksums-multipart.json  |  7 +++
 src/S3/MultipartUploadingTrait.php            | 12 +++--
 tests/S3/MultipartUploaderTest.php            | 47 +++++++++++++++++++
 3 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 .changes/nextrelease/bugfix-flexible-checksums-multipart.json

diff --git a/.changes/nextrelease/bugfix-flexible-checksums-multipart.json b/.changes/nextrelease/bugfix-flexible-checksums-multipart.json
new file mode 100644
index 0000000000..73401793c6
--- /dev/null
+++ b/.changes/nextrelease/bugfix-flexible-checksums-multipart.json
@@ -0,0 +1,7 @@
+[
+  {
+    "type": "feature",
+    "category": "Multipart",
+    "description": "Include ChecksumAlgorithm member in Multipart Upload and Copy for flexible checksums"
+  }
+]
\ No newline at end of file
diff --git a/src/S3/MultipartUploadingTrait.php b/src/S3/MultipartUploadingTrait.php
index baccf58c51..002bd43c46 100644
--- a/src/S3/MultipartUploadingTrait.php
+++ b/src/S3/MultipartUploadingTrait.php
@@ -51,10 +51,14 @@ public static function getStateFromService(
 
     protected function handleResult(CommandInterface $command, ResultInterface $result)
     {
-        $this->getState()->markPartAsUploaded($command['PartNumber'], [
-            'PartNumber' => $command['PartNumber'],
-            'ETag'       => $this->extractETag($result),
-        ]);
+        $partData = [];
+        $partData['PartNumber'] = $command['PartNumber'];
+        $partData['ETag'] = $this->extractETag($result);
+        if (isset($command['ChecksumAlgorithm'])) {
+            $checksumMemberName = 'Checksum' . strtoupper($command['ChecksumAlgorithm']);
+            $partData[$checksumMemberName] = $result[$checksumMemberName];
+        }
+        $this->getState()->markPartAsUploaded($command['PartNumber'], $partData);
     }
 
     abstract protected function extractETag(ResultInterface $result);
diff --git a/tests/S3/MultipartUploaderTest.php b/tests/S3/MultipartUploaderTest.php
index c5beefae92..6abfaa8277 100644
--- a/tests/S3/MultipartUploaderTest.php
+++ b/tests/S3/MultipartUploaderTest.php
@@ -316,4 +316,51 @@ public function testAppliesAmbiguousSuccessParsing()
         );
         $uploader->upload();
     }
+
+    /**
+     * @dataProvider testMultipartSuccessStreams
+     */
+    public function testUploaderAddsFlexibleChecksums($stream, $size)
+    {
+        /** @var \Aws\S3\S3Client $client */
+        $client = $this->getTestClient('s3');
+        $uploadOptions = [
+            'bucket'          => 'foo',
+            'key'             => 'bar',
+            'add_content_md5' => true,
+            'params'          => [
+                'RequestPayer'  => 'test',
+                'ChecksumAlgorithm' => 'Sha256'
+            ],
+            'before_initiate' => function($command) {
+                $this->assertSame('test', $command['RequestPayer']);
+                $this->assertSame('Sha256', $command['ChecksumAlgorithm']);
+            },
+            'before_upload'   => function($command) use ($size) {
+                $this->assertLessThan($size, $command['ContentLength']);
+                $this->assertSame('test', $command['RequestPayer']);
+                $this->assertSame('Sha256', $command['ChecksumAlgorithm']);
+            },
+            'before_complete' => function($command) {
+                $this->assertSame('test', $command['RequestPayer']);
+                $this->assertSame('Sha256', $command['ChecksumAlgorithm']);
+            }
+        ];
+        $url = 'http://foo.s3.amazonaws.com/bar';
+
+        $this->addMockResults($client, [
+            new Result(['UploadId' => 'baz', 'ChecksumSHA256' => 'xyz']),
+            new Result(['ETag' => 'A', 'ChecksumSHA256' => 'xyz']),
+            new Result(['ETag' => 'B', 'ChecksumSHA256' => 'xyz']),
+            new Result(['ETag' => 'C', 'ChecksumSHA256' => 'xyz']),
+            new Result(['Location' => $url, 'ChecksumSHA256' => 'xyz'])
+        ]);
+
+        $uploader = new MultipartUploader($client, $stream, $uploadOptions);
+        $result = $uploader->upload();
+
+        $this->assertTrue($uploader->getState()->isCompleted());
+        $this->assertSame('xyz', $result['ChecksumSHA256']);
+        $this->assertSame($url, $result['ObjectURL']);
+    }
 }