From ad78d15bad837829cd78e263ec7b4f1adeebdfbb Mon Sep 17 00:00:00 2001
From: Gert Hengeveld <info@ghengeveld.nl>
Date: Fri, 20 Oct 2023 11:58:04 +0200
Subject: [PATCH] Abort zip upload if larger than individual files, and log
 size reduction

---
 node-src/lib/compress.test.ts |  6 +++---
 node-src/lib/upload.ts        | 11 ++++++++---
 node-src/tasks/upload.ts      |  5 ++++-
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/node-src/lib/compress.test.ts b/node-src/lib/compress.test.ts
index 6ea486e0b..7172733f2 100644
--- a/node-src/lib/compress.test.ts
+++ b/node-src/lib/compress.test.ts
@@ -14,7 +14,7 @@ afterEach(() => {
 });
 
 const testContext = { sourceDir: '/chromatic-tmp', log: new TestLogger() } as any;
-const fileInfo = { paths: ['file1'] };
+const files = [{ localPath: '/chromatic-tmp/file1', targetPath: 'file1', contentLength: 1 }];
 
 describe('makeZipFile', () => {
   it('adds files to an archive', async () => {
@@ -24,14 +24,14 @@ describe('makeZipFile', () => {
       },
     });
 
-    const result = await makeZipFile(testContext, fileInfo);
+    const result = await makeZipFile(testContext, files);
 
     expect(existsSync(result.path)).toBeTruthy();
     expect(result.size).toBeGreaterThan(0);
   });
 
   it('rejects on error signals', () => {
-    return expect(makeZipFile(testContext, fileInfo)).rejects.toThrow(
+    return expect(makeZipFile(testContext, files)).rejects.toThrow(
       `ENOENT: no such file or directory, open '/chromatic-tmp/file1'`
     );
   });
diff --git a/node-src/lib/upload.ts b/node-src/lib/upload.ts
index 73b024657..c7e679e61 100644
--- a/node-src/lib/upload.ts
+++ b/node-src/lib/upload.ts
@@ -85,8 +85,13 @@ export async function uploadAsZipFile(
     onError?: (error: Error, path?: string) => void;
   } = {}
 ) {
+  const originalSize = files.reduce((acc, { contentLength }) => acc + contentLength, 0);
   const zipped = await makeZipFile(ctx, files);
-  const { path, size: total } = zipped;
+  const { path, size } = zipped;
+
+  if (size > originalSize) throw new Error('Zip file is larger than individual files');
+  ctx.log.debug(`Compression reduced upload size by ${originalSize - size} bytes`);
+
   const { getZipUploadUrl } = await ctx.client.runQuery<GetZipUploadUrlMutationResult>(
     GetZipUploadUrlMutation,
     { buildId: ctx.announcedBuild.id }
@@ -96,12 +101,12 @@ export async function uploadAsZipFile(
   options.onStart?.();
 
   try {
-    await uploadZip(ctx, path, url, total, (progress) => options.onProgress?.(progress, total));
+    await uploadZip(ctx, path, url, size, (progress) => options.onProgress?.(progress, size));
   } catch (e) {
     return options.onError?.(e, path);
   }
 
   await waitForUnpack(ctx, sentinelUrl);
 
-  options.onComplete?.(total, domain);
+  options.onComplete?.(size, domain);
 }
diff --git a/node-src/tasks/upload.ts b/node-src/tasks/upload.ts
index c70ec777a..3175f9ed3 100644
--- a/node-src/tasks/upload.ts
+++ b/node-src/tasks/upload.ts
@@ -217,7 +217,10 @@ export const uploadStorybook = async (ctx: Context, task: Task) => {
     try {
       await uploadAsZipFile(ctx, files, options);
     } catch (err) {
-      ctx.log.debug({ err }, 'Error uploading zip file');
+      ctx.log.debug(
+        { err },
+        'Error uploading zip file, falling back to uploading individual files'
+      );
       await uploadAsIndividualFiles(ctx, files, options);
     }
   } else {