feature(krata): rework api and make ip assignment persistent to database

crates/daemon/src/ip/assignment.rs

@@ -0,0 +1,197 @@
+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<Ipv4Addr, IpReservation>,
+    pub ipv6: HashMap<Ipv6Addr, IpReservation>,
+}
+
+#[derive(Clone)]
+pub struct IpAssignment {
+    ipv4_network: Ipv4Network,
+    ipv6_network: Ipv6Network,
+    gateway_ipv4: Ipv4Addr,
+    gateway_ipv6: Ipv6Addr,
+    gateway_mac: MacAddr6,
+    store: IpReservationStore,
+    state: Arc<RwLock<IpAssignmentState>>,
+}
+
+impl IpAssignment {
+    pub async fn new(
+        host_uuid: Uuid,
+        ipv4_network: Ipv4Network,
+        ipv6_network: Ipv6Network,
+        store: IpReservationStore,
+    ) -> Result<Self> {
+        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<IpAssignmentState> {
+        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<Ipv4Addr>,
+        gateway_ipv6: Option<Ipv6Addr>,
+        gateway_mac: Option<MacAddr6>,
+    ) -> Result<IpReservation> {
+        let mut found_ipv4: Option<Ipv4Addr> = 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<Ipv6Addr> = 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<IpReservation> {
+        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<Option<IpReservation>> {
+        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<IpAssignmentState> {
+        Ok(self.state.read().await.clone())
+    }
+}