fix(lib-storage): add validation for partCount and contentLength for multipart upload (#7363)

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 part(s) but uploaded 1 part(s).");
+    });
+
+    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 byte 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(`The byte size for part number 1, size 0 does not match expected size ${MOCK_PART_SIZE}`);
+    });
+
+    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?: number;
   /**
    * 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 = Math.max(Upload.MIN_PART_SIZE, Math.floor((this.totalBytes || 0) / this.MAX_PARTS));
+    this.expectedPartsCount = this.totalBytes !== undefined ? Math.ceil(this.totalBytes / this.partSize) : undefined;
+
+    this.__validateInput();
   }
 
   async abort(): Promise<void> {
@@ -282,6 +285,8 @@ export class Upload extends EventEmitter {
 
       this.uploadEnqueuedPartsCount += 1;
 
+      this.__validateUploadPart(dataPart);
+
       const partResult = await this.client.send(
         new UploadPartCommand({
           ...this.params,
@@ -364,6 +369,11 @@ export class Upload extends EventEmitter {
 
     let result;
     if (this.isMultiPart) {
+      const { expectedPartsCount, uploadedParts } = this;
+      if (expectedPartsCount !== undefined && uploadedParts.length !== expectedPartsCount) {
+        throw new Error(`Expected ${expectedPartsCount} part(s) but uploaded ${uploadedParts.length} part(s).`);
+      }
+
       this.uploadedParts.sort((a, b) => a.PartNumber! - b.PartNumber!);
 
       const uploadCompleteParams = {
@@ -427,6 +437,28 @@ export class Upload extends EventEmitter {
     });
   }
 
+  private __validateUploadPart(dataPart: RawDataPart): void {
+    const actualPartSize = byteLength(dataPart.data);
+
+    if (actualPartSize === undefined) {
+      throw new Error(
+        `A dataPart was generated without a measurable data chunk size for part number ${dataPart.partNumber}`
+      );
+    }
+
+    // Skip validation for single-part uploads (PUT operations)
+    if (dataPart.partNumber === 1 && dataPart.lastPart) {
+      return;
+    }
+
+    // Validate part size (last part may be smaller)
+    if (!dataPart.lastPart && actualPartSize !== this.partSize) {
+      throw new Error(
+        `The byte size for part number ${dataPart.partNumber}, size ${actualPartSize} does not match expected size ${this.partSize}`
+      );
+    }
+  }
+
   private __validateInput(): void {
     if (!this.params) {
       throw new Error(`InputError: Upload requires params to be passed to upload.`);

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

@@ -139,6 +139,43 @@ describe("@aws-sdk/lib-storage", () => {
             "S3Client AbortMultipartUploadCommand 204",
           ]);
         });
+
+        it("should validate part size constraints", () => {
+          const upload = new Upload({
+            client,
+            params: {
+              Bucket,
+              Key: `validation-test-${Date.now()}`,
+              Body: Buffer.alloc(1024 * 1024 * 10),
+            },
+          });
+
+          const invalidPart = {
+            partNumber: 2,
+            data: Buffer.alloc(1024 * 1024 * 3), // 3MB - too small for non-final part
+            lastPart: false,
+          };
+
+          expect(() => {
+            (upload as any).__validateUploadPart(invalidPart);
+          }).toThrow(/The byte size for part number 2, size \d+ does not match expected size \d+/);
+        });
+
+        it("should validate part count constraints", async () => {
+          const upload = new Upload({
+            client,
+            params: {
+              Bucket,
+              Key: `validation-test-${Date.now()}`,
+              Body: Buffer.alloc(1024 * 1024 * 10),
+            },
+          });
+
+          (upload as any).uploadedParts = [{ PartNumber: 1, ETag: "etag1" }];
+          (upload as any).isMultiPart = true;
+
+          await expect(upload.done()).rejects.toThrow(/Expected \d+ part\(s\) but uploaded \d+ part\(s\)\./);
+        });
       });
     }
   );