chore(server): improve stream read (#10960)

packages/backend/server/src/base/utils/index.ts

@@ -1,4 +1,5 @@
 export * from './promise';
 export * from './request';
+export * from './stream';
 export * from './types';
 export * from './unit';

packages/backend/server/src/base/utils/stream.ts

@@ -0,0 +1,62 @@
+import { Readable } from 'node:stream';
+
+import { BlobQuotaExceeded, StorageQuotaExceeded } from '../error';
+
+export type CheckExceededResult =
+  | {
+      storageQuotaExceeded: boolean;
+      blobQuotaExceeded: boolean;
+    }
+  | undefined;
+
+export async function readBuffer(
+  readable: Readable,
+  checkExceeded: (recvSize: number) => CheckExceededResult
+): Promise<Buffer> {
+  return new Promise<Buffer>((resolve, reject) => {
+    const chunks: Uint8Array[] = [];
+    let totalSize = 0;
+    let result: CheckExceededResult;
+
+    readable.on('data', chunk => {
+      totalSize += chunk.length;
+
+      // check size after receive each chunk to avoid unnecessary memory usage
+      result = checkExceeded(totalSize);
+      if (result?.blobQuotaExceeded) {
+        reject(new BlobQuotaExceeded());
+      } else if (result?.storageQuotaExceeded) {
+        reject(new StorageQuotaExceeded());
+      }
+
+      if (checkExceeded(totalSize)) {
+        reject(new BlobQuotaExceeded());
+        readable.destroy(new BlobQuotaExceeded());
+        return;
+      }
+      chunks.push(chunk);
+    });
+
+    readable.on('error', reject);
+    readable.on('end', () => {
+      const buffer = Buffer.concat(chunks, totalSize);
+
+      if (checkExceeded(buffer.length)) {
+        reject(new BlobQuotaExceeded());
+      } else {
+        resolve(buffer);
+      }
+    });
+  });
+}
+
+export async function readBufferWithLimit(
+  readable: Readable,
+  limit: number
+): Promise<Buffer> {
+  return readBuffer(readable, size =>
+    size > limit
+      ? { blobQuotaExceeded: true, storageQuotaExceeded: false }
+      : undefined
+  );
+}