fix: ban peers when do reorg after client restart

src/protocols/light_client/components/send_last_state_proof.rs

@@ -150,10 +150,10 @@ impl<'a> SendLastStateProofProcess<'a> {
 
         // Check if headers are continuous.
         if reorg_count != 0 {
-            return_if_failed!(check_continuous_headers(&headers[..reorg_count - 1]));
+            return_if_failed!(check_continuous_headers(&headers[..reorg_count]));
         }
         return_if_failed!(check_continuous_headers(
-            &headers[reorg_count + sampled_count..]
+            &headers[(reorg_count + sampled_count)..]
         ));
 
         // Verify MMR proof
@@ -249,6 +249,30 @@ impl<'a> SendLastStateProofProcess<'a> {
                         }
                     } else if reorg_count == 0 {
                         new_last_headers
+                    } else if sampled_count == 0
+                        && check_continuous_headers(&headers[(reorg_count - 1)..=reorg_count])
+                            .is_ok()
+                    {
+                        // At the above part:
+                        // - `reorg_count != 0`, `headers[..reorg_count]` are checked.
+                        // - `headers[(reorg_count + sampled_count)..]` are always checked.
+                        // So, only check `reorg_count - 1` and `reorg_count` is enough to prove
+                        // all header are continuous.
+                        let required_count = last_n_blocks - last_n_count;
+                        let old_last_headers = &headers[..reorg_count];
+                        let old_last_headers_len = old_last_headers.len();
+                        let tmp_headers = old_last_headers
+                            .iter()
+                            .skip(old_last_headers_len.saturating_sub(required_count))
+                            .map(ToOwned::to_owned)
+                            .chain(new_last_headers)
+                            .collect::<Vec<_>>();
+                        debug!("peer {}: reorg but no enough blocks, so all headers should be continuous, \
+                            reorg: {reorg_count}, sampled: {sampled_count}, last_n_calculate: {last_n_count}, \
+                            last_n_real: {}, last_n_param: {last_n_blocks}, original_request: {original_request}",
+                            self.peer_index, tmp_headers.len()
+                        );
+                        tmp_headers
                     } else {
                         // If this branch is reached, the follow conditions must be satisfied:
                         // - No previous prove state.

src/protocols/light_client/peers.rs

@@ -1285,6 +1285,13 @@ impl Peers {
         self.inner.get(index).map(|peer| peer.clone())
     }
 
+    #[cfg(test)]
+    pub(crate) fn mock_initialized(&self, index: PeerIndex) {
+        if let Some(mut peer) = self.inner.get_mut(&index) {
+            _ = peer.state.take();
+        }
+    }
+
     #[cfg(test)]
     pub(crate) fn mock_prove_request(
         &self,

src/tests/protocols/light_client/send_last_state_proof.rs

@@ -1647,6 +1647,18 @@ async fn reorg_rollback_5_blocks() {
     test_with_reorg_blocks(param).await;
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn reorg_rollback_after_restart_and_last_n_blocks_is_not_enough() {
+    let param = ReorgTestParameter {
+        last_number: 30,
+        rollback_blocks_count: 3,
+        last_n_blocks: 20,
+        restart: true,
+        ..Default::default()
+    };
+    test_with_reorg_blocks(param).await;
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn reorg_detect_long_fork_turn_1() {
     let param = ReorgTestParameter {
@@ -1682,6 +1694,8 @@ struct ReorgTestParameter {
     long_fork_detected: bool,
     expected_last_headers_count_opt: Option<BlockNumber>,
     result: StatusCode,
+    // Mock "restart" state: after restart, the first received "last state" is on a forked chain.
+    restart: bool,
 }
 
 async fn test_with_reorg_blocks(param: ReorgTestParameter) {
@@ -1834,6 +1848,11 @@ async fn test_with_reorg_blocks(param: ReorgTestParameter) {
         assert_eq!(chain.shared().snapshot().tip_number(), last_number);
     }
 
+    if param.restart {
+        protocol.peers().mock_initialized(peer_index);
+        protocol.peers().request_last_state(peer_index).unwrap();
+    }
+
     // Run the test.
     {
         let mut prove_request = chain.build_prove_request(