fix: let peers cached their own last headers to avoid be overridden

src/protocols/light_client/peers.rs

@@ -20,8 +20,6 @@ use crate::protocols::{Status, StatusCode, MESSAGE_TIMEOUT};
 
 pub struct Peers {
     inner: DashMap<PeerIndex, Peer>,
-    // verified last N block headers
-    last_headers: RwLock<Vec<HeaderView>>,
     // The headers are fetching, the value is:
     fetching_headers: DashMap<Byte32, FetchInfo>,
     // The transactions are fetching, the value is:
@@ -1124,7 +1122,6 @@ impl Peers {
 
         Self {
             inner: Default::default(),
-            last_headers: Default::default(),
             fetching_headers: DashMap::new(),
             fetching_txs: DashMap::new(),
             matched_blocks: Default::default(),
@@ -1157,10 +1154,6 @@ impl Peers {
         (number.saturating_sub(1) / self.check_point_interval) as u32
     }
 
-    pub(crate) fn last_headers(&self) -> &RwLock<Vec<HeaderView>> {
-        &self.last_headers
-    }
-
     #[cfg(test)]
     pub(crate) fn fetching_headers(&self) -> &DashMap<Byte32, FetchInfo> {
         &self.fetching_headers
@@ -1365,7 +1358,6 @@ impl Peers {
         index: PeerIndex,
         state: ProveState,
     ) -> Result<(), Status> {
-        *self.last_headers.write().expect("poisoned") = state.get_last_headers().to_vec();
         if let Some(mut peer) = self.inner.get_mut(&index) {
             let has_reorg = !state.reorg_last_headers.is_empty();
             peer.state = peer.state.take().receive_last_state_proof(state)?;
@@ -1942,12 +1934,17 @@ impl Peers {
     }
 
     pub(crate) fn find_header_in_proved_state(&self, hash: &Byte32) -> Option<HeaderView> {
-        self.last_headers()
-            .read()
-            .expect("poisoned")
-            .iter()
-            .find(|header| header.hash().eq(hash))
-            .cloned()
+        self.inner.iter().find_map(|item| {
+            let (_, peer) = item.pair();
+            peer.state.get_prove_state().and_then(|prove_state| {
+                // TODO Store last headers in an ordered hashmap could increase performance.
+                prove_state
+                    .get_last_headers()
+                    .iter()
+                    .find(|header| hash == &header.hash())
+                    .cloned()
+            })
+        })
     }
 
     pub(crate) fn get_best_proved_peers(&self, best_tip: &packed::Header) -> Vec<PeerIndex> {

src/tests/protocols/light_client/send_last_state_proof.rs

@@ -2140,14 +2140,6 @@ async fn multi_peers_override_last_headers() {
 
     // Run the test: check last headers which is stored in memory.
     {
-        let last_headers = protocol
-            .peers()
-            .last_headers()
-            .read()
-            .expect("poisoned")
-            .clone();
-        assert_eq!(last_headers.len() as u64, protocol.last_n_blocks());
-        assert_eq!(last_headers.last().expect("checked").number(), num_high - 1);
         assert!(protocol
             .peers()
             .find_header_in_proved_state(&header_hash_for_test)
@@ -2209,18 +2201,9 @@ async fn multi_peers_override_last_headers() {
 
     // Run the test: check last headers which is stored in memory, again.
     {
-        let last_headers = protocol
-            .peers()
-            .last_headers()
-            .read()
-            .expect("poisoned")
-            .clone();
-        assert_eq!(last_headers.len() as u64, protocol.last_n_blocks());
-        assert_eq!(last_headers.last().expect("checked").number(), num_low - 1);
-        // TODO FIXME Last headers from a better chain are overrided by worse data.
         assert!(protocol
             .peers()
             .find_header_in_proved_state(&header_hash_for_test)
-            .is_none());
+            .is_some());
     }
 }

src/tests/service.rs

@@ -1,6 +1,7 @@
 use std::sync::{Arc, RwLock};
 
 use ckb_chain_spec::consensus::Consensus;
+use ckb_network::PeerIndex;
 use ckb_types::{
     bytes::Bytes,
     core::{
@@ -10,11 +11,12 @@ use ckb_types::{
     h256,
     packed::{CellInput, CellOutputBuilder, Header, OutPoint, Script, ScriptBuilder},
     prelude::*,
+    utilities::merkle_mountain_range::VerifiableHeader,
     H256, U256,
 };
 
 use crate::{
-    protocols::{FetchInfo, PendingTxs},
+    protocols::{FetchInfo, LastState, PendingTxs, ProveRequest, ProveState},
     service::{
         BlockFilterRpc, BlockFilterRpcImpl, ChainRpc, ChainRpcImpl, FetchStatus, Order,
         ScriptStatus, ScriptType, SearchKey, SearchKeyFilter, SetScriptsCommand, Status,
@@ -733,11 +735,27 @@ fn rpc() {
 
     // insert fetched headers
     let peers = create_peers();
-    peers
-        .last_headers()
-        .write()
-        .unwrap()
-        .push(extra_header.clone());
+    {
+        let peer_index = PeerIndex::new(3);
+        peers.add_peer(peer_index);
+        let tip_header = VerifiableHeader::new(
+            storage.get_tip_header().into_view(),
+            Default::default(),
+            None,
+            Default::default(),
+        );
+        let last_state = LastState::new(tip_header);
+        let request = ProveRequest::new(last_state.clone(), Default::default());
+        let prove_state = ProveState::new_from_request(
+            request.clone(),
+            Default::default(),
+            vec![extra_header.clone()],
+        );
+        peers.request_last_state(peer_index).unwrap();
+        peers.update_last_state(peer_index, last_state).unwrap();
+        peers.update_prove_request(peer_index, request).unwrap();
+        peers.update_prove_state(peer_index, prove_state).unwrap();
+    }
     peers.fetching_headers().insert(
         h256!("0xaa22").pack(),
         FetchInfo::new(1111, 3344, false, false),