Batch multiple accounts when decrypting

iron-fish · Apr 16, 2024 · 7de7f71 · 7de7f71
1 parent b6e0d37
commit 7de7f71
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 13 deletions.
diff --git a/ironfish/src/rpc/routes/wallet/utils.ts b/ironfish/src/rpc/routes/wallet/utils.ts
@@ -187,6 +187,7 @@ export async function getTransactionNotes(
     }
 
     decryptNotesPayloads.push({
+      accountId: account.id,
       serializedNote: note.serialize(),
       incomingViewKey: account.incomingViewKey,
       outgoingViewKey: account.outgoingViewKey,

diff --git a/ironfish/src/wallet/wallet.ts b/ironfish/src/wallet/wallet.ts
@@ -412,14 +412,14 @@ export class Wallet {
       ))
 
     const batchSize = 20
-    const notePromisesByAccountId: Map<string, Array<Promise<Array<DecryptedNote>>>> = new Map()
+    const notePromises: Array<Promise<Array<DecryptedNote>>> = []
+    let decryptNotesPayloads = []
     for (const account of accountsToCheck) {
-      const notePromises = []
-      let decryptNotesPayloads = []
       let currentNoteIndex = initialNoteIndex
 
       for (const note of transaction.notes) {
         decryptNotesPayloads.push({
+          accountId: account.id,
           serializedNote: note.serialize(),
           incomingViewKey: account.incomingViewKey,
           outgoingViewKey: account.outgoingViewKey,
@@ -437,20 +437,18 @@ export class Wallet {
           decryptNotesPayloads = []
         }
       }
+    }
 
-      if (decryptNotesPayloads.length) {
-        notePromises.push(this.decryptNotesFromTransaction(decryptNotesPayloads))
-      }
-
-      notePromisesByAccountId.set(account.id, notePromises)
+    if (decryptNotesPayloads.length) {
+      notePromises.push(this.decryptNotesFromTransaction(decryptNotesPayloads))
     }
 
     const decryptedNotesByAccountId = new Map<string, Array<DecryptedNote>>()
-    for (const [key, value] of notePromisesByAccountId) {
-      const decryptedNotes = (await Promise.all(value ?? [])).flat()
-      if (decryptedNotes.length) {
-        decryptedNotesByAccountId.set(key, decryptedNotes)
-      }
+    const flatPromises = (await Promise.all(notePromises)).flat()
+    for (const decryptedNote of flatPromises) {
+      const accountNotes = decryptedNotesByAccountId.get(decryptedNote.accountId) ?? []
+      accountNotes.push(decryptedNote)
+      decryptedNotesByAccountId.set(decryptedNote.accountId, accountNotes)
     }
 
     return decryptedNotesByAccountId

diff --git a/ironfish/src/workerPool/tasks/decryptNotes.ts b/ironfish/src/workerPool/tasks/decryptNotes.ts
@@ -11,6 +11,7 @@ import { WorkerTask } from './workerTask'
 
 export interface DecryptNoteOptions {
   serializedNote: Buffer
+  accountId: string
   incomingViewKey: string
   outgoingViewKey: string
   viewKey: string
@@ -19,6 +20,7 @@ export interface DecryptNoteOptions {
 }
 
 export interface DecryptedNote {
+  accountId: string
   index: number | null
   forSpender: boolean
   hash: Buffer
@@ -42,6 +44,7 @@ export class DecryptNotesRequest extends WorkerMessage {
       bw.writeU8(flags)
 
       bw.writeBytes(payload.serializedNote)
+      bw.writeVarString(payload.accountId, 'utf8')
       bw.writeBytes(Buffer.from(payload.incomingViewKey, 'hex'))
       bw.writeBytes(Buffer.from(payload.outgoingViewKey, 'hex'))
       bw.writeBytes(Buffer.from(payload.viewKey, 'hex'))
@@ -61,13 +64,15 @@ export class DecryptNotesRequest extends WorkerMessage {
       const hasCurrentNoteIndex = flags & (1 << 0)
       const decryptForSpender = Boolean(flags & (1 << 1))
       const serializedNote = reader.readBytes(ENCRYPTED_NOTE_LENGTH)
+      const accountId = reader.readVarString('utf8')
       const incomingViewKey = reader.readBytes(ACCOUNT_KEY_LENGTH).toString('hex')
       const outgoingViewKey = reader.readBytes(ACCOUNT_KEY_LENGTH).toString('hex')
       const viewKey = reader.readBytes(VIEW_KEY_LENGTH).toString('hex')
       const currentNoteIndex = hasCurrentNoteIndex ? reader.readU32() : null
 
       payloads.push({
         serializedNote,
+        accountId,
         incomingViewKey,
         outgoingViewKey,
         currentNoteIndex,
@@ -84,6 +89,7 @@ export class DecryptNotesRequest extends WorkerMessage {
     for (const payload of this.payloads) {
       size += 1
       size += ENCRYPTED_NOTE_LENGTH
+      size += bufio.sizeVarString(payload.accountId, 'utf8')
       size += ACCOUNT_KEY_LENGTH
       size += ACCOUNT_KEY_LENGTH
       size += VIEW_KEY_LENGTH
@@ -115,6 +121,7 @@ export class DecryptNotesResponse extends WorkerMessage {
         flags |= Number(note.forSpender) << 2
         bw.writeU8(flags)
         bw.writeHash(note.hash)
+        bw.writeVarString(note.accountId, 'utf8')
         bw.writeBytes(note.serializedNote)
 
         if (note.index) {
@@ -144,6 +151,7 @@ export class DecryptNotesResponse extends WorkerMessage {
       const hasNullifier = flags & (1 << 1)
       const forSpender = Boolean(flags & (1 << 2))
       const hash = reader.readHash()
+      const accountId = reader.readVarString('utf8')
       const serializedNote = reader.readBytes(DECRYPTED_NOTE_LENGTH)
 
       let index = null
@@ -157,6 +165,7 @@ export class DecryptNotesResponse extends WorkerMessage {
       }
 
       notes.push({
+        accountId,
         forSpender,
         index,
         hash,
@@ -176,6 +185,7 @@ export class DecryptNotesResponse extends WorkerMessage {
 
       if (note) {
         size += 1 + 32 + DECRYPTED_NOTE_LENGTH
+        size += bufio.sizeVarString(note.accountId, 'utf8')
 
         if (note.index) {
           size += 4
@@ -205,6 +215,7 @@ export class DecryptNotesTask extends WorkerTask {
     const decryptedNotes = []
 
     for (const {
+      accountId,
       serializedNote,
       incomingViewKey,
       outgoingViewKey,
@@ -218,6 +229,7 @@ export class DecryptNotesTask extends WorkerTask {
       const receivedNote = note.decryptNoteForOwner(incomingViewKey)
       if (receivedNote && receivedNote.value() !== 0n) {
         decryptedNotes.push({
+          accountId,
           index: currentNoteIndex,
           forSpender: false,
           hash: note.hash(),
@@ -235,6 +247,7 @@ export class DecryptNotesTask extends WorkerTask {
         const spentNote = note.decryptNoteForSpender(outgoingViewKey)
         if (spentNote && spentNote.value() !== 0n) {
           decryptedNotes.push({
+            accountId,
             index: currentNoteIndex,
             forSpender: true,
             hash: note.hash(),