From 74cbc3ca7856e37e7e2aa63a5ee2b3b5c9726f53 Mon Sep 17 00:00:00 2001
From: Boyu Yang <yangby@cryptape.com>
Date: Fri, 29 Dec 2023 21:58:01 +0800
Subject: [PATCH] feat: keep the update timestamp if the last state is as the
 same as the previous

---
 .../components/send_last_state.rs             | 64 +++++++++++++------
 src/protocols/light_client/peers.rs           |  1 -
 .../protocols/light_client/send_last_state.rs |  3 +-
 3 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/src/protocols/light_client/components/send_last_state.rs b/src/protocols/light_client/components/send_last_state.rs
index 8f36255..6960b37 100644
--- a/src/protocols/light_client/components/send_last_state.rs
+++ b/src/protocols/light_client/components/send_last_state.rs
@@ -33,32 +33,54 @@ impl<'a> SendLastStateProcess<'a> {
 
         let last_state = LastState::new(last_header);
 
-        return_if_failed!(self
-            .protocol
-            .peers()
-            .update_last_state(self.peer_index, last_state.clone()));
-
         if let Some(prev_last_state) = peer_state.get_last_state() {
-            trace!(
-                "peer {}: update last state from {} to {}",
-                self.peer_index,
-                prev_last_state,
-                last_state,
-            );
-            if prev_last_state.total_difficulty() < last_state.total_difficulty() {
-                if let Some(prove_state) = peer_state.get_prove_state() {
-                    if prove_state.is_parent_of(&last_state) {
-                        trace!("peer {}: new last state could be trusted", self.peer_index);
-                        let last_n_blocks = self.protocol.last_n_blocks() as usize;
-                        let child_prove_state = prove_state.new_child(last_state, last_n_blocks);
-                        return_if_failed!(self
-                            .protocol
-                            .update_prove_state_to_child(self.peer_index, child_prove_state));
+            if last_state.is_same_as(prev_last_state) {
+                trace!(
+                    "peer {}: receive the same last state as previous {}",
+                    self.peer_index,
+                    last_state,
+                );
+                // Do NOT update the timestamp for same last state,
+                // so it could be banned after timeout check.
+            } else {
+                trace!(
+                    "peer {}: update last state from {} to {}",
+                    self.peer_index,
+                    prev_last_state,
+                    last_state,
+                );
+
+                return_if_failed!(self
+                    .protocol
+                    .peers()
+                    .update_last_state(self.peer_index, last_state.clone()));
+
+                if prev_last_state.total_difficulty() < last_state.total_difficulty() {
+                    if let Some(prove_state) = peer_state.get_prove_state() {
+                        if prove_state.is_parent_of(&last_state) {
+                            trace!("peer {}: new last state could be trusted", self.peer_index);
+                            let last_n_blocks = self.protocol.last_n_blocks() as usize;
+                            let child_prove_state =
+                                prove_state.new_child(last_state, last_n_blocks);
+                            return_if_failed!(self
+                                .protocol
+                                .update_prove_state_to_child(self.peer_index, child_prove_state));
+                        }
                     }
                 }
             }
         } else {
-            trace!("peer {}: initialize last state", self.peer_index);
+            trace!(
+                "peer {}: initialize last state {}",
+                self.peer_index,
+                last_state
+            );
+
+            return_if_failed!(self
+                .protocol
+                .peers()
+                .update_last_state(self.peer_index, last_state));
+
             let is_sent =
                 return_if_failed!(self.protocol.get_last_state_proof(self.nc, self.peer_index));
             if !is_sent {
diff --git a/src/protocols/light_client/peers.rs b/src/protocols/light_client/peers.rs
index f56093b..3cbd1b3 100644
--- a/src/protocols/light_client/peers.rs
+++ b/src/protocols/light_client/peers.rs
@@ -254,7 +254,6 @@ impl LastState {
         self.update_ts
     }
 
-    #[cfg(test)]
     pub(crate) fn is_same_as(&self, another: &Self) -> bool {
         if_verifiable_headers_are_same(&self.header, &another.header)
     }
diff --git a/src/tests/protocols/light_client/send_last_state.rs b/src/tests/protocols/light_client/send_last_state.rs
index 94a1a6a..6799d34 100644
--- a/src/tests/protocols/light_client/send_last_state.rs
+++ b/src/tests/protocols/light_client/send_last_state.rs
@@ -215,8 +215,7 @@ async fn update_to_same_last_state() {
         let last_state_after = peer_state_after.get_last_state().expect("has last state");
 
         assert!(last_state_after.is_same_as(&last_state_before));
-        // TODO keep the update timestamp if the last state is not changed
-        assert_ne!(last_state_after.update_ts(), last_state_before.update_ts());
+        assert_eq!(last_state_after.update_ts(), last_state_before.update_ts());
     }
 }