1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::time::{SystemTime, UNIX_EPOCH};
4
5use rand::seq::SliceRandom;
6use ruc::*;
7use serde::{Deserialize, Serialize};
8
9use hotmint_types::validator::ValidatorId;
10use litep2p::PeerId;
11use litep2p::types::multiaddr::Multiaddr;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
15pub enum PeerRole {
16 Validator,
18 Fullnode,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct PeerInfo {
25 pub peer_id: String,
26 pub role: PeerRole,
27 pub validator_id: Option<u64>,
28 pub addresses: Vec<String>,
29 pub last_seen: u64,
30 pub score: i32,
31}
32
33const BAN_THRESHOLD: i32 = -100;
34
35impl PeerInfo {
36 pub fn new(peer_id: PeerId, role: PeerRole, addresses: Vec<Multiaddr>) -> Self {
37 Self {
38 peer_id: peer_id.to_string(),
39 role,
40 validator_id: None,
41 addresses: addresses.iter().map(|a| a.to_string()).collect(),
42 last_seen: now_secs(),
43 score: 0,
44 }
45 }
46
47 pub fn with_validator(mut self, vid: ValidatorId) -> Self {
48 self.validator_id = Some(vid.0);
49 self
50 }
51
52 pub fn is_banned(&self) -> bool {
53 self.score <= BAN_THRESHOLD
54 }
55
56 pub fn touch(&mut self) {
57 self.last_seen = now_secs();
58 }
59}
60
61pub struct PeerBook {
63 peers: HashMap<String, PeerInfo>,
64 path: PathBuf,
65}
66
67impl PeerBook {
68 pub fn new(path: impl AsRef<Path>) -> Self {
69 Self {
70 peers: HashMap::new(),
71 path: path.as_ref().to_path_buf(),
72 }
73 }
74
75 pub fn load(path: impl AsRef<Path>) -> Result<Self> {
76 let p = path.as_ref();
77 if !p.exists() {
78 return Ok(Self::new(p));
79 }
80 let contents = std::fs::read_to_string(p).c(d!("read peer book"))?;
81 let peers: HashMap<String, PeerInfo> =
82 serde_json::from_str(&contents).c(d!("parse peer book"))?;
83 Ok(Self {
84 peers,
85 path: p.to_path_buf(),
86 })
87 }
88
89 pub fn save(&self) -> Result<()> {
90 let contents = serde_json::to_string_pretty(&self.peers).c(d!("serialize peer book"))?;
91 std::fs::write(&self.path, contents).c(d!("write peer book"))
92 }
93
94 pub fn add_peer(&mut self, info: PeerInfo) {
95 self.peers.insert(info.peer_id.clone(), info);
96 }
97
98 pub fn remove_peer(&mut self, peer_id: &str) {
99 self.peers.remove(peer_id);
100 }
101
102 pub fn get(&self, peer_id: &str) -> Option<&PeerInfo> {
103 self.peers.get(peer_id)
104 }
105
106 pub fn get_mut(&mut self, peer_id: &str) -> Option<&mut PeerInfo> {
107 self.peers.get_mut(peer_id)
108 }
109
110 pub fn get_peers_by_role(&self, role: PeerRole) -> Vec<&PeerInfo> {
111 self.peers
112 .values()
113 .filter(|p| p.role == role && !p.is_banned())
114 .collect()
115 }
116
117 pub fn get_random_peers(&self, n: usize) -> Vec<&PeerInfo> {
118 let mut candidates: Vec<&PeerInfo> =
119 self.peers.values().filter(|p| !p.is_banned()).collect();
120 let mut rng = rand::thread_rng();
121 candidates.shuffle(&mut rng);
122 candidates.truncate(n);
123 candidates
124 }
125
126 pub fn adjust_score(&mut self, peer_id: &str, delta: i32) {
127 if let Some(peer) = self.peers.get_mut(peer_id) {
128 peer.score = peer.score.saturating_add(delta);
129 }
130 }
131
132 pub fn prune_stale(&mut self, max_age_secs: u64) {
133 let cutoff = now_secs().saturating_sub(max_age_secs);
134 self.peers
135 .retain(|_, p| p.last_seen >= cutoff || p.role == PeerRole::Validator);
136 }
137
138 pub fn len(&self) -> usize {
139 self.peers.len()
140 }
141
142 pub fn is_empty(&self) -> bool {
143 self.peers.is_empty()
144 }
145
146 pub fn all_peers(&self) -> impl Iterator<Item = &PeerInfo> {
147 self.peers.values()
148 }
149}
150
151fn now_secs() -> u64 {
152 SystemTime::now()
153 .duration_since(UNIX_EPOCH)
154 .unwrap_or_default()
155 .as_secs()
156}