1use crate::{CandidatePeer, ConnectedPeer, ConnectionMode, NodeType, Peer, Resolver};
17
18use snarkos_node_tcp::{ConnectError, P2P, is_bogon_ip, is_unspecified_or_broadcast_ip};
19use snarkvm::prelude::{Address, Network};
20
21use anyhow::Result;
22#[cfg(feature = "locktick")]
23use locktick::parking_lot::RwLock;
24#[cfg(not(feature = "locktick"))]
25use parking_lot::RwLock;
26use std::{
27 cmp,
28 collections::{
29 HashSet,
30 hash_map::{Entry, HashMap},
31 },
32 fs,
33 io::{self, Write},
34 net::{IpAddr, SocketAddr},
35 path::Path,
36 str::FromStr,
37 time::Instant,
38};
39use tokio::task;
40use tracing::*;
41
42#[derive(Debug)]
45pub enum PeeringError {
46 NoExternalPeersAllowed,
47}
48
49impl snarkos_node_tcp::ApplicationError for PeeringError {}
50
51impl std::fmt::Display for PeeringError {
52 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53 match self {
54 Self::NoExternalPeersAllowed => write!(f, "no untrusted peers allowed"),
55 }
56 }
57}
58
59pub trait PeerPoolHandling<N: Network>: P2P {
60 const OWNER: &str;
61
62 const MAXIMUM_POOL_SIZE: usize;
64
65 const PEER_SLASHING_COUNT: usize;
68
69 fn peer_pool(&self) -> &RwLock<HashMap<SocketAddr, Peer<N>>>;
71
72 fn resolver(&self) -> &RwLock<Resolver<N>>;
74
75 fn is_dev(&self) -> bool;
77
78 fn trusted_peers_only(&self) -> bool;
80
81 fn node_type(&self) -> NodeType;
83
84 fn local_ip(&self) -> SocketAddr {
86 self.tcp().listening_addr().expect("The TCP listener is not enabled")
87 }
88
89 fn is_local_ip(&self, addr: SocketAddr) -> bool {
91 addr == self.local_ip()
92 || (addr.ip().is_unspecified() || addr.ip().is_loopback()) && addr.port() == self.local_ip().port()
93 }
94
95 fn is_valid_peer_ip(&self, ip: SocketAddr) -> bool {
97 !self.is_local_ip(ip) && !is_bogon_ip(ip.ip()) && !is_unspecified_or_broadcast_ip(ip.ip())
98 }
99
100 fn max_connected_peers(&self) -> usize {
102 self.tcp().config().max_connections as usize
103 }
104
105 fn check_connection_attempt(&self, listener_addr: SocketAddr) -> Result<(), ConnectError> {
107 if self.is_local_ip(listener_addr) {
109 return Err(ConnectError::SelfConnect { address: listener_addr });
110 }
111 if self.number_of_connected_peers() >= self.max_connected_peers() {
113 return Err(ConnectError::MaximumConnectionsReached { limit: self.max_connected_peers() as u16 });
114 }
115 if self.is_connected(listener_addr) {
117 return Err(ConnectError::AlreadyConnected { address: listener_addr });
118 }
119 if self.is_connecting(listener_addr) {
121 return Err(ConnectError::AlreadyConnecting { address: listener_addr });
122 }
123 if self.is_ip_banned(listener_addr.ip()) {
125 return Err(ConnectError::BannedIp { ip: listener_addr.ip() });
126 }
127 if self.trusted_peers_only() && !self.is_trusted(listener_addr) {
129 return Err(ConnectError::application(PeeringError::NoExternalPeersAllowed));
130 }
131
132 Ok(())
133 }
134
135 fn connect(&self, listener_addr: SocketAddr) -> Result<task::JoinHandle<Result<(), ConnectError>>, ConnectError> {
143 self.check_connection_attempt(listener_addr)?;
145
146 if let Some(Peer::Candidate(peer)) = self.peer_pool().write().get_mut(&listener_addr) {
148 peer.last_connection_attempt = Some(Instant::now());
149 peer.total_connection_attempts += 1;
150 } else {
151 warn!("{} No candidate peer entry exists for '{listener_addr:?}' while connecting.", Self::OWNER);
152 }
153
154 let tcp = self.tcp().clone();
155 Ok(tokio::spawn(async move {
156 debug!("{} Connecting to {listener_addr}...", Self::OWNER);
157 tcp.connect(listener_addr).await
158 }))
159 }
160
161 fn disconnect(&self, listener_addr: SocketAddr) -> task::JoinHandle<bool> {
164 if let Some(connected_addr) = self.resolve_to_ambiguous(listener_addr) {
165 let tcp = self.tcp().clone();
166 tokio::spawn(async move { tcp.disconnect(connected_addr).await })
167 } else {
168 tokio::spawn(async { false })
169 }
170 }
171
172 fn downgrade_peer_to_candidate(&self, listener_addr: SocketAddr) -> bool {
176 let mut peer_pool = self.peer_pool().write();
177 let Some(peer) = peer_pool.get_mut(&listener_addr) else {
178 trace!("{} Downgrade peer to candidate failed - peer not found", Self::OWNER);
179 return false;
180 };
181
182 if let Peer::Connected(conn_peer) = peer {
183 let aleo_addr = if self.node_type() == NodeType::BootstrapClient
188 && conn_peer.connection_mode == ConnectionMode::Router
189 {
190 None
191 } else {
192 Some(conn_peer.aleo_addr)
193 };
194 self.resolver().write().remove_peer(conn_peer.connected_addr, aleo_addr);
195 peer.downgrade_to_candidate(listener_addr);
196 true
197 } else {
198 peer.downgrade_to_candidate(listener_addr);
199 false
200 }
201 }
202
203 fn insert_candidate_peers(&self, mut listener_addrs: Vec<(SocketAddr, Option<u32>)>) {
207 let trusted_peers = self.trusted_peers();
208
209 let mut peer_pool = self.peer_pool().write();
212
213 let mut num_updates: usize = 0;
215 listener_addrs.retain(|&(addr, height)| {
216 !self.is_ip_banned(addr.ip())
217 && if self.is_dev() { !is_bogon_ip(addr.ip()) } else { self.is_valid_peer_ip(addr) }
218 && peer_pool
219 .get(&addr)
220 .map(|peer| peer.is_candidate() && height.is_some())
221 .inspect(|is_valid_update| {
222 if *is_valid_update {
223 num_updates += 1
224 }
225 })
226 .unwrap_or(true)
227 });
228
229 if listener_addrs.is_empty() {
231 return;
232 }
233
234 if peer_pool.len() + listener_addrs.len() - num_updates >= Self::MAXIMUM_POOL_SIZE
236 && Self::PEER_SLASHING_COUNT != 0
237 {
238 let now = Instant::now();
240 let mut peers_to_slash = peer_pool
241 .iter()
242 .filter_map(|(addr, peer)| match peer {
243 Peer::Candidate(candidate) if !trusted_peers.contains(addr) => {
244 Some((*addr, candidate.total_connection_attempts, candidate.last_connection_attempt))
245 }
246 _ => None,
247 })
248 .collect::<Vec<_>>();
249
250 peers_to_slash.sort_unstable_by_key(|(_, attempts, last_attempt)| {
252 (cmp::Reverse(*attempts), last_attempt.unwrap_or(now))
253 });
254
255 peers_to_slash.truncate(Self::PEER_SLASHING_COUNT);
257
258 peer_pool.retain(|addr, _| !peers_to_slash.iter().any(|(slashed, _, _)| slashed == addr));
260 }
261
262 listener_addrs.truncate(Self::MAXIMUM_POOL_SIZE.saturating_sub(peer_pool.len()));
264
265 if listener_addrs.is_empty() {
267 return;
268 }
269
270 for (addr, height) in listener_addrs {
272 match peer_pool.entry(addr) {
273 Entry::Vacant(entry) => {
274 entry.insert(Peer::new_candidate(addr, false));
275 }
276 Entry::Occupied(mut entry) => {
277 if let Peer::Candidate(peer) = entry.get_mut() {
278 peer.last_height_seen = height;
279 }
280 }
281 }
282 }
283 }
284
285 fn remove_peer(&self, listener_addr: SocketAddr) {
287 self.peer_pool().write().remove(&listener_addr);
288 }
289
290 fn resolve_to_ambiguous(&self, listener_addr: SocketAddr) -> Option<SocketAddr> {
292 if let Some(Peer::Connected(peer)) = self.peer_pool().read().get(&listener_addr) {
293 Some(peer.connected_addr)
294 } else {
295 None
296 }
297 }
298
299 fn resolve_to_aleo_addr(&self, listener_addr: SocketAddr) -> Option<Address<N>> {
301 if let Some(Peer::Connected(peer)) = self.peer_pool().read().get(&listener_addr) {
302 Some(peer.aleo_addr)
303 } else {
304 None
305 }
306 }
307
308 fn is_connecting(&self, listener_addr: SocketAddr) -> bool {
310 self.peer_pool().read().get(&listener_addr).is_some_and(|peer| peer.is_connecting())
311 }
312
313 fn is_connected(&self, listener_addr: SocketAddr) -> bool {
315 self.peer_pool().read().get(&listener_addr).is_some_and(|peer| peer.is_connected())
316 }
317
318 fn is_connected_address(&self, aleo_address: Address<N>) -> bool {
320 self.resolver().read().get_peer_ip_for_address(aleo_address).is_some()
322 }
323
324 fn is_connecting_or_connected(&self, listener_addr: SocketAddr) -> bool {
326 self.peer_pool().read().get(&listener_addr).is_some_and(|peer| peer.is_connecting() || peer.is_connected())
327 }
328
329 fn is_trusted(&self, listener_addr: SocketAddr) -> bool {
331 self.peer_pool().read().get(&listener_addr).is_some_and(|peer| peer.is_trusted())
332 }
333
334 fn number_of_peers(&self) -> usize {
336 self.peer_pool().read().len()
337 }
338
339 fn number_of_connected_peers(&self) -> usize {
341 self.peer_pool().read().values().filter(|peer| peer.is_connected()).count()
342 }
343
344 fn number_of_connected_validators(&self) -> Option<usize> {
349 Some(
350 self.peer_pool()
351 .try_read()?
352 .values()
353 .filter(|peer| peer.as_connected().is_some_and(|peer| peer.is_validator()))
354 .count(),
355 )
356 }
357
358 fn number_of_connecting_peers(&self) -> Option<usize> {
363 Some(self.peer_pool().try_read()?.values().filter(|peer| peer.is_connecting()).count())
364 }
365
366 fn number_of_candidate_peers(&self) -> usize {
368 self.peer_pool().read().values().filter(|peer| matches!(peer, Peer::Candidate(_))).count()
369 }
370
371 fn get_connected_peer(&self, listener_addr: SocketAddr) -> Option<ConnectedPeer<N>> {
373 if let Some(Peer::Connected(peer)) = self.peer_pool().read().get(&listener_addr) {
374 Some(peer.clone())
375 } else {
376 None
377 }
378 }
379
380 fn update_connected_peer<F: FnMut(&mut ConnectedPeer<N>)>(
383 &self,
384 listener_addr: &SocketAddr,
385 mut update_fn: F,
386 ) -> bool {
387 if let Some(Peer::Connected(peer)) = self.peer_pool().write().get_mut(listener_addr) {
388 update_fn(peer);
389 true
390 } else {
391 false
392 }
393 }
394
395 fn get_peers(&self) -> Vec<Peer<N>> {
397 self.peer_pool().read().values().cloned().collect()
398 }
399
400 fn get_connected_peers(&self) -> Vec<ConnectedPeer<N>> {
402 self.filter_connected_peers(|_| true)
403 }
404
405 fn get_best_connected_peers(&self, max_entries: Option<usize>) -> Vec<ConnectedPeer<N>> {
408 let mut peers = self.get_connected_peers();
410
411 peers.sort_unstable_by_key(|peer| cmp::Reverse(peer.last_height_seen));
413 if let Some(max) = max_entries {
414 peers.truncate(max);
415 }
416
417 peers
418 }
419
420 fn filter_connected_peers<P: FnMut(&ConnectedPeer<N>) -> bool>(&self, mut predicate: P) -> Vec<ConnectedPeer<N>> {
422 self.peer_pool()
423 .read()
424 .values()
425 .filter_map(|p| {
426 if let Peer::Connected(peer) = p
427 && predicate(peer)
428 {
429 Some(peer)
430 } else {
431 None
432 }
433 })
434 .cloned()
435 .collect()
436 }
437
438 fn connected_peers(&self) -> Vec<SocketAddr> {
440 self.peer_pool().read().iter().filter_map(|(addr, peer)| peer.is_connected().then_some(*addr)).collect()
441 }
442
443 fn trusted_peers(&self) -> Vec<SocketAddr> {
445 self.peer_pool().read().iter().filter_map(|(addr, peer)| peer.is_trusted().then_some(*addr)).collect()
446 }
447
448 fn get_candidate_peers(&self) -> Vec<CandidatePeer<N>> {
450 self.peer_pool()
451 .read()
452 .values()
453 .filter_map(|peer| if let Peer::Candidate(peer) = peer { Some(peer.clone()) } else { None })
454 .collect()
455 }
456
457 fn get_trusted_candidate_peers(&self) -> Vec<CandidatePeer<N>> {
459 self.peer_pool()
460 .read()
461 .values()
462 .filter_map(|peer| {
463 if let Peer::Candidate(peer) = peer
464 && peer.trusted
465 {
466 Some(peer.clone())
467 } else {
468 None
469 }
470 })
471 .collect()
472 }
473
474 fn load_cached_peers(path: &Path) -> Result<Vec<SocketAddr>> {
477 let peers = match fs::read_to_string(path) {
478 Ok(cached_peers_str) => {
479 let mut cached_peers = Vec::new();
480 for peer_addr_str in cached_peers_str.lines() {
481 match SocketAddr::from_str(peer_addr_str) {
482 Ok(addr) => cached_peers.push(addr),
483 Err(error) => warn!("Couldn't parse the cached peer address '{peer_addr_str}': {error}"),
484 }
485 }
486 cached_peers
487 }
488 Err(error) if error.kind() == io::ErrorKind::NotFound => {
489 Vec::new()
491 }
492 Err(error) => {
493 warn!("{} Couldn't load cached peers at {}: {error}", Self::OWNER, path.display());
494 Vec::new()
495 }
496 };
497
498 Ok(peers)
499 }
500
501 fn save_best_peers(&self, path: &Path, max_entries: Option<usize>, store_ports: bool) -> Result<()> {
509 let mut peers = self.get_peers();
511
512 peers.sort_unstable_by_key(|peer| (cmp::Reverse(peer.last_height_seen()), peer.failed_connection_attempts()));
514 if let Some(max) = max_entries {
515 peers.truncate(max);
516 }
517
518 let addrs: HashSet<_> = peers
520 .iter()
521 .map(
522 |peer| {
523 if store_ports { peer.listener_addr().to_string() } else { peer.listener_addr().ip().to_string() }
524 },
525 )
526 .collect();
527
528 let mut file = fs::File::create(path)?;
529 for addr in addrs {
530 writeln!(file, "{addr}")?;
531 }
532
533 Ok(())
534 }
535
536 fn add_connecting_peer(&self, listener_addr: SocketAddr) -> Result<(), ConnectError> {
541 match self.peer_pool().write().entry(listener_addr) {
542 Entry::Vacant(entry) => {
543 entry.insert(Peer::new_connecting(listener_addr, false));
544 Ok(())
545 }
546 Entry::Occupied(mut entry) => match entry.get() {
547 Peer::Candidate(_) => {
550 entry.get_mut().promote_to_connecting();
551 Ok(())
552 }
553 Peer::Connecting(_) => Err(ConnectError::AlreadyConnecting { address: listener_addr }),
554 Peer::Connected(_) => Err(ConnectError::AlreadyConnected { address: listener_addr }),
555 },
556 }
557 }
558
559 fn ip_ban_peer(&self, listener_addr: SocketAddr, reason: Option<&str>) {
562 if self.is_dev() {
564 return;
565 }
566
567 let ip = listener_addr.ip();
568 debug!("IP-banning {ip}{}", reason.map(|r| format!(" reason: {r}")).unwrap_or_default());
569
570 self.tcp().banned_peers().update_ip_ban(ip);
572
573 self.disconnect(listener_addr);
575 self.remove_peer(listener_addr);
577 }
578
579 fn is_ip_banned(&self, ip: IpAddr) -> bool {
581 self.tcp().banned_peers().is_ip_banned(&ip)
582 }
583
584 fn update_ip_ban(&self, ip: IpAddr) {
586 self.tcp().banned_peers().update_ip_ban(ip);
587 }
588}
589
590#[cfg(test)]
591mod tests {
592 use super::*;
593 use crate::Peer;
594 use snarkos_node_tcp::{Config, P2P, Tcp};
595 use snarkvm::{prelude::Rng, utilities::TestRng};
596
597 use std::{collections::HashMap, net::SocketAddr, time::Instant};
598
599 type CurrentNetwork = snarkvm::prelude::MainnetV0;
600
601 struct MockPeerPool<N: Network> {
602 tcp: Tcp,
603 peer_pool: RwLock<HashMap<SocketAddr, Peer<N>>>,
604 resolver: RwLock<Resolver<N>>,
605 }
606
607 impl<N: Network> MockPeerPool<N> {
608 fn new() -> Self {
609 let config = Config { listener_ip: None, ..Default::default() };
610 Self { tcp: Tcp::new(config), peer_pool: Default::default(), resolver: Default::default() }
611 }
612 }
613
614 impl<N: Network> P2P for MockPeerPool<N> {
615 fn tcp(&self) -> &Tcp {
616 &self.tcp
617 }
618 }
619
620 impl<N: Network> PeerPoolHandling<N> for MockPeerPool<N> {
621 const MAXIMUM_POOL_SIZE: usize = 100;
622 const OWNER: &str = "MockPeerPool";
623 const PEER_SLASHING_COUNT: usize = 10;
624
625 fn peer_pool(&self) -> &RwLock<HashMap<SocketAddr, Peer<N>>> {
626 &self.peer_pool
627 }
628
629 fn resolver(&self) -> &RwLock<Resolver<N>> {
630 &self.resolver
631 }
632
633 fn is_dev(&self) -> bool {
634 false
635 }
636
637 fn trusted_peers_only(&self) -> bool {
638 false
639 }
640
641 fn node_type(&self) -> NodeType {
642 NodeType::Client
643 }
644 }
645
646 fn make_connected_peer(port: u16, node_type: NodeType, rng: &mut TestRng) -> (SocketAddr, Peer<CurrentNetwork>) {
647 use snarkvm::prelude::Address;
648 let listener_addr = SocketAddr::from(([127, 0, 0, 1], port));
649 let connected_addr = SocketAddr::from(([127, 0, 0, 1], port + 10000));
650 let now = Instant::now();
651 let peer = Peer::Connected(ConnectedPeer {
652 listener_addr,
653 connected_addr,
654 connection_mode: ConnectionMode::Router,
655 trusted: false,
656 aleo_addr: Address::<CurrentNetwork>::new(rng.random()),
657 node_type,
658 version: 1,
659 snarkos_sha: None,
660 last_height_seen: None,
661 first_seen: now,
662 last_seen: now,
663 });
664 (listener_addr, peer)
665 }
666
667 #[test]
668 fn test_peer_state_transitions() {
669 use snarkvm::prelude::Address;
670
671 let pool = MockPeerPool::<CurrentNetwork>::new();
672 let mut rng = TestRng::default();
673
674 let listener_addr = SocketAddr::from(([192, 0, 2, 1], 4000));
675 let connected_addr = SocketAddr::from(([192, 0, 2, 1], 14000));
676 let aleo_addr = Address::<CurrentNetwork>::new(rng.random());
677
678 pool.peer_pool().write().insert(listener_addr, Peer::new_candidate(listener_addr, false));
680
681 assert_eq!(pool.number_of_candidate_peers(), 1);
682 assert_eq!(pool.number_of_connecting_peers(), Some(0));
683 assert_eq!(pool.number_of_connected_peers(), 0);
684 assert!(!pool.is_connecting(listener_addr));
685 assert!(!pool.is_connected(listener_addr));
686
687 assert!(pool.add_connecting_peer(listener_addr).is_ok());
689
690 assert_eq!(pool.number_of_candidate_peers(), 0);
691 assert_eq!(pool.number_of_connecting_peers(), Some(1));
692 assert_eq!(pool.number_of_connected_peers(), 0);
693 assert!(pool.is_connecting(listener_addr));
694 assert!(!pool.is_connected(listener_addr));
695
696 pool.peer_pool().write().get_mut(&listener_addr).unwrap().upgrade_to_connected(
698 connected_addr,
699 listener_addr.port(),
700 aleo_addr,
701 NodeType::Validator,
702 1,
703 None,
704 ConnectionMode::Router,
705 );
706
707 assert_eq!(pool.number_of_candidate_peers(), 0);
708 assert_eq!(pool.number_of_connecting_peers(), Some(0));
709 assert_eq!(pool.number_of_connected_peers(), 1);
710 assert!(!pool.is_connecting(listener_addr));
711 assert!(pool.is_connected(listener_addr));
712 assert_eq!(pool.number_of_connected_validators(), Some(1));
713
714 let connected = pool.get_connected_peer(listener_addr).expect("peer should be connected");
716 assert_eq!(connected.listener_addr, listener_addr);
717 assert_eq!(connected.connected_addr, connected_addr);
718 assert_eq!(connected.aleo_addr, aleo_addr);
719 assert_eq!(connected.node_type, NodeType::Validator);
720 }
721
722 #[test]
723 fn test_number_of_connected_validators() {
724 let pool = MockPeerPool::<CurrentNetwork>::new();
725 let mut rng = TestRng::default();
726
727 assert_eq!(pool.number_of_connected_validators(), Some(0));
729
730 let (addr1, peer1) = make_connected_peer(3000, NodeType::Validator, &mut rng);
732 let (addr2, peer2) = make_connected_peer(3001, NodeType::Validator, &mut rng);
733 let (addr3, peer3) = make_connected_peer(3002, NodeType::Client, &mut rng);
734 {
735 let mut pool_write = pool.peer_pool().write();
736 pool_write.insert(addr1, peer1);
737 pool_write.insert(addr2, peer2);
738 pool_write.insert(addr3, peer3);
739 }
740
741 assert_eq!(pool.number_of_connected_validators(), Some(2));
742 assert_eq!(pool.number_of_connected_peers(), 3);
743
744 let candidate_addr = SocketAddr::from(([127, 0, 0, 1], 3003));
746 pool.peer_pool().write().insert(candidate_addr, Peer::new_candidate(candidate_addr, false));
747
748 assert_eq!(pool.number_of_connected_validators(), Some(2));
749 assert_eq!(pool.number_of_connected_peers(), 3);
750 }
751}