kaspa-addressmanager 0.13.4

Kaspa address manager
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

use kaspa_database::{
    prelude::{CachePolicy, StoreError, StoreResult},
    prelude::{CachedDbAccess, DirectDbWriter, DB},
    registry::DatabaseStorePrefixes,
};
use kaspa_utils::mem_size::MemSizeEstimator;
use serde::{Deserialize, Serialize};
use std::net::{IpAddr, Ipv6Addr};
use std::{error::Error, fmt::Display, sync::Arc};

#[derive(Clone, Copy, Serialize, Deserialize)]
pub struct ConnectionBanTimestamp(pub u64);

impl MemSizeEstimator for ConnectionBanTimestamp {}

pub trait BannedAddressesStoreReader {
    fn get(&self, address: IpAddr) -> Result<ConnectionBanTimestamp, StoreError>;
}

pub trait BannedAddressesStore: BannedAddressesStoreReader {
    fn set(&mut self, ip: IpAddr, timestamp: ConnectionBanTimestamp) -> StoreResult<()>;
    fn remove(&mut self, ip: IpAddr) -> StoreResult<()>;
}

const IPV6_LEN: usize = 16;
const ADDRESS_KEY_SIZE: usize = IPV6_LEN;

#[derive(Eq, Hash, PartialEq, Debug, Copy, Clone)]
struct AddressKey([u8; ADDRESS_KEY_SIZE]);

impl AsRef<[u8]> for AddressKey {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

impl Display for AddressKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let ip: Ipv6Addr = (*self).into();
        write!(f, "{ip}")
    }
}

impl From<IpAddr> for AddressKey {
    fn from(ip: IpAddr) -> Self {
        Self(match ip {
            IpAddr::V4(ip) => ip.to_ipv6_mapped().octets(),
            IpAddr::V6(ip) => ip.octets(),
        })
    }
}

impl From<AddressKey> for Ipv6Addr {
    fn from(k: AddressKey) -> Self {
        k.0.into()
    }
}

impl From<AddressKey> for IpAddr {
    fn from(k: AddressKey) -> Self {
        let ipv6: Ipv6Addr = k.0.into();
        match ipv6.to_ipv4_mapped() {
            Some(ipv4) => IpAddr::V4(ipv4),
            None => IpAddr::V6(ipv6),
        }
    }
}

#[derive(Clone)]
pub struct DbBannedAddressesStore {
    db: Arc<DB>,
    access: CachedDbAccess<AddressKey, ConnectionBanTimestamp>,
}

impl DbBannedAddressesStore {
    pub fn new(db: Arc<DB>, cache_policy: CachePolicy) -> Self {
        Self { db: Arc::clone(&db), access: CachedDbAccess::new(db, cache_policy, DatabaseStorePrefixes::BannedAddresses.into()) }
    }

    pub fn iterator(&self) -> impl Iterator<Item = Result<(IpAddr, ConnectionBanTimestamp), Box<dyn Error>>> + '_ {
        self.access.iterator().map(|iter_result| match iter_result {
            Ok((key_bytes, connection_ban_timestamp)) => match <[u8; ADDRESS_KEY_SIZE]>::try_from(&key_bytes[..]) {
                Ok(address_key_slice) => {
                    let addr_key = AddressKey(address_key_slice);
                    let address: IpAddr = addr_key.into();
                    Ok((address, connection_ban_timestamp))
                }
                Err(e) => Err(e.into()),
            },
            Err(e) => Err(e),
        })
    }
}

impl BannedAddressesStoreReader for DbBannedAddressesStore {
    fn get(&self, ip: IpAddr) -> Result<ConnectionBanTimestamp, StoreError> {
        self.access.read(ip.into())
    }
}

impl BannedAddressesStore for DbBannedAddressesStore {
    fn set(&mut self, ip: IpAddr, timestamp: ConnectionBanTimestamp) -> StoreResult<()> {
        self.access.write(DirectDbWriter::new(&self.db), ip.into(), timestamp)
    }

    fn remove(&mut self, ip: IpAddr) -> StoreResult<()> {
        self.access.delete(DirectDbWriter::new(&self.db), ip.into())
    }
}