fix: fallback to multipart, test coverage

Signed-off-by: ferhat elmas <elmas.ferhat@gmail.com>

Author: ferhatelmas

src/config.ts

@@ -66,6 +66,7 @@ type StorageConfigType = {
   storageS3InternalTracesEnabled?: boolean
   storageS3MaxSockets: number
   storageS3DisableChecksum: boolean
+  storageS3UploadQueueSize: number
   storageS3Bucket: string
   storageS3Endpoint?: string
   storageS3ForcePathStyle?: boolean
@@ -364,6 +365,8 @@ export function getConfig(options?: { reload?: boolean }): StorageConfigType {
       10
     ),
     storageS3DisableChecksum: getOptionalConfigFromEnv('STORAGE_S3_DISABLE_CHECKSUM') === 'true',
+    storageS3UploadQueueSize:
+      envNumber(getOptionalConfigFromEnv('STORAGE_S3_UPLOAD_QUEUE_SIZE')) ?? 2,
     storageS3InternalTracesEnabled:
       getOptionalConfigFromEnv('STORAGE_S3_ENABLED_METRICS') === 'true',
     storageS3Bucket: getOptionalConfigFromEnv('STORAGE_S3_BUCKET', 'GLOBAL_S3_BUCKET'),

src/storage/backend/s3/adapter.ts

@@ -35,13 +35,14 @@ import {
 } from './../adapter'
 
 const {
+  storageS3UploadQueueSize,
   tracingFeatures,
   storageS3MaxSockets,
   tracingEnabled,
   storageS3DisableChecksum,
 } = getConfig()
 
-const MAX_PUT_OBJECT_SIZE = 5 * 1024 * 1024 * 1024 // 5GB
+export const MAX_PUT_OBJECT_SIZE = 5 * 1024 * 1024 * 1024 // 5GB
 
 export interface S3ClientOptions {
   endpoint?: string
@@ -152,13 +153,17 @@ export class S3Backend implements StorageBackendAdapter {
       throw ERRORS.Aborted('Upload was aborted')
     }
 
-    if (typeof contentLength !== 'number') {
-      // If content length is unknown, use streaming multipart upload which does not require buffering the entire stream in memory.
-      return this.bufferedMultipartUpload(bucketName, key, version, body, contentType, cacheControl, signal)
-    }
-
-    if (contentLength > MAX_PUT_OBJECT_SIZE) {
-      throw ERRORS.EntityTooLarge(undefined)
+    if (typeof contentLength !== 'number' || contentLength > MAX_PUT_OBJECT_SIZE) {
+      // Use multipart when the length is unknown or exceeds S3's 5GB single-request limit.
+      return this.bufferedMultipartUpload(
+        bucketName,
+        key,
+        version,
+        body,
+        contentType,
+        cacheControl,
+        signal
+      )
     }
 
     // Use PutObject directly when content-length is known and within S3's single-object limit (5GB).
@@ -232,7 +237,7 @@ export class S3Backend implements StorageBackendAdapter {
 
     const upload = new Upload({
       client: this.client,
-      queueSize: 1,
+      queueSize: storageS3UploadQueueSize,
       params: {
         Bucket: bucketName,
         Key: withOptionalVersion(key, version),

src/storage/uploader.ts

@@ -411,6 +411,10 @@ export async function fileUploadFromRequest(
     }
 
     fileContentLength = getKnownRequestContentLength(request)
+    if (typeof fileContentLength === 'number' && fileContentLength > maxFileSize) {
+      throw ERRORS.EntityTooLarge()
+    }
+
     isTruncated = () => {
       // @todo more secure to get this from the stream or from s3 in the next step
       return typeof fileContentLength === 'number' && fileContentLength > maxFileSize

src/test/s3-adapter.test.ts

@@ -1,8 +1,11 @@
 'use strict'
 
-import { S3Client } from '@aws-sdk/client-s3'
+import { HeadObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3'
+import { Upload } from '@aws-sdk/lib-storage'
 import { Readable } from 'stream'
-import { S3Backend } from '../storage/backend/s3/adapter'
+import { getConfig } from '../config'
+import { ErrorCode, isStorageError } from '../internal/errors'
+import { MAX_PUT_OBJECT_SIZE, S3Backend } from '../storage/backend/s3/adapter'
 
 jest.mock('@aws-sdk/client-s3', () => {
   const originalModule = jest.requireActual('@aws-sdk/client-s3')
@@ -14,17 +17,76 @@ jest.mock('@aws-sdk/client-s3', () => {
   }
 })
 
+jest.mock('@aws-sdk/lib-storage', () => {
+  const originalModule = jest.requireActual('@aws-sdk/lib-storage')
+  return {
+    ...originalModule,
+    Upload: jest.fn(),
+  }
+})
+
+type MockUploadInstance = {
+  options: any
+  abort: jest.Mock
+  done: jest.Mock<Promise<any>, []>
+  on: jest.Mock
+  off: jest.Mock
+  emit: (event: string, payload: unknown) => void
+}
+
 describe('S3Backend', () => {
   let mockSend: jest.Mock
+  let mockUploadDone: jest.Mock<Promise<any>, [MockUploadInstance]>
+  let uploadInstances: MockUploadInstance[]
 
   beforeEach(() => {
     jest.clearAllMocks()
     mockSend = jest.fn()
+    mockUploadDone = jest.fn().mockResolvedValue({
+      ETag: '"multipart-etag"',
+      $metadata: {
+        httpStatusCode: 200,
+      },
+    })
+    uploadInstances = []
+
     ;(S3Client as jest.Mock).mockImplementation(() => ({
       send: mockSend,
     }))
+
+    ;(Upload as unknown as jest.Mock).mockImplementation((options) => {
+      const handlers = new Map<string, Set<(payload: unknown) => void>>()
+      const instance = {} as MockUploadInstance
+
+      instance.options = options
+      instance.abort = jest.fn()
+      instance.done = jest.fn(() => mockUploadDone(instance))
+      instance.on = jest.fn((event: string, handler: (payload: unknown) => void) => {
+        const eventHandlers = handlers.get(event) ?? new Set()
+        eventHandlers.add(handler)
+        handlers.set(event, eventHandlers)
+        return instance
+      })
+      instance.off = jest.fn((event: string, handler: (payload: unknown) => void) => {
+        handlers.get(event)?.delete(handler)
+        return instance
+      })
+      instance.emit = (event: string, payload: unknown) => {
+        handlers.get(event)?.forEach((handler) => handler(payload))
+      }
+
+      uploadInstances.push(instance)
+      return instance
+    })
   })
 
+  function createBackend() {
+    return new S3Backend({
+      region: 'us-east-1',
+      endpoint: 'http://localhost:9000',
+    })
+  }
+
   describe('getObject', () => {
     test('should return correct default MIME type when S3 returns no ContentType', async () => {
       mockSend.mockResolvedValue({
@@ -38,10 +100,7 @@ describe('S3Backend', () => {
         },
       })
 
-      const backend = new S3Backend({
-        region: 'us-east-1',
-        endpoint: 'http://localhost:9000',
-      })
+      const backend = createBackend()
 
       const result = await backend.getObject('test-bucket', 'test-key', undefined)
 
@@ -64,14 +123,151 @@ describe('S3Backend', () => {
         },
       })
 
-      const backend = new S3Backend({
-        region: 'us-east-1',
-        endpoint: 'http://localhost:9000',
-      })
+      const backend = createBackend()
 
       const result = await backend.getObject('test-bucket', 'test-key', undefined)
 
       expect(result.metadata.mimetype).toBe('image/png')
     })
   })
+
+  describe('uploadObject', () => {
+    test('uses PutObject for known-size uploads within the single-request limit', async () => {
+      mockSend.mockResolvedValue({
+        ETag: '"put-etag"',
+        $metadata: {
+          httpStatusCode: 200,
+        },
+      })
+
+      const backend = createBackend()
+      const result = await backend.uploadObject(
+        'test-bucket',
+        'test-key',
+        undefined,
+        Readable.from(['hello']),
+        'text/plain',
+        'max-age=60',
+        undefined,
+        5
+      )
+
+      expect(mockSend).toHaveBeenCalledTimes(1)
+      expect(mockSend.mock.calls[0][0]).toBeInstanceOf(PutObjectCommand)
+      expect(mockSend.mock.calls[0][0].input).toMatchObject({
+        Bucket: 'test-bucket',
+        Key: 'test-key',
+        ContentType: 'text/plain',
+        CacheControl: 'max-age=60',
+        ContentLength: 5,
+      })
+      expect(Upload).not.toHaveBeenCalled()
+      expect(result).toMatchObject({
+        httpStatusCode: 200,
+        cacheControl: 'max-age=60',
+        eTag: '"put-etag"',
+        mimetype: 'text/plain',
+        contentLength: 5,
+        size: 5,
+      })
+    })
+
+    test('falls back to multipart upload when content length exceeds the single-request limit', async () => {
+      const overLimit = MAX_PUT_OBJECT_SIZE + 1
+      const lastModified = new Date('2024-01-01T00:00:00.000Z')
+
+      mockUploadDone.mockImplementationOnce(async (instance) => {
+        instance.emit('httpUploadProgress', { loaded: 1 })
+        return {
+          ETag: '"multipart-etag"',
+          $metadata: {
+            httpStatusCode: 200,
+          },
+        }
+      })
+      mockSend.mockResolvedValueOnce({
+        CacheControl: 'max-age=60',
+        ContentType: 'text/plain',
+        ContentLength: overLimit,
+        ETag: '"head-etag"',
+        LastModified: lastModified,
+        $metadata: {
+          httpStatusCode: 200,
+        },
+      })
+
+      const backend = createBackend()
+      const result = await backend.uploadObject(
+        'test-bucket',
+        'test-key',
+        undefined,
+        Readable.from(['hello']),
+        'text/plain',
+        'max-age=60',
+        undefined,
+        overLimit
+      )
+
+      expect(Upload).toHaveBeenCalledTimes(1)
+      expect(uploadInstances[0].options.queueSize).toBe(getConfig().storageS3UploadQueueSize)
+      expect(mockSend).toHaveBeenCalledTimes(1)
+      expect(mockSend.mock.calls[0][0]).toBeInstanceOf(HeadObjectCommand)
+      expect(result).toMatchObject({
+        httpStatusCode: 200,
+        cacheControl: 'max-age=60',
+        eTag: '"head-etag"',
+        mimetype: 'text/plain',
+        contentLength: overLimit,
+        size: overLimit,
+        lastModified,
+      })
+    })
+
+    test('uses multipart upload when content length is unknown', async () => {
+      const backend = createBackend()
+      const result = await backend.uploadObject(
+        'test-bucket',
+        'test-key',
+        undefined,
+        Readable.from(['hello']),
+        'text/plain',
+        'max-age=60'
+      )
+
+      expect(Upload).toHaveBeenCalledTimes(1)
+      expect(uploadInstances[0].options.queueSize).toBe(getConfig().storageS3UploadQueueSize)
+      expect(mockSend).not.toHaveBeenCalled()
+      expect(result).toMatchObject({
+        httpStatusCode: 200,
+        cacheControl: 'max-age=60',
+        eTag: '"multipart-etag"',
+        mimetype: 'text/plain',
+        contentLength: 0,
+        size: 0,
+      })
+    })
+
+    test('maps PutObject abort errors to AbortedTerminate', async () => {
+      mockSend.mockRejectedValueOnce(Object.assign(new Error('aborted'), { name: 'AbortError' }))
+
+      const backend = createBackend()
+
+      try {
+        await backend.uploadObject(
+          'test-bucket',
+          'test-key',
+          undefined,
+          Readable.from(['hello']),
+          'text/plain',
+          'max-age=60',
+          undefined,
+          5
+        )
+        throw new Error('Expected uploadObject to throw')
+      } catch (error) {
+        expect(isStorageError(ErrorCode.AbortedTerminate, error)).toBe(true)
+        expect((error as Error).message).toBe('Upload was aborted')
+      }
+    })
+  })
 })