reddwarf/crates/reddwarf-runtime/src/network/ipam.rs

use crate::error::{Result, RuntimeError};
use reddwarf_storage::KVStore;
use std::collections::BTreeMap;
use std::net::Ipv4Addr;
use std::sync::Arc;
use tracing::debug;

/// Parsed CIDR configuration
#[derive(Debug, Clone)]
pub struct CidrConfig {
    /// Base network address
    pub network: Ipv4Addr,
    /// CIDR prefix length
    pub prefix_len: u8,
    /// Gateway address (network + 1)
    pub gateway: Ipv4Addr,
    /// First allocatable host address (network + 2)
    pub first_host: Ipv4Addr,
    /// Broadcast address (last in range)
    pub broadcast: Ipv4Addr,
}

/// An allocated IP for a pod
#[derive(Debug, Clone)]
pub struct IpAllocation {
    pub ip_address: Ipv4Addr,
    pub gateway: Ipv4Addr,
    pub prefix_len: u8,
}

/// IPAM (IP Address Management) backed by a KVStore
///
/// Storage keys (with default "ipam" prefix):
/// - `{prefix}/_cidr` → the CIDR string (e.g. "10.88.0.0/16")
/// - `{prefix}/alloc/{ip}` → `"{namespace}/{name}"`
pub struct Ipam {
    storage: Arc<dyn KVStore>,
    cidr: CidrConfig,
    /// Storage key for CIDR config
    #[allow(dead_code)]
    cidr_key: Vec<u8>,
    /// Storage key prefix for allocations
    alloc_prefix: Vec<u8>,
}

impl Ipam {
    /// Create a new IPAM instance with default "ipam" prefix, persisting the CIDR config
    pub fn new(storage: Arc<dyn KVStore>, cidr_str: &str) -> Result<Self> {
        Self::with_prefix(storage, cidr_str, "ipam")
    }

    /// Create a new IPAM instance with a custom storage key prefix
    pub fn with_prefix(storage: Arc<dyn KVStore>, cidr_str: &str, prefix: &str) -> Result<Self> {
        let cidr = parse_cidr(cidr_str)?;

        let cidr_key = format!("{}/_cidr", prefix).into_bytes();
        let alloc_prefix = format!("{}/alloc/", prefix).into_bytes();

        // Persist the CIDR configuration
        storage.put(&cidr_key, cidr_str.as_bytes())?;

        debug!(
            "IPAM ({}) initialized: network={}, gateway={}, first_host={}, broadcast={}, prefix_len={}",
            prefix, cidr.network, cidr.gateway, cidr.first_host, cidr.broadcast, cidr.prefix_len
        );

        Ok(Self {
            storage,
            cidr,
            cidr_key,
            alloc_prefix,
        })
    }

    /// Allocate an IP for a resource. Idempotent: returns existing allocation if one exists.
    pub fn allocate(&self, namespace: &str, name: &str) -> Result<IpAllocation> {
        let resource_key = format!("{}/{}", namespace, name);
        let prefix_len = self.alloc_prefix.len();

        // Check if this resource already has an allocation
        let allocations = self.storage.scan(&self.alloc_prefix)?;
        for (key, value) in &allocations {
            let existing = String::from_utf8_lossy(value);
            if existing == resource_key {
                let key_str = String::from_utf8_lossy(key);
                let ip_str = &key_str[prefix_len..];
                if let Ok(ip) = ip_str.parse::<Ipv4Addr>() {
                    debug!(
                        "IPAM: returning existing allocation {} for {}",
                        ip, resource_key
                    );
                    return Ok(IpAllocation {
                        ip_address: ip,
                        gateway: self.cidr.gateway,
                        prefix_len: self.cidr.prefix_len,
                    });
                }
            }
        }

        // Collect already-allocated IPs
        let allocated: std::collections::HashSet<Ipv4Addr> = allocations
            .iter()
            .filter_map(|(key, _)| {
                let key_str = String::from_utf8_lossy(key);
                let ip_str = &key_str[prefix_len..];
                ip_str.parse::<Ipv4Addr>().ok()
            })
            .collect();

        // Find next free IP starting from first_host
        let mut candidate = self.cidr.first_host;
        loop {
            if candidate >= self.cidr.broadcast {
                return Err(RuntimeError::IpamPoolExhausted {
                    cidr: format!("{}/{}", self.cidr.network, self.cidr.prefix_len),
                });
            }

            if !allocated.contains(&candidate) {
                // Allocate this IP
                let alloc_key_str = String::from_utf8_lossy(&self.alloc_prefix);
                let alloc_key = format!("{}{}", alloc_key_str, candidate);
                self.storage
                    .put(alloc_key.as_bytes(), resource_key.as_bytes())?;

                debug!("IPAM: allocated {} for {}", candidate, resource_key);
                return Ok(IpAllocation {
                    ip_address: candidate,
                    gateway: self.cidr.gateway,
                    prefix_len: self.cidr.prefix_len,
                });
            }

            candidate = next_ip(candidate);
        }
    }

    /// Release the IP allocated to a resource
    pub fn release(&self, namespace: &str, name: &str) -> Result<Option<Ipv4Addr>> {
        let resource_key = format!("{}/{}", namespace, name);
        let prefix_len = self.alloc_prefix.len();

        let allocations = self.storage.scan(&self.alloc_prefix)?;
        for (key, value) in &allocations {
            let existing = String::from_utf8_lossy(value);
            if existing == resource_key {
                let key_str = String::from_utf8_lossy(key);
                let ip_str = &key_str[prefix_len..];
                let ip = ip_str.parse::<Ipv4Addr>().ok();

                self.storage.delete(key)?;
                debug!("IPAM: released {:?} for {}", ip, resource_key);
                return Ok(ip);
            }
        }

        debug!("IPAM: no allocation found for {}", resource_key);
        Ok(None)
    }

    /// Get all current allocations
    pub fn get_all_allocations(&self) -> Result<BTreeMap<Ipv4Addr, String>> {
        let prefix_len = self.alloc_prefix.len();
        let allocations = self.storage.scan(&self.alloc_prefix)?;
        let mut result = BTreeMap::new();

        for (key, value) in &allocations {
            let key_str = String::from_utf8_lossy(key);
            let ip_str = &key_str[prefix_len..];
            if let Ok(ip) = ip_str.parse::<Ipv4Addr>() {
                result.insert(ip, String::from_utf8_lossy(value).into_owned());
            }
        }

        Ok(result)
    }
}

/// Parse a CIDR string like "10.88.0.0/16" into a CidrConfig
pub fn parse_cidr(cidr_str: &str) -> Result<CidrConfig> {
    let parts: Vec<&str> = cidr_str.split('/').collect();
    if parts.len() != 2 {
        return Err(RuntimeError::invalid_config(
            format!("Invalid CIDR format: '{}'", cidr_str),
            "Use format like '10.88.0.0/16'",
        ));
    }

    let network: Ipv4Addr = parts[0].parse().map_err(|_| {
        RuntimeError::invalid_config(
            format!("Invalid network address: '{}'", parts[0]),
            "Use a valid IPv4 address like '10.88.0.0'",
        )
    })?;

    let prefix_len: u8 = parts[1].parse().map_err(|_| {
        RuntimeError::invalid_config(
            format!("Invalid prefix length: '{}'", parts[1]),
            "Use a number between 0 and 32",
        )
    })?;

    if prefix_len > 32 {
        return Err(RuntimeError::invalid_config(
            format!("Prefix length {} is out of range", prefix_len),
            "Use a number between 0 and 32",
        ));
    }

    let network_u32 = u32::from(network);
    let host_bits = 32 - prefix_len;
    let mask = if prefix_len == 0 {
        0u32
    } else {
        !((1u32 << host_bits) - 1)
    };
    let broadcast_u32 = network_u32 | !mask;

    let gateway = Ipv4Addr::from(network_u32 + 1);
    let first_host = Ipv4Addr::from(network_u32 + 2);
    let broadcast = Ipv4Addr::from(broadcast_u32);

    Ok(CidrConfig {
        network,
        prefix_len,
        gateway,
        first_host,
        broadcast,
    })
}

/// Increment an IPv4 address by one
fn next_ip(ip: Ipv4Addr) -> Ipv4Addr {
    Ipv4Addr::from(u32::from(ip) + 1)
}

#[cfg(test)]
mod tests {
    use super::*;
    use reddwarf_storage::RedbBackend;
    use tempfile::tempdir;

    fn make_test_ipam(cidr: &str) -> Ipam {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("test-ipam.redb");
        let storage = Arc::new(RedbBackend::new(&db_path).unwrap());
        // We need to keep tempdir alive for the duration, but for tests
        // we leak it to avoid dropping the temp dir too early
        std::mem::forget(dir);
        Ipam::new(storage, cidr).unwrap()
    }

    #[test]
    fn test_parse_cidr_valid() {
        let cidr = parse_cidr("10.88.0.0/16").unwrap();
        assert_eq!(cidr.network, Ipv4Addr::new(10, 88, 0, 0));
        assert_eq!(cidr.prefix_len, 16);
        assert_eq!(cidr.gateway, Ipv4Addr::new(10, 88, 0, 1));
        assert_eq!(cidr.first_host, Ipv4Addr::new(10, 88, 0, 2));
        assert_eq!(cidr.broadcast, Ipv4Addr::new(10, 88, 255, 255));
    }

    #[test]
    fn test_parse_cidr_slash24() {
        let cidr = parse_cidr("192.168.1.0/24").unwrap();
        assert_eq!(cidr.network, Ipv4Addr::new(192, 168, 1, 0));
        assert_eq!(cidr.gateway, Ipv4Addr::new(192, 168, 1, 1));
        assert_eq!(cidr.first_host, Ipv4Addr::new(192, 168, 1, 2));
        assert_eq!(cidr.broadcast, Ipv4Addr::new(192, 168, 1, 255));
    }

    #[test]
    fn test_parse_cidr_invalid() {
        assert!(parse_cidr("not-a-cidr").is_err());
        assert!(parse_cidr("10.88.0.0").is_err());
        assert!(parse_cidr("10.88.0.0/33").is_err());
        assert!(parse_cidr("bad/16").is_err());
    }

    #[test]
    fn test_allocate_sequential() {
        let ipam = make_test_ipam("10.88.0.0/16");

        let alloc1 = ipam.allocate("default", "pod-a").unwrap();
        assert_eq!(alloc1.ip_address, Ipv4Addr::new(10, 88, 0, 2));
        assert_eq!(alloc1.gateway, Ipv4Addr::new(10, 88, 0, 1));
        assert_eq!(alloc1.prefix_len, 16);

        let alloc2 = ipam.allocate("default", "pod-b").unwrap();
        assert_eq!(alloc2.ip_address, Ipv4Addr::new(10, 88, 0, 3));
    }

    #[test]
    fn test_allocate_idempotent() {
        let ipam = make_test_ipam("10.88.0.0/16");

        let alloc1 = ipam.allocate("default", "pod-a").unwrap();
        let alloc2 = ipam.allocate("default", "pod-a").unwrap();
        assert_eq!(alloc1.ip_address, alloc2.ip_address);
    }

    #[test]
    fn test_release_and_reallocate() {
        let ipam = make_test_ipam("10.88.0.0/16");

        let alloc1 = ipam.allocate("default", "pod-a").unwrap();
        let first_ip = alloc1.ip_address;

        // Allocate a second pod
        let _alloc2 = ipam.allocate("default", "pod-b").unwrap();

        // Release first pod
        let released = ipam.release("default", "pod-a").unwrap();
        assert_eq!(released, Some(first_ip));

        // New pod should reuse the freed IP
        let alloc3 = ipam.allocate("default", "pod-c").unwrap();
        assert_eq!(alloc3.ip_address, first_ip);
    }

    #[test]
    fn test_pool_exhaustion() {
        // /30 gives us network .0, gateway .1, one host .2, broadcast .3
        let ipam = make_test_ipam("10.0.0.0/30");

        // First allocation should succeed (.2)
        let alloc = ipam.allocate("default", "pod-a").unwrap();
        assert_eq!(alloc.ip_address, Ipv4Addr::new(10, 0, 0, 2));

        // Second allocation should fail (only .2 is usable, .3 is broadcast)
        let result = ipam.allocate("default", "pod-b");
        assert!(matches!(
            result.unwrap_err(),
            RuntimeError::IpamPoolExhausted { .. }
        ));
    }

    #[test]
    fn test_get_all_allocations() {
        let ipam = make_test_ipam("10.88.0.0/16");

        ipam.allocate("default", "pod-a").unwrap();
        ipam.allocate("kube-system", "pod-b").unwrap();

        let allocs = ipam.get_all_allocations().unwrap();
        assert_eq!(allocs.len(), 2);
        assert_eq!(allocs[&Ipv4Addr::new(10, 88, 0, 2)], "default/pod-a");
        assert_eq!(allocs[&Ipv4Addr::new(10, 88, 0, 3)], "kube-system/pod-b");
    }

    #[test]
    fn test_release_nonexistent() {
        let ipam = make_test_ipam("10.88.0.0/16");
        let released = ipam.release("default", "nonexistent").unwrap();
        assert_eq!(released, None);
    }
}