From a5c64058503888bc38cd13755db2f717a139f7e1 Mon Sep 17 00:00:00 2001
From: David Salami <31099392+Wizdave97@users.noreply.github.com>
Date: Mon, 17 Jul 2023 13:02:01 +0100
Subject: [PATCH] Change return type for `verify_consensus` (#63)

---
 ismp-testsuite/src/mocks.rs    |  6 ++---
 ismp/src/consensus.rs          |  5 +++-
 ismp/src/handlers/consensus.rs | 45 +++++++++++++++++++---------------
 3 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/ismp-testsuite/src/mocks.rs b/ismp-testsuite/src/mocks.rs
index d6abb0c..c43c7b5 100644
--- a/ismp-testsuite/src/mocks.rs
+++ b/ismp-testsuite/src/mocks.rs
@@ -1,11 +1,11 @@
 use ismp::{
     consensus::{
         ConsensusClient, ConsensusClientId, ConsensusStateId, StateCommitment, StateMachineClient,
-        StateMachineHeight, StateMachineId,
+        StateMachineHeight, StateMachineId, VerifiedCommitments,
     },
     error::Error,
     host::{IsmpHost, StateMachine},
-    messaging::{Proof, StateCommitmentHeight},
+    messaging::Proof,
     module::IsmpModule,
     router::{
         DispatchRequest, Get, IsmpDispatcher, IsmpRouter, Post, PostResponse, Request,
@@ -39,7 +39,7 @@ impl ConsensusClient for MockClient {
         _consensus_state_id: ConsensusStateId,
         _trusted_consensus_state: Vec<u8>,
         _proof: Vec<u8>,
-    ) -> Result<(Vec<u8>, BTreeMap<StateMachine, StateCommitmentHeight>), Error> {
+    ) -> Result<(Vec<u8>, VerifiedCommitments), Error> {
         Ok(Default::default())
     }
 
diff --git a/ismp/src/consensus.rs b/ismp/src/consensus.rs
index 37e2fd9..be794ac 100644
--- a/ismp/src/consensus.rs
+++ b/ismp/src/consensus.rs
@@ -89,6 +89,9 @@ pub struct StateMachineHeight {
     pub height: u64,
 }
 
+/// A map of state machine to verified state commitments
+pub type VerifiedCommitments = BTreeMap<StateMachine, Vec<StateCommitmentHeight>>;
+
 /// We define the consensus client as a module that handles logic for consensus proof verification,
 /// and State-Proof verification as well.
 pub trait ConsensusClient {
@@ -99,7 +102,7 @@ pub trait ConsensusClient {
         consensus_state_id: ConsensusStateId,
         trusted_consensus_state: Vec<u8>,
         proof: Vec<u8>,
-    ) -> Result<(Vec<u8>, BTreeMap<StateMachine, StateCommitmentHeight>), Error>;
+    ) -> Result<(Vec<u8>, VerifiedCommitments), Error>;
 
     /// Given two distinct consensus proofs, verify that they're both valid and represent
     /// conflicting views of the network. returns Ok(()) if they're both valid.
diff --git a/ismp/src/handlers/consensus.rs b/ismp/src/handlers/consensus.rs
index 4c773ca..65ea2db 100644
--- a/ismp/src/handlers/consensus.rs
+++ b/ismp/src/handlers/consensus.rs
@@ -63,31 +63,36 @@ where
     let timestamp = host.timestamp();
     host.store_consensus_update_time(msg.consensus_state_id, timestamp)?;
     let mut state_updates = BTreeSet::new();
-    for (id, commitment_height) in intermediate_states {
+    for (id, mut commitment_heights) in intermediate_states {
+        commitment_heights.sort_unstable_by(|a, b| a.height.cmp(&b.height));
         let id = StateMachineId { state_id: id, consensus_state_id: msg.consensus_state_id };
-        let state_height = StateMachineHeight { id, height: commitment_height.height };
-        // If a state machine is frozen, we skip it
-        if host.is_state_machine_frozen(state_height).is_err() {
-            continue
-        }
-
         let previous_latest_height = host.latest_commitment_height(id)?;
-
-        // Only allow heights greater than latest height
-        if previous_latest_height > commitment_height.height {
-            continue
+        for commitment_height in commitment_heights.iter() {
+            let state_height = StateMachineHeight { id, height: commitment_height.height };
+            // If a state machine is frozen, we skip it
+            if host.is_state_machine_frozen(state_height).is_err() {
+                continue
+            }
+
+            // Only allow heights greater than latest height
+            if previous_latest_height > commitment_height.height {
+                continue
+            }
+
+            // Skip duplicate states
+            if host.state_machine_commitment(state_height).is_ok() {
+                continue
+            }
+
+            host.store_state_machine_commitment(state_height, commitment_height.commitment)?;
         }
 
-        // Skip duplicate states
-        if host.state_machine_commitment(state_height).is_ok() {
-            continue
+        if let Some(latest_height) = commitment_heights.last() {
+            let latest_height = StateMachineHeight { id, height: latest_height.height };
+            state_updates
+                .insert((StateMachineHeight { id, height: previous_latest_height }, latest_height));
+            host.store_latest_commitment_height(latest_height)?;
         }
-
-        host.store_state_machine_commitment(state_height, commitment_height.commitment)?;
-
-        state_updates
-            .insert((StateMachineHeight { id, height: previous_latest_height }, state_height));
-        host.store_latest_commitment_height(state_height)?;
     }
 
     let result = ConsensusUpdateResult {