snarkos_node_tcp/helpers/
banned_peers.rs

1// Copyright (c) 2019-2025 Provable Inc.
2// This file is part of the snarkOS library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use std::{collections::HashMap, net::IpAddr, time::Instant};
17
18#[cfg(feature = "locktick")]
19use locktick::parking_lot::RwLock;
20#[cfg(not(feature = "locktick"))]
21use parking_lot::RwLock;
22
23/// Contains the ban details for a banned peer.
24#[derive(Clone)]
25pub struct BanDetails {
26    /// The time when the ban was created.
27    banned_at: Instant,
28}
29
30impl BanDetails {
31    /// Creates a new ban at the given time.
32    pub fn new() -> Self {
33        Self { banned_at: Instant::now() }
34    }
35}
36
37impl Default for BanDetails {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43/// Contains the set of peers currently banned by IP.
44#[derive(Default)]
45pub struct BannedPeers(RwLock<HashMap<IpAddr, BanDetails>>);
46
47impl BannedPeers {
48    /// Check whether the given IP address is currently banned.
49    pub fn is_ip_banned(&self, ip: &IpAddr) -> bool {
50        self.0.read().contains_key(ip)
51    }
52
53    /// Get all banned IPs.
54    pub fn get_banned_ips(&self) -> Vec<IpAddr> {
55        self.0.read().keys().cloned().collect()
56    }
57
58    /// Get ban config for the given IP address.
59    pub fn get_ban_config(&self, ip: IpAddr) -> Option<BanDetails> {
60        self.0.read().get(&ip).cloned()
61    }
62
63    /// Insert or update a banned IP.
64    pub fn update_ip_ban(&self, ip: IpAddr) {
65        self.0.write().insert(ip, BanDetails::default());
66    }
67
68    /// Remove the expired entries
69    pub fn remove_old_bans(&self, ban_time_in_secs: u64) {
70        self.0.write().retain(|_, ban_config| ban_config.banned_at.elapsed().as_secs() < ban_time_in_secs);
71    }
72}