Merge pull request #170 from yangby-cryptape/feature/ban-peers-if-sta…

…te-not-changed-after-timeout

feat: ban a peer if its last state isn't changed after timeout

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 {

src/protocols/light_client/peers.rs

@@ -249,6 +249,14 @@ impl LastState {
     pub(crate) fn header(&self) -> &HeaderView {
         self.as_ref().header()
     }
+
+    pub(crate) fn update_ts(&self) -> u64 {
+        self.update_ts
+    }
+
+    pub(crate) fn is_same_as(&self, another: &Self) -> bool {
+        if_verifiable_headers_are_same(&self.header, &another.header)
+    }
 }
 
 impl fmt::Display for ProveRequest {
@@ -1025,7 +1033,7 @@ impl PeerState {
     fn require_new_last_state(&self, before_ts: u64) -> bool {
         match self {
             Self::Initialized => true,
-            Self::Ready { ref last_state, .. } => last_state.update_ts < before_ts,
+            Self::Ready { ref last_state, .. } => last_state.update_ts() < before_ts,
             Self::OnlyHasLastState { .. }
             | Self::RequestFirstLastState { .. }
             | Self::RequestFirstLastStateProof { .. }
@@ -1848,6 +1856,15 @@ impl Peers {
                             None
                         }
                     })
+                    .or_else(|| {
+                        peer.state.get_last_state().and_then(|state| {
+                            if now > state.update_ts + MESSAGE_TIMEOUT {
+                                Some(*peer_index)
+                            } else {
+                                None
+                            }
+                        })
+                    })
                     .or_else(|| {
                         peer.get_blocks_proof_request().and_then(|req| {
                             if now > req.when_sent + MESSAGE_TIMEOUT {

src/tests/protocols/light_client/send_last_state.rs

@@ -160,6 +160,65 @@ async fn initialize_last_state() {
     assert_eq!(content.last_hash().as_slice(), last_hash.as_slice());
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn update_to_same_last_state() {
+    let chain = MockChain::new_with_dummy_pow("test-light-client").start();
+    let nc = MockNetworkContext::new(SupportProtocols::LightClient);
+
+    let peer_index = PeerIndex::new(1);
+    let peers = {
+        let peers = chain.create_peers();
+        peers.add_peer(peer_index);
+        peers.request_last_state(peer_index).unwrap();
+        peers
+    };
+    let mut protocol = chain.create_light_client_protocol(peers);
+
+    let num = 12;
+    chain.mine_to(num);
+
+    let snapshot = chain.shared().snapshot();
+    let last_header = snapshot
+        .get_verifiable_header_by_number(num)
+        .expect("block stored");
+    let data = {
+        let content = packed::SendLastState::new_builder()
+            .last_header(last_header)
+            .build();
+        packed::LightClientMessage::new_builder()
+            .set(content)
+            .build()
+    }
+    .as_bytes();
+
+    // Setup the test fixture:
+    // - Update last state.
+    {
+        protocol
+            .received(nc.context(), peer_index, data.clone())
+            .await;
+    }
+
+    // Run the test.
+    {
+        let peer_state_before = protocol
+            .get_peer_state(&peer_index)
+            .expect("has peer state");
+        let last_state_before = peer_state_before.get_last_state().expect("has last state");
+
+        tokio::time::sleep(tokio::time::Duration::from_millis(2000)).await;
+        protocol.received(nc.context(), peer_index, data).await;
+
+        let peer_state_after = protocol
+            .get_peer_state(&peer_index)
+            .expect("has peer 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));
+        assert_eq!(last_state_after.update_ts(), last_state_before.update_ts());
+    }
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn update_to_continuous_last_state() {
     let chain = MockChain::new_with_dummy_pow("test-light-client").start();