From a37836a91d8c0ca8ce20d4eb1fb6954a1d4b7b52 Mon Sep 17 00:00:00 2001
From: boqiu <82121246@qq.com>
Date: Fri, 22 Nov 2024 14:50:22 +0800
Subject: [PATCH] Add shard config in status messgae

---
 node/network/src/rpc/codec/ssz_snappy.rs | 9 ++-------
 node/network/src/rpc/methods.rs          | 6 +++++-
 node/network/tests/rpc_tests.rs          | 8 ++------
 node/router/src/libp2p_event_handler.rs  | 8 ++++++++
 4 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/node/network/src/rpc/codec/ssz_snappy.rs b/node/network/src/rpc/codec/ssz_snappy.rs
index d0e9370..e35c197 100644
--- a/node/network/src/rpc/codec/ssz_snappy.rs
+++ b/node/network/src/rpc/codec/ssz_snappy.rs
@@ -399,9 +399,7 @@ mod tests {
     use std::io::Write;
 
     fn status_message() -> StatusMessage {
-        StatusMessage {
-            data: Default::default(),
-        }
+        Default::default()
     }
 
     fn ping_message() -> Ping {
@@ -570,10 +568,7 @@ mod tests {
         assert_eq!(stream_identifier.len(), 10);
 
         // Status message is 84 bytes uncompressed. `max_compressed_len` is 32 + 84 + 84/6 = 130.
-        let status_message_bytes = StatusMessage {
-            data: Default::default(),
-        }
-        .as_ssz_bytes();
+        let status_message_bytes = StatusMessage::default().as_ssz_bytes();
 
         let mut uvi_codec: Uvi<usize> = Uvi::default();
         let mut dst = BytesMut::with_capacity(1024);
diff --git a/node/network/src/rpc/methods.rs b/node/network/src/rpc/methods.rs
index 554eb8b..8322492 100644
--- a/node/network/src/rpc/methods.rs
+++ b/node/network/src/rpc/methods.rs
@@ -69,9 +69,13 @@ impl ToString for ErrorType {
 /* Requests */
 
 /// The STATUS request/response handshake message.
-#[derive(Encode, Decode, Clone, Debug, PartialEq, Eq)]
+#[derive(Encode, Decode, Clone, Debug, PartialEq, Eq, Default)]
 pub struct StatusMessage {
     pub data: NetworkIdentity,
+
+    // shard config
+    pub num_shard: usize,
+    pub shard_id: usize,
 }
 
 /// The PING request/response message.
diff --git a/node/network/tests/rpc_tests.rs b/node/network/tests/rpc_tests.rs
index 6c7feba..cb14c35 100644
--- a/node/network/tests/rpc_tests.rs
+++ b/node/network/tests/rpc_tests.rs
@@ -23,14 +23,10 @@ fn test_status_rpc() {
         let (mut sender, mut receiver) = common::build_node_pair(Arc::downgrade(&rt)).await;
 
         // Dummy STATUS RPC message
-        let rpc_request = Request::Status(StatusMessage {
-            data: Default::default(),
-        });
+        let rpc_request = Request::Status(Default::default());
 
         // Dummy STATUS RPC message
-        let rpc_response = Response::Status(StatusMessage {
-            data: Default::default(),
-        });
+        let rpc_response = Response::Status(Default::default());
 
         // build the sender future
         let sender_future = async {
diff --git a/node/router/src/libp2p_event_handler.rs b/node/router/src/libp2p_event_handler.rs
index 9a69dee..6458fed 100644
--- a/node/router/src/libp2p_event_handler.rs
+++ b/node/router/src/libp2p_event_handler.rs
@@ -172,8 +172,11 @@ impl Libp2pEventHandler {
     }
 
     pub fn send_status(&self, peer_id: PeerId) {
+        let shard_config = self.store.get_store().get_shard_config();
         let status_message = StatusMessage {
             data: self.network_globals.network_id(),
+            num_shard: shard_config.num_shard,
+            shard_id: shard_config.shard_id,
         };
         debug!(%peer_id, ?status_message, "Sending Status request");
 
@@ -254,8 +257,11 @@ impl Libp2pEventHandler {
         debug!(%peer_id, ?status, "Received Status request");
 
         let network_id = self.network_globals.network_id();
+        let shard_config = self.store.get_store().get_shard_config();
         let status_message = StatusMessage {
             data: network_id.clone(),
+            num_shard: shard_config.num_shard,
+            shard_id: shard_config.shard_id,
         };
         debug!(%peer_id, ?status_message, "Sending Status response");
 
@@ -1216,6 +1222,8 @@ mod tests {
         let req_id = (ConnectionId::new(4), SubstreamId(12));
         let request = Request::Status(StatusMessage {
             data: Default::default(),
+            num_shard: 1,
+            shard_id: 0,
         });
         handler.on_rpc_request(alice, req_id, request).await;