use advmac::MacAddr6; use anyhow::{anyhow, Result}; use ipnetwork::{Ipv4Network, Ipv6Network}; use std::{ collections::HashMap, net::{Ipv4Addr, Ipv6Addr}, sync::Arc, }; use tokio::sync::RwLock; use uuid::Uuid; use crate::db::ip::{IpReservation, IpReservationStore}; #[derive(Default, Clone)] pub struct IpAssignmentState { pub ipv4: HashMap, pub ipv6: HashMap, } #[derive(Clone)] pub struct IpAssignment { ipv4_network: Ipv4Network, ipv6_network: Ipv6Network, gateway_ipv4: Ipv4Addr, gateway_ipv6: Ipv6Addr, gateway_mac: MacAddr6, store: IpReservationStore, state: Arc>, } impl IpAssignment { pub async fn new( host_uuid: Uuid, ipv4_network: Ipv4Network, ipv6_network: Ipv6Network, store: IpReservationStore, ) -> Result { let mut state = IpAssignment::fetch_current_state(&store).await?; let reservation = if let Some(reservation) = store.read(host_uuid).await? { reservation } else { IpAssignment::allocate( &mut state, &store, host_uuid, ipv4_network, ipv6_network, None, None, None, ) .await? }; let assignment = IpAssignment { ipv4_network, ipv6_network, gateway_ipv4: reservation.ipv4, gateway_ipv6: reservation.ipv6, gateway_mac: reservation.gateway_mac, store, state: Arc::new(RwLock::new(state)), }; Ok(assignment) } async fn fetch_current_state(store: &IpReservationStore) -> Result { let reservations = store.list().await?; let mut state = IpAssignmentState::default(); for reservation in reservations.values() { state.ipv4.insert(reservation.ipv4, reservation.clone()); state.ipv6.insert(reservation.ipv6, reservation.clone()); } Ok(state) } #[allow(clippy::too_many_arguments)] async fn allocate( state: &mut IpAssignmentState, store: &IpReservationStore, uuid: Uuid, ipv4_network: Ipv4Network, ipv6_network: Ipv6Network, gateway_ipv4: Option, gateway_ipv6: Option, gateway_mac: Option, ) -> Result { let mut found_ipv4: Option = None; for ip in ipv4_network.iter() { if ip.is_loopback() || ip.is_multicast() || ip.is_broadcast() { continue; } if !ip.is_private() { continue; } let last = ip.octets()[3]; if last == 0 || last > 250 { continue; } if state.ipv4.contains_key(&ip) { continue; } found_ipv4 = Some(ip); break; } let mut found_ipv6: Option = None; for ip in ipv6_network.iter() { if ip.is_loopback() || ip.is_multicast() { continue; } if state.ipv6.contains_key(&ip) { continue; } found_ipv6 = Some(ip); break; } let Some(ipv4) = found_ipv4 else { return Err(anyhow!( "unable to allocate ipv4 address, assigned network is exhausted" )); }; let Some(ipv6) = found_ipv6 else { return Err(anyhow!( "unable to allocate ipv6 address, assigned network is exhausted" )); }; let mut mac = MacAddr6::random(); mac.set_local(false); mac.set_multicast(false); let reservation = IpReservation { uuid: uuid.to_string(), ipv4, ipv6, mac, ipv4_prefix: ipv4_network.prefix(), ipv6_prefix: ipv6_network.prefix(), gateway_ipv4: gateway_ipv4.unwrap_or(ipv4), gateway_ipv6: gateway_ipv6.unwrap_or(ipv6), gateway_mac: gateway_mac.unwrap_or(mac), }; state.ipv4.insert(ipv4, reservation.clone()); state.ipv6.insert(ipv6, reservation.clone()); store.update(uuid, reservation.clone()).await?; Ok(reservation) } pub async fn assign(&self, uuid: Uuid) -> Result { let mut state = self.state.write().await; let reservation = IpAssignment::allocate( &mut state, &self.store, uuid, self.ipv4_network, self.ipv6_network, Some(self.gateway_ipv4), Some(self.gateway_ipv6), Some(self.gateway_mac), ) .await?; Ok(reservation) } pub async fn recall(&self, uuid: Uuid) -> Result<()> { let mut state = self.state.write().await; self.store.remove(uuid).await?; state .ipv4 .retain(|_, reservation| reservation.uuid != uuid.to_string()); state .ipv6 .retain(|_, reservation| reservation.uuid != uuid.to_string()); Ok(()) } pub async fn retrieve(&self, uuid: Uuid) -> Result> { self.store.read(uuid).await } pub async fn reload(&self) -> Result<()> { let mut state = self.state.write().await; let intermediate = IpAssignment::fetch_current_state(&self.store).await?; *state = intermediate; Ok(()) } pub async fn read(&self) -> Result { Ok(self.state.read().await.clone()) } }