zlayer_overlay/

allocator.rs

//! IP address allocation for overlay networks
//!
//! Manages allocation and tracking of overlay IP addresses within a CIDR range.
//! Supports both IPv4 and IPv6 (dual-stack) networks.

use crate::error::{OverlayError, Result};
use ipnet::{IpNet, Ipv4Net, Ipv6Net};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::path::Path;

/// IP allocator for overlay network addresses
///
/// Tracks allocated IP addresses and provides next-available allocation
/// from a configured CIDR range. Supports both IPv4 and IPv6 networks.
#[derive(Debug, Clone)]
pub struct IpAllocator {
    /// Network CIDR range (IPv4 or IPv6)
    network: IpNet,
    /// Set of allocated IP addresses
    allocated: HashSet<IpAddr>,
}

/// Persistent state for IP allocator
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IpAllocatorState {
    /// CIDR string
    pub cidr: String,
    /// List of allocated IPs (serializes as strings, backward-compatible)
    pub allocated: Vec<IpAddr>,
}

/// Increment an IPv6 address by a u128 offset from a base address.
///
/// Returns `None` if the result would overflow.
fn ipv6_add(base: Ipv6Addr, offset: u128) -> Option<Ipv6Addr> {
    let base_u128 = u128::from(base);
    base_u128.checked_add(offset).map(Ipv6Addr::from)
}

/// Compute the number of host addresses for a given address family and prefix length.
///
/// For IPv4: `2^(32 - prefix) - 2` (excludes network and broadcast).
/// For IPv6: `2^(128 - prefix) - 1` (excludes the network address).
///
/// Returns `None` if the result overflows u128 (only for /0 edge cases).
fn host_count(is_ipv6: bool, prefix_len: u8) -> u128 {
    if is_ipv6 {
        let bits = 128 - u32::from(prefix_len);
        if bits == 128 {
            // /0 network — saturate
            u128::MAX
        } else if bits == 0 {
            // /128 — single host, no usable addresses (it IS the network address)
            0
        } else {
            // 2^bits - 1 (skip network address)
            (1u128 << bits) - 1
        }
    } else {
        let bits = 32 - u32::from(prefix_len);
        if bits <= 1 {
            // /31 or /32 — no usable hosts in classical networking
            0
        } else {
            // 2^bits - 2 (skip network and broadcast)
            (1u128 << bits) - 2
        }
    }
}

impl IpAllocator {
    /// Create a new IP allocator for the given CIDR range
    ///
    /// Supports both IPv4 (e.g., "10.200.0.0/16") and IPv6 (e.g., `fd00::/48`).
    ///
    /// # Arguments
    /// * `cidr` - Network CIDR notation
    ///
    /// # Errors
    ///
    /// Returns `OverlayError::InvalidCidr` if the CIDR string cannot be parsed.
    ///
    /// # Example
    /// ```
    /// use zlayer_overlay::allocator::IpAllocator;
    ///
    /// let v4 = IpAllocator::new("10.200.0.0/16").unwrap();
    /// let v6 = IpAllocator::new("fd00::/48").unwrap();
    /// ```
    pub fn new(cidr: &str) -> Result<Self> {
        let network: IpNet = cidr
            .parse()
            .map_err(|e| OverlayError::InvalidCidr(format!("{cidr}: {e}")))?;

        Ok(Self {
            network,
            allocated: HashSet::new(),
        })
    }

    /// Create an allocator from persisted state
    ///
    /// # Errors
    ///
    /// Returns an error if the CIDR is invalid or any IP is out of range.
    pub fn from_state(state: IpAllocatorState) -> Result<Self> {
        let mut allocator = Self::new(&state.cidr)?;
        for ip in state.allocated {
            allocator.mark_allocated(ip)?;
        }
        Ok(allocator)
    }

    /// Get the current state for persistence
    #[must_use]
    pub fn to_state(&self) -> IpAllocatorState {
        IpAllocatorState {
            cidr: self.network.to_string(),
            allocated: self.allocated.iter().copied().collect(),
        }
    }

    /// Load allocator state from a file
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read or the state is invalid.
    pub async fn load(path: &Path) -> Result<Self> {
        let contents = tokio::fs::read_to_string(path).await?;
        let state: IpAllocatorState = serde_json::from_str(&contents)?;
        Self::from_state(state)
    }

    /// Save allocator state to a file
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be written or serialization fails.
    pub async fn save(&self, path: &Path) -> Result<()> {
        let state = self.to_state();
        let contents = serde_json::to_string_pretty(&state)?;
        tokio::fs::write(path, contents).await?;
        Ok(())
    }

    /// Allocate the next available IP address
    ///
    /// For IPv4, skips the network and broadcast addresses.
    /// For IPv6, skips the network address.
    ///
    /// Returns `None` if all addresses in the CIDR range are allocated.
    ///
    /// # Example
    /// ```
    /// use zlayer_overlay::allocator::IpAllocator;
    ///
    /// let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
    /// let ip = allocator.allocate().unwrap();
    /// assert_eq!(ip.to_string(), "10.200.0.1");
    /// ```
    pub fn allocate(&mut self) -> Option<IpAddr> {
        match self.network {
            IpNet::V4(v4net) => {
                // IPv4: iterate hosts() which skips network and broadcast
                for ip in v4net.hosts() {
                    let addr = IpAddr::V4(ip);
                    if !self.allocated.contains(&addr) {
                        self.allocated.insert(addr);
                        return Some(addr);
                    }
                }
                None
            }
            IpNet::V6(v6net) => {
                // IPv6: counter-based allocation starting from base+1
                // We skip the network address itself (offset 0) and allocate from offset 1.
                let base = v6net.network();
                let total = host_count(true, v6net.prefix_len());

                for offset in 1..=total {
                    if let Some(candidate) = ipv6_add(base, offset) {
                        let addr = IpAddr::V6(candidate);
                        if !self.allocated.contains(&addr) {
                            self.allocated.insert(addr);
                            return Some(addr);
                        }
                    } else {
                        break;
                    }
                }
                None
            }
        }
    }

    /// Allocate a specific IP address
    ///
    /// # Errors
    ///
    /// Returns an error if the IP is already allocated or not in the CIDR range.
    pub fn allocate_specific(&mut self, ip: IpAddr) -> Result<()> {
        if !self.network.contains(&ip) {
            return Err(OverlayError::IpNotInRange(ip, self.network.to_string()));
        }

        if self.allocated.contains(&ip) {
            return Err(OverlayError::IpAlreadyAllocated(ip));
        }

        self.allocated.insert(ip);
        Ok(())
    }

    /// Allocate the first usable IP in the range (typically for the leader)
    ///
    /// # Example
    /// ```
    /// use zlayer_overlay::allocator::IpAllocator;
    ///
    /// let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
    /// let ip = allocator.allocate_first().unwrap();
    /// assert_eq!(ip.to_string(), "10.200.0.1");
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an error if no IPs are available or the first IP is already allocated.
    pub fn allocate_first(&mut self) -> Result<IpAddr> {
        let first_ip = self.first_host().ok_or(OverlayError::NoAvailableIps)?;

        if self.allocated.contains(&first_ip) {
            return Err(OverlayError::IpAlreadyAllocated(first_ip));
        }

        self.allocated.insert(first_ip);
        Ok(first_ip)
    }

    /// Get the first usable host address in the network.
    ///
    /// For IPv4: first host from `hosts()` (skips network address).
    /// For IPv6: network address + 1 (skips the network address).
    fn first_host(&self) -> Option<IpAddr> {
        match self.network {
            IpNet::V4(v4net) => v4net.hosts().next().map(IpAddr::V4),
            IpNet::V6(v6net) => {
                let base = v6net.network();
                ipv6_add(base, 1).map(IpAddr::V6)
            }
        }
    }

    /// Mark an IP address as allocated (for restoring state)
    ///
    /// # Errors
    ///
    /// Returns an error if the IP is not in the CIDR range.
    pub fn mark_allocated(&mut self, ip: IpAddr) -> Result<()> {
        if !self.network.contains(&ip) {
            return Err(OverlayError::IpNotInRange(ip, self.network.to_string()));
        }
        self.allocated.insert(ip);
        Ok(())
    }

    /// Release an IP address back to the pool
    ///
    /// Returns `true` if the IP was released, `false` if it wasn't allocated.
    pub fn release(&mut self, ip: IpAddr) -> bool {
        self.allocated.remove(&ip)
    }

    /// Check if an IP address is allocated
    #[must_use]
    pub fn is_allocated(&self, ip: IpAddr) -> bool {
        self.allocated.contains(&ip)
    }

    /// Check if an IP address is within the CIDR range
    #[must_use]
    pub fn contains(&self, ip: IpAddr) -> bool {
        self.network.contains(&ip)
    }

    /// Get the number of allocated addresses
    #[must_use]
    pub fn allocated_count(&self) -> usize {
        self.allocated.len()
    }

    /// Get the total number of usable addresses in the range
    ///
    /// For IPv6 networks with large host spaces, this saturates at `u32::MAX`.
    #[must_use]
    #[allow(clippy::cast_possible_truncation)]
    pub fn total_hosts(&self) -> u32 {
        let is_v6 = matches!(self.network, IpNet::V6(_));
        let count = host_count(is_v6, self.network.prefix_len());
        // Saturate to u32::MAX for enormous IPv6 subnets
        if count > u128::from(u32::MAX) {
            u32::MAX
        } else {
            count as u32
        }
    }

    /// Get the number of available addresses
    #[must_use]
    #[allow(clippy::cast_possible_truncation)]
    pub fn available_count(&self) -> u32 {
        self.total_hosts()
            .saturating_sub(self.allocated.len() as u32)
    }

    /// Get the CIDR string
    #[must_use]
    pub fn cidr(&self) -> String {
        self.network.to_string()
    }

    /// Get the network address
    #[must_use]
    pub fn network_addr(&self) -> IpAddr {
        self.network.network()
    }

    /// Get the broadcast address
    ///
    /// For IPv6, returns the last address in the range (all host bits set to 1).
    #[must_use]
    pub fn broadcast_addr(&self) -> IpAddr {
        self.network.broadcast()
    }

    /// Get the prefix length
    #[must_use]
    pub fn prefix_len(&self) -> u8 {
        self.network.prefix_len()
    }

    /// Get the host prefix length (32 for IPv4, 128 for IPv6)
    #[must_use]
    pub fn host_prefix_len(&self) -> u8 {
        self.network.max_prefix_len()
    }

    /// Get all allocated IPs
    #[must_use]
    pub fn allocated_ips(&self) -> Vec<IpAddr> {
        self.allocated.iter().copied().collect()
    }
}

/// Leader-side allocator that carves per-node slices out of a cluster CIDR.
///
/// Used to fix the latent IP-collision bug where every agent independently
/// allocated container IPs from the full cluster `/16`. With a `NodeSliceAllocator`
/// the leader hands each joining node its own non-overlapping slice, and the
/// agent-local `IpAllocator` is bounded to that slice.
///
/// Slice assignment is deterministic within a leader process: the node ID hashes
/// to a candidate slice index; collisions are resolved by linear probing forward
/// until a free slot is found. Existing assignments are preserved across leader
/// restart via `snapshot()` / `restore()`.
#[derive(Debug, Clone)]
pub struct NodeSliceAllocator {
    cluster_cidr: IpNet,
    slice_prefix: u8,
    assigned: HashMap<String, IpNet>,
}

/// Persistent snapshot of a `NodeSliceAllocator` for raft/disk persistence.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeSliceAllocatorSnapshot {
    pub cluster_cidr: String,
    pub slice_prefix: u8,
    pub assigned: Vec<(String, String)>,
}

/// Deterministic FNV-1a 64-bit hash for a node ID string.
///
/// Chosen over `DefaultHasher` because `DefaultHasher` is seeded per-process
/// and slice assignments should be reproducible from a snapshot.
fn hash_node_id(node_id: &str) -> u64 {
    const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
    let mut hash = FNV_OFFSET;
    for &b in node_id.as_bytes() {
        hash ^= u64::from(b);
        hash = hash.wrapping_mul(FNV_PRIME);
    }
    hash
}

impl NodeSliceAllocator {
    /// Create a new slice allocator that carves `/slice_prefix`-sized slices
    /// out of `cluster_cidr`.
    ///
    /// # Errors
    ///
    /// Returns `OverlayError::InvalidCidr` if `slice_prefix` is not strictly
    /// more specific than `cluster_cidr.prefix_len()`, or if it exceeds the
    /// address family's maximum prefix length.
    pub fn new(cluster_cidr: IpNet, slice_prefix: u8) -> Result<Self> {
        if slice_prefix <= cluster_cidr.prefix_len() {
            return Err(OverlayError::InvalidCidr(format!(
                "slice prefix /{} must be more specific than cluster prefix /{}",
                slice_prefix,
                cluster_cidr.prefix_len()
            )));
        }
        if slice_prefix > cluster_cidr.max_prefix_len() {
            return Err(OverlayError::InvalidCidr(format!(
                "slice prefix /{} exceeds address family max /{}",
                slice_prefix,
                cluster_cidr.max_prefix_len()
            )));
        }
        Ok(Self {
            cluster_cidr,
            slice_prefix,
            assigned: HashMap::new(),
        })
    }

    /// Assign (or return an existing) slice for `node_id`.
    ///
    /// Idempotent: calling `assign` with a node ID that already has a slice
    /// returns the existing slice without re-assigning.
    ///
    /// # Errors
    ///
    /// Returns `OverlayError::NoAvailableIps` if every slice in the cluster
    /// CIDR is already assigned.
    pub fn assign(&mut self, node_id: &str) -> Result<IpNet> {
        if let Some(existing) = self.assigned.get(node_id) {
            return Ok(*existing);
        }

        let num_slices = self.num_slices();
        if num_slices == 0 {
            return Err(OverlayError::NoAvailableIps);
        }

        let taken: HashSet<IpNet> = self.assigned.values().copied().collect();
        let start = hash_node_id(node_id) % num_slices;

        for i in 0..num_slices {
            let idx = (start + i) % num_slices;
            let slice = self.slice_at_index(idx);
            if !taken.contains(&slice) {
                self.assigned.insert(node_id.to_string(), slice);
                return Ok(slice);
            }
        }

        Err(OverlayError::NoAvailableIps)
    }

    /// Release `node_id`'s slice back to the free pool.
    ///
    /// Returns `true` if a slice was released, `false` if the node was not assigned.
    pub fn release(&mut self, node_id: &str) -> bool {
        self.assigned.remove(node_id).is_some()
    }

    /// Look up a node's assigned slice without mutating state.
    #[must_use]
    pub fn slice_for(&self, node_id: &str) -> Option<IpNet> {
        self.assigned.get(node_id).copied()
    }

    /// Number of currently-assigned slices.
    #[must_use]
    pub fn assigned_count(&self) -> usize {
        self.assigned.len()
    }

    /// Total number of slices the cluster CIDR can hold at the configured slice prefix.
    #[must_use]
    pub fn capacity(&self) -> u64 {
        self.num_slices()
    }

    /// Cluster CIDR the allocator operates over.
    #[must_use]
    pub fn cluster_cidr(&self) -> IpNet {
        self.cluster_cidr
    }

    /// Slice prefix length (e.g. `28` for `/28` slices).
    #[must_use]
    pub fn slice_prefix(&self) -> u8 {
        self.slice_prefix
    }

    /// Build a persistable snapshot for durable leader state.
    #[must_use]
    pub fn snapshot(&self) -> NodeSliceAllocatorSnapshot {
        NodeSliceAllocatorSnapshot {
            cluster_cidr: self.cluster_cidr.to_string(),
            slice_prefix: self.slice_prefix,
            assigned: self
                .assigned
                .iter()
                .map(|(k, v)| (k.clone(), v.to_string()))
                .collect(),
        }
    }

    /// Rebuild an allocator from a snapshot.
    ///
    /// # Errors
    ///
    /// Returns `OverlayError::InvalidCidr` if the snapshot's CIDR or any
    /// assigned slice fails to parse, or if the slice prefix is inconsistent.
    pub fn restore(snapshot: NodeSliceAllocatorSnapshot) -> Result<Self> {
        let cluster_cidr: IpNet = snapshot
            .cluster_cidr
            .parse()
            .map_err(|e| OverlayError::InvalidCidr(format!("{}: {e}", snapshot.cluster_cidr)))?;
        let mut allocator = Self::new(cluster_cidr, snapshot.slice_prefix)?;
        for (node_id, slice_str) in snapshot.assigned {
            let slice: IpNet = slice_str
                .parse()
                .map_err(|e| OverlayError::InvalidCidr(format!("{slice_str}: {e}")))?;
            if slice.prefix_len() != snapshot.slice_prefix {
                return Err(OverlayError::InvalidCidr(format!(
                    "assigned slice {slice} does not match configured prefix /{}",
                    snapshot.slice_prefix
                )));
            }
            if !cluster_cidr.contains(&slice.network()) {
                return Err(OverlayError::InvalidCidr(format!(
                    "assigned slice {slice} is not contained in cluster CIDR {cluster_cidr}"
                )));
            }
            allocator.assigned.insert(node_id, slice);
        }
        Ok(allocator)
    }

    fn num_slices(&self) -> u64 {
        let bits = self.slice_prefix - self.cluster_cidr.prefix_len();
        // bits is in 1..=32 for v4 or 1..=128 for v6. For a /16 cluster with /28
        // slices, bits = 12 → 4096 slices, safely inside u64 range.
        if bits >= 64 {
            u64::MAX
        } else {
            1u64 << bits
        }
    }

    fn slice_at_index(&self, idx: u64) -> IpNet {
        let shift = u32::from(self.cluster_cidr.max_prefix_len() - self.slice_prefix);
        match self.cluster_cidr {
            IpNet::V4(v4) => {
                let base = u32::from(v4.network());
                // idx fits in 32 bits whenever slice_prefix − cluster_prefix ≤ 32.
                #[allow(clippy::cast_possible_truncation)]
                let offset = (idx as u32).wrapping_shl(shift);
                let slice_addr = Ipv4Addr::from(base.wrapping_add(offset));
                IpNet::V4(
                    Ipv4Net::new(slice_addr, self.slice_prefix)
                        .expect("slice_prefix validated in constructor"),
                )
            }
            IpNet::V6(v6) => {
                let base = u128::from(v6.network());
                let offset = u128::from(idx).wrapping_shl(shift);
                let slice_addr = Ipv6Addr::from(base.wrapping_add(offset));
                IpNet::V6(
                    Ipv6Net::new(slice_addr, self.slice_prefix)
                        .expect("slice_prefix validated in constructor"),
                )
            }
        }
    }
}

/// Helper function to get the first usable IP from a CIDR
///
/// Supports both IPv4 and IPv6 CIDR notation.
///
/// # Errors
///
/// Returns an error if the CIDR is invalid or has no usable hosts.
pub fn first_ip_from_cidr(cidr: &str) -> Result<IpAddr> {
    let network: IpNet = cidr
        .parse()
        .map_err(|e| OverlayError::InvalidCidr(format!("{cidr}: {e}")))?;

    match network {
        IpNet::V4(v4net) => v4net
            .hosts()
            .next()
            .map(IpAddr::V4)
            .ok_or(OverlayError::NoAvailableIps),
        IpNet::V6(v6net) => {
            let base = v6net.network();
            ipv6_add(base, 1)
                .map(IpAddr::V6)
                .ok_or(OverlayError::NoAvailableIps)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{Ipv4Addr, Ipv6Addr};

    /// Increment an IPv4 address by a u32 offset from a base address.
    ///
    /// Returns `None` if the result would overflow.
    fn ipv4_add(base: Ipv4Addr, offset: u32) -> Option<Ipv4Addr> {
        let base_u32 = u32::from(base);
        base_u32.checked_add(offset).map(Ipv4Addr::from)
    }

    // ========================
    // IPv4 Tests (existing, updated for IpAddr)
    // ========================

    #[test]
    fn test_allocator_new() {
        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
        assert_eq!(allocator.cidr(), "10.200.0.0/24");
        assert_eq!(allocator.allocated_count(), 0);
    }

    #[test]
    fn test_allocator_invalid_cidr() {
        let result = IpAllocator::new("invalid");
        assert!(result.is_err());
    }

    #[test]
    fn test_allocate_sequential() {
        let mut allocator = IpAllocator::new("10.200.0.0/30").unwrap();

        // /30 has 2 usable hosts (excluding network and broadcast)
        let ip1 = allocator.allocate().unwrap();
        let ip2 = allocator.allocate().unwrap();

        assert_eq!(ip1.to_string(), "10.200.0.1");
        assert_eq!(ip2.to_string(), "10.200.0.2");

        // Should be exhausted
        assert!(allocator.allocate().is_none());
    }

    #[test]
    fn test_allocate_first() {
        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();

        let first = allocator.allocate_first().unwrap();
        assert_eq!(first.to_string(), "10.200.0.1");

        // Can't allocate first again
        assert!(allocator.allocate_first().is_err());
    }

    #[test]
    fn test_allocate_specific() {
        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();

        let specific_ip: IpAddr = "10.200.0.50".parse().unwrap();
        allocator.allocate_specific(specific_ip).unwrap();

        assert!(allocator.is_allocated(specific_ip));

        // Can't allocate same IP again
        assert!(allocator.allocate_specific(specific_ip).is_err());
    }

    #[test]
    fn test_allocate_specific_out_of_range() {
        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();

        let out_of_range: IpAddr = "192.168.1.1".parse().unwrap();
        assert!(allocator.allocate_specific(out_of_range).is_err());
    }

    #[test]
    fn test_release() {
        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();

        let ip = allocator.allocate().unwrap();
        assert!(allocator.is_allocated(ip));

        assert!(allocator.release(ip));
        assert!(!allocator.is_allocated(ip));

        // Can allocate same IP again
        let ip2 = allocator.allocate().unwrap();
        assert_eq!(ip, ip2);
    }

    #[test]
    fn test_mark_allocated() {
        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();

        let ip: IpAddr = "10.200.0.100".parse().unwrap();
        allocator.mark_allocated(ip).unwrap();

        assert!(allocator.is_allocated(ip));
    }

    #[test]
    fn test_contains() {
        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();

        assert!(allocator.contains("10.200.0.50".parse().unwrap()));
        assert!(!allocator.contains("10.201.0.50".parse().unwrap()));
    }

    #[test]
    fn test_total_hosts() {
        // /24 has 254 usable hosts
        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
        assert_eq!(allocator.total_hosts(), 254);

        // /30 has 2 usable hosts
        let allocator = IpAllocator::new("10.200.0.0/30").unwrap();
        assert_eq!(allocator.total_hosts(), 2);
    }

    #[test]
    fn test_available_count() {
        let mut allocator = IpAllocator::new("10.200.0.0/30").unwrap();

        assert_eq!(allocator.available_count(), 2);

        allocator.allocate();
        assert_eq!(allocator.available_count(), 1);

        allocator.allocate();
        assert_eq!(allocator.available_count(), 0);
    }

    #[test]
    fn test_state_roundtrip() {
        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
        allocator.allocate();
        allocator.allocate();

        let state = allocator.to_state();
        let restored = IpAllocator::from_state(state).unwrap();

        assert_eq!(allocator.cidr(), restored.cidr());
        assert_eq!(allocator.allocated_count(), restored.allocated_count());
    }

    #[test]
    fn test_first_ip_from_cidr() {
        let ip = first_ip_from_cidr("10.200.0.0/24").unwrap();
        assert_eq!(ip.to_string(), "10.200.0.1");
    }

    #[test]
    fn test_network_addr_v4() {
        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
        assert_eq!(
            allocator.network_addr(),
            IpAddr::V4("10.200.0.0".parse().unwrap())
        );
    }

    #[test]
    fn test_broadcast_addr_v4() {
        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
        assert_eq!(
            allocator.broadcast_addr(),
            IpAddr::V4("10.200.0.255".parse().unwrap())
        );
    }

    #[test]
    fn test_host_prefix_len_v4() {
        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
        assert_eq!(allocator.host_prefix_len(), 32);
    }

    // ========================
    // IPv6 Tests
    // ========================

    #[test]
    fn test_allocator_new_v6() {
        let allocator = IpAllocator::new("fd00::/48").unwrap();
        assert_eq!(allocator.cidr(), "fd00::/48");
        assert_eq!(allocator.allocated_count(), 0);
    }

    #[test]
    fn test_allocate_sequential_v6() {
        let mut allocator = IpAllocator::new("fd00::/126").unwrap();

        // /126 has 3 usable hosts (4 addresses total, minus the network address)
        let ip1 = allocator.allocate().unwrap();
        let ip2 = allocator.allocate().unwrap();
        let ip3 = allocator.allocate().unwrap();

        assert_eq!(ip1.to_string(), "fd00::1");
        assert_eq!(ip2.to_string(), "fd00::2");
        assert_eq!(ip3.to_string(), "fd00::3");

        // Should be exhausted
        assert!(allocator.allocate().is_none());
    }

    #[test]
    fn test_allocate_first_v6() {
        let mut allocator = IpAllocator::new("fd00::/48").unwrap();

        let first = allocator.allocate_first().unwrap();
        assert_eq!(first.to_string(), "fd00::1");

        // Can't allocate first again
        assert!(allocator.allocate_first().is_err());
    }

    #[test]
    fn test_allocate_specific_v6() {
        let mut allocator = IpAllocator::new("fd00::/48").unwrap();

        let specific_ip: IpAddr = "fd00::beef".parse().unwrap();
        allocator.allocate_specific(specific_ip).unwrap();

        assert!(allocator.is_allocated(specific_ip));

        // Can't allocate same IP again
        assert!(allocator.allocate_specific(specific_ip).is_err());
    }

    #[test]
    fn test_allocate_specific_out_of_range_v6() {
        let mut allocator = IpAllocator::new("fd00::/48").unwrap();

        let out_of_range: IpAddr = "fe80::1".parse().unwrap();
        assert!(allocator.allocate_specific(out_of_range).is_err());
    }

    #[test]
    fn test_release_v6() {
        let mut allocator = IpAllocator::new("fd00::/48").unwrap();

        let ip = allocator.allocate().unwrap();
        assert!(allocator.is_allocated(ip));

        assert!(allocator.release(ip));
        assert!(!allocator.is_allocated(ip));

        // Can allocate same IP again
        let ip2 = allocator.allocate().unwrap();
        assert_eq!(ip, ip2);
    }

    #[test]
    fn test_mark_allocated_v6() {
        let mut allocator = IpAllocator::new("fd00::/48").unwrap();

        let ip: IpAddr = "fd00::ff".parse().unwrap();
        allocator.mark_allocated(ip).unwrap();

        assert!(allocator.is_allocated(ip));
    }

    #[test]
    fn test_contains_v6() {
        let allocator = IpAllocator::new("fd00::/48").unwrap();

        assert!(allocator.contains("fd00::50".parse().unwrap()));
        assert!(!allocator.contains("fe80::1".parse().unwrap()));
    }

    #[test]
    fn test_total_hosts_v6_small() {
        // /126 has 3 usable hosts (skip network addr)
        let allocator = IpAllocator::new("fd00::/126").unwrap();
        assert_eq!(allocator.total_hosts(), 3);

        // /127 has 1 usable host
        let allocator = IpAllocator::new("fd00::/127").unwrap();
        assert_eq!(allocator.total_hosts(), 1);
    }

    #[test]
    fn test_total_hosts_v6_large() {
        // /48 has 2^80 - 1 usable hosts, which saturates to u32::MAX
        let allocator = IpAllocator::new("fd00::/48").unwrap();
        assert_eq!(allocator.total_hosts(), u32::MAX);
    }

    #[test]
    fn test_available_count_v6() {
        let mut allocator = IpAllocator::new("fd00::/126").unwrap();

        assert_eq!(allocator.available_count(), 3);

        allocator.allocate();
        assert_eq!(allocator.available_count(), 2);

        allocator.allocate();
        assert_eq!(allocator.available_count(), 1);

        allocator.allocate();
        assert_eq!(allocator.available_count(), 0);
    }

    #[test]
    fn test_state_roundtrip_v6() {
        let mut allocator = IpAllocator::new("fd00::/48").unwrap();
        allocator.allocate();
        allocator.allocate();

        let state = allocator.to_state();

        // Verify IpAddr serializes as strings (backward-compatible)
        let json = serde_json::to_string_pretty(&state).unwrap();
        assert!(json.contains("fd00::1"));
        assert!(json.contains("fd00::2"));

        let restored = IpAllocator::from_state(state).unwrap();

        assert_eq!(allocator.cidr(), restored.cidr());
        assert_eq!(allocator.allocated_count(), restored.allocated_count());
    }

    #[test]
    fn test_first_ip_from_cidr_v6() {
        let ip = first_ip_from_cidr("fd00::/48").unwrap();
        assert_eq!(ip.to_string(), "fd00::1");
    }

    #[test]
    fn test_network_addr_v6() {
        let allocator = IpAllocator::new("fd00::/48").unwrap();
        assert_eq!(
            allocator.network_addr(),
            IpAddr::V6("fd00::".parse().unwrap())
        );
    }

    #[test]
    fn test_broadcast_addr_v6() {
        let allocator = IpAllocator::new("fd00::/126").unwrap();
        assert_eq!(
            allocator.broadcast_addr(),
            IpAddr::V6("fd00::3".parse().unwrap())
        );
    }

    #[test]
    fn test_host_prefix_len_v6() {
        let allocator = IpAllocator::new("fd00::/48").unwrap();
        assert_eq!(allocator.host_prefix_len(), 128);
    }

    // ========================
    // Cross-protocol tests
    // ========================

    #[test]
    fn test_v4_and_v6_allocators_independent() {
        let mut v4 = IpAllocator::new("10.200.0.0/30").unwrap();
        let mut v6 = IpAllocator::new("fd00::/126").unwrap();

        let v4_ip = v4.allocate().unwrap();
        let v6_ip = v6.allocate().unwrap();

        assert!(v4_ip.is_ipv4());
        assert!(v6_ip.is_ipv6());
        assert_eq!(v4_ip.to_string(), "10.200.0.1");
        assert_eq!(v6_ip.to_string(), "fd00::1");
    }

    #[test]
    fn test_ipv6_does_not_contain_ipv4() {
        let allocator = IpAllocator::new("fd00::/48").unwrap();
        assert!(!allocator.contains("10.200.0.1".parse().unwrap()));
    }

    #[test]
    fn test_ipv4_does_not_contain_ipv6() {
        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
        assert!(!allocator.contains("fd00::1".parse().unwrap()));
    }

    #[test]
    fn test_allocate_specific_wrong_family() {
        let mut v4_alloc = IpAllocator::new("10.200.0.0/24").unwrap();
        let v6_ip: IpAddr = "fd00::1".parse().unwrap();
        assert!(v4_alloc.allocate_specific(v6_ip).is_err());

        let mut v6_alloc = IpAllocator::new("fd00::/48").unwrap();
        let v4_ip: IpAddr = "10.200.0.1".parse().unwrap();
        assert!(v6_alloc.allocate_specific(v4_ip).is_err());
    }

    // ========================
    // Helper function tests
    // ========================

    #[test]
    fn test_ipv4_add() {
        let base: Ipv4Addr = "10.0.0.0".parse().unwrap();
        assert_eq!(ipv4_add(base, 1), Some("10.0.0.1".parse().unwrap()));
        assert_eq!(ipv4_add(base, 256), Some("10.0.1.0".parse().unwrap()));
    }

    #[test]
    fn test_ipv4_add_overflow() {
        let base: Ipv4Addr = "255.255.255.255".parse().unwrap();
        assert_eq!(ipv4_add(base, 1), None);
    }

    #[test]
    fn test_ipv6_add() {
        let base: Ipv6Addr = "fd00::".parse().unwrap();
        assert_eq!(ipv6_add(base, 1), Some("fd00::1".parse().unwrap()));
        assert_eq!(ipv6_add(base, 0xffff), Some("fd00::ffff".parse().unwrap()));
    }

    #[test]
    fn test_ipv6_add_overflow() {
        let base: Ipv6Addr = "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff".parse().unwrap();
        assert_eq!(ipv6_add(base, 1), None);
    }

    #[test]
    fn test_host_count_v4() {
        assert_eq!(host_count(false, 24), 254); // 2^8 - 2
        assert_eq!(host_count(false, 30), 2); // 2^2 - 2
        assert_eq!(host_count(false, 16), 65534); // 2^16 - 2
        assert_eq!(host_count(false, 31), 0); // /31 — no classical hosts
        assert_eq!(host_count(false, 32), 0); // /32 — single address
    }

    #[test]
    fn test_host_count_v6() {
        assert_eq!(host_count(true, 126), 3); // 2^2 - 1
        assert_eq!(host_count(true, 127), 1); // 2^1 - 1
        assert_eq!(host_count(true, 128), 0); // /128 — single address (is network addr)
        assert_eq!(host_count(true, 64), (1u128 << 64) - 1); // 2^64 - 1
    }

    // ========================
    // NodeSliceAllocator tests
    // ========================

    fn cluster() -> IpNet {
        "10.200.0.0/16".parse().unwrap()
    }

    #[test]
    fn test_slice_new_rejects_equal_prefix() {
        let err = NodeSliceAllocator::new(cluster(), 16).unwrap_err();
        assert!(matches!(err, OverlayError::InvalidCidr(_)));
    }

    #[test]
    fn test_slice_new_rejects_smaller_prefix() {
        let err = NodeSliceAllocator::new(cluster(), 8).unwrap_err();
        assert!(matches!(err, OverlayError::InvalidCidr(_)));
    }

    #[test]
    fn test_slice_new_rejects_over_max() {
        let err = NodeSliceAllocator::new(cluster(), 33).unwrap_err();
        assert!(matches!(err, OverlayError::InvalidCidr(_)));
    }

    #[test]
    fn test_slice_capacity_28_in_16() {
        let allocator = NodeSliceAllocator::new(cluster(), 28).unwrap();
        // /16 → /28 ⇒ 2^12 = 4096 slices
        assert_eq!(allocator.capacity(), 4096);
    }

    #[test]
    fn test_slice_capacity_24_in_16() {
        let allocator = NodeSliceAllocator::new(cluster(), 24).unwrap();
        // /16 → /24 ⇒ 2^8 = 256 slices
        assert_eq!(allocator.capacity(), 256);
    }

    #[test]
    fn test_slice_assign_is_within_cluster() {
        let mut allocator = NodeSliceAllocator::new(cluster(), 28).unwrap();
        let slice = allocator.assign("node-a").unwrap();
        assert_eq!(slice.prefix_len(), 28);
        assert!(cluster().contains(&slice.network()));
    }

    #[test]
    fn test_slice_assign_is_idempotent() {
        let mut allocator = NodeSliceAllocator::new(cluster(), 28).unwrap();
        let first = allocator.assign("node-a").unwrap();
        let second = allocator.assign("node-a").unwrap();
        assert_eq!(first, second);
        assert_eq!(allocator.assigned_count(), 1);
    }

    #[test]
    fn test_slice_assign_different_nodes_get_different_slices() {
        let mut allocator = NodeSliceAllocator::new(cluster(), 28).unwrap();
        let a = allocator.assign("node-a").unwrap();
        let b = allocator.assign("node-b").unwrap();
        let c = allocator.assign("node-c").unwrap();
        assert_ne!(a, b);
        assert_ne!(b, c);
        assert_ne!(a, c);
    }

    #[test]
    fn test_slice_release() {
        let mut allocator = NodeSliceAllocator::new(cluster(), 28).unwrap();
        let slice = allocator.assign("node-a").unwrap();
        assert_eq!(allocator.slice_for("node-a"), Some(slice));

        assert!(allocator.release("node-a"));
        assert_eq!(allocator.slice_for("node-a"), None);

        // Release of unknown node returns false.
        assert!(!allocator.release("node-a"));
    }

    #[test]
    fn test_slice_collision_probes_forward() {
        // Use a very small cluster → few slices → high probability that two
        // arbitrary IDs hash to the same candidate index. Force a true collision
        // by manually occupying the slot a second node's hash maps to.
        let small: IpNet = "10.200.0.0/28".parse().unwrap();
        let mut allocator = NodeSliceAllocator::new(small, 30).unwrap();
        // /28 → /30 ⇒ 2^2 = 4 slices
        assert_eq!(allocator.capacity(), 4);

        // Assign 4 nodes — all must succeed and all must land on distinct slices.
        let ids = ["a", "b", "c", "d"];
        let mut slices: Vec<IpNet> = Vec::new();
        for id in ids {
            let slice = allocator.assign(id).unwrap();
            assert!(
                !slices.contains(&slice),
                "slice {slice} re-assigned; all slices must be distinct"
            );
            slices.push(slice);
        }
        assert_eq!(allocator.assigned_count(), 4);
    }

    #[test]
    fn test_slice_exhaustion_4096() {
        let mut allocator = NodeSliceAllocator::new(cluster(), 28).unwrap();
        // Fill every one of the 4096 slices.
        for i in 0..4096u32 {
            let id = format!("node-{i}");
            allocator.assign(&id).unwrap();
        }
        assert_eq!(allocator.assigned_count(), 4096);

        // The next assignment must fail with NoAvailableIps.
        let err = allocator.assign("node-4096").unwrap_err();
        assert!(matches!(err, OverlayError::NoAvailableIps));
    }

    #[test]
    fn test_slice_snapshot_roundtrip() {
        let mut allocator = NodeSliceAllocator::new(cluster(), 28).unwrap();
        let slice_a = allocator.assign("node-a").unwrap();
        let slice_b = allocator.assign("node-b").unwrap();
        let slice_c = allocator.assign("node-c").unwrap();

        let snapshot = allocator.snapshot();

        // Round-trip through JSON serialization too.
        let json = serde_json::to_string(&snapshot).unwrap();
        let snapshot_restored: NodeSliceAllocatorSnapshot = serde_json::from_str(&json).unwrap();

        let restored = NodeSliceAllocator::restore(snapshot_restored).unwrap();
        assert_eq!(restored.slice_for("node-a"), Some(slice_a));
        assert_eq!(restored.slice_for("node-b"), Some(slice_b));
        assert_eq!(restored.slice_for("node-c"), Some(slice_c));
        assert_eq!(restored.capacity(), 4096);
        assert_eq!(restored.slice_prefix(), 28);
        assert_eq!(restored.cluster_cidr(), cluster());
    }

    #[test]
    fn test_slice_restore_rejects_mismatched_prefix() {
        let snapshot = NodeSliceAllocatorSnapshot {
            cluster_cidr: "10.200.0.0/16".to_string(),
            slice_prefix: 28,
            assigned: vec![("node-a".to_string(), "10.200.0.0/24".to_string())],
        };
        let err = NodeSliceAllocator::restore(snapshot).unwrap_err();
        assert!(matches!(err, OverlayError::InvalidCidr(_)));
    }

    #[test]
    fn test_slice_restore_rejects_out_of_cluster() {
        let snapshot = NodeSliceAllocatorSnapshot {
            cluster_cidr: "10.200.0.0/16".to_string(),
            slice_prefix: 28,
            assigned: vec![("node-a".to_string(), "10.201.0.0/28".to_string())],
        };
        let err = NodeSliceAllocator::restore(snapshot).unwrap_err();
        assert!(matches!(err, OverlayError::InvalidCidr(_)));
    }

    #[test]
    fn test_slice_hash_is_deterministic() {
        // Two allocators built fresh should produce the same first-assignment
        // index for the same node ID — critical for consistency across leader
        // restart on a fresh cluster (before any snapshot exists).
        let mut a = NodeSliceAllocator::new(cluster(), 28).unwrap();
        let mut b = NodeSliceAllocator::new(cluster(), 28).unwrap();
        let slice_a = a.assign("my-node-id").unwrap();
        let slice_b = b.assign("my-node-id").unwrap();
        assert_eq!(slice_a, slice_b);
    }

    #[test]
    fn test_slice_allocator_v6() {
        let cluster_v6: IpNet = "fd00:200::/48".parse().unwrap();
        let mut allocator = NodeSliceAllocator::new(cluster_v6, 64).unwrap();
        // /48 → /64 ⇒ 2^16 = 65536 slices
        assert_eq!(allocator.capacity(), 65536);

        let slice = allocator.assign("node-a").unwrap();
        assert_eq!(slice.prefix_len(), 64);
        assert!(cluster_v6.contains(&slice.network()));
    }
}