use crate::traceroute::AsnInfo;
use ip_network::Ipv4Network;
use ip_network_table::IpNetworkTable;
use std::sync::{Arc, RwLock};
pub struct AsnCache {
cache: Arc<RwLock<IpNetworkTable<AsnInfo>>>,
}
impl std::fmt::Debug for AsnCache {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AsnCache").finish()
}
}
impl Default for AsnCache {
fn default() -> Self {
Self::new()
}
}
impl AsnCache {
pub fn new() -> Self {
Self {
cache: Arc::new(RwLock::new(IpNetworkTable::new())),
}
}
pub fn get(&self, ip: &std::net::Ipv4Addr) -> Option<AsnInfo> {
let cache = self.cache.read().expect("rwlock poisoned");
cache
.longest_match(*ip)
.map(|(_, asn_info)| asn_info.clone())
}
pub fn insert(&self, prefix: Ipv4Network, asn_info: AsnInfo) {
let mut cache = self.cache.write().expect("rwlock poisoned");
cache.insert(prefix, asn_info);
}
pub fn len(&self) -> usize {
let cache = self.cache.read().expect("rwlock poisoned");
let (ipv4_len, ipv6_len) = cache.len();
ipv4_len + ipv6_len
}
pub fn is_empty(&self) -> bool {
let cache = self.cache.read().expect("rwlock poisoned");
cache.is_empty()
}
pub fn clear(&self) {
let mut cache = self.cache.write().expect("rwlock poisoned");
*cache = IpNetworkTable::new();
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::Ipv4Addr;
#[test]
fn test_asn_cache() {
let cache = AsnCache::new();
assert!(cache.is_empty());
let asn_info = AsnInfo {
asn: 13335,
prefix: "104.16.0.0/12".to_string(),
country_code: "US".to_string(),
registry: "ARIN".to_string(),
name: "CLOUDFLARENET".to_string(),
};
let prefix: Ipv4Network = "104.16.0.0/12".parse().unwrap();
cache.insert(prefix, asn_info.clone());
assert_eq!(cache.len(), 1);
assert!(!cache.is_empty());
let ip: Ipv4Addr = "104.16.1.1".parse().unwrap();
let result = cache.get(&ip);
assert!(result.is_some());
assert_eq!(result.unwrap().asn, 13335);
let ip: Ipv4Addr = "8.8.8.8".parse().unwrap();
let result = cache.get(&ip);
assert!(result.is_none());
cache.clear();
assert!(cache.is_empty());
assert_eq!(cache.len(), 0);
}
#[test]
fn test_overlapping_prefixes() {
let cache = AsnCache::new();
let specific_prefix = "192.168.1.0/24".parse().unwrap();
let specific_info = AsnInfo {
asn: 1,
prefix: "192.168.1.0/24".to_string(),
country_code: "US".to_string(),
registry: "ARIN".to_string(),
name: "Specific".to_string(),
};
cache.insert(specific_prefix, specific_info);
let broader_prefix = "192.168.0.0/16".parse().unwrap();
let broader_info = AsnInfo {
asn: 2,
prefix: "192.168.0.0/16".to_string(),
country_code: "US".to_string(),
registry: "ARIN".to_string(),
name: "Broader".to_string(),
};
cache.insert(broader_prefix, broader_info);
let ip: Ipv4Addr = "192.168.1.1".parse().unwrap();
let result = cache.get(&ip);
assert!(result.is_some());
assert_eq!(result.unwrap().asn, 1);
}
}