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}