added part count and content range validation for download

Author: Fred1155

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartDownloaderSubscriber.java

@@ -23,6 +23,7 @@
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
@@ -76,6 +77,16 @@ public class MultipartDownloaderSubscriber implements Subscriber<AsyncResponseTr
      */
     private volatile String eTag;
 
+    /**
+     * The size of each part of the object being downloaded.
+     */
+    private volatile Long partSize;
+
+    /**
+     * The total size of the object being downloaded.
+     */
+    private volatile Long totalContentLength;
+
     /**
      * The Subscription lock
      */
@@ -117,6 +128,7 @@ public void onNext(AsyncResponseTransformer<GetObjectResponse, GetObjectResponse
 
         synchronized (lock) {
             if (totalParts != null && nextPartToGet > totalParts) {
+                validatePartsCount(completedParts.get());
                 log.debug(() -> String.format("Completing multipart download after a total of %d parts downloaded.", totalParts));
                 subscription.cancel();
                 return;
@@ -162,10 +174,20 @@ private void requestMoreIfNeeded(GetObjectResponse response) {
             totalParts = partCount;
         }
 
+        String actualContentRange = response.contentRange();
+        if (actualContentRange != null && partSize == null) {
+            getRangeInfo(actualContentRange);
+            log.debug(() -> String.format("Part size of the object to download: " + partSize));
+            log.debug(() -> String.format("Total Content Length of the object to download: " + totalContentLength));
+        }
+
+        validateContentRange(totalComplete, actualContentRange);
+
         synchronized (lock) {
             if (totalParts != null && totalParts > 1 && totalComplete < totalParts) {
                 subscription.request(1);
             } else {
+                validatePartsCount(completedParts.get());
                 log.debug(() -> String.format("Completing multipart download after a total of %d parts downloaded.", totalParts));
                 subscription.cancel();
             }
@@ -198,4 +220,45 @@ private GetObjectRequest nextRequest(int nextPartToGet) {
             }
         });
     }
+
+    private void validatePartsCount(int currentGetCount) {
+        if (totalParts != null && currentGetCount != totalParts) {
+            String errorMessage = "PartsCount validation failed. Expected " + totalParts + ", downloaded"
+                                  + " " + currentGetCount + " parts.";
+            log.error(() -> errorMessage);
+            subscription.cancel();
+            SdkClientException exception = SdkClientException.create(errorMessage);
+            onError(exception);
+        }
+    }
+
+    private void validateContentRange(int partNumber, String contentRange) {
+        if (contentRange == null) {
+            return;
+        }
+
+        long expectedStart = (partNumber - 1) * partSize;
+        long expectedEnd = partNumber == totalParts ? totalContentLength - 1 : expectedStart + partSize - 1;
+
+        String expectedContentRange = String.format("bytes %d-%d/%d", expectedStart, expectedEnd, totalContentLength);
+
+        if (!expectedContentRange.equals(contentRange)) {
+            String errorMessage = String.format(
+                "Content-Range validation failed for part %d. Expected: %s, Actual: %s",
+                partNumber, expectedContentRange, contentRange);
+            log.error(() -> errorMessage);
+            onError(SdkClientException.create(errorMessage));
+        }
+    }
+
+    private void getRangeInfo(String contentRange) {
+        String rangeInfo = contentRange.substring(6);
+        String[] parts = rangeInfo.split("/");
+
+        this.totalContentLength = Long.parseLong(parts[1]);
+        String[] rangeParts = parts[0].split("-");
+        long startByte = Long.parseLong(rangeParts[0]);
+        long endByte = Long.parseLong(rangeParts[1]);
+        this.partSize = endByte - startByte + 1;
+    }
 }

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartDownloadTestUtil.java

@@ -70,6 +70,21 @@ public byte[] stubForPart(String testBucket, String testKey,int part, int totalP
             aResponse()
                 .withHeader("x-amz-mp-parts-count", totalPart + "")
                 .withHeader("ETag", eTag)
+                .withHeader("Content-Length", String.valueOf(body.length))
+                .withHeader("Content-Range", contentRange(part, totalPart, partSize))
+                .withBody(body)));
+        return body;
+    }
+
+    public byte[] stubForPartwithWrongContentRange(String testBucket, String testKey,int part, int totalPart, int partSize) {
+        byte[] body = new byte[partSize];
+        random.nextBytes(body);
+        stubFor(get(urlEqualTo(String.format("/%s/%s?partNumber=%d", testBucket, testKey, part))).willReturn(
+            aResponse()
+                .withHeader("x-amz-mp-parts-count", totalPart + "")
+                .withHeader("ETag", eTag)
+                .withHeader("Content-Length", String.valueOf(body.length))
+                .withHeader("Content-Range", contentRange(part, totalPart, partSize + 1))
                 .withBody(body)));
         return body;
     }
@@ -95,4 +110,16 @@ public byte[] stubForPartSuccess(int part, int totalPart, int partSize) {
                             .withBody(body)));
         return body;
     }
+
+    private String contentRange(int part, int totalPart, int partSize) {
+        long totalObjectSize = (long) totalPart * partSize;
+        long startByte = (long) (part - 1) * partSize;
+        long endByte = startByte + partSize - 1;
+
+        if (part == totalPart) {
+            endByte = totalObjectSize - 1;
+        }
+
+        return String.format("bytes %d-%d/%d", startByte, endByte, totalObjectSize);
+    }
 }

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartDownloaderSubscriberWiremockTest.java

@@ -160,6 +160,32 @@ <T> void errorOnThirdRequest_shouldCompleteExceptionallyOnlyPartsGreaterThanTwo(
         }
     }
 
+    @ParameterizedTest
+    @MethodSource("argumentsProvider")
+    <T> void wrongContentRangeOnSecondRequest_should(AsyncResponseTransformerTestSupplier<T> supplier,
+                                                     int amountOfPartToTest,
+                                                     int partSize) {
+        util.stubForPart(testBucket, testKey, 1, 3, partSize);
+        util.stubForPartwithWrongContentRange(testBucket, testKey, 2, 3, partSize);
+        util.stubForPart(testBucket, testKey, 3, 3, partSize);
+        //byte[] expectedBody = util.stubAllParts(testBucket, testKey, amountOfPartToTest, partSize);
+        AsyncResponseTransformer<GetObjectResponse, T> transformer = supplier.transformer();
+        AsyncResponseTransformer.SplitResult<GetObjectResponse, T> split = transformer.split(
+            SplittingTransformerConfiguration.builder()
+                                             .bufferSizeInBytes(1024 * 32L)
+                                             .build());
+        Subscriber<AsyncResponseTransformer<GetObjectResponse, GetObjectResponse>> subscriber = new MultipartDownloaderSubscriber(
+            s3AsyncClient,
+            GetObjectRequest.builder()
+                            .bucket(testBucket)
+                            .key(testKey)
+                            .build());
+
+        split.publisher().subscribe(subscriber);
+        T response = split.resultFuture().join();
+
+    }
+
     private static Stream<Arguments> argumentsProvider() {
         // amount of part, individual part size
         List<Pair<Integer, Integer>> partSizes = Arrays.asList(