use super::cache::AsnCache;
use super::lookup::{lookup_asn_with_cache, AsnLookupError};
use crate::traceroute::AsnInfo;
use std::net::{IpAddr, Ipv4Addr};
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Clone, Debug)]
pub struct AsnLookup {
cache: Arc<RwLock<AsnCache>>,
}
impl AsnLookup {
pub fn new() -> Self {
Self {
cache: Arc::new(RwLock::new(AsnCache::new())),
}
}
pub fn with_cache(cache: AsnCache) -> Self {
Self {
cache: Arc::new(RwLock::new(cache)),
}
}
pub async fn lookup(&self, ip: IpAddr) -> Result<AsnInfo, AsnLookupError> {
match ip {
IpAddr::V4(ipv4) => lookup_asn_with_cache(ipv4, &self.cache).await,
IpAddr::V6(_) => Err(AsnLookupError::NotFound),
}
}
pub async fn lookup_ipv4(&self, ip: Ipv4Addr) -> Result<AsnInfo, AsnLookupError> {
self.lookup(IpAddr::V4(ip)).await
}
pub async fn clear_cache(&self) {
let cache = self.cache.write().await;
cache.clear();
}
pub async fn cache_stats(&self) -> CacheStats {
let cache = self.cache.read().await;
CacheStats {
entries: cache.len(),
is_empty: cache.is_empty(),
}
}
pub async fn is_cached(&self, ip: &Ipv4Addr) -> bool {
let cache = self.cache.read().await;
cache.get(ip).is_some()
}
}
impl Default for AsnLookup {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub entries: usize,
pub is_empty: bool,
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::IpAddr;
#[tokio::test]
async fn test_asn_lookup_service() {
let service = AsnLookup::new();
let private_ip: IpAddr = "192.168.1.1".parse().expect("valid IP");
let result = service.lookup(private_ip).await;
assert!(result.is_ok());
let info = result.expect("should succeed");
assert_eq!(info.asn, 0);
assert_eq!(info.name, "Private Network");
}
#[tokio::test]
async fn test_cache_operations() {
let service = AsnLookup::new();
let stats = service.cache_stats().await;
assert!(stats.is_empty);
let ip: IpAddr = "10.0.0.1".parse().expect("valid IP");
let _ = service.lookup(ip).await;
if let IpAddr::V4(ipv4) = ip {
assert!(service.is_cached(&ipv4).await);
}
service.clear_cache().await;
let stats = service.cache_stats().await;
assert!(stats.is_empty);
}
#[tokio::test]
async fn test_public_ip_lookup() {
let service = AsnLookup::new();
let ip: IpAddr = "8.8.8.8".parse().expect("valid IP");
let result = service.lookup(ip).await;
if let Ok(info) = result {
assert_eq!(info.asn, 15169);
assert!(info.name.contains("GOOGLE"));
}
}
}