From 85a7810cc3158cd81d55f3e8d7747767bf62fc37 Mon Sep 17 00:00:00 2001
From: boqiu <82121246@qq.com>
Date: Fri, 22 Nov 2024 15:13:16 +0800
Subject: [PATCH] verify shard config for status message

---
 node/router/src/libp2p_event_handler.rs | 51 +++++++++++++++++++++----
 1 file changed, 44 insertions(+), 7 deletions(-)

diff --git a/node/router/src/libp2p_event_handler.rs b/node/router/src/libp2p_event_handler.rs
index 6458fed..f418e53 100644
--- a/node/router/src/libp2p_event_handler.rs
+++ b/node/router/src/libp2p_event_handler.rs
@@ -258,8 +258,13 @@ impl Libp2pEventHandler {
 
         let network_id = self.network_globals.network_id();
         let shard_config = self.store.get_store().get_shard_config();
+
+        if !self.verify_status_message(peer_id, status, network_id.clone(), &shard_config) {
+            return;
+        }
+
         let status_message = StatusMessage {
-            data: network_id.clone(),
+            data: network_id,
             num_shard: shard_config.num_shard,
             shard_id: shard_config.shard_id,
         };
@@ -270,12 +275,12 @@ impl Libp2pEventHandler {
             id: request_id,
             response: Response::Status(status_message),
         });
-        self.on_status_message(peer_id, status, network_id);
     }
 
     fn on_status_response(&self, peer_id: PeerId, status: StatusMessage) {
         let network_id = self.network_globals.network_id();
-        self.on_status_message(peer_id, status, network_id);
+        let shard_config = self.store.get_store().get_shard_config();
+        self.verify_status_message(peer_id, status, network_id, &shard_config);
     }
 
     pub async fn on_rpc_response(
@@ -956,21 +961,53 @@ impl Libp2pEventHandler {
         MessageAcceptance::Accept
     }
 
-    fn on_status_message(
+    fn verify_status_message(
         &self,
         peer_id: PeerId,
         status: StatusMessage,
         network_id: NetworkIdentity,
-    ) {
+        shard_config: &ShardConfig,
+    ) -> bool {
         if status.data != network_id {
             warn!(%peer_id, ?network_id, ?status.data, "Report peer with incompatible network id");
             self.send_to_network(NetworkMessage::ReportPeer {
                 peer_id,
                 action: PeerAction::Fatal,
-                source: ReportSource::Gossipsub,
+                source: ReportSource::RPC,
                 msg: "Incompatible network id in StatusMessage",
-            })
+            });
+            return false;
         }
+
+        let peer_shard_config = match ShardConfig::new(status.shard_id, status.num_shard) {
+            Ok(v) => v,
+            Err(err) => {
+                warn!(%peer_id, ?status, ?err, "Report peer with invalid shard config");
+                self.send_to_network(NetworkMessage::ReportPeer {
+                    peer_id,
+                    action: PeerAction::Fatal,
+                    source: ReportSource::RPC,
+                    msg: "Invalid shard config in StatusMessage",
+                });
+                return false;
+            }
+        };
+
+        if !peer_shard_config.intersect(shard_config) {
+            info!(%peer_id, ?shard_config, ?status, "Report peer with mismatched shard config");
+            self.send_to_network(NetworkMessage::ReportPeer {
+                peer_id,
+                action: PeerAction::Fatal,
+                source: ReportSource::RPC,
+                msg: "Shard config mismatch in StatusMessage",
+            });
+            return false;
+        }
+
+        self.file_location_cache
+            .insert_peer_config(peer_id, peer_shard_config);
+
+        true
     }
 
     async fn publish_file(&self, tx_id: TxID) -> Option<bool> {