diff --git a/lib/lib-storage/src/Upload.spec.ts b/lib/lib-storage/src/Upload.spec.ts index 0d36c71250d2..383549a58edf 100644 --- a/lib/lib-storage/src/Upload.spec.ts +++ b/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(); + }); + }); }); diff --git a/lib/lib-storage/src/Upload.ts b/lib/lib-storage/src/Upload.ts index 06de3cccdaef..99098f80c8d8 100644 --- a/lib/lib-storage/src/Upload.ts +++ b/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 { @@ -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.`); diff --git a/lib/lib-storage/src/lib-storage.e2e.spec.ts b/lib/lib-storage/src/lib-storage.e2e.spec.ts index d3c6fb089d51..e615e615f49a 100644 --- a/lib/lib-storage/src/lib-storage.e2e.spec.ts +++ b/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\)\./); + }); }); } );