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 = 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;