peer_book.rs (13810B)
1 use std::collections::BTreeMap; 2 3 use serde::{Deserialize, Serialize}; 4 5 use super::{ 6 PEER_CLOCK_OFFSET_ACCEPTANCE_MS, PEER_CLOCK_OFFSET_STALE_MS, PEER_MISBEHAVIOR_BAN_MS, 7 PEER_MISBEHAVIOR_BAN_SCORE, now_ms, 8 }; 9 10 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] 11 pub struct PeerBook { 12 peers: BTreeMap<String, PeerInfo>, 13 } 14 15 impl PeerBook { 16 pub fn from_addresses(addresses: Vec<String>) -> Self { 17 let mut book = Self::default(); 18 for address in addresses { 19 book.add_peer(address); 20 } 21 book 22 } 23 24 pub fn add_peer(&mut self, address: impl Into<String>) { 25 let address = address.into(); 26 let peer = self 27 .peers 28 .entry(address.clone()) 29 .or_insert_with(|| PeerInfo::new(address, PeerDirection::Outbound)); 30 if peer.direction != PeerDirection::Outbound { 31 peer.direction = PeerDirection::Outbound; 32 } 33 } 34 35 pub fn add_discovered_peer(&mut self, address: impl Into<String>) { 36 let address = address.into(); 37 let peer = self 38 .peers 39 .entry(address.clone()) 40 .or_insert_with(|| PeerInfo::new(address, PeerDirection::Discovered)); 41 if peer.direction == PeerDirection::Inbound { 42 peer.direction = PeerDirection::Discovered; 43 } 44 } 45 46 pub fn observe_inbound_peer(&mut self, address: impl Into<String>) { 47 let address = address.into(); 48 self.peers 49 .entry(address.clone()) 50 .or_insert_with(|| PeerInfo::new(address, PeerDirection::Inbound)); 51 } 52 53 pub fn replace_peer_address(&mut self, from: &str, to: impl Into<String>) { 54 let to = to.into(); 55 if from == to { 56 if !self.peers.contains_key(from) { 57 self.add_peer(to); 58 } 59 return; 60 } 61 62 let Some(from_peer) = self.peers.remove(from) else { 63 self.add_peer(to); 64 return; 65 }; 66 67 let to_peer = self 68 .peers 69 .entry(to.clone()) 70 .or_insert_with(|| PeerInfo::new(to, from_peer.direction.clone())); 71 if from_peer.direction == PeerDirection::Outbound { 72 to_peer.direction = PeerDirection::Outbound; 73 } else if from_peer.direction == PeerDirection::Discovered 74 && to_peer.direction == PeerDirection::Inbound 75 { 76 to_peer.direction = PeerDirection::Discovered; 77 } 78 to_peer.messages_sent = to_peer 79 .messages_sent 80 .saturating_add(from_peer.messages_sent); 81 to_peer.messages_received = to_peer 82 .messages_received 83 .saturating_add(from_peer.messages_received); 84 to_peer.last_known_height = to_peer.last_known_height.or(from_peer.last_known_height); 85 to_peer.last_known_tip_hash = to_peer 86 .last_known_tip_hash 87 .clone() 88 .or(from_peer.last_known_tip_hash); 89 if from_peer.last_clock_observed_ms > to_peer.last_clock_observed_ms { 90 to_peer.last_clock_offset_ms = from_peer.last_clock_offset_ms; 91 to_peer.last_clock_offset_accepted = from_peer.last_clock_offset_accepted; 92 to_peer.last_clock_observed_ms = from_peer.last_clock_observed_ms; 93 } 94 to_peer.last_contact_ms = to_peer.last_contact_ms.max(from_peer.last_contact_ms); 95 to_peer.last_success_ms = to_peer.last_success_ms.max(from_peer.last_success_ms); 96 to_peer.last_error_ms = to_peer.last_error_ms.max(from_peer.last_error_ms); 97 if to_peer.last_error.is_none() { 98 to_peer.last_error = from_peer.last_error; 99 } 100 to_peer.misbehavior_score = to_peer 101 .misbehavior_score 102 .saturating_add(from_peer.misbehavior_score); 103 to_peer.banned_until_ms = to_peer.banned_until_ms.max(from_peer.banned_until_ms); 104 if to_peer.ban_reason.is_none() { 105 to_peer.ban_reason = from_peer.ban_reason; 106 } 107 } 108 109 pub fn remove_peer(&mut self, address: &str) -> bool { 110 if self 111 .peers 112 .get(address) 113 .is_some_and(|peer| peer.direction != PeerDirection::Inbound) 114 { 115 self.peers.remove(address); 116 true 117 } else { 118 false 119 } 120 } 121 122 pub fn is_connectable_peer(&self, address: &str) -> bool { 123 self.peers 124 .get(address) 125 .is_some_and(|peer| peer.direction != PeerDirection::Inbound) 126 } 127 128 pub fn addresses(&self) -> Vec<String> { 129 self.peers 130 .values() 131 .filter(|peer| peer.direction != PeerDirection::Inbound) 132 .map(|peer| peer.address.clone()) 133 .collect() 134 } 135 136 pub fn connectable_addresses_at(&self, now_ms: u64) -> Vec<String> { 137 self.peers 138 .values() 139 .filter(|peer| peer.direction != PeerDirection::Inbound) 140 .filter(|peer| !peer.is_banned_at(now_ms)) 141 .map(|peer| peer.address.clone()) 142 .collect() 143 } 144 145 pub fn addresses_except(&self, excluded: &str) -> Vec<String> { 146 self.connectable_addresses_at(now_ms()) 147 .into_iter() 148 .filter(|address| address != excluded) 149 .collect() 150 } 151 152 pub fn list(&self) -> Vec<PeerInfo> { 153 self.peers.values().cloned().collect() 154 } 155 156 pub fn prune_stale_inbound_peers_at(&mut self, now_ms: u64, max_age_ms: u64) -> usize { 157 let before = self.peers.len(); 158 self.peers.retain(|_, peer| { 159 if peer.direction != PeerDirection::Inbound || peer.is_banned_at(now_ms) { 160 return true; 161 } 162 peer.last_contact_ms 163 .is_some_and(|last_contact| now_ms.saturating_sub(last_contact) <= max_age_ms) 164 }); 165 before.saturating_sub(self.peers.len()) 166 } 167 168 pub fn record_sent(&mut self, address: &str, count: u64) { 169 let now = now_ms(); 170 let peer = self.ensure(address, PeerDirection::Outbound); 171 peer.messages_sent += count; 172 peer.last_contact_ms = Some(now); 173 peer.last_success_ms = Some(now); 174 if !peer.is_banned_at(now) { 175 peer.last_error = None; 176 peer.clear_misbehavior(); 177 } 178 } 179 180 pub fn record_status(&mut self, address: &str, height: u64, tip_hash: String) { 181 let now = now_ms(); 182 let peer = self.ensure(address, PeerDirection::Outbound); 183 peer.last_known_height = Some(height); 184 peer.last_known_tip_hash = Some(tip_hash); 185 peer.last_contact_ms = Some(now); 186 peer.last_success_ms = Some(now); 187 if !peer.is_banned_at(now) { 188 peer.last_error = None; 189 peer.clear_misbehavior(); 190 } 191 } 192 193 pub fn record_clock_observation( 194 &mut self, 195 address: &str, 196 direction: PeerDirection, 197 remote_time_ms: u64, 198 local_receive_time_ms: u64, 199 ) { 200 if remote_time_ms == 0 { 201 return; 202 } 203 let offset = remote_time_ms as i128 - local_receive_time_ms as i128; 204 let offset = offset.clamp(i64::MIN as i128, i64::MAX as i128) as i64; 205 let accepted = offset.abs() <= PEER_CLOCK_OFFSET_ACCEPTANCE_MS; 206 let peer = self.ensure(address, direction); 207 peer.last_clock_offset_ms = Some(offset); 208 peer.last_clock_offset_accepted = Some(accepted); 209 peer.last_clock_observed_ms = Some(local_receive_time_ms); 210 } 211 212 pub fn network_time_offset_ms_at(&self, now_ms: u64) -> Option<i64> { 213 median_i64( 214 self.peers 215 .values() 216 .filter(|peer| !peer.is_banned_at(now_ms)) 217 .filter(|peer| peer.last_error.is_none()) 218 .filter(|peer| peer.last_clock_offset_accepted == Some(true)) 219 .filter(|peer| { 220 peer.last_clock_observed_ms.is_some_and(|observed_ms| { 221 now_ms.saturating_sub(observed_ms) <= PEER_CLOCK_OFFSET_STALE_MS 222 }) 223 }) 224 .filter_map(|peer| peer.last_clock_offset_ms) 225 .collect(), 226 ) 227 } 228 229 pub fn adjusted_time_ms_at(&self, now_ms: u64) -> u64 { 230 match self.network_time_offset_ms_at(now_ms) { 231 Some(offset) if offset >= 0 => now_ms.saturating_add(offset as u64), 232 Some(offset) => now_ms.saturating_sub(offset.unsigned_abs()), 233 None => now_ms, 234 } 235 } 236 237 pub fn bad_clock_peer_count_at(&self, now_ms: u64) -> usize { 238 self.peers 239 .values() 240 .filter(|peer| !peer.is_banned_at(now_ms)) 241 .filter(|peer| { 242 peer.last_clock_observed_ms.is_some_and(|observed_ms| { 243 now_ms.saturating_sub(observed_ms) <= PEER_CLOCK_OFFSET_STALE_MS 244 }) 245 }) 246 .filter(|peer| peer.last_clock_offset_accepted == Some(false)) 247 .count() 248 } 249 250 pub fn record_error(&mut self, address: &str, error: impl Into<String>) { 251 let now = now_ms(); 252 let peer = self.ensure(address, PeerDirection::Outbound); 253 peer.last_contact_ms = Some(now); 254 peer.last_error_ms = Some(now); 255 peer.last_error = Some(error.into()); 256 } 257 258 pub fn record_inbound_error(&mut self, address: &str, error: impl Into<String>) { 259 let now = now_ms(); 260 let peer = self.ensure(address, PeerDirection::Inbound); 261 peer.last_contact_ms = Some(now); 262 peer.last_error_ms = Some(now); 263 peer.last_error = Some(error.into()); 264 } 265 266 pub fn record_received(&mut self, address: &str, count: u64) { 267 let now = now_ms(); 268 let peer = self.ensure(address, PeerDirection::Inbound); 269 peer.messages_received += count; 270 peer.last_contact_ms = Some(now); 271 peer.last_success_ms = Some(now); 272 if !peer.is_banned_at(now) { 273 peer.last_error = None; 274 peer.clear_misbehavior(); 275 } 276 } 277 278 pub fn record_misbehavior(&mut self, address: &str, reason: impl Into<String>) { 279 self.record_misbehavior_at(address, reason, now_ms()); 280 } 281 282 pub fn record_misbehavior_at(&mut self, address: &str, reason: impl Into<String>, now_ms: u64) { 283 self.record_misbehavior_with_direction(address, reason, now_ms, PeerDirection::Outbound); 284 } 285 286 pub fn record_inbound_misbehavior(&mut self, address: &str, reason: impl Into<String>) { 287 self.record_misbehavior_with_direction(address, reason, now_ms(), PeerDirection::Inbound); 288 } 289 290 fn record_misbehavior_with_direction( 291 &mut self, 292 address: &str, 293 reason: impl Into<String>, 294 now_ms: u64, 295 direction: PeerDirection, 296 ) { 297 let reason = reason.into(); 298 let peer = self.ensure(address, direction); 299 peer.last_contact_ms = Some(now_ms); 300 peer.last_error_ms = Some(now_ms); 301 peer.last_error = Some(reason.clone()); 302 peer.misbehavior_score = peer.misbehavior_score.saturating_add(1); 303 peer.ban_reason = Some(reason); 304 if peer.misbehavior_score >= PEER_MISBEHAVIOR_BAN_SCORE { 305 peer.banned_until_ms = Some(now_ms.saturating_add(PEER_MISBEHAVIOR_BAN_MS)); 306 } 307 } 308 309 pub fn is_banned(&self, address: &str) -> bool { 310 self.is_banned_at(address, now_ms()) 311 } 312 313 pub fn is_banned_at(&self, address: &str, now_ms: u64) -> bool { 314 self.peers 315 .get(address) 316 .is_some_and(|peer| peer.is_banned_at(now_ms)) 317 } 318 319 fn ensure(&mut self, address: &str, direction: PeerDirection) -> &mut PeerInfo { 320 self.peers 321 .entry(address.to_string()) 322 .or_insert_with(|| PeerInfo::new(address.to_string(), direction)) 323 } 324 } 325 326 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] 327 pub struct PeerInfo { 328 pub address: String, 329 pub direction: PeerDirection, 330 pub messages_sent: u64, 331 pub messages_received: u64, 332 pub last_known_height: Option<u64>, 333 pub last_known_tip_hash: Option<String>, 334 #[serde(default)] 335 pub last_clock_offset_ms: Option<i64>, 336 #[serde(default)] 337 pub last_clock_offset_accepted: Option<bool>, 338 #[serde(default)] 339 pub last_clock_observed_ms: Option<u64>, 340 pub last_error: Option<String>, 341 pub last_contact_ms: Option<u64>, 342 pub last_success_ms: Option<u64>, 343 pub last_error_ms: Option<u64>, 344 pub misbehavior_score: u32, 345 pub banned_until_ms: Option<u64>, 346 pub ban_reason: Option<String>, 347 } 348 349 impl PeerInfo { 350 fn new(address: String, direction: PeerDirection) -> Self { 351 Self { 352 address, 353 direction, 354 messages_sent: 0, 355 messages_received: 0, 356 last_known_height: None, 357 last_known_tip_hash: None, 358 last_clock_offset_ms: None, 359 last_clock_offset_accepted: None, 360 last_clock_observed_ms: None, 361 last_error: None, 362 last_contact_ms: None, 363 last_success_ms: None, 364 last_error_ms: None, 365 misbehavior_score: 0, 366 banned_until_ms: None, 367 ban_reason: None, 368 } 369 } 370 371 pub fn is_banned_at(&self, now_ms: u64) -> bool { 372 self.banned_until_ms 373 .is_some_and(|banned_until| banned_until > now_ms) 374 } 375 376 fn clear_misbehavior(&mut self) { 377 self.misbehavior_score = 0; 378 self.banned_until_ms = None; 379 self.ban_reason = None; 380 } 381 } 382 383 fn median_i64(mut values: Vec<i64>) -> Option<i64> { 384 if values.is_empty() { 385 return None; 386 } 387 values.sort_unstable(); 388 Some(values[values.len() / 2]) 389 } 390 391 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] 392 #[serde(rename_all = "snake_case")] 393 pub enum PeerDirection { 394 Outbound, 395 Discovered, 396 Inbound, 397 }