feat: Add validation around bytes received vs bytes expected (#2078)

sydney-munro · web-flow · commit 45d142a7ba45 · 2023-06-21T16:39:22.000-07:00
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/transfermanager/ChunkedDownloadCallable.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/transfermanager/ChunkedDownloadCallable.java
@@ -20,6 +20,7 @@
 import com.google.cloud.storage.BlobInfo;
 import com.google.cloud.storage.Storage;
 import com.google.cloud.storage.Storage.BlobSourceOption;
+import com.google.cloud.storage.StorageException;
 import com.google.common.io.ByteStreams;
 import java.nio.channels.FileChannel;
 import java.nio.file.Path;
@@ -64,6 +65,20 @@ public DownloadSegment call() {
       rc.limit(endPosition);
       wc.position(startPosition);
       bytesCopied = ByteStreams.copy(rc, wc);
+      long bytesExpected = endPosition - startPosition;
+      if (bytesCopied != bytesExpected) {
+        return DownloadSegment.newBuilder(originalBlob, TransferStatus.FAILED_TO_FINISH)
+            .setException(
+                new StorageException(
+                    0,
+                    "Unexpected end of stream, read "
+                        + bytesCopied
+                        + " expected "
+                        + bytesExpected
+                        + " from object "
+                        + originalBlob.getBlobId().toGsUtilUriWithGeneration()))
+            .build();
+      }
     } catch (Exception e) {
       if (bytesCopied == -1) {
         return DownloadSegment.newBuilder(originalBlob, TransferStatus.FAILED_TO_START)
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/transfermanager/DirectDownloadCallable.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/transfermanager/DirectDownloadCallable.java
@@ -21,6 +21,7 @@
 import com.google.cloud.storage.BlobInfo;
 import com.google.cloud.storage.Storage;
 import com.google.cloud.storage.Storage.BlobSourceOption;
+import com.google.cloud.storage.StorageException;
 import com.google.common.io.ByteStreams;
 import java.nio.channels.FileChannel;
 import java.nio.file.Path;
@@ -60,6 +61,21 @@ public DownloadResult call() {
               StandardOpenOption.CREATE,
               StandardOpenOption.TRUNCATE_EXISTING);
       bytesCopied = ByteStreams.copy(rc, wc);
+      if (originalBlob.getSize() != null) {
+        if (bytesCopied != originalBlob.getSize()) {
+          return DownloadResult.newBuilder(originalBlob, TransferStatus.FAILED_TO_FINISH)
+              .setException(
+                  new StorageException(
+                      0,
+                      "Unexpected end of stream, read "
+                          + bytesCopied
+                          + " expected "
+                          + originalBlob.getSize()
+                          + " from object "
+                          + originalBlob.getBlobId().toGsUtilUriWithGeneration()))
+              .build();
+        }
+      }
     } catch (Exception e) {
       if (bytesCopied == -1) {
         return DownloadResult.newBuilder(originalBlob, TransferStatus.FAILED_TO_START)
