use libp2p::PeerId;
use log::trace;
use parking_lot::Mutex;
use partial_sort::PartialSort;
use sc_network_common::types::ReputationChange;
use std::{
cmp::{Ord, Ordering, PartialOrd},
collections::{hash_map::Entry, HashMap, HashSet},
fmt::Debug,
sync::Arc,
time::{Duration, Instant},
};
use wasm_timer::Delay;
use crate::protocol_controller::ProtocolHandle;
pub const LOG_TARGET: &str = "peerset";
pub const BANNED_THRESHOLD: i32 = 82 * (i32::MIN / 100);
const DISCONNECT_REPUTATION_CHANGE: i32 = -256;
const INVERSE_DECREMENT: i32 = 50;
const FORGET_AFTER: Duration = Duration::from_secs(3600);
pub trait PeerStoreProvider: Debug + Send {
fn is_banned(&self, peer_id: &PeerId) -> bool;
fn register_protocol(&self, protocol_handle: ProtocolHandle);
fn report_disconnect(&mut self, peer_id: PeerId);
fn report_peer(&mut self, peer_id: PeerId, change: ReputationChange);
fn peer_reputation(&self, peer_id: &PeerId) -> i32;
fn outgoing_candidates(&self, count: usize, ignored: HashSet<&PeerId>) -> Vec<PeerId>;
}
#[derive(Debug, Clone)]
pub struct PeerStoreHandle {
inner: Arc<Mutex<PeerStoreInner>>,
}
impl PeerStoreProvider for PeerStoreHandle {
fn is_banned(&self, peer_id: &PeerId) -> bool {
self.inner.lock().is_banned(peer_id)
}
fn register_protocol(&self, protocol_handle: ProtocolHandle) {
self.inner.lock().register_protocol(protocol_handle);
}
fn report_disconnect(&mut self, peer_id: PeerId) {
self.inner.lock().report_disconnect(peer_id)
}
fn report_peer(&mut self, peer_id: PeerId, change: ReputationChange) {
self.inner.lock().report_peer(peer_id, change)
}
fn peer_reputation(&self, peer_id: &PeerId) -> i32 {
self.inner.lock().peer_reputation(peer_id)
}
fn outgoing_candidates(&self, count: usize, ignored: HashSet<&PeerId>) -> Vec<PeerId> {
self.inner.lock().outgoing_candidates(count, ignored)
}
}
impl PeerStoreHandle {
pub fn num_known_peers(&self) -> usize {
self.inner.lock().peers.len()
}
pub fn add_known_peer(&mut self, peer_id: PeerId) {
self.inner.lock().add_known_peer(peer_id);
}
}
#[derive(Debug, Clone, Copy)]
struct PeerInfo {
reputation: i32,
last_updated: Instant,
}
impl Default for PeerInfo {
fn default() -> Self {
Self { reputation: 0, last_updated: Instant::now() }
}
}
impl PartialEq for PeerInfo {
fn eq(&self, other: &Self) -> bool {
self.reputation == other.reputation
}
}
impl Eq for PeerInfo {}
impl Ord for PeerInfo {
fn cmp(&self, other: &Self) -> Ordering {
self.reputation.cmp(&other.reputation).reverse()
}
}
impl PartialOrd for PeerInfo {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl PeerInfo {
fn is_banned(&self) -> bool {
self.reputation < BANNED_THRESHOLD
}
fn add_reputation(&mut self, increment: i32) {
self.reputation = self.reputation.saturating_add(increment);
self.bump_last_updated();
}
fn decay_reputation(&mut self, seconds_passed: u64) {
for _ in 0..seconds_passed {
let mut diff = self.reputation / INVERSE_DECREMENT;
if diff == 0 && self.reputation < 0 {
diff = -1;
} else if diff == 0 && self.reputation > 0 {
diff = 1;
}
self.reputation = self.reputation.saturating_sub(diff);
if self.reputation == 0 {
break
}
}
}
fn bump_last_updated(&mut self) {
self.last_updated = Instant::now();
}
}
#[derive(Debug)]
struct PeerStoreInner {
peers: HashMap<PeerId, PeerInfo>,
protocols: Vec<ProtocolHandle>,
}
impl PeerStoreInner {
fn is_banned(&self, peer_id: &PeerId) -> bool {
self.peers.get(peer_id).map_or(false, |info| info.is_banned())
}
fn register_protocol(&mut self, protocol_handle: ProtocolHandle) {
self.protocols.push(protocol_handle);
}
fn report_disconnect(&mut self, peer_id: PeerId) {
let peer_info = self.peers.entry(peer_id).or_default();
peer_info.add_reputation(DISCONNECT_REPUTATION_CHANGE);
log::trace!(
target: LOG_TARGET,
"Peer {} disconnected, reputation: {:+} to {}",
peer_id,
DISCONNECT_REPUTATION_CHANGE,
peer_info.reputation,
);
}
fn report_peer(&mut self, peer_id: PeerId, change: ReputationChange) {
let peer_info = self.peers.entry(peer_id).or_default();
peer_info.add_reputation(change.value);
if peer_info.reputation < BANNED_THRESHOLD {
self.protocols.iter().for_each(|handle| handle.disconnect_peer(peer_id));
log::trace!(
target: LOG_TARGET,
"Report {}: {:+} to {}. Reason: {}. Banned, disconnecting.",
peer_id,
change.value,
peer_info.reputation,
change.reason,
);
} else {
log::trace!(
target: LOG_TARGET,
"Report {}: {:+} to {}. Reason: {}.",
peer_id,
change.value,
peer_info.reputation,
change.reason,
);
}
}
fn peer_reputation(&self, peer_id: &PeerId) -> i32 {
self.peers.get(peer_id).map_or(0, |info| info.reputation)
}
fn outgoing_candidates(&self, count: usize, ignored: HashSet<&PeerId>) -> Vec<PeerId> {
let mut candidates = self
.peers
.iter()
.filter_map(|(peer_id, info)| {
(!info.is_banned() && !ignored.contains(peer_id)).then_some((*peer_id, *info))
})
.collect::<Vec<_>>();
let count = std::cmp::min(count, candidates.len());
candidates.partial_sort(count, |(_, info1), (_, info2)| info1.cmp(info2));
candidates.iter().take(count).map(|(peer_id, _)| *peer_id).collect()
}
fn progress_time(&mut self, seconds_passed: u64) {
if seconds_passed == 0 {
return
}
self.peers
.iter_mut()
.for_each(|(_, info)| info.decay_reputation(seconds_passed));
let now = Instant::now();
self.peers
.retain(|_, info| info.reputation != 0 || info.last_updated + FORGET_AFTER > now);
}
fn add_known_peer(&mut self, peer_id: PeerId) {
match self.peers.entry(peer_id) {
Entry::Occupied(mut e) => {
trace!(
target: LOG_TARGET,
"Trying to add an already known peer {peer_id}, bumping `last_updated`.",
);
e.get_mut().bump_last_updated();
},
Entry::Vacant(e) => {
trace!(target: LOG_TARGET, "Adding a new known peer {peer_id}.");
e.insert(PeerInfo::default());
},
}
}
}
#[derive(Debug)]
pub struct PeerStore {
inner: Arc<Mutex<PeerStoreInner>>,
}
impl PeerStore {
pub fn new(bootnodes: Vec<PeerId>) -> Self {
PeerStore {
inner: Arc::new(Mutex::new(PeerStoreInner {
peers: bootnodes
.into_iter()
.map(|peer_id| (peer_id, PeerInfo::default()))
.collect(),
protocols: Vec::new(),
})),
}
}
pub fn handle(&self) -> PeerStoreHandle {
PeerStoreHandle { inner: self.inner.clone() }
}
pub async fn run(self) {
let started = Instant::now();
let mut latest_time_update = started;
loop {
let now = Instant::now();
let seconds_passed = {
let elapsed_latest = latest_time_update - started;
let elapsed_now = now - started;
latest_time_update = now;
elapsed_now.as_secs() - elapsed_latest.as_secs()
};
self.inner.lock().progress_time(seconds_passed);
let _ = Delay::new(Duration::from_secs(1)).await;
}
}
}
#[cfg(test)]
mod tests {
use super::PeerInfo;
#[test]
fn decaying_zero_reputation_yields_zero() {
let mut peer_info = PeerInfo::default();
assert_eq!(peer_info.reputation, 0);
peer_info.decay_reputation(1);
assert_eq!(peer_info.reputation, 0);
peer_info.decay_reputation(100_000);
assert_eq!(peer_info.reputation, 0);
}
#[test]
fn decaying_positive_reputation_decreases_it() {
const INITIAL_REPUTATION: i32 = 100;
let mut peer_info = PeerInfo::default();
peer_info.reputation = INITIAL_REPUTATION;
peer_info.decay_reputation(1);
assert!(peer_info.reputation >= 0);
assert!(peer_info.reputation < INITIAL_REPUTATION);
}
#[test]
fn decaying_negative_reputation_increases_it() {
const INITIAL_REPUTATION: i32 = -100;
let mut peer_info = PeerInfo::default();
peer_info.reputation = INITIAL_REPUTATION;
peer_info.decay_reputation(1);
assert!(peer_info.reputation <= 0);
assert!(peer_info.reputation > INITIAL_REPUTATION);
}
#[test]
fn decaying_max_reputation_finally_yields_zero() {
const INITIAL_REPUTATION: i32 = i32::MAX;
const SECONDS: u64 = 1000;
let mut peer_info = PeerInfo::default();
peer_info.reputation = INITIAL_REPUTATION;
peer_info.decay_reputation(SECONDS / 2);
assert!(peer_info.reputation > 0);
peer_info.decay_reputation(SECONDS / 2);
assert_eq!(peer_info.reputation, 0);
}
#[test]
fn decaying_min_reputation_finally_yields_zero() {
const INITIAL_REPUTATION: i32 = i32::MIN;
const SECONDS: u64 = 1000;
let mut peer_info = PeerInfo::default();
peer_info.reputation = INITIAL_REPUTATION;
peer_info.decay_reputation(SECONDS / 2);
assert!(peer_info.reputation < 0);
peer_info.decay_reputation(SECONDS / 2);
assert_eq!(peer_info.reputation, 0);
}
}