feat(lib-storage): add validation for partCount and contentLength for multipart upload

Author: smilkuri

lib/lib-storage/src/Upload.spec.ts

@@ -792,4 +792,91 @@ describe(Upload.name, () => {
       new Error("@aws-sdk/lib-storage: this instance of Upload has already executed .done(). Create a new instance.")
     );
   });
+
+  describe("Upload Part and parts count validation", () => {
+    const MOCK_PART_SIZE = 1024 * 1024 * 5; // 5MB
+
+    it("should throw error when uploaded parts count doesn't match expected parts count", async () => {
+      const largeBuffer = Buffer.from("#".repeat(MOCK_PART_SIZE * 2 + 100));
+      const upload = new Upload({
+        params: { ...params, Body: largeBuffer },
+        client: new S3({}),
+      });
+
+      (upload as any).__doConcurrentUpload = vi.fn().mockResolvedValue(undefined);
+
+      (upload as any).uploadedParts = [{ PartNumber: 1, ETag: "etag1" }];
+      (upload as any).isMultiPart = true;
+
+      await expect(upload.done()).rejects.toThrow("Expected 3 number of parts but uploaded only 1 part");
+    });
+
+    it("should throw error when part size doesn't match expected size except for laast part", () => {
+      const upload = new Upload({
+        params,
+        client: new S3({}),
+      });
+
+      const invalidPart = {
+        partNumber: 1,
+        data: Buffer.from("small"),
+        lastPart: false,
+      };
+
+      expect(() => {
+        (upload as any).__validateUploadPart(invalidPart, MOCK_PART_SIZE);
+      }).toThrow(`The Part size for part number 1, size 5 does not match expected size ${MOCK_PART_SIZE}`);
+    });
+
+    it("should allow smaller size for last part", () => {
+      const upload = new Upload({
+        params,
+        client: new S3({}),
+      });
+
+      const lastPart = {
+        partNumber: 2,
+        data: Buffer.from("small"),
+        lastPart: true,
+      };
+
+      expect(() => {
+        (upload as any).__validateUploadPart(lastPart, MOCK_PART_SIZE);
+      }).not.toThrow();
+    });
+
+    it("should throw error when part has zero content length", () => {
+      const upload = new Upload({
+        params,
+        client: new S3({}),
+      });
+
+      const emptyPart = {
+        partNumber: 1,
+        data: Buffer.from(""),
+        lastPart: false,
+      };
+
+      expect(() => {
+        (upload as any).__validateUploadPart(emptyPart, MOCK_PART_SIZE);
+      }).toThrow("Content length is missing on the data for part number 1");
+    });
+
+    it("should skip validation for single-part uploads", () => {
+      const upload = new Upload({
+        params,
+        client: new S3({}),
+      });
+
+      const singlePart = {
+        partNumber: 1,
+        data: Buffer.from("small"),
+        lastPart: true,
+      };
+
+      expect(() => {
+        (upload as any).__validateUploadPart(singlePart, MOCK_PART_SIZE);
+      }).not.toThrow();
+    });
+  });
 });

lib/lib-storage/src/Upload.ts

@@ -47,7 +47,7 @@ export class Upload extends EventEmitter {
 
   // Defaults.
   private readonly queueSize: number = 4;
-  private readonly partSize = Upload.MIN_PART_SIZE;
+  private readonly partSize: number;
   private readonly leavePartsOnError: boolean = false;
   private readonly tags: Tag[] = [];
 
@@ -66,6 +66,7 @@ export class Upload extends EventEmitter {
 
   private uploadedParts: CompletedPart[] = [];
   private uploadEnqueuedPartsCount = 0;
+  private expectedPartsCount = 0;
   /**
    * Last UploadId if the upload was done with MultipartUpload and not PutObject.
    */
@@ -81,19 +82,21 @@ export class Upload extends EventEmitter {
 
     // set defaults from options.
     this.queueSize = options.queueSize || this.queueSize;
-    this.partSize = options.partSize || this.partSize;
     this.leavePartsOnError = options.leavePartsOnError || this.leavePartsOnError;
     this.tags = options.tags || this.tags;
 
     this.client = options.client;
     this.params = options.params;
 
-    this.__validateInput();
-
     // set progress defaults
     this.totalBytes = byteLength(this.params.Body);
     this.bytesUploadedSoFar = 0;
     this.abortController = options.abortController ?? new AbortController();
+
+    this.partSize = this.__calculatePartSize(this.totalBytes ?? 0, Upload.MIN_PART_SIZE);
+    this.expectedPartsCount = this.totalBytes ? Math.ceil(this.totalBytes / this.partSize) : 1;
+
+    this.__validateInput();
   }
 
   async abort(): Promise<void> {
@@ -234,6 +237,8 @@ export class Upload extends EventEmitter {
         return;
       }
 
+      this.__validateUploadPart(dataPart, this.partSize);
+
       // Use put instead of multipart for one chunk uploads.
       if (dataPart.partNumber === 1 && dataPart.lastPart) {
         return await this.__uploadUsingPut(dataPart);
@@ -364,6 +369,12 @@ export class Upload extends EventEmitter {
 
     let result;
     if (this.isMultiPart) {
+      if (this.totalBytes && this.uploadedParts.length !== this.expectedPartsCount) {
+        throw new Error(
+          `Expected ${this.expectedPartsCount} number of parts but uploaded only ${this.uploadedParts.length} parts`
+        );
+      }
+
       this.uploadedParts.sort((a, b) => a.PartNumber! - b.PartNumber!);
 
       const uploadCompleteParams = {
@@ -427,6 +438,31 @@ export class Upload extends EventEmitter {
     });
   }
 
+  private __calculatePartSize(targetPartSizeBytes: number, minPartSize: number): number {
+    const calculatedPartSize = Math.floor(targetPartSizeBytes / this.MAX_PARTS);
+    return Math.max(minPartSize, calculatedPartSize);
+  }
+
+  private __validateUploadPart(dataPart: RawDataPart, partSize: number): void {
+    const actualPartSize = byteLength(dataPart.data) || 0;
+
+    // Skip validation for single-part uploads (PUT operations)
+    if (dataPart.partNumber === 1 && dataPart.lastPart) {
+      return;
+    }
+
+    if (actualPartSize === 0) {
+      throw new Error(`Content length is missing on the data for part number ${dataPart.partNumber}`);
+    }
+
+    // Validate part size (last part may be smaller)
+    if (!dataPart.lastPart && actualPartSize !== partSize) {
+      throw new Error(
+        `The Part size for part number ${dataPart.partNumber}, size ${actualPartSize} does not match expected size ${partSize}`
+      );
+    }
+  }
+
   private __validateInput(): void {
     if (!this.params) {
       throw new Error(`InputError: Upload requires params to be passed to upload.`);