zlayer_overlay/

bootstrap.rs

//! Overlay network bootstrap functionality
//!
//! Provides initialization and joining capabilities for overlay networks,
//! including keypair generation, interface creation, and peer management.

use crate::allocator::IpAllocator;
use crate::allocator::{NodeSliceAllocator, NodeSliceAllocatorSnapshot};
use crate::config::PeerInfo;
use crate::dns::{peer_hostname, DnsConfig, DnsHandle, DnsServer, DEFAULT_DNS_PORT};
use crate::error::{OverlayError, Result};
#[cfg(feature = "nat")]
use crate::nat::{Candidate, ConnectionType, NatTraversal, RelayServer};
use crate::transport::OverlayTransport;
use ipnet::IpNet;
use serde::{Deserialize, Serialize};
use std::net::{IpAddr, SocketAddr};
use std::path::{Path, PathBuf};
use std::time::Duration;
use tracing::{debug, info, warn};

/// Default overlay interface name for `ZLayer`
///
/// On macOS, this is `"utun"` which tells boringtun to let the kernel
/// auto-assign a `utunN` device. On Linux, a custom name is used.
#[cfg(target_os = "macos")]
pub const DEFAULT_INTERFACE_NAME: &str = "utun";
#[cfg(not(target_os = "macos"))]
pub const DEFAULT_INTERFACE_NAME: &str = "zl-overlay0";

/// Default overlay listen port (re-exported from `zlayer-core`).
pub use zlayer_core::DEFAULT_WG_PORT;

/// Default overlay network CIDR (IPv4)
pub const DEFAULT_OVERLAY_CIDR: &str = "10.200.0.0/16";

/// Default overlay network CIDR (IPv6)
///
/// Uses a ULA (Unique Local Address) prefix in the `fd00::/8` range.
/// The `fd00:200::/48` prefix mirrors the IPv4 `10.200.0.0/16` convention.
pub const DEFAULT_OVERLAY_CIDR_V6: &str = "fd00:200::/48";

/// Default persistent keepalive interval (seconds)
pub const DEFAULT_KEEPALIVE_SECS: u16 = 25;

/// Default per-node slice prefix length for carving the cluster CIDR into
/// per-node allocation slices. A `/28` gives each node 16 addresses
/// (14 usable hosts) from the cluster `/16`, supporting up to 4096 nodes.
pub const DEFAULT_SLICE_PREFIX: u8 = 28;

/// Serde helper for `Option<IpNet>`: serialize as a CIDR string (e.g.
/// `"10.200.7.0/28"`), so we don't depend on `ipnet`'s optional `serde`
/// feature flag. Keeps persisted state compact and human-readable.
mod option_ipnet_str {
    use ipnet::IpNet;
    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    #[allow(clippy::ref_option)]
    pub fn serialize<S>(value: &Option<IpNet>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        value.map(|v| v.to_string()).serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Option<IpNet>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let opt = Option::<String>::deserialize(deserializer)?;
        match opt {
            None => Ok(None),
            Some(s) => s
                .parse::<IpNet>()
                .map(Some)
                .map_err(serde::de::Error::custom),
        }
    }
}

/// Overlay network bootstrap configuration
///
/// Contains all configuration needed to initialize and manage
/// an overlay network on a node.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BootstrapConfig {
    /// Network CIDR (e.g., "10.200.0.0/16")
    pub cidr: String,

    /// This node's overlay IP address (IPv4 or IPv6)
    pub node_ip: IpAddr,

    /// Overlay interface name
    pub interface: String,

    /// Overlay listen port
    pub port: u16,

    /// This node's overlay private key
    pub private_key: String,

    /// This node's overlay public key
    pub public_key: String,

    /// Whether this node is the cluster leader
    pub is_leader: bool,

    /// Creation timestamp (Unix epoch seconds)
    pub created_at: u64,

    /// Per-node slice of the cluster CIDR that this node draws container
    /// IPs from. `None` for pre-slice-aware configs (back-compat). When set,
    /// this replaces `node_ip/32` in `allowed_ip()`.
    #[serde(default, with = "option_ipnet_str")]
    pub slice_cidr: Option<IpNet>,
}

impl BootstrapConfig {
    /// Get the overlay IP with host prefix for allowed IPs
    ///
    /// When `slice_cidr` is set, returns the slice CIDR string unchanged so
    /// the overlay accepts the full per-node slice of container IPs. Otherwise
    /// falls back to the legacy `/32` (IPv4) or `/128` (IPv6) single-host
    /// representation of `node_ip` for back-compat with pre-slice configs.
    #[must_use]
    pub fn allowed_ip(&self) -> String {
        if let Some(slice) = self.slice_cidr {
            return slice.to_string();
        }
        let prefix = match self.node_ip {
            IpAddr::V4(_) => 32,
            IpAddr::V6(_) => 128,
        };
        format!("{}/{}", self.node_ip, prefix)
    }
}

/// Peer configuration for overlay network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerConfig {
    /// Peer's node ID (for identification)
    pub node_id: String,

    /// Peer's overlay public key
    pub public_key: String,

    /// Peer's public endpoint (host:port)
    pub endpoint: String,

    /// Peer's overlay IP address (IPv4 or IPv6)
    pub overlay_ip: IpAddr,

    /// Optional persistent keepalive interval in seconds
    #[serde(default)]
    pub keepalive: Option<u16>,

    /// Optional custom DNS hostname for this peer (without zone suffix)
    /// If provided, the peer will be registered with this name in addition
    /// to the auto-generated IP-based hostname.
    #[serde(default)]
    pub hostname: Option<String>,

    /// NAT traversal candidates for this peer
    #[serde(default)]
    #[cfg(feature = "nat")]
    pub candidates: Vec<Candidate>,

    /// How this peer is currently connected
    #[serde(default)]
    #[cfg(feature = "nat")]
    pub connection_type: ConnectionType,

    /// Peer's per-node slice CIDR. `None` on old configs; when set, used as
    /// the peer's `allowed_ips` instead of `overlay_ip/32`.
    #[serde(default, with = "option_ipnet_str")]
    pub slice_cidr: Option<IpNet>,
}

impl PeerConfig {
    /// Create a new peer configuration
    #[must_use]
    pub fn new(node_id: String, public_key: String, endpoint: String, overlay_ip: IpAddr) -> Self {
        Self {
            node_id,
            public_key,
            endpoint,
            overlay_ip,
            keepalive: Some(DEFAULT_KEEPALIVE_SECS),
            hostname: None,
            #[cfg(feature = "nat")]
            candidates: Vec::new(),
            #[cfg(feature = "nat")]
            connection_type: ConnectionType::default(),
            slice_cidr: None,
        }
    }

    /// Set a custom DNS hostname for this peer
    #[must_use]
    pub fn with_hostname(mut self, hostname: impl Into<String>) -> Self {
        self.hostname = Some(hostname.into());
        self
    }

    /// Set the per-node slice CIDR for this peer.
    ///
    /// When set, `to_peer_info()` uses the slice CIDR as the peer's
    /// `allowed_ips` (so the overlay accepts the peer's full slice of
    /// container IPs) instead of the legacy single-host `overlay_ip/32`.
    #[must_use]
    pub fn with_slice_cidr(mut self, cidr: IpNet) -> Self {
        self.slice_cidr = Some(cidr);
        self
    }

    /// Convert to `PeerInfo` for overlay transport configuration
    ///
    /// When `slice_cidr` is set, the peer's `allowed_ips` is the slice CIDR
    /// string. Otherwise, falls back to the legacy `overlay_ip/{32,128}`
    /// host-prefix representation for back-compat.
    ///
    /// # Errors
    ///
    /// Returns an error if the endpoint address cannot be parsed.
    pub fn to_peer_info(&self) -> std::result::Result<PeerInfo, Box<dyn std::error::Error>> {
        let endpoint: SocketAddr = self.endpoint.parse()?;
        let keepalive =
            Duration::from_secs(u64::from(self.keepalive.unwrap_or(DEFAULT_KEEPALIVE_SECS)));

        let allowed_ips = if let Some(slice) = self.slice_cidr {
            slice.to_string()
        } else {
            let prefix = match self.overlay_ip {
                IpAddr::V4(_) => 32,
                IpAddr::V6(_) => 128,
            };
            format!("{}/{}", self.overlay_ip, prefix)
        };

        Ok(PeerInfo::new(
            self.public_key.clone(),
            endpoint,
            &allowed_ips,
            keepalive,
        ))
    }
}

/// Persistent state for the overlay bootstrap
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BootstrapState {
    /// Bootstrap configuration
    pub config: BootstrapConfig,

    /// List of configured peers
    pub peers: Vec<PeerConfig>,

    /// IP allocator state (only for leader)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub allocator_state: Option<crate::allocator::IpAllocatorState>,

    /// Per-node slice allocator state (only for leader).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub slice_allocator_state: Option<NodeSliceAllocatorSnapshot>,
}

/// Bootstrap manager for overlay network
///
/// Handles overlay network initialization, peer management,
/// and overlay transport interface configuration.
pub struct OverlayBootstrap {
    /// Bootstrap configuration
    config: BootstrapConfig,

    /// Configured peers
    peers: Vec<PeerConfig>,

    /// Data directory for persistent state
    data_dir: PathBuf,

    /// IP allocator (only for leader nodes)
    allocator: Option<IpAllocator>,

    /// Per-node slice allocator (only for leader nodes) — carves the cluster
    /// CIDR into `/28` slices assigned to each joining node. Added to fix the
    /// latent IP-collision bug where every agent independently allocated
    /// container IPs from the full cluster `/16`.
    slice_allocator: Option<NodeSliceAllocator>,

    /// DNS configuration (opt-in)
    dns_config: Option<DnsConfig>,

    /// DNS handle for managing records (available after `start()` if DNS enabled)
    dns_handle: Option<DnsHandle>,

    /// Overlay transport (boringtun device handle).
    ///
    /// Must be kept alive for the overlay network lifetime; dropping the
    /// transport destroys the TUN device.
    transport: Option<OverlayTransport>,

    /// NAT traversal orchestrator (available after `start()` if NAT is enabled)
    #[cfg(feature = "nat")]
    nat_traversal: Option<NatTraversal>,

    /// Built-in relay server (available after `start()` if relay is configured)
    #[cfg(feature = "nat")]
    relay_server: Option<RelayServer>,

    /// NAT traversal configuration. When `None`, [`NatConfig::default()`] is
    /// used at `start()` time so callers that never set a value still pick up
    /// the defaults (NAT enabled, public STUN servers).
    #[cfg(feature = "nat")]
    nat_config: Option<crate::nat::NatConfig>,
}

impl OverlayBootstrap {
    /// Initialize as cluster leader (first node in the overlay)
    ///
    /// This generates a new overlay keypair, allocates the first IP
    /// in the CIDR range, and prepares the node as the overlay leader.
    ///
    /// # Arguments
    /// * `cidr` - Overlay network CIDR (e.g., "10.200.0.0/16")
    /// * `port` - Overlay listen port
    /// * `data_dir` - Directory for persistent state
    ///
    /// # Example
    /// ```ignore
    /// let bootstrap = OverlayBootstrap::init_leader(
    ///     "10.200.0.0/16",
    ///     51820,
    ///     Path::new("/var/lib/zlayer"),
    /// ).await?;
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an error if already initialized, key generation fails, or state cannot be saved.
    pub async fn init_leader(cidr: &str, port: u16, data_dir: &Path) -> Result<Self> {
        // Check if already initialized
        let config_path = data_dir.join("overlay_bootstrap.json");
        if config_path.exists() {
            return Err(OverlayError::AlreadyInitialized(
                config_path.display().to_string(),
            ));
        }

        // Ensure data directory exists
        tokio::fs::create_dir_all(data_dir).await?;

        // Generate overlay keypair
        info!("Generating overlay keypair for leader");
        let (private_key, public_key) = OverlayTransport::generate_keys()
            .await
            .map_err(|e| OverlayError::TransportCommand(e.to_string()))?;

        // Initialize IP allocator (legacy flat allocator, kept in sync for
        // back-compat with callers that still read it) and allocate its
        // first IP so its state tracks at least the leader.
        let mut allocator = IpAllocator::new(cidr)?;
        let _legacy_first = allocator.allocate_first()?;

        // Initialize the per-node slice allocator and carve the leader's
        // own slice. The leader's `node_ip` is taken as the first usable
        // host of its slice so it always lives inside the slice it owns.
        let cluster_cidr: IpNet = cidr
            .parse()
            .map_err(|e: ipnet::AddrParseError| OverlayError::InvalidCidr(e.to_string()))?;
        let mut slice_allocator = NodeSliceAllocator::new(cluster_cidr, DEFAULT_SLICE_PREFIX)?;
        let leader_slice = slice_allocator.assign("leader")?;
        let node_ip = leader_slice.hosts().next().unwrap_or_else(|| {
            // Fall back to network() + 1 for very small slices where
            // `hosts()` can return empty. We still need a valid address
            // to assign to the interface.
            match leader_slice.network() {
                IpAddr::V4(v4) => {
                    let bits = u32::from(v4).saturating_add(1);
                    IpAddr::V4(std::net::Ipv4Addr::from(bits))
                }
                IpAddr::V6(v6) => {
                    let bits = u128::from(v6).saturating_add(1);
                    IpAddr::V6(std::net::Ipv6Addr::from(bits))
                }
            }
        });

        info!(
            node_ip = %node_ip,
            cidr = cidr,
            slice = %leader_slice,
            "Allocated leader IP from slice"
        );

        // Create config
        let config = BootstrapConfig {
            cidr: cidr.to_string(),
            node_ip,
            interface: DEFAULT_INTERFACE_NAME.to_string(),
            port,
            private_key,
            public_key,
            is_leader: true,
            created_at: current_timestamp(),
            slice_cidr: Some(leader_slice),
        };

        let bootstrap = Self {
            config,
            peers: Vec::new(),
            data_dir: data_dir.to_path_buf(),
            allocator: Some(allocator),
            slice_allocator: Some(slice_allocator),
            dns_config: None,
            dns_handle: None,
            transport: None,
            #[cfg(feature = "nat")]
            nat_traversal: None,
            #[cfg(feature = "nat")]
            relay_server: None,
            #[cfg(feature = "nat")]
            nat_config: None,
        };

        // Persist state
        bootstrap.save().await?;

        Ok(bootstrap)
    }

    /// Join an existing overlay network
    ///
    /// Generates a new overlay keypair and configures this node
    /// to connect to an existing overlay network.
    ///
    /// # Arguments
    /// * `leader_cidr` - Leader's overlay network CIDR
    /// * `leader_endpoint` - Leader's public endpoint (host:port)
    /// * `leader_public_key` - Leader's overlay public key
    /// * `leader_overlay_ip` - Leader's overlay IP address
    /// * `allocated_ip` - IP address allocated for this node by the leader
    /// * `port` - Overlay listen port for this node
    /// * `slice_cidr` - Per-node slice assigned by the leader (if any). When
    ///   set, the worker uses this slice as its `allowed_ip`; `None` preserves
    ///   the legacy `allocated_ip/32` behavior for pre-slice-aware clusters.
    /// * `data_dir` - Directory for persistent state
    ///
    /// # Errors
    ///
    /// Returns an error if already initialized, key generation fails, or state cannot be saved.
    #[allow(clippy::too_many_arguments)]
    pub async fn join(
        leader_cidr: &str,
        leader_endpoint: &str,
        leader_public_key: &str,
        leader_overlay_ip: IpAddr,
        allocated_ip: IpAddr,
        port: u16,
        slice_cidr: Option<IpNet>,
        data_dir: &Path,
    ) -> Result<Self> {
        // Check if already initialized
        let config_path = data_dir.join("overlay_bootstrap.json");
        if config_path.exists() {
            return Err(OverlayError::AlreadyInitialized(
                config_path.display().to_string(),
            ));
        }

        // Ensure data directory exists
        tokio::fs::create_dir_all(data_dir).await?;

        // Generate overlay keypair for this node
        info!("Generating overlay keypair for joining node");
        let (private_key, public_key) = OverlayTransport::generate_keys()
            .await
            .map_err(|e| OverlayError::TransportCommand(e.to_string()))?;

        // Create config
        let config = BootstrapConfig {
            cidr: leader_cidr.to_string(),
            node_ip: allocated_ip,
            interface: DEFAULT_INTERFACE_NAME.to_string(),
            port,
            private_key,
            public_key,
            is_leader: false,
            created_at: current_timestamp(),
            slice_cidr,
        };

        // Add leader as the first peer
        let leader_peer = PeerConfig {
            node_id: "leader".to_string(),
            public_key: leader_public_key.to_string(),
            endpoint: leader_endpoint.to_string(),
            overlay_ip: leader_overlay_ip,
            keepalive: Some(DEFAULT_KEEPALIVE_SECS),
            hostname: None, // Leader gets its own DNS alias "leader.zone"
            #[cfg(feature = "nat")]
            candidates: Vec::new(),
            #[cfg(feature = "nat")]
            connection_type: ConnectionType::default(),
            slice_cidr: None,
        };

        info!(
            leader_endpoint = leader_endpoint,
            overlay_ip = %allocated_ip,
            "Configured leader as peer"
        );

        let bootstrap = Self {
            config,
            peers: vec![leader_peer],
            data_dir: data_dir.to_path_buf(),
            allocator: None, // Workers don't manage IP allocation
            slice_allocator: None,
            dns_config: None,
            dns_handle: None,
            transport: None,
            #[cfg(feature = "nat")]
            nat_traversal: None,
            #[cfg(feature = "nat")]
            relay_server: None,
            #[cfg(feature = "nat")]
            nat_config: None,
        };

        // Persist state
        bootstrap.save().await?;

        Ok(bootstrap)
    }

    /// Load existing bootstrap state from disk
    ///
    /// # Errors
    ///
    /// Returns an error if the state file is missing, unreadable, or invalid.
    pub async fn load(data_dir: &Path) -> Result<Self> {
        let config_path = data_dir.join("overlay_bootstrap.json");

        if !config_path.exists() {
            return Err(OverlayError::NotInitialized);
        }

        let contents = tokio::fs::read_to_string(&config_path).await?;
        let state: BootstrapState = serde_json::from_str(&contents)?;

        let allocator = if let Some(alloc_state) = state.allocator_state {
            Some(IpAllocator::from_state(alloc_state)?)
        } else {
            None
        };

        let slice_allocator = if let Some(snapshot) = state.slice_allocator_state {
            Some(NodeSliceAllocator::restore(snapshot)?)
        } else {
            None
        };

        Ok(Self {
            config: state.config,
            peers: state.peers,
            data_dir: data_dir.to_path_buf(),
            allocator,
            slice_allocator,
            dns_config: None, // DNS config must be re-enabled after load
            dns_handle: None,
            transport: None,
            #[cfg(feature = "nat")]
            nat_traversal: None,
            #[cfg(feature = "nat")]
            relay_server: None,
            #[cfg(feature = "nat")]
            nat_config: None,
        })
    }

    /// Save bootstrap state to disk
    ///
    /// # Errors
    ///
    /// Returns an error if serialization or file writing fails.
    pub async fn save(&self) -> Result<()> {
        let config_path = self.data_dir.join("overlay_bootstrap.json");

        let state = BootstrapState {
            config: self.config.clone(),
            peers: self.peers.clone(),
            allocator_state: self
                .allocator
                .as_ref()
                .map(super::allocator::IpAllocator::to_state),
            slice_allocator_state: self
                .slice_allocator
                .as_ref()
                .map(NodeSliceAllocator::snapshot),
        };

        let contents = serde_json::to_string_pretty(&state)?;
        tokio::fs::write(&config_path, contents).await?;

        debug!(path = %config_path.display(), "Saved bootstrap state");
        Ok(())
    }

    /// Enable DNS service discovery for the overlay network
    ///
    /// When DNS is enabled, peers are automatically registered with both:
    /// - An IP-based hostname: `node-X-Y.zone` (e.g., `node-0-5.overlay.local`)
    /// - A custom hostname if provided in `PeerConfig`
    ///
    /// The leader node additionally gets a `leader.zone` alias.
    ///
    /// # Arguments
    /// * `zone` - DNS zone (e.g., "overlay.local.")
    /// * `port` - DNS server port (default: 15353 to avoid conflicts)
    ///
    /// # Example
    /// ```ignore
    /// let bootstrap = OverlayBootstrap::init_leader(cidr, port, data_dir)
    ///     .await?
    ///     .with_dns("overlay.local.", 15353)?;
    /// bootstrap.start().await?;
    /// ```
    ///
    /// # Errors
    ///
    /// This method currently always succeeds but returns `Result` for API consistency.
    pub fn with_dns(mut self, zone: &str, port: u16) -> Result<Self> {
        self.dns_config = Some(DnsConfig {
            zone: zone.to_string(),
            port,
            bind_addr: self.config.node_ip,
        });
        Ok(self)
    }

    /// Enable DNS with default port (15353)
    ///
    /// # Errors
    ///
    /// This method currently always succeeds but returns `Result` for API consistency.
    pub fn with_dns_default(self, zone: &str) -> Result<Self> {
        self.with_dns(zone, DEFAULT_DNS_PORT)
    }

    /// Override the NAT traversal configuration used by `start()`.
    ///
    /// When unset, `start()` falls back to [`crate::nat::NatConfig::default()`]
    /// which enables NAT traversal with public STUN servers.
    #[cfg(feature = "nat")]
    #[must_use]
    pub fn with_nat_config(mut self, nat: crate::nat::NatConfig) -> Self {
        self.nat_config = Some(nat);
        self
    }

    /// Get the DNS handle for managing records
    ///
    /// Returns None if DNS is not enabled or `start()` hasn't been called yet.
    #[must_use]
    pub fn dns_handle(&self) -> Option<&DnsHandle> {
        self.dns_handle.as_ref()
    }

    /// Check if DNS is enabled
    #[must_use]
    pub fn dns_enabled(&self) -> bool {
        self.dns_config.is_some()
    }

    /// Start the overlay network (create and configure overlay transport)
    ///
    /// This creates the boringtun TUN interface, assigns the overlay IP,
    /// configures all known peers, and starts the DNS server if enabled.
    ///
    /// # Errors
    ///
    /// Returns an error if interface creation, peer configuration, or DNS startup fails.
    pub async fn start(&mut self) -> Result<()> {
        info!(
            interface = %self.config.interface,
            overlay_ip = %self.config.node_ip,
            port = self.config.port,
            dns_enabled = self.dns_config.is_some(),
            "Starting overlay network"
        );

        // Convert our config to OverlayConfig
        let overlay_config = crate::config::OverlayConfig {
            local_endpoint: SocketAddr::new(
                std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
                self.config.port,
            ),
            private_key: self.config.private_key.clone(),
            public_key: self.config.public_key.clone(),
            overlay_cidr: self.config.allowed_ip(),
            cluster_cidr: Some(self.config.cidr.clone()),
            peer_discovery_interval: Duration::from_secs(30),
            #[cfg(feature = "nat")]
            nat: self.nat_config.clone().unwrap_or_default(),
        };

        #[cfg(feature = "nat")]
        let nat_config = overlay_config.nat.clone();

        // Create overlay transport
        let mut transport = OverlayTransport::new(overlay_config, self.config.interface.clone());

        // Create the interface
        transport
            .create_interface()
            .await
            .map_err(|e| OverlayError::TransportCommand(e.to_string()))?;

        // On macOS, the kernel assigns a utunN name that may differ from
        // the requested name. Update our config to reflect the actual name.
        let actual_name = transport.interface_name().to_string();
        if actual_name != self.config.interface {
            info!(
                requested = %self.config.interface,
                actual = %actual_name,
                "Interface name resolved by kernel"
            );
            self.config.interface = actual_name;
        }

        // Convert peers to PeerInfo
        let peer_infos: Vec<PeerInfo> = self
            .peers
            .iter()
            .filter_map(|p| match p.to_peer_info() {
                Ok(info) => Some(info),
                Err(e) => {
                    warn!(peer = %p.node_id, error = %e, "Failed to parse peer info");
                    None
                }
            })
            .collect();

        // Configure transport with peers
        transport
            .configure(&peer_infos)
            .await
            .map_err(|e| OverlayError::TransportCommand(e.to_string()))?;

        // Store the transport so the TUN device stays alive for the overlay
        // lifetime. Dropping the OverlayTransport destroys the boringtun device.
        self.transport = Some(transport);

        // NAT traversal: gather candidates and connect to peers
        #[cfg(feature = "nat")]
        self.start_nat_traversal(nat_config).await;

        // Start DNS server if configured
        self.start_dns().await?;

        info!("Overlay network started successfully");
        Ok(())
    }

    /// Start the DNS server and register all known peers.
    async fn start_dns(&mut self) -> Result<()> {
        let Some(dns_config) = &self.dns_config else {
            return Ok(());
        };

        info!(
            zone = %dns_config.zone,
            port = dns_config.port,
            "Starting DNS server for overlay"
        );

        let dns_server =
            DnsServer::from_config(dns_config).map_err(|e| OverlayError::Dns(e.to_string()))?;

        // Register self with IP-based hostname
        let self_hostname = peer_hostname(self.config.node_ip);
        dns_server
            .add_record(&self_hostname, self.config.node_ip)
            .await
            .map_err(|e| OverlayError::Dns(e.to_string()))?;

        // If leader, also register "leader" alias
        if self.config.is_leader {
            dns_server
                .add_record("leader", self.config.node_ip)
                .await
                .map_err(|e| OverlayError::Dns(e.to_string()))?;
            debug!(ip = %self.config.node_ip, "Registered leader.{}", dns_config.zone);
        }

        // Register existing peers
        for peer in &self.peers {
            // Always register IP-based hostname
            let hostname = peer_hostname(peer.overlay_ip);
            dns_server
                .add_record(&hostname, peer.overlay_ip)
                .await
                .map_err(|e| OverlayError::Dns(e.to_string()))?;

            // Also register custom hostname if provided
            if let Some(custom) = &peer.hostname {
                dns_server
                    .add_record(custom, peer.overlay_ip)
                    .await
                    .map_err(|e| OverlayError::Dns(e.to_string()))?;
                debug!(
                    hostname = custom,
                    ip = %peer.overlay_ip,
                    "Registered custom hostname"
                );
            }
        }

        // Start the DNS server and store the handle
        let handle = dns_server
            .start()
            .await
            .map_err(|e| OverlayError::Dns(e.to_string()))?;
        self.dns_handle = Some(handle);

        info!("DNS server started successfully");
        Ok(())
    }

    /// Initialize NAT traversal, gather candidates, and connect to known peers.
    #[cfg(feature = "nat")]
    async fn start_nat_traversal(&mut self, nat_config: crate::nat::NatConfig) {
        if !nat_config.enabled {
            return;
        }

        // Optionally start built-in relay server
        if let Some(ref relay_config) = nat_config.relay_server {
            let relay = RelayServer::new(relay_config, &self.config.private_key);
            match relay.start().await {
                Ok(()) => {
                    info!("Built-in relay server started");
                    self.relay_server = Some(relay);
                }
                Err(e) => {
                    warn!(error = %e, "Failed to start relay server");
                }
            }
        }

        let mut nat = NatTraversal::new(nat_config, self.config.port);
        match nat.gather_candidates().await {
            Ok(candidates) => {
                info!(count = candidates.len(), "Gathered NAT candidates");
                if let Some(ref transport) = self.transport {
                    for peer in &mut self.peers {
                        if !peer.candidates.is_empty() {
                            match nat
                                .connect_to_peer(transport, &peer.public_key, &peer.candidates)
                                .await
                            {
                                Ok(ct) => {
                                    peer.connection_type = ct;
                                    info!(
                                        peer = %peer.node_id,
                                        connection = %ct,
                                        "NAT traversal succeeded"
                                    );
                                }
                                Err(e) => warn!(
                                    peer = %peer.node_id,
                                    error = %e,
                                    "NAT traversal failed"
                                ),
                            }
                        }
                    }
                }
                self.nat_traversal = Some(nat);
            }
            Err(e) => warn!(error = %e, "NAT candidate gathering failed"),
        }
    }

    /// Stop the overlay network (shut down the boringtun transport)
    ///
    /// # Errors
    ///
    /// This method currently always succeeds but returns `Result` for API consistency.
    #[allow(clippy::unused_async)]
    pub async fn stop(&mut self) -> Result<()> {
        info!(interface = %self.config.interface, "Stopping overlay network");

        if let Some(mut transport) = self.transport.take() {
            transport.shutdown();
        }

        Ok(())
    }

    /// Replay assigned slices from persisted `PeerConfig.slice_cidr` values
    /// back into the slice allocator. Call on leader startup after `load()`
    /// so the in-memory allocator matches what's on disk.
    ///
    /// No-op for workers (who have no allocator).
    ///
    /// # Errors
    ///
    /// Returns an error if the snapshot rebuild fails.
    pub fn reconcile_existing_peers(&mut self) -> Result<()> {
        let Some(ref mut allocator) = self.slice_allocator else {
            return Ok(());
        };
        // Collect into a snapshot so we can restore cleanly.
        let mut assigned: Vec<(String, String)> = Vec::new();
        for peer in &self.peers {
            if let Some(slice) = peer.slice_cidr {
                assigned.push((peer.node_id.clone(), slice.to_string()));
            }
        }
        if assigned.is_empty() {
            return Ok(());
        }
        let snapshot = NodeSliceAllocatorSnapshot {
            cluster_cidr: allocator.cluster_cidr().to_string(),
            slice_prefix: allocator.slice_prefix(),
            assigned,
        };
        *allocator = NodeSliceAllocator::restore(snapshot)?;
        Ok(())
    }

    /// Add a new peer to the overlay network
    ///
    /// For leader nodes, this also allocates an IP address for the peer.
    ///
    /// # Errors
    ///
    /// Returns an error if no IPs are available, DNS registration fails, or state cannot be saved.
    pub async fn add_peer(&mut self, mut peer: PeerConfig) -> Result<IpAddr> {
        // If we're the leader, allocate an IP for this peer
        let overlay_ip = if let Some(ref mut allocator) = self.allocator {
            let ip = allocator.allocate().ok_or(OverlayError::NoAvailableIps)?;
            peer.overlay_ip = ip;
            ip
        } else {
            peer.overlay_ip
        };

        // If we're the leader and have a slice allocator, assign a /28 slice
        // to this peer. Peers track their slice so `to_peer_info()` can emit the
        // correct `allowed_ips`.
        if let Some(ref mut slice_allocator) = self.slice_allocator {
            let slice = slice_allocator.assign(&peer.node_id)?;
            peer.slice_cidr = Some(slice);
            // The peer's overlay_ip (reallocated above from the flat allocator)
            // is left as-is for compatibility with callers that still read it;
            // the authoritative routing info is `slice_cidr`.
        }

        // Add peer to overlay transport via UAPI
        if let Ok(peer_info) = peer.to_peer_info() {
            // Prefer the stored transport; fall back to a temporary instance
            // (UAPI calls work via the Unix socket regardless of DeviceHandle)
            let transport_ref: Option<&OverlayTransport> = self.transport.as_ref();

            let result = if let Some(t) = transport_ref {
                t.add_peer(&peer_info).await
            } else {
                let overlay_config = crate::config::OverlayConfig {
                    local_endpoint: SocketAddr::new(
                        std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
                        self.config.port,
                    ),
                    private_key: self.config.private_key.clone(),
                    public_key: self.config.public_key.clone(),
                    overlay_cidr: self.config.allowed_ip(),
                    cluster_cidr: Some(self.config.cidr.clone()),
                    peer_discovery_interval: Duration::from_secs(30),
                    #[cfg(feature = "nat")]
                    nat: crate::nat::NatConfig::default(),
                };
                let tmp = OverlayTransport::new(overlay_config, self.config.interface.clone());
                tmp.add_peer(&peer_info).await
            };

            match result {
                Ok(()) => debug!(peer = %peer.node_id, "Added peer to overlay"),
                Err(e) => {
                    warn!(peer = %peer.node_id, error = %e, "Failed to add peer to overlay (interface may not be up)");
                }
            }
        }

        // Register peer in DNS if enabled
        if let Some(ref dns_handle) = self.dns_handle {
            // IP-based hostname
            let hostname = peer_hostname(overlay_ip);
            dns_handle
                .add_record(&hostname, overlay_ip)
                .await
                .map_err(|e| OverlayError::Dns(e.to_string()))?;
            debug!(hostname = %hostname, ip = %overlay_ip, "Registered peer in DNS");

            // Custom hostname alias if provided
            if let Some(ref custom) = peer.hostname {
                dns_handle
                    .add_record(custom, overlay_ip)
                    .await
                    .map_err(|e| OverlayError::Dns(e.to_string()))?;
                debug!(hostname = %custom, ip = %overlay_ip, "Registered custom hostname in DNS");
            }
        }

        // NAT traversal for new peer
        #[cfg(feature = "nat")]
        {
            if let (Some(ref nat), Some(ref transport)) = (&self.nat_traversal, &self.transport) {
                if !peer.candidates.is_empty() {
                    match nat
                        .connect_to_peer(transport, &peer.public_key, &peer.candidates)
                        .await
                    {
                        Ok(ct) => {
                            peer.connection_type = ct;
                            info!(
                                peer = %peer.node_id,
                                connection = %ct,
                                "NAT traversal for new peer"
                            );
                        }
                        Err(e) => warn!(
                            peer = %peer.node_id,
                            error = %e,
                            "NAT failed for new peer"
                        ),
                    }
                }
            }
        }

        // Add to peer list
        self.peers.push(peer);

        // Persist state
        self.save().await?;

        info!(peer_ip = %overlay_ip, "Added peer to overlay");
        Ok(overlay_ip)
    }

    /// Remove a peer from the overlay network
    ///
    /// # Errors
    ///
    /// Returns an error if the peer is not found, DNS removal fails, or state cannot be saved.
    pub async fn remove_peer(&mut self, public_key: &str) -> Result<()> {
        // Find the peer
        let peer_idx = self
            .peers
            .iter()
            .position(|p| p.public_key == public_key)
            .ok_or_else(|| OverlayError::PeerNotFound(public_key.to_string()))?;

        let peer = &self.peers[peer_idx];

        // Capture peer info for DNS removal before we lose the reference
        let peer_overlay_ip = peer.overlay_ip;
        let peer_custom_hostname = peer.hostname.clone();

        // Release IP if we're managing allocation
        if let Some(ref mut allocator) = self.allocator {
            allocator.release(peer_overlay_ip);
        }

        // Remove from DNS if enabled
        if let Some(ref dns_handle) = self.dns_handle {
            // Remove IP-based hostname
            let hostname = peer_hostname(peer_overlay_ip);
            dns_handle
                .remove_record(&hostname)
                .await
                .map_err(|e| OverlayError::Dns(e.to_string()))?;
            debug!(hostname = %hostname, "Removed peer from DNS");

            // Remove custom hostname if it was set
            if let Some(ref custom) = peer_custom_hostname {
                dns_handle
                    .remove_record(custom)
                    .await
                    .map_err(|e| OverlayError::Dns(e.to_string()))?;
                debug!(hostname = %custom, "Removed custom hostname from DNS");
            }
        }

        // Remove peer from overlay transport via UAPI
        let transport_ref: Option<&OverlayTransport> = self.transport.as_ref();

        let result = if let Some(t) = transport_ref {
            t.remove_peer(public_key).await
        } else {
            let overlay_config = crate::config::OverlayConfig {
                local_endpoint: SocketAddr::new(
                    std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
                    self.config.port,
                ),
                private_key: self.config.private_key.clone(),
                public_key: self.config.public_key.clone(),
                overlay_cidr: self.config.allowed_ip(),
                cluster_cidr: Some(self.config.cidr.clone()),
                peer_discovery_interval: Duration::from_secs(30),
                #[cfg(feature = "nat")]
                nat: crate::nat::NatConfig::default(),
            };
            let tmp = OverlayTransport::new(overlay_config, self.config.interface.clone());
            tmp.remove_peer(public_key).await
        };

        match result {
            Ok(()) => debug!(public_key = public_key, "Removed peer from overlay"),
            Err(e) => {
                warn!(public_key = public_key, error = %e, "Failed to remove peer from overlay");
            }
        }

        // Remove from peer list
        self.peers.remove(peer_idx);

        // Persist state
        self.save().await?;

        info!(public_key = public_key, "Removed peer from overlay");
        Ok(())
    }

    /// Get this node's public key
    #[must_use]
    pub fn public_key(&self) -> &str {
        &self.config.public_key
    }

    /// Get this node's overlay IP (IPv4 or IPv6)
    #[must_use]
    pub fn node_ip(&self) -> IpAddr {
        self.config.node_ip
    }

    /// Get the overlay CIDR
    #[must_use]
    pub fn cidr(&self) -> &str {
        &self.config.cidr
    }

    /// Get the overlay interface name
    #[must_use]
    pub fn interface(&self) -> &str {
        &self.config.interface
    }

    /// Get the overlay listen port
    #[must_use]
    pub fn port(&self) -> u16 {
        self.config.port
    }

    /// Check if this node is the leader
    #[must_use]
    pub fn is_leader(&self) -> bool {
        self.config.is_leader
    }

    /// Get configured peers
    #[must_use]
    pub fn peers(&self) -> &[PeerConfig] {
        &self.peers
    }

    /// Get the bootstrap config
    #[must_use]
    pub fn config(&self) -> &BootstrapConfig {
        &self.config
    }

    /// Allocate an IP for a new peer (leader only)
    ///
    /// This is used by the control plane when processing join requests.
    ///
    /// # Errors
    ///
    /// Returns an error if this node is not a leader or no IPs are available.
    pub fn allocate_peer_ip(&mut self) -> Result<IpAddr> {
        let allocator = self
            .allocator
            .as_mut()
            .ok_or(OverlayError::Config("Not a leader node".to_string()))?;

        allocator.allocate().ok_or(OverlayError::NoAvailableIps)
    }

    /// Get IP allocation statistics (leader only)
    #[must_use]
    #[allow(clippy::cast_possible_truncation)]
    pub fn allocation_stats(&self) -> Option<(u32, u32)> {
        self.allocator
            .as_ref()
            .map(|a| (a.allocated_count() as u32, a.total_hosts()))
    }

    /// Perform NAT maintenance: refresh STUN, attempt relay upgrades.
    ///
    /// Call this periodically from the runtime's main loop. Re-probes
    /// STUN servers to detect reflexive address changes and attempts
    /// to upgrade relayed connections to direct or hole-punched.
    ///
    /// # Errors
    ///
    /// Returns an error if STUN refresh fails.
    #[cfg(feature = "nat")]
    pub async fn nat_maintenance_tick(&mut self) -> Result<()> {
        let (Some(nat), Some(transport)) = (&mut self.nat_traversal, &self.transport) else {
            return Ok(());
        };

        if nat.refresh().await? {
            info!("Reflexive address changed");
        }

        for peer in &mut self.peers {
            if peer.connection_type == ConnectionType::Relayed && !peer.candidates.is_empty() {
                if let Ok(Some(upgraded)) = nat
                    .attempt_upgrade(transport, &peer.public_key, &peer.candidates)
                    .await
                {
                    peer.connection_type = upgraded;
                    info!(
                        peer = %peer.node_id,
                        connection = %upgraded,
                        "Upgraded relayed connection"
                    );
                }
            }
        }

        Ok(())
    }

    /// Get this node's NAT candidates for sharing with peers.
    ///
    /// Returns an empty vec if NAT traversal has not been initialized
    /// or no candidates were gathered.
    #[cfg(feature = "nat")]
    #[must_use]
    pub fn nat_candidates(&self) -> Vec<Candidate> {
        self.nat_traversal
            .as_ref()
            .map(|n| n.local_candidates().to_vec())
            .unwrap_or_default()
    }
}

/// Get current Unix timestamp
fn current_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

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

    #[test]
    fn test_bootstrap_config_allowed_ip_v4() {
        let config = BootstrapConfig {
            cidr: "10.200.0.0/16".to_string(),
            node_ip: IpAddr::V4(Ipv4Addr::new(10, 200, 0, 1)),
            interface: DEFAULT_INTERFACE_NAME.to_string(),
            port: DEFAULT_WG_PORT,
            private_key: "test_private".to_string(),
            public_key: "test_public".to_string(),
            is_leader: true,
            created_at: 0,
            slice_cidr: None,
        };

        assert_eq!(config.allowed_ip(), "10.200.0.1/32");
    }

    #[test]
    fn test_bootstrap_config_allowed_ip_v6() {
        let config = BootstrapConfig {
            cidr: "fd00:200::/48".to_string(),
            node_ip: "fd00:200::1".parse::<IpAddr>().unwrap(),
            interface: DEFAULT_INTERFACE_NAME.to_string(),
            port: DEFAULT_WG_PORT,
            private_key: "test_private".to_string(),
            public_key: "test_public".to_string(),
            is_leader: true,
            created_at: 0,
            slice_cidr: None,
        };

        assert_eq!(config.allowed_ip(), "fd00:200::1/128");
    }

    #[test]
    fn test_peer_config_new_v4() {
        let peer = PeerConfig::new(
            "node-1".to_string(),
            "pubkey123".to_string(),
            "192.168.1.100:51820".to_string(),
            IpAddr::V4(Ipv4Addr::new(10, 200, 0, 5)),
        );

        assert_eq!(peer.node_id, "node-1");
        assert_eq!(peer.keepalive, Some(DEFAULT_KEEPALIVE_SECS));
        assert_eq!(peer.hostname, None);
    }

    #[test]
    fn test_peer_config_new_v6() {
        let peer = PeerConfig::new(
            "node-1".to_string(),
            "pubkey123".to_string(),
            "[::1]:51820".to_string(),
            "fd00:200::5".parse::<IpAddr>().unwrap(),
        );

        assert_eq!(peer.node_id, "node-1");
        assert_eq!(peer.keepalive, Some(DEFAULT_KEEPALIVE_SECS));
        assert_eq!(peer.hostname, None);
    }

    #[test]
    fn test_peer_config_with_hostname() {
        let peer = PeerConfig::new(
            "node-1".to_string(),
            "pubkey123".to_string(),
            "192.168.1.100:51820".to_string(),
            IpAddr::V4(Ipv4Addr::new(10, 200, 0, 5)),
        )
        .with_hostname("web-server");

        assert_eq!(peer.hostname, Some("web-server".to_string()));
    }

    #[test]
    fn test_peer_config_to_peer_info_v4() {
        let peer = PeerConfig::new(
            "node-1".to_string(),
            "pubkey123".to_string(),
            "192.168.1.100:51820".to_string(),
            IpAddr::V4(Ipv4Addr::new(10, 200, 0, 5)),
        );

        let peer_info = peer.to_peer_info().unwrap();
        assert_eq!(peer_info.public_key, "pubkey123");
        assert_eq!(peer_info.allowed_ips, "10.200.0.5/32");
    }

    #[test]
    fn test_peer_config_to_peer_info_v6() {
        let peer = PeerConfig::new(
            "node-1".to_string(),
            "pubkey123".to_string(),
            "[::1]:51820".to_string(),
            "fd00:200::5".parse::<IpAddr>().unwrap(),
        );

        let peer_info = peer.to_peer_info().unwrap();
        assert_eq!(peer_info.public_key, "pubkey123");
        assert_eq!(peer_info.allowed_ips, "fd00:200::5/128");
    }

    #[test]
    fn test_bootstrap_state_serialization_v4() {
        let config = BootstrapConfig {
            cidr: "10.200.0.0/16".to_string(),
            node_ip: IpAddr::V4(Ipv4Addr::new(10, 200, 0, 1)),
            interface: DEFAULT_INTERFACE_NAME.to_string(),
            port: DEFAULT_WG_PORT,
            private_key: "private".to_string(),
            public_key: "public".to_string(),
            is_leader: true,
            created_at: 1_234_567_890,
            slice_cidr: None,
        };

        let state = BootstrapState {
            config,
            peers: vec![],
            allocator_state: None,
            slice_allocator_state: None,
        };

        let json = serde_json::to_string_pretty(&state).unwrap();
        let deserialized: BootstrapState = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized.config.cidr, "10.200.0.0/16");
        assert_eq!(deserialized.config.node_ip.to_string(), "10.200.0.1");
    }

    #[test]
    fn test_bootstrap_state_serialization_v6() {
        let config = BootstrapConfig {
            cidr: "fd00:200::/48".to_string(),
            node_ip: "fd00:200::1".parse::<IpAddr>().unwrap(),
            interface: DEFAULT_INTERFACE_NAME.to_string(),
            port: DEFAULT_WG_PORT,
            private_key: "private".to_string(),
            public_key: "public".to_string(),
            is_leader: true,
            created_at: 1_234_567_890,
            slice_cidr: None,
        };

        let state = BootstrapState {
            config,
            peers: vec![],
            allocator_state: None,
            slice_allocator_state: None,
        };

        let json = serde_json::to_string_pretty(&state).unwrap();
        let deserialized: BootstrapState = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized.config.cidr, "fd00:200::/48");
        assert_eq!(deserialized.config.node_ip.to_string(), "fd00:200::1");
    }

    #[test]
    fn test_default_overlay_cidr_v6_constant() {
        // Verify the IPv6 CIDR constant is valid
        let net: ipnet::IpNet = DEFAULT_OVERLAY_CIDR_V6.parse().unwrap();
        assert!(matches!(net, ipnet::IpNet::V6(_)));
        assert_eq!(net.prefix_len(), 48);
    }

    #[test]
    fn test_to_peer_info_uses_slice_when_set() {
        let peer = PeerConfig::new(
            "node-42".to_string(),
            "pubkey-xyz".to_string(),
            "192.168.1.100:51820".to_string(),
            IpAddr::V4(Ipv4Addr::new(10, 200, 42, 1)),
        )
        .with_slice_cidr("10.200.42.0/28".parse().unwrap());

        let peer_info = peer.to_peer_info().unwrap();
        assert_eq!(peer_info.allowed_ips, "10.200.42.0/28");
    }

    #[test]
    fn test_to_peer_info_falls_back_to_node_ip_when_no_slice() {
        let peer = PeerConfig::new(
            "node-5".to_string(),
            "pubkey-abc".to_string(),
            "192.168.1.100:51820".to_string(),
            "10.200.0.5".parse().unwrap(),
        );

        let peer_info = peer.to_peer_info().unwrap();
        assert_eq!(peer_info.allowed_ips, "10.200.0.5/32");
    }

    #[test]
    fn test_bootstrap_config_allowed_ip_prefers_slice() {
        let config = BootstrapConfig {
            cidr: "10.200.0.0/16".to_string(),
            node_ip: IpAddr::V4(Ipv4Addr::new(10, 200, 7, 1)),
            interface: DEFAULT_INTERFACE_NAME.to_string(),
            port: DEFAULT_WG_PORT,
            private_key: "private".to_string(),
            public_key: "public".to_string(),
            is_leader: false,
            created_at: 0,
            slice_cidr: Some("10.200.7.0/28".parse().unwrap()),
        };

        assert_eq!(config.allowed_ip(), "10.200.7.0/28");
    }

    #[test]
    fn test_bootstrap_config_allowed_ip_falls_back_to_node_ip() {
        let config = BootstrapConfig {
            cidr: "10.200.0.0/16".to_string(),
            node_ip: IpAddr::V4(Ipv4Addr::new(10, 200, 0, 9)),
            interface: DEFAULT_INTERFACE_NAME.to_string(),
            port: DEFAULT_WG_PORT,
            private_key: "private".to_string(),
            public_key: "public".to_string(),
            is_leader: false,
            created_at: 0,
            slice_cidr: None,
        };

        assert_eq!(config.allowed_ip(), "10.200.0.9/32");
    }
}