Roll back merkle tree if note size mismatch

iron-fish · Sep 13, 2023 · 738a832 · 738a832
1 parent d8210d9
commit 738a832
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/example/src/Client/utils/BlockProcessor.ts b/example/src/Client/utils/BlockProcessor.ts
@@ -8,7 +8,11 @@ import {
   LightBlock,
   LightStreamerClient,
 } from "../../../../src/models/lightstreamer";
-import { addNotesToMerkleTree, revertToNoteSize } from "./merkle";
+import {
+  addNotesToMerkleTree,
+  getNotesTreeSize,
+  revertToNoteSize,
+} from "./merkle";
 import { logThrottled } from "./logThrottled";
 import { AccountsManager } from "./AccountsManager";
 
@@ -158,8 +162,6 @@ export class BlockProcessor {
       return;
     }
 
-    await this.blockCache.cacheBlock(block);
-
     const notes: Buffer[] = [];
 
     for (const transaction of block.transactions) {
@@ -168,7 +170,16 @@ export class BlockProcessor {
       }
     }
 
+    const prevBlockNoteSize = block.noteSize - notes.length;
+    const notesTreeSize = await getNotesTreeSize();
+
+    if (notesTreeSize > prevBlockNoteSize) {
+      await revertToNoteSize(prevBlockNoteSize);
+    }
+
     await addNotesToMerkleTree(notes);
+
+    await this.blockCache.cacheBlock(block);
   }
 
   private async _checkForReorg(block: LightBlock) {

diff --git a/example/src/Client/utils/merkle.ts b/example/src/Client/utils/merkle.ts
@@ -28,3 +28,5 @@ export const addNotesToMerkleTree = async (notes: Buffer[]) => {
 export const revertToNoteSize = async (noteSize: number) => {
   await notesTree.truncate(noteSize);
 };
+
+export const getNotesTreeSize = () => notesTree.size();