From 3597e7789baab9e35d755e1e96da252677dad05c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=B2=81=E6=A0=91=E4=BA=BA?= <lu.shuren@um-react.app>
Date: Sat, 10 Jun 2023 16:22:36 +0100
Subject: [PATCH] refactor: extend `transformBlob` to take cleanup function

---
 src/decrypt-worker/crypto/qmc/qmc_v2.ts  | 20 +++++++++++---------
 src/decrypt-worker/util/transformBlob.ts | 15 +++++++++------
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/src/decrypt-worker/crypto/qmc/qmc_v2.ts b/src/decrypt-worker/crypto/qmc/qmc_v2.ts
index af95bc5..116a2a8 100644
--- a/src/decrypt-worker/crypto/qmc/qmc_v2.ts
+++ b/src/decrypt-worker/crypto/qmc/qmc_v2.ts
@@ -9,11 +9,11 @@ export class QMC2Crypto implements CryptoBase {
   checkByDecryptHeader = false;
 
   async decrypt(buffer: ArrayBuffer): Promise<Blob> {
-    // FIXME: Move the cleanup to transformBlob
-    const mod = await fetchParakeet();
-    const footerParser = mod.make.QMCv2FooterParser(SEED, ENC_V2_KEY_1, ENC_V2_KEY_2);
-    return transformBlob(buffer, (p) => p.make.QMCv2(footerParser)).finally(() => {
-      footerParser.delete();
+    const parakeet = await fetchParakeet();
+    const footerParser = parakeet.make.QMCv2FooterParser(SEED, ENC_V2_KEY_1, ENC_V2_KEY_2);
+    return transformBlob(buffer, (p) => p.make.QMCv2(footerParser), {
+      parakeet,
+      cleanup: () => footerParser.delete(),
     });
   }
 
@@ -35,12 +35,14 @@ export class QMC2CryptoWithKey implements CryptoBase {
       throw new Error('key was not provided');
     }
 
-    // FIXME: Move the cleanup to transformBlob
-    const mod = await fetchParakeet();
+    const parakeet = await fetchParakeet();
     const textEncoder = new TextEncoder();
     const key = textEncoder.encode(options.qmc2Key);
-    const keyCrypto = mod.make.QMCv2KeyCrypto(SEED, ENC_V2_KEY_1, ENC_V2_KEY_2);
-    return transformBlob(buffer, (p) => p.make.QMCv2EKey(key, keyCrypto));
+    const keyCrypto = parakeet.make.QMCv2KeyCrypto(SEED, ENC_V2_KEY_1, ENC_V2_KEY_2);
+    return transformBlob(buffer, (p) => p.make.QMCv2EKey(key, keyCrypto), {
+      parakeet,
+      cleanup: () => keyCrypto.delete(),
+    });
   }
 
   public static make() {
diff --git a/src/decrypt-worker/util/transformBlob.ts b/src/decrypt-worker/util/transformBlob.ts
index 2fd1bb5..ba25371 100644
--- a/src/decrypt-worker/util/transformBlob.ts
+++ b/src/decrypt-worker/util/transformBlob.ts
@@ -5,21 +5,24 @@ import { UnsupportedSourceFile } from './DecryptError';
 export async function transformBlob(
   blob: Blob | ArrayBuffer,
   transformerFactory: (p: Parakeet) => Transformer | Promise<Transformer>,
-  parakeet?: Parakeet
+  { cleanup, parakeet }: { cleanup?: () => void; parakeet?: Parakeet } = {}
 ) {
-  const cleanup: (() => void)[] = [];
+  const registeredCleanupFns: (() => void)[] = [];
+  if (cleanup) {
+    registeredCleanupFns.push(cleanup);
+  }
 
   try {
     const mod = parakeet ?? (await fetchParakeet());
     const transformer = await transformerFactory(mod);
-    cleanup.push(() => transformer.delete());
+    registeredCleanupFns.push(() => transformer.delete());
 
     const reader = mod.make.Reader(await toArrayBuffer(blob));
-    cleanup.push(() => reader.delete());
+    registeredCleanupFns.push(() => reader.delete());
 
     const sink = mod.make.WriterSink();
     const writer = sink.getWriter();
-    cleanup.push(() => writer.delete());
+    registeredCleanupFns.push(() => writer.delete());
 
     const result = transformer.Transform(writer, reader);
     if (result === TransformResult.ERROR_INVALID_FORMAT) {
@@ -30,6 +33,6 @@ export async function transformBlob(
 
     return sink.collectBlob();
   } finally {
-    cleanup.forEach((clean) => clean());
+    registeredCleanupFns.forEach((cleanup) => cleanup());
   }
 }