Fix handling of empty S3 objects in S3ChecksumValidatingInputStream

Author: RanVaknin

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/checksums/S3ChecksumValidatingInputStream.java

@@ -58,6 +58,18 @@ public S3ChecksumValidatingInputStream(InputStream in, SdkChecksum cksum, long s
      */
     @Override
     public int read() throws IOException {
+        if (strippedLength == 0 && lengthRead == 0) {
+            for (int i = 0; i < CHECKSUM_SIZE; i++) {
+                int b = inputStream.read();
+                if (b != -1) {
+                    streamChecksum[i] = (byte) b;
+                }
+            }
+            lengthRead = CHECKSUM_SIZE;
+            validateAndThrow();
+            return -1;
+        }
+
         int read = inputStream.read();
 
         if (read != -1 && lengthRead < strippedLength) {

services/s3/src/test/java/software/amazon/awssdk/services/s3/checksums/S3ChecksumValidatingInputStreamTest.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.services.s3.checksums;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.fail;
 
 import java.io.ByteArrayInputStream;
@@ -56,20 +57,51 @@ public static void populateData() {
 
     @Test
     public void validChecksumSucceeds() throws IOException {
-        InputStream validatingInputStream = newValidatingStream(testData);
+        InputStream validatingInputStream = newValidatingStream(testData, TEST_DATA_SIZE, CHECKSUM_SIZE);
         byte[] dataFromValidatingStream = IoUtils.toByteArray(validatingInputStream);
 
         assertArrayEquals(testDataWithoutChecksum, dataFromValidatingStream);
     }
 
+    @Test
+    public void emptyObjectSingleByteReadReturnsEOF() throws IOException {
+        Md5Checksum checksum = new Md5Checksum();
+        byte[] checksumBytes = checksum.getChecksumBytes();
+        byte[] emptyWithChecksum = new byte[CHECKSUM_SIZE];
+
+        for (int i = 0; i < CHECKSUM_SIZE; i++) {
+            emptyWithChecksum[i] = checksumBytes[i];
+        }
+
+        InputStream validatingInputStream = newValidatingStream(emptyWithChecksum, 0, CHECKSUM_SIZE);
+
+        assertEquals(-1, validatingInputStream.read());
+    }
+
+    @Test
+    public void emptyObjectByteArrayReadReturnsEOF() throws IOException {
+        Md5Checksum checksum = new Md5Checksum();
+        byte[] checksumBytes = checksum.getChecksumBytes();
+        byte[] emptyWithChecksum = new byte[CHECKSUM_SIZE];
+
+        for (int i = 0; i < CHECKSUM_SIZE; i++) {
+            emptyWithChecksum[i] = checksumBytes[i];
+        }
+
+        InputStream validatingInputStream = newValidatingStream(emptyWithChecksum, 0, CHECKSUM_SIZE);
+        byte[] buffer = new byte[1];
+
+        assertEquals(-1, validatingInputStream.read(buffer));
+    }
+
     @Test
     public void invalidChecksumFails() throws IOException {
         for (int i = 0; i < testData.length; i++) {
             // Make sure that corruption of any byte in the test data causes a checksum validation failure.
             byte[] corruptedChecksumData = Arrays.copyOf(testData, testData.length);
             corruptedChecksumData[i] = (byte) ~corruptedChecksumData[i];
 
-            InputStream validatingInputStream = newValidatingStream(corruptedChecksumData);
+            InputStream validatingInputStream = newValidatingStream(corruptedChecksumData, TEST_DATA_SIZE, CHECKSUM_SIZE);
 
             try {
                 IoUtils.toByteArray(validatingInputStream);
@@ -80,9 +112,9 @@ public void invalidChecksumFails() throws IOException {
         }
     }
 
-    private InputStream newValidatingStream(byte[] dataFromS3) {
+    private InputStream newValidatingStream(byte[] dataFromS3, int testDataSize, int checkSumSize) {
         return new S3ChecksumValidatingInputStream(new ByteArrayInputStream(dataFromS3),
                                                    new Md5Checksum(),
-                                                 TEST_DATA_SIZE + CHECKSUM_SIZE);
+                                                   testDataSize + checkSumSize);
     }
 }

services/s3/src/test/java/software/amazon/awssdk/services/s3/checksums/S3ChecksumValidatingPublisherTest.java

@@ -169,6 +169,28 @@ public void checksumValidationFailure_throwsSdkClientException_NotNPE() {
     assertFalse(s.hasCompleted());
   }
 
+  @Test
+  public void emptyObjectReturnsNoData() {
+    Md5Checksum checksum = new Md5Checksum();
+    byte[] checksumBytes = checksum.getChecksumBytes();
+    byte[] emptyWithChecksum = new byte[CHECKSUM_SIZE];
+    for (int i = 0; i < CHECKSUM_SIZE; i++) {
+      emptyWithChecksum[i] = checksumBytes[i];
+    }
+
+    final TestPublisher driver = new TestPublisher();
+    final TestSubscriber s = new TestSubscriber();
+    final S3ChecksumValidatingPublisher p = new S3ChecksumValidatingPublisher(driver, new Md5Checksum(), CHECKSUM_SIZE);
+    p.subscribe(s);
+
+    driver.doOnNext(ByteBuffer.wrap(emptyWithChecksum));
+    driver.doOnComplete();
+
+    assertArrayEquals(new byte[0], s.receivedData());
+    assertTrue(s.hasCompleted());
+    assertFalse(s.isOnErrorCalled());
+  }
+
   private class TestSubscriber implements Subscriber<ByteBuffer> {
     final List<ByteBuffer> received;
     boolean completed;