From 738a832e7e710cd9b5d30e268b3946133b9606fd Mon Sep 17 00:00:00 2001
From: Daniel Cortes <danielgcortes@gmail.com>
Date: Wed, 13 Sep 2023 16:23:53 -0600
Subject: [PATCH] Roll back merkle tree if note size mismatch

---
 example/src/Client/utils/BlockProcessor.ts | 17 ++++++++++++++---
 example/src/Client/utils/merkle.ts         |  2 ++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/example/src/Client/utils/BlockProcessor.ts b/example/src/Client/utils/BlockProcessor.ts
index 13357f1..75d44df 100644
--- 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
index 74d16f5..bc30a00 100644
--- 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();