Merge pull request #639 from internxt/feat/fix-upload-file-streams

[_]: fix/upload-file-streams

Author: larryrider

src/commands/upload-file.ts

@@ -11,7 +11,6 @@ import { EncryptionVersion } from '@internxt/sdk/dist/drive/storage/types';
 import { ThumbnailService } from '../services/thumbnail.service';
 import { AuthService } from '../services/auth.service';
 import { UploadUtils } from '../utils/upload.utils';
-import { BufferStream } from '../utils/stream.utils';
 
 export default class UploadFile extends Command {
   static readonly args = {};
@@ -81,14 +80,11 @@ export default class UploadFile extends Command {
     progressBar?.start(100, 0);
 
     let fileId: string | undefined;
-    let thumbnailStream: BufferStream | undefined;
     const fileSize = stats.size ?? 0;
 
     if (fileSize > 0) {
       // Upload file to the Network
       const readStream = createReadStream(filePath);
-      const preparedStreams = UploadUtils.prepareUploadStreams(readStream, fileType);
-      thumbnailStream = preparedStreams.thumbnailStream;
 
       const progressCallback = (progress: number) => {
         progressBar?.update(progress * 100 * 0.99);
@@ -97,7 +93,7 @@ export default class UploadFile extends Command {
       const abortable = new AbortController();
 
       fileId = await networkFacade.uploadFile({
-        from: preparedStreams.fileStream,
+        from: readStream,
         size: fileSize,
         bucketId: bucket,
         progressCallback,
@@ -128,7 +124,7 @@ export default class UploadFile extends Command {
 
     const thumbnailTimer = CLIUtils.timer();
     await ThumbnailService.instance.tryUploadThumbnail({
-      bufferStream: thumbnailStream,
+      input: filePath,
       fileType,
       bucket,
       fileUuid: createdDriveFile.uuid,

src/services/network/upload/upload-file.service.ts

@@ -11,11 +11,11 @@ import { dirname, extname } from 'node:path';
 import { ErrorUtils } from '../../../utils/errors.utils';
 import { stat } from 'node:fs/promises';
 import { EncryptionVersion } from '@internxt/sdk/dist/drive/storage/types';
-import { BufferStream } from '../../../utils/stream.utils';
 import { DriveFileItem } from '../../../types/drive.types';
 import { CLIUtils } from '../../../utils/cli.utils';
 import { ThumbnailService } from '../../thumbnail.service';
 import { FormatUtils } from '../../../utils/format.utils';
+import { createReadStream } from 'node:fs';
 
 export class UploadFileService {
   static readonly instance = new UploadFileService();
@@ -84,7 +84,6 @@ export class UploadFileService {
         const fileType = extname(file.absolutePath).replaceAll('.', '');
 
         let fileId: string | undefined;
-        let thumbnailStream: BufferStream | undefined;
 
         const timings = {
           networkUpload: 0,
@@ -93,16 +92,12 @@ export class UploadFileService {
         };
 
         if (fileSize > 0) {
-          const { fileStream, bufferStream } = ThumbnailService.instance.createFileStreamWithBuffer({
-            path: file.absolutePath,
-            fileType,
-          });
+          const readStream = createReadStream(file.absolutePath);
 
           const uploadTimer = CLIUtils.timer();
-          thumbnailStream = bufferStream;
 
           fileId = await network.uploadFile({
-            from: fileStream,
+            from: readStream,
             size: fileSize,
             bucketId: bucket,
             progressCallback: () => {},
@@ -125,9 +120,9 @@ export class UploadFileService {
         timings.driveUpload = driveTimer.stop();
 
         const thumbnailTimer = CLIUtils.timer();
-        if (thumbnailStream && fileSize > 0) {
+        if (fileSize > 0) {
           await ThumbnailService.instance.tryUploadThumbnail({
-            bufferStream: thumbnailStream,
+            input: file.absolutePath,
             fileType,
             bucket,
             fileUuid: createdDriveFile.uuid,

src/services/thumbnail.service.ts

@@ -1,10 +1,8 @@
 import { Readable } from 'node:stream';
-import { createReadStream } from 'node:fs';
 import { DriveFileService } from './drive/drive-file.service';
 import { StorageTypes } from '@internxt/sdk/dist/drive';
 import { NetworkFacade } from './network/network-facade.service';
 import { ThumbnailConfig, ThumbnailUtils } from '../utils/thumbnail.utils';
-import { BufferStream } from '../utils/stream.utils';
 import { ErrorUtils } from '../utils/errors.utils';
 import { AsyncUtils } from '../utils/async.utils';
 
@@ -26,7 +24,7 @@ export class ThumbnailService {
   private static readonly MAX_THUMBNAIL_TIMEOUT = 30000;
 
   public uploadThumbnail = async (
-    fileContent: Buffer,
+    input: string | Buffer,
     fileType: string,
     userBucket: string,
     file_id: string,
@@ -35,7 +33,7 @@ export class ThumbnailService {
   ): Promise<StorageTypes.Thumbnail | undefined> => {
     let thumbnailBuffer: Buffer | undefined;
     if (ThumbnailUtils.isImageThumbnailable(fileType, fileSize)) {
-      thumbnailBuffer = await this.getThumbnailFromImageBuffer(fileContent);
+      thumbnailBuffer = await this.generateThumbnail(input);
     }
     if (thumbnailBuffer) {
       const size = thumbnailBuffer.length;
@@ -61,10 +59,10 @@ export class ThumbnailService {
     }
   };
 
-  private readonly getThumbnailFromImageBuffer = async (buffer: Buffer): Promise<Buffer | undefined> => {
+  private readonly generateThumbnail = async (input: string | Buffer): Promise<Buffer | undefined> => {
     const sharp = await getSharp();
     if (sharp) {
-      return sharp(buffer, { failOn: 'none' })
+      return sharp(input, { failOn: 'none' })
         .resize({
           height: ThumbnailConfig.MaxHeight,
           width: ThumbnailConfig.MaxWidth,
@@ -78,25 +76,24 @@ export class ThumbnailService {
   };
 
   public tryUploadThumbnail = async ({
-    bufferStream,
+    input,
     fileType,
     bucket,
     fileUuid,
     networkFacade,
     size,
   }: {
-    bufferStream?: BufferStream;
+    input?: string | Buffer;
     fileType: string;
     bucket: string;
     fileUuid: string;
     networkFacade: NetworkFacade;
     size: number;
   }) => {
     try {
-      const thumbnailBuffer = bufferStream?.getBuffer();
-      if (thumbnailBuffer && size > 0) {
+      if (input && size > 0) {
         await AsyncUtils.withTimeout(
-          ThumbnailService.instance.uploadThumbnail(thumbnailBuffer, fileType, bucket, fileUuid, networkFacade, size),
+          ThumbnailService.instance.uploadThumbnail(input, fileType, bucket, fileUuid, networkFacade, size),
           ThumbnailService.MAX_THUMBNAIL_TIMEOUT,
           'Thumbnail upload timeout',
         );
@@ -105,25 +102,4 @@ export class ThumbnailService {
       ErrorUtils.report(error);
     }
   };
-
-  public createFileStreamWithBuffer = ({
-    path,
-    fileType,
-  }: {
-    path: string;
-    fileType: string;
-  }): {
-    bufferStream?: BufferStream;
-    fileStream: Readable;
-  } => {
-    const readable: Readable = createReadStream(path);
-    if (ThumbnailUtils.isFileThumbnailable(fileType)) {
-      const bufferStream = new BufferStream();
-      return {
-        bufferStream,
-        fileStream: readable.pipe(bufferStream),
-      };
-    }
-    return { fileStream: readable };
-  };
 }

src/utils/stream.utils.ts

@@ -67,16 +67,15 @@ export class StreamUtils {
 }
 
 export class BufferStream extends Transform {
-  public buffer: Buffer | null;
+  private chunks: Buffer[];
 
   constructor(opts?: TransformOptions) {
     super(opts);
-    this.buffer = null;
+    this.chunks = [];
   }
 
   _transform(chunk: Buffer, _: BufferEncoding, callback: TransformCallback) {
-    const currentBuffer = this.buffer ?? Buffer.alloc(0);
-    this.buffer = Buffer.concat([currentBuffer, chunk]);
+    this.chunks.push(chunk);
     callback(null, chunk);
   }
 
@@ -85,10 +84,11 @@ export class BufferStream extends Transform {
   }
 
   reset() {
-    this.buffer = null;
+    this.chunks = [];
   }
 
-  getBuffer(): Buffer | null {
-    return this.buffer;
+  getBuffer(): Buffer | undefined {
+    if (this.chunks.length === 0) return undefined;
+    return Buffer.concat(this.chunks);
   }
 }

src/utils/thumbnail.utils.ts

@@ -17,32 +17,21 @@ const imageExtensions: FileExtensionMap = {
   raw: ['raw', 'cr2', 'nef', 'orf', 'sr2'],
   webp: ['webp'],
 };
-const pdfExtensions: FileExtensionMap = {
-  pdf: ['pdf'],
-};
 const thumbnailableImageExtension: Set<string> = new Set([
   ...imageExtensions['jpg'],
   ...imageExtensions['png'],
   ...imageExtensions['webp'],
   ...imageExtensions['gif'],
   ...imageExtensions['tiff'],
 ]);
-const thumbnailablePdfExtension: Set<string> = new Set(pdfExtensions['pdf']);
-const thumbnailableExtension: Set<string> = new Set(thumbnailableImageExtension);
 
 export class ThumbnailUtils {
-  static readonly MAX_IMAGE_THUMBNAILABLE_SIZE_IN_MB = 500 * 1024 * 1024;
-
-  static readonly isFileThumbnailable = (fileType: string) => {
-    return fileType.trim().length > 0 && thumbnailableExtension.has(fileType.trim().toLowerCase());
-  };
-
-  static readonly isPDFThumbnailable = (fileType: string) => {
-    return fileType.trim().length > 0 && thumbnailablePdfExtension.has(fileType.trim().toLowerCase());
-  };
+  static readonly MAX_IMAGE_THUMBNAILABLE_SIZE_IN_BYTES = 128 * 1024 * 1024;
 
   static readonly isImageThumbnailable = (fileType: string, size: number) => {
-    if (size > ThumbnailUtils.MAX_IMAGE_THUMBNAILABLE_SIZE_IN_MB) return false;
+    if (size <= 0 || size > ThumbnailUtils.MAX_IMAGE_THUMBNAILABLE_SIZE_IN_BYTES) {
+      return false;
+    }
     return fileType.trim().length > 0 && thumbnailableImageExtension.has(fileType.trim().toLowerCase());
   };
 }

src/utils/upload.utils.ts

@@ -25,12 +25,13 @@ export class UploadUtils {
   static readonly prepareUploadStreams = (
     readable: Readable,
     fileType: string,
+    size: number,
   ): {
     fileStream: Readable;
     thumbnailStream: BufferStream | undefined;
     isThumbnailable: boolean;
   } => {
-    const isThumbnailable = ThumbnailUtils.isFileThumbnailable(fileType);
+    const isThumbnailable = ThumbnailUtils.isImageThumbnailable(fileType, size);
     if (!isThumbnailable) {
       return { fileStream: readable, thumbnailStream: undefined, isThumbnailable };
     }

src/webdav/handlers/PUT.handler.ts

@@ -57,7 +57,7 @@ export class PUTRequestHandler implements WebDavMethodHandler {
     const { user } = await AuthService.instance.getAuthDetails();
     const fileType = resource.path.ext.replace('.', '');
 
-    const { fileStream, thumbnailStream } = UploadUtils.prepareUploadStreams(req, fileType);
+    const { fileStream, thumbnailStream } = UploadUtils.prepareUploadStreams(req, fileType, contentLength);
 
     const { networkFacade, bucket } = await CLIUtils.prepareNetwork(user);
 
@@ -144,7 +144,7 @@ export class PUTRequestHandler implements WebDavMethodHandler {
     const thumbnailTimer = CLIUtils.timer();
     await ThumbnailService.instance.tryUploadThumbnail({
       fileUuid: file.uuid,
-      bufferStream: thumbnailStream,
+      input: thumbnailStream?.getBuffer(),
       fileType,
       bucket,
       networkFacade,

test/services/network/upload/upload-file.service.test.ts

@@ -13,7 +13,6 @@ import {
   createProgressFixtures,
 } from './upload.service.helpers';
 import { newFileItem } from '../../../fixtures/drive.fixture';
-import { ThumbnailUtils } from '../../../../src/utils/thumbnail.utils';
 import { ThumbnailService } from '../../../../src/services/thumbnail.service';
 
 vi.mock('fs', () => ({
@@ -37,12 +36,7 @@ describe('Upload File Service', () => {
     vi.mocked(stat).mockResolvedValue(createMockStats(1024) as Awaited<ReturnType<typeof stat>>);
     vi.mocked(createReadStream).mockReturnValue(createMockReadStream() as ReturnType<typeof createReadStream>);
     vi.spyOn(ErrorUtils, 'isAlreadyExistsError').mockReturnValue(false);
-    vi.spyOn(ThumbnailUtils, 'isFileThumbnailable').mockReturnValue(false);
     vi.spyOn(ThumbnailService.instance, 'tryUploadThumbnail').mockResolvedValue(undefined);
-    vi.spyOn(ThumbnailService.instance, 'createFileStreamWithBuffer').mockReturnValue({
-      fileStream: createMockReadStream() as ReturnType<typeof createReadStream>,
-      bufferStream: undefined,
-    });
     vi.spyOn(DriveFileService.instance, 'createFile').mockResolvedValue(mockFile);
   });
 
@@ -300,14 +294,6 @@ describe('Upload File Service', () => {
     });
 
     test('when a thumbnailable file is uploaded, then a thumbnail is generated', async () => {
-      const mockBufferStream = { getBuffer: vi.fn() };
-      vi.spyOn(ThumbnailService.instance, 'createFileStreamWithBuffer').mockReturnValue({
-        fileStream: createMockReadStream() as ReturnType<typeof createReadStream>,
-        bufferStream: mockBufferStream as unknown as ReturnType<
-          typeof ThumbnailService.instance.createFileStreamWithBuffer
-        >['bufferStream'],
-      });
-
       const file = createFileSystemNodeFixture({
         type: 'file',
         name: 'image.png',
@@ -325,7 +311,7 @@ describe('Upload File Service', () => {
       });
 
       expect(ThumbnailService.instance.tryUploadThumbnail).toHaveBeenCalledWith({
-        bufferStream: mockBufferStream,
+        input: file.absolutePath,
         fileType: 'png',
         bucket,
         fileUuid: mockFile.uuid,