OpenPXE/crates/pxeforge/src/main.rs

//! PXEForge entry point. Wires the three protocol servers (DHCP proxy,
//! TFTP, HTTP) to the shared ISO store and client registry, then runs
//! them concurrently.

use clap::{Parser, Subcommand};
use pxeforge_core::{
    ClientRegistry, Config, DhcpMode, GateQueue, HostBindings, LogBus, LogBusLayer, Metrics,
    SettingsStore,
};
use pxeforge_dhcp_proxy::DhcpProxyServer;
use pxeforge_http_api::{build_router, AppState};
use pxeforge_iso_store::{IsoStore, NfsManager, SmbManager};
use std::sync::Arc;
use pxeforge_tftp::TftpServer;
use std::net::{Ipv4Addr, SocketAddr};
use std::path::PathBuf;
use tokio::io::AsyncReadExt;

#[derive(Debug, Parser)]
#[command(name = "pxeforge", about = "Container-native PXE boot server", version)]
struct Cli {
    /// Path to a TOML config file. All fields have sensible defaults and can
    /// also be overridden with env vars (PXEFORGE_*).
    #[arg(long, env = "PXEFORGE_CONFIG")]
    config: Option<PathBuf>,

    #[command(subcommand)]
    command: Option<Command>,
}

#[derive(Debug, Subcommand)]
enum Command {
    /// Import every `.iso` from a host directory into the ISO store, running
    /// the same introspection + boot-entry generation pass the web upload
    /// does. Useful for pre-populating the store before starting the server
    /// (e.g. in a CI setup or after copying ISOs to a fresh PVC).
    ///
    /// Example:
    ///   docker run --rm \
    ///     -v /my/isos:/seed:ro \
    ///     -v pxeforge-data:/var/lib/pxeforge/isos \
    ///     pxeforge:0.1.0 seed --from /seed
    Seed {
        /// Source directory containing one or more `.iso` files.
        #[arg(long)]
        from: PathBuf,
        /// Don't actually import — print what would happen.
        #[arg(long)]
        dry_run: bool,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // The LogBus has to exist before we install tracing layers, since one
    // of those layers fans out into it. The web UI's Terminal tab
    // subscribes to this bus over SSE.
    let log_bus = LogBus::new(500);
    init_tracing(log_bus.clone());

    let cli = Cli::parse();
    let mut config = match &cli.config {
        Some(p) if p.exists() => Config::from_toml_file(p)?,
        _ => Config::default(),
    };
    config.apply_env();

    // Dispatch subcommands before bringing up the server.
    if let Some(cmd) = cli.command {
        return run_command(cmd, config).await;
    }

    pxeforge_ipxe_assets::log_availability();

    let our_ip = match config.server.public_ip {
        Some(ip) => {
            tracing::info!("using configured public IP: {ip}");
            ip
        }
        None => match detect_primary_ipv4() {
            Some(ip) => {
                tracing::info!("auto-detected public IPv4: {ip}");
                ip
            }
            None => {
                // Without a reachable IP, every generated iPXE URL would
                // be unreachable from real clients. Exit with a clear
                // message instead of serving a broken deployment.
                anyhow::bail!(
                    "could not detect a non-loopback IPv4 address for this host. \
                     Set PXEFORGE_PUBLIC_IP=<your-ip> (e.g. `-e PXEFORGE_PUBLIC_IP=10.0.0.5` \
                     in docker, or the env block in OpenShift Deployment) to advertise \
                     a specific IP to PXE clients."
                );
            }
        },
    };
    let public_base_url = format!("http://{our_ip}");

    let iso_store = IsoStore::new(config.paths.iso_dir.clone());
    iso_store.load_from_disk().await?;
    let clients = ClientRegistry::new();
    let gates = GateQueue::new();
    let settings = SettingsStore::load_or_default(&config.paths.work_dir);
    let hosts = HostBindings::load_or_default(&config.paths.work_dir);
    let metrics = Metrics::new();

    // Build the SMB manager unconditionally — it starts/stops on the
    // Windows toggle, not at process start. If the `smb_dir` isn't
    // writable (e.g. read-only filesystem), the manager will surface that
    // as `SmbState::Failed` when the operator flips the toggle.
    let smb = Arc::new(SmbManager::new(config.paths.smb_dir.clone()));
    if settings.snapshot().windows_enabled {
        let _ = smb.start();
    }

    // NFS manager. The mount root has to be set on the IsoStore *before*
    // we replay any persisted mounts, otherwise an in-memory IsoMeta
    // pointing at an NFS source can't resolve to a path.
    let nfs = NfsManager::new(&config.paths.work_dir, iso_store.clone());
    iso_store.set_nfs_root(nfs.mount_root());
    if let Err(e) = nfs.load_and_remount().await {
        tracing::warn!(target: "pxeforge::nfs", "could not reload NFS mounts: {e}");
    }

    // Sniff network details for the Network tab. None of these are
    // required for PXE to work — they're informational, surfaced in the
    // UI so an operator doesn't have to drop to a shell to find their
    // own gateway.
    let net = detect_network_info(our_ip);
    tracing::info!(
        target: "pxeforge::net",
        nic = %net.nic_name, mask = %net.subnet_mask, gateway = %net.gateway,
        "network info"
    );

    let state = AppState {
        iso_store: iso_store.clone(),
        clients: clients.clone(),
        settings: settings.clone(),
        gates: gates.clone(),
        hosts: hosts.clone(),
        metrics: metrics.clone(),
        smb: Some(smb.clone()),
        nfs: nfs.clone(),
        log_bus: log_bus.clone(),
        started_at: time::OffsetDateTime::now_utc(),
        public_base_url: public_base_url.clone(),
        nic_name: net.nic_name,
        subnet_mask: net.subnet_mask,
        gateway: net.gateway,
    };

    let http_addr = SocketAddr::new(config.server.http_bind, config.server.http_port);
    let router = build_router(state);
    let http_task = tokio::spawn(async move {
        let listener = tokio::net::TcpListener::bind(http_addr).await?;
        tracing::info!(target: "pxeforge::http", "HTTP listening on {http_addr}");
        axum::serve(listener, router).await?;
        Ok::<_, anyhow::Error>(())
    });

    let tftp = TftpServer::new(
        config.server.tftp_bind,
        config.server.tftp_port,
        clients.clone(),
        metrics.clone(),
    );
    let tftp_task = tokio::spawn(tftp.run());

    let dhcp_task: tokio::task::JoinHandle<anyhow::Result<()>> = match config.network.dhcp_mode {
        DhcpMode::Proxy => {
            let s = DhcpProxyServer::new(
                config.network.dhcp_bind,
                config.network.dhcp_port,
                config.network.pxe_port,
                our_ip,
                public_base_url.clone(),
                clients.clone(),
                metrics.clone(),
            );
            tokio::spawn(s.run())
        }
        DhcpMode::Disabled => {
            tracing::info!(target: "pxeforge::dhcp", "DHCP disabled — external DHCP must set next-server + filename");
            tokio::spawn(async { futures_forever().await })
        }
    };

    tokio::select! {
        r = http_task => { tracing::error!("http task exited: {:?}", r); r??; }
        r = tftp_task => { tracing::error!("tftp task exited: {:?}", r); r??; }
        r = dhcp_task => { tracing::error!("dhcp task exited: {:?}", r); r??; }
    }
    Ok(())
}

async fn futures_forever() -> anyhow::Result<()> {
    std::future::pending::<()>().await;
    Ok(())
}

async fn run_command(cmd: Command, config: Config) -> anyhow::Result<()> {
    match cmd {
        Command::Seed { from, dry_run } => seed_from_dir(&from, &config, dry_run).await,
    }
}

/// Walk `src` for `*.iso`, stream each file through the normal upload path.
/// Reuses `IsoStore::begin_upload` / `finish` so the resulting meta on disk
/// is identical to a web upload — same slug rules, same introspection, same
/// sha256.
async fn seed_from_dir(src: &std::path::Path, config: &Config, dry_run: bool) -> anyhow::Result<()> {
    let store = IsoStore::new(config.paths.iso_dir.clone());
    store.load_from_disk().await?;
    let mut entries = tokio::fs::read_dir(src).await?;
    let mut imported = 0u32;
    let mut skipped = 0u32;
    while let Some(entry) = entries.next_entry().await? {
        let p = entry.path();
        if p.extension().and_then(|e| e.to_str()).map(str::to_ascii_lowercase).as_deref() != Some("iso") {
            continue;
        }
        let filename = p
            .file_name()
            .and_then(|s| s.to_str())
            .ok_or_else(|| anyhow::anyhow!("non-utf8 filename: {}", p.display()))?
            .to_string();
        println!("  {} ({} bytes)", filename, tokio::fs::metadata(&p).await?.len());
        if dry_run { continue; }

        let mut handle = match store.begin_upload(&filename).await {
            Ok(h) => h,
            Err(pxeforge_core::Error::Invalid(e)) => {
                eprintln!("    skip: {e}");
                skipped += 1;
                continue;
            }
            Err(e) => return Err(e.into()),
        };
        let mut file = tokio::fs::File::open(&p).await?;
        let mut buf = vec![0u8; 1024 * 1024];
        loop {
            let n = file.read(&mut buf).await?;
            if n == 0 { break; }
            let chunk: bytes::Bytes = buf[..n].to_vec().into();
            handle.write_chunk(&chunk).await?;
        }
        let meta = handle.finish(&store).await?;
        println!("    -> id={} family={:?}", meta.id, meta.introspection.family);
        imported += 1;
    }
    println!("\nimported={imported}  skipped={skipped}  {}", if dry_run { "(dry run)" } else { "" });
    Ok(())
}

/// Pick the first non-loopback IPv4 address on this host. Returns `None` if
/// detection fails — callers should fail startup rather than silently using
/// a loopback address (which would give every PXE client an unreachable
/// `http://127.0.0.1/...`). Users in multi-homed setups should set
/// `PXEFORGE_PUBLIC_IP` explicitly.
fn detect_primary_ipv4() -> Option<Ipv4Addr> {
    // First try: route to the public internet. `UdpSocket::connect` to a
    // well-known external address causes the OS to populate `local_addr`
    // with the source IP it would use — this is the standard "which of my
    // interfaces is the internet-facing one" idiom.
    if let Ok(sock) = std::net::UdpSocket::bind("0.0.0.0:0") {
        if sock.connect("8.8.8.8:80").is_ok() {
            if let Ok(std::net::SocketAddr::V4(addr)) = sock.local_addr() {
                let v4 = *addr.ip();
                if !v4.is_loopback() && !v4.is_unspecified() {
                    return Some(v4);
                }
            }
        }
    }
    // Fallback: hostname resolution.
    if let Ok(hostname) = hostname() {
        if let Ok(addrs) = std::net::ToSocketAddrs::to_socket_addrs(&format!("{hostname}:0")) {
            for a in addrs {
                if let std::net::IpAddr::V4(v4) = a.ip() {
                    if !v4.is_loopback() && !v4.is_unspecified() {
                        return Some(v4);
                    }
                }
            }
        }
    }
    None
}

fn hostname() -> std::io::Result<String> {
    // Tiny shim: read /proc/sys/kernel/hostname on Linux, fall back to `uname -n` via env.
    if let Ok(h) = std::fs::read_to_string("/proc/sys/kernel/hostname") {
        return Ok(h.trim().to_string());
    }
    std::env::var("HOSTNAME").map_err(|_| std::io::Error::new(
        std::io::ErrorKind::NotFound, "no hostname",
    ))
}

fn init_tracing(bus: Arc<LogBus>) {
    use tracing_subscriber::{fmt, prelude::*, EnvFilter};
    let filter = EnvFilter::try_from_env("PXEFORGE_LOG")
        .unwrap_or_else(|_| EnvFilter::new("info,pxeforge=debug"));
    tracing_subscriber::registry()
        .with(filter)
        .with(fmt::layer().with_target(true))
        .with(LogBusLayer::new(bus))
        .init();
}

#[derive(Debug, Default)]
struct NetworkInfo {
    nic_name: String,
    subnet_mask: String,
    gateway: String,
}

/// Best-effort population of the Network tab's read-only fields. We shell
/// out to standard Linux tools (`ip route`) instead of pulling in a
/// netlink crate — these calls run once at startup, so the cost of
/// spawning a process is negligible. Empty strings are perfectly fine
/// fallbacks; the UI shows a `?` placeholder.
fn detect_network_info(our_ip: Ipv4Addr) -> NetworkInfo {
    use std::process::Command;
    let mut info = NetworkInfo::default();

    // `ip -o -f inet addr show` lists every interface with its
    // `inet a.b.c.d/mask`. We match the line that mentions our IP.
    if let Ok(out) = Command::new("ip").args(["-o", "-f", "inet", "addr", "show"]).output() {
        if let Ok(text) = String::from_utf8(out.stdout) {
            for line in text.lines() {
                if !line.contains(&our_ip.to_string()) {
                    continue;
                }
                // Format: "2: enp1s0    inet 10.0.0.5/24 brd ..."
                let mut parts = line.split_whitespace();
                let _idx = parts.next();
                if let Some(name) = parts.next() {
                    info.nic_name = name.trim_end_matches(':').to_string();
                }
                if let Some(addr) = line.split_whitespace().find(|p| p.contains('/')) {
                    if let Some((_, prefix_str)) = addr.split_once('/') {
                        if let Ok(prefix) = prefix_str.parse::<u8>() {
                            info.subnet_mask = prefix_to_dotted(prefix);
                        }
                    }
                }
                break;
            }
        }
    }

    // `ip route show default` -> "default via 10.0.0.1 dev enp1s0 ..."
    if let Ok(out) = Command::new("ip").args(["route", "show", "default"]).output() {
        if let Ok(text) = String::from_utf8(out.stdout) {
            if let Some(line) = text.lines().next() {
                let mut parts = line.split_whitespace();
                while let Some(p) = parts.next() {
                    if p == "via" {
                        if let Some(gw) = parts.next() {
                            info.gateway = gw.to_string();
                        }
                        break;
                    }
                }
            }
        }
    }

    info
}

fn prefix_to_dotted(prefix: u8) -> String {
    let prefix = prefix.min(32);
    let mask: u32 = if prefix == 0 { 0 } else { u32::MAX << (32 - prefix) };
    format!(
        "{}.{}.{}.{}",
        (mask >> 24) & 0xff,
        (mask >> 16) & 0xff,
        (mask >> 8) & 0xff,
        mask & 0xff
    )
}