Fixed an issue that could cause checksum mismatch errors in S3 uploads. (#5836)

* Fixed an issue that could cause checksum mismatch errors in S3 uploads.

This error is sometimes encountered when a customer uses:
1. The AsyncS3Client
2. A ChecksumAlgorithm of SHA1 or SHA256 (instead of the default CRC32)
3. Parallel uploads

The root cause was the SDK using thread locals to cache the SHA1 or SHA256 message digest implementations. This meant that if a single event loop thread was processing multiple requests, those two requests would use the same digest implementation to calculate the checksum.

* Fixed an issue where SignerUtils was using the internal class of another module.

* Update clone() failure message. Move to LinkedBlockingDeque for DigestAlgorithm cache.

* Added direct testing of the DigestAlgorithm cache.

* Fixed a transient issue that could occur when the first attempt fully sends the payload, and then a retry happens with checksums enabled.

Author: millems

.changes/next-release/bugfix-AmazonS3-6cefd9c.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "Amazon S3",
+    "contributor": "",
+    "description": "Fixed an issue that could cause checksum mismatch errors when performing parallel uploads with the async S3 client and the SHA1 or SHA256 checksum algorithms selected."
+}

core/checksums/pom.xml

@@ -63,6 +63,11 @@
             <artifactId>junit-jupiter</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

core/checksums/src/main/java/software/amazon/awssdk/checksums/SdkChecksum.java

@@ -21,9 +21,8 @@
 import software.amazon.awssdk.checksums.internal.Crc32Checksum;
 import software.amazon.awssdk.checksums.internal.Crc64NvmeChecksum;
 import software.amazon.awssdk.checksums.internal.CrcChecksumProvider;
-import software.amazon.awssdk.checksums.internal.Md5Checksum;
-import software.amazon.awssdk.checksums.internal.Sha1Checksum;
-import software.amazon.awssdk.checksums.internal.Sha256Checksum;
+import software.amazon.awssdk.checksums.internal.DigestAlgorithm;
+import software.amazon.awssdk.checksums.internal.DigestAlgorithmChecksum;
 import software.amazon.awssdk.checksums.spi.ChecksumAlgorithm;
 
 /**
@@ -43,11 +42,11 @@ static SdkChecksum forAlgorithm(ChecksumAlgorithm algorithm) {
             case "CRC32":
                 return new Crc32Checksum();
             case "SHA1":
-                return new Sha1Checksum();
+                return new DigestAlgorithmChecksum(DigestAlgorithm.SHA1);
             case "SHA256":
-                return new Sha256Checksum();
+                return new DigestAlgorithmChecksum(DigestAlgorithm.SHA256);
             case "MD5":
-                return new Md5Checksum();
+                return new DigestAlgorithmChecksum(DigestAlgorithm.MD5);
             case "CRC64NVME":
                 return new Crc64NvmeChecksum();
             default:

core/checksums/src/main/java/software/amazon/awssdk/checksums/internal/DigestAlgorithm.java

@@ -17,51 +17,115 @@
 
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.Deque;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.function.Supplier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
+import software.amazon.awssdk.utils.SdkAutoCloseable;
 
 @SdkInternalApi
 public enum DigestAlgorithm {
-
     SHA1("SHA-1"),
+
     MD5("MD5"),
     SHA256("SHA-256")
     ;
 
+    private static final Supplier<MessageDigest> CLOSED_DIGEST = () -> {
+        throw new IllegalStateException("This message digest is closed.");
+    };
+
+    private static final int MAX_CACHED_DIGESTS = 10_000;
     private final String algorithmName;
-    private final DigestThreadLocal digestReference;
+    private final Deque<MessageDigest> digestCache = new LinkedBlockingDeque<>(MAX_CACHED_DIGESTS); // LIFO
 
     DigestAlgorithm(String algorithmName) {
         this.algorithmName = algorithmName;
-        digestReference = new DigestThreadLocal(algorithmName);
     }
 
     public String getAlgorithmName() {
         return algorithmName;
     }
 
     /**
-     * Returns the thread local reference for the {@link MessageDigest} algorithm
+     * Returns a {@link CloseableMessageDigest} to use for this algorithm.
      */
-    public MessageDigest getDigest() {
-        MessageDigest digest = digestReference.get();
-        digest.reset();
-        return digest;
+    public CloseableMessageDigest getDigest() {
+        MessageDigest digest = digestCache.pollFirst();
+        if (digest != null) {
+            digest.reset();
+            return new CloseableMessageDigest(digest);
+        }
+        return new CloseableMessageDigest(newDigest());
     }
 
-    private static class DigestThreadLocal extends ThreadLocal<MessageDigest> {
-        private final String algorithmName;
+    private MessageDigest newDigest() {
+        try {
+            return MessageDigest.getInstance(algorithmName);
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException("Unable to fetch message digest instance for Algorithm "
+                                       + algorithmName + ": " + e.getMessage(), e);
+        }
+    }
+
+    @SdkTestInternalApi
+    static void clearCaches() {
+        for (DigestAlgorithm value : values()) {
+            value.digestCache.clear();
+        }
+    }
+
+    public final class CloseableMessageDigest implements SdkAutoCloseable, Cloneable {
+
+        private Supplier<MessageDigest> digest;
+        private byte[] messageDigest;
+
+        private CloseableMessageDigest(MessageDigest digest) {
+            this.digest = () -> digest;
+        }
+
+        /**
+         * Retrieve the message digest instance.
+         */
+        public MessageDigest messageDigest() {
+            return digest.get();
+        }
+
+        /**
+         * Retrieve the message digest bytes. This will close the message digest when invoked. This is because the underlying
+         * message digest is reset on read, and we'd rather fail future interactions with the digest than act on the wrong data.
+         */
+        public byte[] digest() {
+            if (messageDigest != null) {
+                return messageDigest;
+            }
+            messageDigest = messageDigest().digest();
+            close();
+            return messageDigest;
+        }
+
+        /**
+         * Release this message digest back to the cache. Once released, you must not use the digest anymore.
+         */
+        @Override
+        public void close() {
+            if (digest == CLOSED_DIGEST) {
+                return;
+            }
+
+            // Drop this digest is the cache is full.
+            digestCache.offerFirst(digest.get());
 
-        DigestThreadLocal(String algorithmName) {
-            this.algorithmName = algorithmName;
+            digest = CLOSED_DIGEST;
         }
 
         @Override
-        protected MessageDigest initialValue() {
+        public CloseableMessageDigest clone() {
             try {
-                return MessageDigest.getInstance(algorithmName);
-            } catch (NoSuchAlgorithmException e) {
-                throw new RuntimeException("Unable to fetch message digest instance for Algorithm "
-                                           + algorithmName + ": " + e.getMessage(), e);
+                return new CloseableMessageDigest((MessageDigest) digest.get().clone());
+            } catch (CloneNotSupportedException e) {
+                throw new IllegalStateException("Clone was not supported by this digest type.", e);
             }
         }
     }

core/checksums/src/main/java/software/amazon/awssdk/checksums/internal/Md5Checksum.java renamed to core/checksums/src/main/java/software/amazon/awssdk/checksums/internal/DigestAlgorithmChecksum.java

@@ -15,32 +15,39 @@
 
 package software.amazon.awssdk.checksums.internal;
 
-import java.security.MessageDigest;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.checksums.SdkChecksum;
+import software.amazon.awssdk.checksums.internal.DigestAlgorithm.CloseableMessageDigest;
 
 /**
- * Implementation of {@link SdkChecksum} to calculate an MD5 checksum.
+ * An implementation of {@link SdkChecksum} that uses a {@link DigestAlgorithm}.
  */
 @SdkInternalApi
-public class Md5Checksum implements SdkChecksum {
+public class DigestAlgorithmChecksum implements SdkChecksum {
 
-    private MessageDigest digest;
+    private final DigestAlgorithm algorithm;
 
-    private MessageDigest digestLastMarked;
+    private CloseableMessageDigest digest;
 
-    public Md5Checksum() {
-        this.digest = getDigest();
+    private CloseableMessageDigest digestLastMarked;
+
+    public DigestAlgorithmChecksum(DigestAlgorithm algorithm) {
+        this.algorithm = algorithm;
+        this.digest = newDigest();
+    }
+
+    private CloseableMessageDigest newDigest() {
+        return algorithm.getDigest();
     }
 
     @Override
     public void update(int b) {
-        digest.update((byte) b);
+        digest.messageDigest().update((byte) b);
     }
 
     @Override
     public void update(byte[] b, int off, int len) {
-        digest.update(b, off, len);
+        digest.messageDigest().update(b, off, len);
     }
 
     @Override
@@ -50,15 +57,12 @@ public long getValue() {
 
     @Override
     public void reset() {
-        digest = (digestLastMarked == null)
-                 // This is necessary so that should there be a reset without a
-                 // preceding mark, the MD5 would still be computed correctly.
-                 ? getDigest()
-                 : cloneFrom(digestLastMarked);
-    }
-
-    private MessageDigest getDigest() {
-        return DigestAlgorithm.MD5.getDigest();
+        digest.close();
+        if (digestLastMarked == null) {
+            digest = newDigest();
+        } else {
+            digest = digestLastMarked;
+        }
     }
 
     @Override
@@ -68,14 +72,6 @@ public byte[] getChecksumBytes() {
 
     @Override
     public void mark(int readLimit) {
-        digestLastMarked = cloneFrom(digest);
-    }
-
-    private MessageDigest cloneFrom(MessageDigest from) {
-        try {
-            return (MessageDigest) from.clone();
-        } catch (CloneNotSupportedException e) { // should never occur
-            throw new IllegalStateException("unexpected", e);
-        }
+        digestLastMarked = digest.clone();
     }
 }