Draft for testing: Supply Content-MD5 to S3 uploads

This is just to kick off the s3 third party test suite.

Author: bcully

modules/repository-s3/src/main/java/org/elasticsearch/repositories/s3/S3BlobContainer.java

@@ -69,8 +69,11 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Base64;
 import java.util.Date;
 import java.util.Iterator;
 import java.util.List;
@@ -189,13 +192,18 @@ private void flushBuffer(boolean lastPart) throws IOException {
                     }
                 }
                 assert lastPart == false || successful : "must only write last part if successful";
+                final var byteStream = buffer.bytes().streamInput();
+                byteStream.mark(0);
+                final var md5 = md5DigestOfInputStream(byteStream, Long.MAX_VALUE);
+                byteStream.reset();
                 final UploadPartRequest uploadRequest = createPartUploadRequest(
                     purpose,
-                    buffer.bytes().streamInput(),
+                    byteStream,
                     uploadId.get(),
                     parts.size() + 1,
                     absoluteBlobKey,
                     buffer.size(),
+                    md5,
                     lastPart
                 );
                 final UploadPartResult uploadResponse;
@@ -260,6 +268,7 @@ private UploadPartRequest createPartUploadRequest(
         int number,
         String blobName,
         long size,
+        byte[] partMd5,
         boolean lastPart
     ) {
         final UploadPartRequest uploadRequest = new UploadPartRequest();
@@ -270,6 +279,7 @@ private UploadPartRequest createPartUploadRequest(
         uploadRequest.setInputStream(stream);
         S3BlobStore.configureRequestForMetrics(uploadRequest, blobStore, Operation.PUT_MULTIPART_OBJECT, purpose);
         uploadRequest.setPartSize(size);
+        uploadRequest.setMd5Digest(Base64.getEncoder().encodeToString(partMd5));
         uploadRequest.setLastPart(lastPart);
         return uploadRequest;
     }
@@ -457,8 +467,18 @@ void executeSingleUpload(
         if (blobSize > s3BlobStore.bufferSizeInBytes()) {
             throw new IllegalArgumentException("Upload request size [" + blobSize + "] can't be larger than buffer size");
         }
+        // required to reset the stream for MD5 calculation
+        if (input.markSupported() == false) {
+            throw new IllegalArgumentException("input stream mark not supported");
+        }
 
         final ObjectMetadata md = new ObjectMetadata();
+
+        input.mark(0);
+        final byte[] md5 = md5DigestOfInputStream(input, Long.MAX_VALUE);
+        input.reset();
+        md.setContentMD5(Base64.getEncoder().encodeToString(md5));
+
         md.setContentLength(blobSize);
         if (s3BlobStore.serverSideEncryption()) {
             md.setSSEAlgorithm(ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION);
@@ -485,7 +505,6 @@ void executeMultipartUpload(
         final InputStream input,
         final long blobSize
     ) throws IOException {
-
         ensureMultiPartUploadSize(blobSize);
         final long partSize = s3BlobStore.bufferSizeInBytes();
         final Tuple<Long, Long> multiparts = numberOfMultiparts(blobSize, partSize);
@@ -498,6 +517,19 @@ void executeMultipartUpload(
         final long lastPartSize = multiparts.v2();
         assert blobSize == (((nbParts - 1) * partSize) + lastPartSize) : "blobSize does not match multipart sizes";
 
+        // required to reset the stream for MD5 calculation
+        if (input.markSupported() == false) {
+            throw new IllegalArgumentException("input stream mark not supported");
+        }
+
+        final byte[][] md5s = new byte[nbParts][];
+        input.mark(0);
+        for (int i = 0; i < nbParts - 1; i++) {
+            md5s[i] = md5DigestOfInputStream(input, partSize);
+        }
+        md5s[nbParts - 1] = md5DigestOfInputStream(input, lastPartSize);
+        input.reset();
+
         final SetOnce<String> uploadId = new SetOnce<>();
         final String bucketName = s3BlobStore.bucket();
         boolean success = false;
@@ -525,6 +557,7 @@ void executeMultipartUpload(
                     i,
                     blobName,
                     lastPart ? lastPartSize : partSize,
+                    md5s[i - 1],
                     lastPart
                 );
                 bytesCount += uploadRequest.getPartSize();
@@ -564,6 +597,28 @@ void executeMultipartUpload(
         }
     }
 
+    // Calculate the MD5 of up to remaining bytes of the given InputStream
+    private byte[] md5DigestOfInputStream(final InputStream inputStream, long remaining) throws IOException {
+        try {
+            final MessageDigest md5 = MessageDigest.getInstance("MD5");
+            // update in chunks to bound memory usage while amortizing read cost
+            byte[] buffer = new byte[65536];
+            int bytesRead;
+            do {
+                final int toRead = (int) Math.min(remaining, buffer.length);
+                bytesRead = inputStream.read(buffer, 0, toRead);
+                if (bytesRead > 0) {
+                    md5.update(buffer, 0, bytesRead);
+                    remaining -= bytesRead;
+                }
+            } while (bytesRead > 0);
+
+            return md5.digest();
+        } catch (NoSuchAlgorithmException e) {
+            throw new IOException(e);
+        }
+    }
+
     // non-static, package private for testing
     void ensureMultiPartUploadSize(final long blobSize) {
         if (blobSize > MAX_FILE_SIZE_USING_MULTIPART.getBytes()) {

modules/repository-s3/src/test/java/org/elasticsearch/repositories/s3/S3BlobStoreContainerTests.java

@@ -213,7 +213,7 @@ public void testExecuteMultipartUpload() throws IOException {
         final ArgumentCaptor<CompleteMultipartUploadRequest> compArgCaptor = ArgumentCaptor.forClass(CompleteMultipartUploadRequest.class);
         when(client.completeMultipartUpload(compArgCaptor.capture())).thenReturn(new CompleteMultipartUploadResult());
 
-        final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[0]);
+        final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[1]);
         final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);
         blobContainer.executeMultipartUpload(randomPurpose(), blobStore, blobName, inputStream, blobSize);