s3: Improve AWS checksum validation

Fixes: #244
PiperOrigin-RevId: 786337455
Change-Id: Ie2e88cb54ddd2d4917cd496aa540dad9f024f508

Author: laramiel

tensorstore/kvstore/s3/s3_key_value_store.cc

@@ -437,8 +437,12 @@ struct ReadTask : public RateLimiterNode,
 
     if (options.byte_range.size() != 0) {
       request_builder.MaybeAddRangeHeader(options.byte_range);
+      if (options.byte_range.IsFull()) {
+        // Checksum validation is only supported when reading the entire object,
+        // and when the object was not written with a chunked upload.
+        request_builder.AddHeader("x-amz-checksum-mode", "ENABLED");
+      }
     }
-    request_builder.AddHeader("x-amz-checksum-mode", "ENABLED");
 
     const auto& ehr = endpoint_region_.value();
     start_time_ = absl::Now();
@@ -515,13 +519,22 @@ struct ReadTask : public RateLimiterNode,
             options.generation_conditions.if_not_equal, start_time_});
     }
 
-    // Validate checksum if present.
-    if (auto it = httpresponse.headers.find("x-amz-checksum-sha256");
-        it != httpresponse.headers.end()) {
-      auto payload_sha256 = PayloadSha256Base64(httpresponse.payload);
-      if (payload_sha256 != it->second) {
-        return absl::DataLossError(absl::StrFormat(
-            "Content hash mismatch: %s vs %s", payload_sha256, it->second));
+    // Validate checksum if possible.
+    if (auto checksum_sha256 =
+            httpresponse.headers.find("x-amz-checksum-sha256");
+        checksum_sha256 != httpresponse.headers.end()) {
+      if (auto checksum_type = httpresponse.headers.find("x-amz-checksum-type");
+          checksum_type == httpresponse.headers.end() ||
+          checksum_type->second == "FULL_OBJECT") {
+        // Only validate the payload checksum if the header indicates a full
+        // object checksum was sent; multi-part uploads need to be validated
+        // using the etag and multi-part configuration.
+        auto payload_sha256 = PayloadSha256Base64(httpresponse.payload);
+        if (payload_sha256 != checksum_sha256->second) {
+          return absl::DataLossError(
+              absl::StrFormat("Content hash mismatch: %s vs %s", payload_sha256,
+                              checksum_sha256->second));
+        }
       }
     }
 

tensorstore/kvstore/s3/s3_key_value_store_test.cc

@@ -155,26 +155,46 @@ TEST(S3KeyValueStoreTest, SimpleMock_VirtualHost) {
        HttpResponse{200, absl::Cord(),
                     HeaderMap{{"x-amz-bucket-region", "us-east-1"}}}},
 
-      {"GET https://my-bucket.s3.us-east-1.amazonaws.com/tmp:1/key_read",
+      {"GET https://my-bucket.s3.us-east-1.amazonaws.com/tmp:1/key_read1",
        HttpResponse{
            200, absl::Cord("abcd"),
            HeaderMap{{"etag", "\"900150983cd24fb0d6963f7d28e17f72\""},
                      {"x-amz-checksum-sha256",
                       "iNQmb9TmM40TuEX88olXnSCciXgjuSF9o+Fhk28DFYk="}}}},
 
-      {"GET https://my-bucket.s3.us-east-1.amazonaws.com/tmp:1/empty_read",
+      {"GET https://my-bucket.s3.us-east-1.amazonaws.com/tmp:1/key_read2",
+       HttpResponse{200, absl::Cord("abcd"),
+                    HeaderMap{
+                        {"etag", "\"900150983cd24fb0d6963f7d28e17f72\""},
+                        {"x-amz-checksum-sha256",
+                         "iNQmb9TmM40TuEX88olXnSCciXgjuSF9o+Fhk28DFYk="},
+                        {"x-amz-checksum-type", "FULL_OBJECT"},
+                    }}},
+
+      {"GET https://my-bucket.s3.us-east-1.amazonaws.com/tmp:1/key_read3",
        HttpResponse{
-           200, absl::Cord(),
-           HeaderMap{{"etag", "\"900150983cd24fb0d6963f7d28e17f73\""},
-                     {"x-amz-checksum-sha256",
-                      "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU="}}}},
+           200, absl::Cord("abcd"),
+           HeaderMap{
+               {"etag", "\"900150983cd24fb0d6963f7d28e17f72\""},
+               {"x-amz-checksum-sha256",
+                "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU="},  // mismatch
+               {"x-amz-checksum-type", "COMPOSITE"},  // not FULL_OBJECT
+           }}},
+
+      {"GET https://my-bucket.s3.us-east-1.amazonaws.com/tmp:1/empty_read",
+       HttpResponse{200, absl::Cord(),
+                    HeaderMap{
+                        {"etag", "\"900150983cd24fb0d6963f7d28e17f73\""},
+                        {"x-amz-checksum-sha256",
+                         "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU="},
+                    }}},
 
       {"GET https://my-bucket.s3.us-east-1.amazonaws.com/tmp:1/sha_mismatch",
-       HttpResponse{
-           200, absl::Cord("xyz"),
-           HeaderMap{{"etag", "\"900150983cd24fb0d6963f7d28e17f73\""},
-                     {"x-amz-checksum-sha256",
-                      "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU="}}}},
+       HttpResponse{200, absl::Cord("xyz"),
+                    HeaderMap{{"etag", "\"900150983cd24fb0d6963f7d28e17f73\""},
+                              {"x-amz-checksum-sha256",
+                               "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU="},
+                              {"x-amz-checksum-type", "FULL_OBJECT"}}}},
 
       {"PUT https://my-bucket.s3.us-east-1.amazonaws.com/tmp:1/key_write",
        HttpResponse{
@@ -197,7 +217,19 @@ TEST(S3KeyValueStoreTest, SimpleMock_VirtualHost) {
           .result());
 
   EXPECT_THAT(
-      kvstore::Read(store, "key_read").result(),
+      kvstore::Read(store, "key_read1").result(),
+      MatchesKvsReadResult(absl::Cord("abcd"),
+                           StorageGeneration::FromString(
+                               "\"900150983cd24fb0d6963f7d28e17f72\"")));
+
+  EXPECT_THAT(
+      kvstore::Read(store, "key_read2").result(),
+      MatchesKvsReadResult(absl::Cord("abcd"),
+                           StorageGeneration::FromString(
+                               "\"900150983cd24fb0d6963f7d28e17f72\"")));
+
+  EXPECT_THAT(
+      kvstore::Read(store, "key_read3").result(),
       MatchesKvsReadResult(absl::Cord("abcd"),
                            StorageGeneration::FromString(
                                "\"900150983cd24fb0d6963f7d28e17f72\"")));