From 3ae506878d02df06ed62a6ad06fbb2127f56238b Mon Sep 17 00:00:00 2001
From: TejasvOnly <tejasvonly@gmail.com>
Date: Thu, 5 Sep 2024 04:24:48 -0700
Subject: [PATCH] feat: add chunking for decryption

---
 .../decryptMessagesWithPin/index.ts           | 27 +++++++++++++++----
 .../encryptMessagesWithPin/index.ts           |  2 +-
 .../src/utils/operationHelper.ts              |  2 --
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/packages/app-inheritance/src/operations/decryptMessagesWithPin/index.ts b/packages/app-inheritance/src/operations/decryptMessagesWithPin/index.ts
index aa6d924e..35938dfa 100644
--- a/packages/app-inheritance/src/operations/decryptMessagesWithPin/index.ts
+++ b/packages/app-inheritance/src/operations/decryptMessagesWithPin/index.ts
@@ -6,6 +6,10 @@ import {
 } from '@cypherock/sdk-utils';
 import { WALLET_ID_LENGTH } from '../../constants';
 import { APP_VERSION } from '../../constants/appId';
+import {
+  DecryptDataWithPinDecryptedDataStructure,
+  DecryptDataWithPinEncryptedDataStructure,
+} from '../../proto/generated/inheritance/decrypt_data_with_pin';
 import { DecryptDataStatus } from '../../types';
 import {
   assertOrThrowInvalidResult,
@@ -59,18 +63,31 @@ export const decryptMessagesWithPin = async (
   await helper.sendQuery({
     initiate: {
       walletId: params.walletId,
-      encryptedData: params.encryptedData,
     },
   });
 
-  const result = await helper.waitForResult();
+  await helper.waitForResult();
+  logger.verbose('decryptMessages confirmed');
+
+  const rawData = DecryptDataWithPinEncryptedDataStructure.encode(
+    DecryptDataWithPinEncryptedDataStructure.create({
+      data: params.encryptedData,
+    }),
+  ).finish();
+  await helper.sendInChunks(rawData, 'encryptedData', 'dataAccepted');
+
+  const decryptedData = await helper.receiveInChunks(
+    'decryptedDataRequest',
+    'decryptedData',
+  );
+  const result = DecryptDataWithPinDecryptedDataStructure.decode(decryptedData);
   logger.verbose('decryptMessages response', result);
 
-  assertOrThrowInvalidResult(result.messages?.plainData);
+  assertOrThrowInvalidResult(result.decryptedData);
 
   const output = {
-    decryptedData: result.messages.plainData.map(data => data.message),
-    decryptedDataAsStrings: result.messages.plainData.map(data =>
+    decryptedData: result.decryptedData.map(data => data.message),
+    decryptedDataAsStrings: result.decryptedData.map(data =>
       Buffer.from(data.message).toString(),
     ),
   };
diff --git a/packages/app-inheritance/src/operations/encryptMessagesWithPin/index.ts b/packages/app-inheritance/src/operations/encryptMessagesWithPin/index.ts
index c171314f..4e3c2464 100644
--- a/packages/app-inheritance/src/operations/encryptMessagesWithPin/index.ts
+++ b/packages/app-inheritance/src/operations/encryptMessagesWithPin/index.ts
@@ -65,8 +65,8 @@ export const encryptMessageWithPin = async (
     },
   });
 
-  // Wait for confirmation
   await helper.waitForResult();
+  logger.verbose('encryptMessages confirmed');
 
   const rawData = EncryptDataWithPinPlainDataStructure.encode(
     EncryptDataWithPinPlainDataStructure.create({
diff --git a/packages/app-inheritance/src/utils/operationHelper.ts b/packages/app-inheritance/src/utils/operationHelper.ts
index 011be3a9..81cd8f38 100644
--- a/packages/app-inheritance/src/utils/operationHelper.ts
+++ b/packages/app-inheritance/src/utils/operationHelper.ts
@@ -130,8 +130,6 @@ export class OperationHelper<Q extends QueryKey, R extends ResultKey> {
   >(queryKey: QK, resultKey: RK): Promise<Uint8Array> {
     const chunks: Uint8Array[] = [];
 
-    await this.sendQuery({ [queryKey]: {} });
-
     let index = 0;
 
     while (1) {