From dfe69c8d1f0d42e0a0000dc55926e2cc1cdec9be Mon Sep 17 00:00:00 2001
From: nullptr-0 <959220793@qq.com>
Date: Tue, 29 Aug 2023 00:29:05 +0800
Subject: [PATCH] bugfix for #36 #50

---
 src/decrypt/kgm_wasm.ts | 18 +++++++++---------
 src/decrypt/qmc_wasm.ts |  7 ++++---
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/src/decrypt/kgm_wasm.ts b/src/decrypt/kgm_wasm.ts
index 039c3be..1a9b021 100644
--- a/src/decrypt/kgm_wasm.ts
+++ b/src/decrypt/kgm_wasm.ts
@@ -2,7 +2,7 @@ import { KgmCrypto } from '@xhacker/kgmwasm/KgmWasmBundle';
 import KgmCryptoModule from '@xhacker/kgmwasm/KgmWasmBundle';
 import { MergeUint8Array } from '@/utils/MergeUint8Array';
 
-// 每次处理 2M 的数据
+// 每次可以处理 2M 的数据
 const DECRYPTION_BUF_SIZE = 2 *1024 * 1024;
 
 export interface KGMDecryptionResult {
@@ -36,12 +36,12 @@ export async function DecryptKgmWasm(kgmBlob: ArrayBuffer, ext: string): Promise
 
   // 申请内存块，并文件末端数据到 WASM 的内存堆
   let kgmBuf = new Uint8Array(kgmBlob);
-  const pQmcBuf = KgmCryptoObj._malloc(DECRYPTION_BUF_SIZE);
-  KgmCryptoObj.writeArrayToMemory(kgmBuf.slice(0, DECRYPTION_BUF_SIZE), pQmcBuf);
+  const pKgmBuf = KgmCryptoObj._malloc(DECRYPTION_BUF_SIZE);
+  const preDecDataSize = Math.min(DECRYPTION_BUF_SIZE, kgmBlob.byteLength); // 初始化缓冲区大小
+  KgmCryptoObj.writeArrayToMemory(kgmBuf.slice(0, preDecDataSize), pKgmBuf);
 
   // 进行解密初始化
-  const headerSize = KgmCryptoObj.preDec(pQmcBuf, DECRYPTION_BUF_SIZE, ext);
-  console.log(headerSize);
+  const headerSize = KgmCryptoObj.preDec(pKgmBuf, preDecDataSize, ext);
   kgmBuf = kgmBuf.slice(headerSize);
 
   const decryptedParts = [];
@@ -52,14 +52,14 @@ export async function DecryptKgmWasm(kgmBlob: ArrayBuffer, ext: string): Promise
 
     // 解密一些片段
     const blockData = new Uint8Array(kgmBuf.slice(offset, offset + blockSize));
-    KgmCryptoObj.writeArrayToMemory(blockData, pQmcBuf);
-    KgmCryptoObj.decBlob(pQmcBuf, blockSize, offset);
-    decryptedParts.push(KgmCryptoObj.HEAPU8.slice(pQmcBuf, pQmcBuf + blockSize));
+    KgmCryptoObj.writeArrayToMemory(blockData, pKgmBuf);
+    KgmCryptoObj.decBlob(pKgmBuf, blockSize, offset);
+    decryptedParts.push(KgmCryptoObj.HEAPU8.slice(pKgmBuf, pKgmBuf + blockSize));
 
     offset += blockSize;
     bytesToDecrypt -= blockSize;
   }
-  KgmCryptoObj._free(pQmcBuf);
+  KgmCryptoObj._free(pKgmBuf);
 
   result.data = MergeUint8Array(decryptedParts);
   result.success = true;
diff --git a/src/decrypt/qmc_wasm.ts b/src/decrypt/qmc_wasm.ts
index f747bf4..0ab2dad 100644
--- a/src/decrypt/qmc_wasm.ts
+++ b/src/decrypt/qmc_wasm.ts
@@ -2,7 +2,7 @@ import { QmcCrypto } from '@xhacker/qmcwasm/QmcWasmBundle';
 import QmcCryptoModule from '@xhacker/qmcwasm/QmcWasmBundle';
 import { MergeUint8Array } from '@/utils/MergeUint8Array';
 
-// 每次处理 2M 的数据
+// 每次可以处理 2M 的数据
 const DECRYPTION_BUF_SIZE = 2 *1024 * 1024;
 
 export interface QMCDecryptionResult {
@@ -38,11 +38,12 @@ export async function DecryptQmcWasm(qmcBlob: ArrayBuffer, ext: string): Promise
   // 申请内存块，并文件末端数据到 WASM 的内存堆
   const qmcBuf = new Uint8Array(qmcBlob);
   const pQmcBuf = QmcCryptoObj._malloc(DECRYPTION_BUF_SIZE);
-  QmcCryptoObj.writeArrayToMemory(qmcBuf.slice(-DECRYPTION_BUF_SIZE), pQmcBuf);
+  const preDecDataSize = Math.min(DECRYPTION_BUF_SIZE, qmcBlob.byteLength); // 初始化缓冲区大小
+  QmcCryptoObj.writeArrayToMemory(qmcBuf.slice(-preDecDataSize), pQmcBuf);
 
   // 进行解密初始化
   ext = '.' + ext;
-  const tailSize = QmcCryptoObj.preDec(pQmcBuf, DECRYPTION_BUF_SIZE, ext);
+  const tailSize = QmcCryptoObj.preDec(pQmcBuf, preDecDataSize, ext);
   if (tailSize == -1) {
     result.error = QmcCryptoObj.getErr();
     return result;