use std::collections::BTreeSet;
use std::net::Ipv4Addr;
use ipnet::Ipv4Net;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IpamError {
BadCidr(String),
Exhausted,
}
impl std::fmt::Display for IpamError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::BadCidr(c) => write!(f, "invalid IPAM CIDR: {c}"),
Self::Exhausted => write!(f, "IPAM pool exhausted"),
}
}
}
impl std::error::Error for IpamError {}
#[derive(Debug, Clone)]
pub struct IpPool {
net: Ipv4Net,
gateway: Ipv4Addr,
allocated: BTreeSet<u32>,
}
impl IpPool {
pub fn new(cidr: &str) -> Result<Self, IpamError> {
let net: Ipv4Net = cidr
.parse()
.map_err(|_| IpamError::BadCidr(cidr.to_string()))?;
let gateway = net.hosts().next().unwrap_or(net.network());
Ok(Self {
net,
gateway,
allocated: BTreeSet::new(),
})
}
pub fn gateway(&self) -> Ipv4Addr {
self.gateway
}
pub fn reserve(&mut self, ip: Ipv4Addr) {
self.allocated.insert(ip.into());
}
pub fn allocate(&mut self) -> Result<Ipv4Addr, IpamError> {
for ip in self.net.hosts() {
if ip == self.gateway {
continue;
}
let key = u32::from(ip);
if !self.allocated.contains(&key) {
self.allocated.insert(key);
return Ok(ip);
}
}
Err(IpamError::Exhausted)
}
pub fn release(&mut self, ip: Ipv4Addr) {
self.allocated.remove(&u32::from(ip));
}
pub fn allocated_count(&self) -> usize {
self.allocated.len()
}
pub fn mac_for(ip: Ipv4Addr) -> String {
let o = ip.octets();
format!("02:00:{:02x}:{:02x}:{:02x}:{:02x}", o[0], o[1], o[2], o[3])
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn allocates_sequentially_skipping_gateway() {
let mut pool = IpPool::new("10.0.0.0/24").unwrap();
assert_eq!(pool.gateway(), Ipv4Addr::new(10, 0, 0, 1));
assert_eq!(pool.allocate().unwrap(), Ipv4Addr::new(10, 0, 0, 2));
assert_eq!(pool.allocate().unwrap(), Ipv4Addr::new(10, 0, 0, 3));
assert_eq!(pool.allocated_count(), 2);
}
#[test]
fn release_makes_an_address_reusable() {
let mut pool = IpPool::new("10.0.0.0/24").unwrap();
let a = pool.allocate().unwrap();
let b = pool.allocate().unwrap();
pool.release(a);
assert_eq!(pool.allocate().unwrap(), a);
assert_ne!(a, b);
}
#[test]
fn reserve_marks_in_use() {
let mut pool = IpPool::new("10.0.0.0/24").unwrap();
pool.reserve(Ipv4Addr::new(10, 0, 0, 2));
assert_eq!(pool.allocate().unwrap(), Ipv4Addr::new(10, 0, 0, 3));
}
#[test]
fn tiny_pool_exhausts() {
let mut pool = IpPool::new("10.0.0.0/30").unwrap();
assert_eq!(pool.allocate().unwrap(), Ipv4Addr::new(10, 0, 0, 2));
assert_eq!(pool.allocate(), Err(IpamError::Exhausted));
}
#[test]
fn mac_is_locally_administered_and_stable() {
let mac = IpPool::mac_for(Ipv4Addr::new(10, 0, 0, 5));
assert_eq!(mac, "02:00:0a:00:00:05");
assert_eq!(mac, IpPool::mac_for(Ipv4Addr::new(10, 0, 0, 5)));
}
#[test]
fn bad_cidr_errors() {
assert!(matches!(
IpPool::new("not-a-cidr"),
Err(IpamError::BadCidr(_))
));
}
}