use crate::protocol::{
BalancingIpDetails,
routeplanner::{FailingAddress, IpBlock, RotatingIpDetails, RoutePlannerStatus},
};
use async_trait::async_trait;
use ipnet::IpNet;
use rand::Rng;
use std::{
collections::HashMap,
net::{IpAddr, Ipv4Addr, Ipv6Addr},
str::FromStr,
sync::Mutex,
};
use tracing::{debug, info};
#[async_trait]
pub trait RoutePlanner: Send + Sync {
fn get_status(&self) -> RoutePlannerStatus;
fn free_address(&self, address: &str);
fn free_all_addresses(&self);
fn mark_failed(&self, address: &str);
fn get_address(&self) -> Option<std::net::IpAddr>;
}
pub struct BalancingIpRoutePlanner {
ip_blocks: Vec<IpBlock>,
parsed_blocks: Vec<IpNet>,
failing_addresses: Mutex<HashMap<String, u64>>,
block_index: Mutex<usize>,
ip_indices: Mutex<Vec<u128>>,
}
impl BalancingIpRoutePlanner {
pub fn new(cidrs: Vec<String>) -> Result<Self, String> {
let mut ip_blocks = Vec::with_capacity(cidrs.len());
let mut parsed_blocks = Vec::with_capacity(cidrs.len());
let mut total_ips = 0;
for cidr in cidrs.iter() {
let parsed = IpNet::from_str(cidr)
.or_else(|_| {
let suffix = if cidr.contains(':') { "/128" } else { "/32" };
IpNet::from_str(&format!("{}{}", cidr, suffix))
})
.map_err(|e| format!("Invalid CIDR or IP '{}': {}", cidr, e))?;
let block_type = match parsed {
IpNet::V4(_) => "Inet4Address",
IpNet::V6(_) => "Inet6Address",
}
.to_string();
let size = match parsed {
IpNet::V4(net) => {
let prefix_len = net.prefix_len();
if prefix_len >= 32 {
1
} else {
2u128.pow(32 - prefix_len as u32)
}
}
IpNet::V6(net) => {
let prefix_len = net.prefix_len();
if prefix_len >= 128 {
1
} else if prefix_len <= 64 {
u128::MAX
} else {
2u128.pow(128 - prefix_len as u32)
}
}
};
if size == u128::MAX {
info!(
"Added {} block: {} (virtually unlimited addresses)",
block_type, cidr
);
} else {
info!("Added {} block: {} ({} addresses)", block_type, cidr, size);
total_ips += size;
}
ip_blocks.push(IpBlock {
block_type,
size: cidr.clone(),
});
parsed_blocks.push(parsed);
}
if total_ips > 0 && total_ips != u128::MAX {
info!(
"Route planner initialized with {} total addresses",
total_ips
);
} else if total_ips == u128::MAX {
info!("Route planner initialized with virtually unlimited addresses");
}
Ok(Self {
ip_blocks,
parsed_blocks,
failing_addresses: Mutex::new(HashMap::new()),
block_index: Mutex::new(0),
ip_indices: Mutex::new(vec![0; cidrs.len()]),
})
}
fn calculate_ip(block: &IpNet, index: u128) -> IpAddr {
let prefix_len = block.prefix_len();
match block {
IpNet::V4(net) => {
let addr_u32 = u32::from(net.addr());
let offset = if prefix_len >= 32 {
0
} else if prefix_len == 0 {
index as u32
} else {
(index as u32) & (!0u32 >> prefix_len)
};
IpAddr::V4(Ipv4Addr::from(addr_u32 + offset))
}
IpNet::V6(net) => {
let addr_u128 = u128::from(net.addr());
let offset = if prefix_len >= 128 {
0
} else if prefix_len == 0 {
index
} else {
index & (!0u128 >> prefix_len)
};
IpAddr::V6(Ipv6Addr::from(addr_u128 + offset))
}
}
}
fn next_ip(&self) -> IpAddr {
let mut b_idx = self.block_index.lock().unwrap_or_else(|e| e.into_inner());
let mut indices = self.ip_indices.lock().unwrap_or_else(|e| e.into_inner());
let block_idx = *b_idx % self.parsed_blocks.len();
let block = &self.parsed_blocks[block_idx];
let prefix_len = block.prefix_len();
let max_bits = match block {
IpNet::V4(_) => 32,
IpNet::V6(_) => 128,
};
let size_bits = max_bits - prefix_len;
let increment = if size_bits > 7 {
rand::thread_rng().gen_range(10..20) as u128
} else {
1
};
let current_index = indices[block_idx];
indices[block_idx] = current_index.wrapping_add(increment);
let final_index = indices[block_idx];
*b_idx = (*b_idx + 1) % self.parsed_blocks.len();
let ip = Self::calculate_ip(block, final_index);
debug!("Route planner picked IP: {}", ip);
ip
}
fn get_address_internal(&self) -> Option<IpAddr> {
for _ in 0..100 {
let ip = self.next_ip();
let ip_str = ip.to_string();
let mut failing = self
.failing_addresses
.lock()
.unwrap_or_else(|e| e.into_inner());
if let Some(×tamp) = failing.get(&ip_str) {
if crate::common::utils::now_ms()
> timestamp + crate::audio::constants::ROUTE_PLANNER_FAIL_EXPIRE_MS
{
failing.remove(&ip_str);
} else {
continue;
}
}
return Some(ip);
}
None
}
}
#[async_trait]
impl RoutePlanner for BalancingIpRoutePlanner {
fn get_status(&self) -> RoutePlannerStatus {
let failing = self
.failing_addresses
.lock()
.unwrap_or_else(|e| e.into_inner());
let failing_vec: Vec<FailingAddress> = failing
.iter()
.map(|(addr, &ts)| FailingAddress {
failing_address: addr.clone(),
failing_timestamp: ts,
failing_time: String::new(),
})
.collect();
if self.ip_blocks.len() == 1 {
let index = self.ip_indices.lock().unwrap_or_else(|e| e.into_inner())[0];
let current_ip = Self::calculate_ip(&self.parsed_blocks[0], index).to_string();
RoutePlannerStatus::RotatingIpRoutePlanner(RotatingIpDetails {
ip_block: self.ip_blocks[0].clone(),
failing_addresses: failing_vec,
rotate_index: "0".to_string(),
ip_index: index.to_string(),
current_address: current_ip,
})
} else {
RoutePlannerStatus::BalancingIpRoutePlanner(BalancingIpDetails {
ip_block: self.ip_blocks[0].clone(),
failing_addresses: failing_vec,
})
}
}
fn free_address(&self, address: &str) {
self.failing_addresses
.lock()
.unwrap_or_else(|e| e.into_inner())
.remove(address);
}
fn free_all_addresses(&self) {
self.failing_addresses
.lock()
.unwrap_or_else(|e| e.into_inner())
.clear();
}
fn mark_failed(&self, address: &str) {
self.failing_addresses
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(address.to_string(), crate::common::utils::now_ms());
}
fn get_address(&self) -> Option<std::net::IpAddr> {
self.get_address_internal()
}
}