From 9852302d68ec4050f0c1e7139dd705cf6f332cd3 Mon Sep 17 00:00:00 2001
From: David Salami <31099392+Wizdave97@users.noreply.github.com>
Date: Sat, 22 Apr 2023 09:10:54 +0100
Subject: [PATCH] Finalize mmr only when leaves are added (#31)

---
 pallet-ismp/rpc/src/lib.rs |  2 +-
 pallet-ismp/src/lib.rs     | 29 ++++++++++++++++++++++++++---
 pallet-ismp/src/router.rs  | 31 +++++++++++++++----------------
 3 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/pallet-ismp/rpc/src/lib.rs b/pallet-ismp/rpc/src/lib.rs
index 92dd735ac..d286f54f1 100644
--- a/pallet-ismp/rpc/src/lib.rs
+++ b/pallet-ismp/rpc/src/lib.rs
@@ -108,7 +108,7 @@ where
 
     /// Query ISMP Events that were deposited in a series of blocks
     /// Using String keys because HashMap fails to deserialize when key is not a String
-    #[method(name = "ibc_queryEvents")]
+    #[method(name = "ismp_queryEvents")]
     fn query_events(
         &self,
         block_numbers: Vec<BlockNumberOrHash<Hash>>,
diff --git a/pallet-ismp/src/lib.rs b/pallet-ismp/src/lib.rs
index adf0b6aa6..3f8d0c8ed 100644
--- a/pallet-ismp/src/lib.rs
+++ b/pallet-ismp/src/lib.rs
@@ -21,7 +21,7 @@ extern crate alloc;
 mod errors;
 pub mod events;
 pub mod host;
-pub mod mmr;
+mod mmr;
 pub mod primitives;
 pub mod router;
 
@@ -182,6 +182,11 @@ pub mod pallet {
         OptionQuery,
     >;
 
+    /// State variable that tells us if at least one new leaf was added to the mmr
+    #[pallet::storage]
+    #[pallet::getter(fn new_leaves)]
+    pub type NewLeavesAdded<T> = StorageValue<_, LeafIndex, OptionQuery>;
+
     // Pallet implements [`Hooks`] trait to define some logic to execute in some context.
     #[pallet::hooks]
     impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T>
@@ -194,6 +199,10 @@ pub mod pallet {
         }
 
         fn on_finalize(_n: T::BlockNumber) {
+            // Only finalize if mmr was modified
+            if !NewLeavesAdded::<T>::exists() {
+                return
+            }
             let leaves = Self::number_of_leaves();
 
             let mmr: Mmr<mmr::storage::RuntimeStorage, T> = Mmr::new(leaves);
@@ -214,6 +223,7 @@ pub mod pallet {
 
             let digest = sp_runtime::generic::DigestItem::Consensus(ISMP_ID, log.encode());
             <frame_system::Pallet<T>>::deposit_log(digest);
+            NewLeavesAdded::<T>::kill();
         }
 
         fn offchain_worker(_n: T::BlockNumber) {}
@@ -419,7 +429,10 @@ pub struct RequestResponseLog<T: Config> {
     mmr_root_hash: <T as frame_system::Config>::Hash,
 }
 
-impl<T: Config> Pallet<T> {
+impl<T: Config> Pallet<T>
+where
+    <T as frame_system::Config>::Hash: From<H256>,
+{
     pub fn request_leaf_index_offchain_key(
         source_chain: StateMachine,
         dest_chain: StateMachine,
@@ -436,7 +449,7 @@ impl<T: Config> Pallet<T> {
         (T::INDEXING_PREFIX, "Responses/leaf_indices", source_chain, dest_chain, nonce).encode()
     }
 
-    fn store_leaf_index_offchain(key: Vec<u8>, leaf_index: LeafIndex) {
+    pub fn store_leaf_index_offchain(key: Vec<u8>, leaf_index: LeafIndex) {
         sp_io::offchain_index::set(&key, &leaf_index.encode());
     }
 
@@ -531,4 +544,14 @@ impl<T: Config> Pallet<T> {
     pub fn get_responses(leaf_indices: Vec<LeafIndex>) -> Vec<Response> {
         leaf_indices.into_iter().filter_map(|leaf_index| Self::get_response(leaf_index)).collect()
     }
+
+    pub fn mmr_push(leaf: Leaf) -> Option<NodeIndex> {
+        let leaves = Self::number_of_leaves();
+        let mut mmr: Mmr<mmr::storage::RuntimeStorage, T> = Mmr::new(leaves);
+        let index = mmr.push(leaf);
+        if !NewLeavesAdded::<T>::exists() && index.is_some() {
+            NewLeavesAdded::<T>::put(index.unwrap())
+        }
+        index
+    }
 }
diff --git a/pallet-ismp/src/router.rs b/pallet-ismp/src/router.rs
index 21ebb7d81..50c912a1a 100644
--- a/pallet-ismp/src/router.rs
+++ b/pallet-ismp/src/router.rs
@@ -1,4 +1,4 @@
-use crate::{host::Host, mmr, mmr::mmr::Mmr, Config, Event, Pallet, RequestAcks, ResponseAcks};
+use crate::{host::Host, Config, Event, Pallet, RequestAcks, ResponseAcks};
 use alloc::{boxed::Box, string::ToString};
 use codec::{Decode, Encode};
 use core::marker::PhantomData;
@@ -58,13 +58,12 @@ where
                 })?
             }
 
-            let leaves = Pallet::<T>::number_of_leaves();
             let (dest_chain, source_chain, nonce) =
                 (request.dest_chain(), request.source_chain(), request.nonce());
-            let mut mmr: Mmr<mmr::storage::RuntimeStorage, T> = Mmr::new(leaves);
             let offchain_key =
                 Pallet::<T>::request_leaf_index_offchain_key(source_chain, dest_chain, nonce);
-            let leaf_index = if let Some(leaf_index) = mmr.push(Leaf::Request(request)) {
+            let leaf_index = if let Some(leaf_index) = Pallet::<T>::mmr_push(Leaf::Request(request))
+            {
                 leaf_index
             } else {
                 Err(DispatchError {
@@ -124,25 +123,25 @@ where
                 })?
             }
 
-            let leaves = Pallet::<T>::number_of_leaves();
             let (dest_chain, source_chain, nonce) = (
                 response.request.source_chain(),
                 response.request.dest_chain(),
                 response.request.nonce(),
             );
-            let mut mmr: Mmr<mmr::storage::RuntimeStorage, T> = Mmr::new(leaves);
+
             let offchain_key =
                 Pallet::<T>::response_leaf_index_offchain_key(source_chain, dest_chain, nonce);
-            let leaf_index = if let Some(leaf_index) = mmr.push(Leaf::Response(response)) {
-                leaf_index
-            } else {
-                Err(DispatchError {
-                    msg: "Failed to push response into mmr".to_string(),
-                    nonce,
-                    source: source_chain,
-                    dest: dest_chain,
-                })?
-            };
+            let leaf_index =
+                if let Some(leaf_index) = Pallet::<T>::mmr_push(Leaf::Response(response)) {
+                    leaf_index
+                } else {
+                    Err(DispatchError {
+                        msg: "Failed to push response into mmr".to_string(),
+                        nonce,
+                        source: source_chain,
+                        dest: dest_chain,
+                    })?
+                };
 
             Pallet::<T>::deposit_event(Event::Response {
                 request_nonce: nonce,