add PeerActor handling for received msg

saketh-are · Nov 2, 2023 · 3144647 · 3144647
1 parent 96a86d1
commit 3144647
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git a/chain/network/src/peer/peer_actor.rs b/chain/network/src/peer/peer_actor.rs
@@ -17,6 +17,7 @@ use crate::private_actix::{RegisterPeerError, SendMessage};
 use crate::routing::edge::verify_nonce;
 use crate::routing::NetworkTopologyChange;
 use crate::shards_manager::ShardsManagerRequestFromNetwork;
+use crate::snapshot_hosts::SnapshotHostInfoError;
 use crate::stats::metrics;
 use crate::tcp;
 use crate::types::{
@@ -1261,6 +1262,26 @@ impl PeerActor {
                     message_processed_event();
                 }));
             }
+            PeerMessage::SyncSnapshotHosts(msg) => {
+                metrics::SYNC_SNAPSHOT_HOSTS.with_label_values(&["received"]).inc();
+                // Early exit, if there is no data in the message.
+                if msg.hosts.is_empty() {
+                    message_processed_event();
+                    return;
+                }
+                let network_state = self.network_state.clone();
+                ctx.spawn(wrap_future(async move {
+                    if let Some(err) = network_state.add_snapshot_hosts(msg.hosts).await {
+                        conn.stop(Some(match err {
+                            SnapshotHostInfoError::InvalidSignature => {
+                                ReasonForBan::InvalidSignature
+                            }
+                            SnapshotHostInfoError::DuplicatePeerId => ReasonForBan::Abusive,
+                        }));
+                    }
+                    message_processed_event();
+                }));
+            }
             PeerMessage::Routed(mut msg) => {
                 tracing::trace!(
                     target: "network",

diff --git a/chain/network/src/stats/metrics.rs b/chain/network/src/stats/metrics.rs
@@ -200,6 +200,14 @@ pub(crate) static SYNC_ACCOUNTS_DATA: Lazy<IntCounterVec> = Lazy::new(|| {
     )
     .unwrap()
 });
+pub(crate) static SYNC_SNAPSHOT_HOSTS: Lazy<IntCounterVec> = Lazy::new(|| {
+    try_create_int_counter_vec(
+        "near_sync_snapshot_hosts",
+        "Number of SyncSnapshotHost messages sent/received",
+        &["direction"],
+    )
+    .unwrap()
+});
 
 pub(crate) static REQUEST_COUNT_BY_TYPE_TOTAL: Lazy<IntCounterVec> = Lazy::new(|| {
     try_create_int_counter_vec(