OpenPXE/crates/iso-store/src/store.rs

//! On-disk ISO store with sidecar metadata files.

use crate::entry::{BootEntry, BootKind, KernelArgs};
use crate::introspect::{introspect, DistroFamily, IntrospectionReport};
use bytes::Bytes;
use parking_lot::RwLock;
use pxeforge_core::{Error, Result};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use time::OffsetDateTime;
use tokio::io::AsyncWriteExt;

/// Where the bytes for an ISO actually live.
///
/// The default is `Local` — uploaded ISOs sit in `<iso_dir>/<id>.iso`.
/// `Nfs` entries point at a file inside a remote share that the
/// `NfsManager` is keeping mounted. We resolve the on-disk path lazily
/// in [`IsoStore::iso_path_for`] using the `nfs_root` set at startup.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum IsoSource {
    #[default]
    Local,
    Nfs {
        mount_id: String,
        /// Path relative to the mount point — typically just the filename.
        relative_path: String,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IsoMeta {
    /// Stable slug used in URLs (derived from the uploaded filename).
    pub id: String,
    pub filename: String,
    pub size_bytes: u64,
    pub sha256_hex: Option<String>,
    #[serde(with = "time::serde::rfc3339")]
    pub uploaded_at: OffsetDateTime,
    pub introspection: IntrospectionReport,
    /// Boot entries this ISO currently exposes in the PXE menu. Usually one,
    /// occasionally two (BIOS + UEFI variant for some RHEL ISOs).
    pub boot_entries: Vec<BootEntry>,
    /// Source of the bytes — local upload (default) or NFS mount.
    /// Old `meta.json` files without this field deserialize as `Local`.
    #[serde(default)]
    pub source: IsoSource,
}

pub struct UploadHandle {
    pub id: String,
    pub partial_path: PathBuf,
    final_path: PathBuf,
    filename: String,
    hasher: Sha256,
    bytes_written: u64,
    file: tokio::fs::File,
}

impl UploadHandle {
    pub async fn write_chunk(&mut self, chunk: &Bytes) -> Result<()> {
        self.file.write_all(chunk).await?;
        self.hasher.update(chunk);
        self.bytes_written += chunk.len() as u64;
        Ok(())
    }

    /// Flush, close, and atomically rename to the final path. Returns the
    /// final `IsoMeta` including introspection results.
    pub async fn finish(mut self, store: &IsoStore) -> Result<IsoMeta> {
        self.file.flush().await?;
        self.file.sync_all().await?;
        drop(self.file);
        tokio::fs::rename(&self.partial_path, &self.final_path).await?;

        let hash = hex::encode(self.hasher.finalize());
        let introspection = {
            let p = self.final_path.clone();
            tokio::task::spawn_blocking(move || introspect(&p))
                .await
                .map_err(|e| Error::Other(e.into()))?
        };

        let boot_entries = generate_boot_entries(&self.id, &self.filename, &introspection);
        let meta = IsoMeta {
            id: self.id.clone(),
            filename: self.filename,
            size_bytes: self.bytes_written,
            sha256_hex: Some(hash),
            uploaded_at: OffsetDateTime::now_utc(),
            introspection,
            boot_entries,
            source: IsoSource::Local,
        };
        store.persist_meta(&meta).await?;
        store.insert(meta.clone());
        Ok(meta)
    }

    pub async fn abort(self) -> Result<()> {
        drop(self.file);
        let _ = tokio::fs::remove_file(&self.partial_path).await;
        Ok(())
    }
}

#[derive(Debug, Default)]
struct Inner {
    isos: HashMap<String, IsoMeta>,
}

#[derive(Debug, Clone)]
pub struct IsoStore {
    iso_dir: Arc<PathBuf>,
    /// Where NFS mounts land on disk. Set at startup via
    /// [`IsoStore::set_nfs_root`]; required for resolving any
    /// `IsoSource::Nfs` entry.
    nfs_root: Arc<RwLock<Option<PathBuf>>>,
    inner: Arc<RwLock<Inner>>,
}

impl IsoStore {
    pub fn new(iso_dir: PathBuf) -> Self {
        Self {
            iso_dir: Arc::new(iso_dir),
            nfs_root: Arc::new(RwLock::new(None)),
            inner: Arc::new(RwLock::new(Inner::default())),
        }
    }

    /// Tell the store where NFS mounts live. Without this set,
    /// `IsoSource::Nfs` entries cannot be resolved to a file path.
    pub fn set_nfs_root(&self, root: PathBuf) {
        *self.nfs_root.write() = Some(root);
    }

    pub async fn ensure_dirs(&self) -> Result<()> {
        tokio::fs::create_dir_all(self.iso_dir.as_path()).await?;
        Ok(())
    }

    /// Scan the ISO directory on startup and load any sidecar `.meta.json`
    /// files. ISOs without metadata are introspected lazily — we don't block
    /// startup on potentially many GB of scanning.
    pub async fn load_from_disk(&self) -> Result<()> {
        self.ensure_dirs().await?;
        let mut entries = tokio::fs::read_dir(self.iso_dir.as_path()).await?;
        while let Some(e) = entries.next_entry().await? {
            let p = e.path();
            if p.extension().and_then(|s| s.to_str()) != Some("json") { continue; }
            if !p
                .file_name()
                .and_then(|s| s.to_str())
                .is_some_and(|n| n.ends_with(".meta.json"))
            {
                continue;
            }
            if let Ok(text) = tokio::fs::read_to_string(&p).await {
                if let Ok(meta) = serde_json::from_str::<IsoMeta>(&text) {
                    self.insert(meta);
                }
            }
        }
        Ok(())
    }

    fn insert(&self, meta: IsoMeta) {
        self.inner.write().isos.insert(meta.id.clone(), meta);
    }

    async fn persist_meta(&self, meta: &IsoMeta) -> Result<()> {
        let path = self.meta_path(&meta.id);
        let text = serde_json::to_string_pretty(meta).map_err(|e| Error::Other(e.into()))?;
        tokio::fs::write(path, text).await?;
        Ok(())
    }

    fn meta_path(&self, id: &str) -> PathBuf {
        self.iso_dir.join(format!("{id}.meta.json"))
    }

    fn iso_path(&self, id: &str) -> PathBuf {
        self.iso_dir.join(format!("{id}.iso"))
    }

    pub async fn begin_upload(&self, filename: &str) -> Result<UploadHandle> {
        self.ensure_dirs().await?;
        let id = slugify(filename);
        let final_path = self.iso_path(&id);
        if final_path.exists() {
            return Err(Error::Invalid(format!("iso '{id}' already exists")));
        }
        let partial_path = self.iso_dir.join(format!("{id}.partial"));
        let file = tokio::fs::File::create(&partial_path).await?;
        Ok(UploadHandle {
            id,
            partial_path,
            final_path,
            filename: filename.to_string(),
            hasher: Sha256::new(),
            bytes_written: 0,
            file,
        })
    }

    /// Readiness probe — is the backing directory reachable? Distinct from
    /// "is there content in it", to avoid an empty store failing health.
    #[must_use]
    pub fn list_ok(&self) -> bool {
        std::fs::read_dir(self.iso_dir.as_path()).is_ok()
    }

    #[must_use]
    pub fn list(&self) -> Vec<IsoMeta> {
        let g = self.inner.read();
        let mut v: Vec<_> = g.isos.values().cloned().collect();
        // Newest-first by upload time.
        v.sort_by_key(|m| std::cmp::Reverse(m.uploaded_at));
        v
    }

    #[must_use]
    pub fn get(&self, id: &str) -> Option<IsoMeta> {
        self.inner.read().isos.get(id).cloned()
    }

    /// Resolve an ISO id to its on-disk path, if any. For local entries
    /// this is `<iso_dir>/<id>.iso`; for NFS entries it's
    /// `<nfs_root>/<mount_id>/<relative_path>`. Returns None if the file
    /// is missing or the source isn't resolvable (e.g. NFS share
    /// unmounted).
    pub fn iso_path_for(&self, id: &str) -> Option<PathBuf> {
        let meta = self.get(id)?;
        let path = match &meta.source {
            IsoSource::Local => self.iso_path(id),
            IsoSource::Nfs {
                mount_id,
                relative_path,
            } => {
                let root = self.nfs_root.read().clone()?;
                root.join(mount_id).join(relative_path)
            }
        };
        if path.exists() {
            Some(path)
        } else {
            None
        }
    }

    /// Delete an ISO and its sidecar metadata. Only acts on local ISOs;
    /// for NFS-backed ISOs the operator must remove the file from the
    /// share or unmount the NFS share entirely.
    pub async fn delete(&self, id: &str) -> Result<()> {
        let meta = self.get(id);
        let is_local = matches!(meta.as_ref().map(|m| &m.source), Some(IsoSource::Local) | None);
        if is_local {
            let iso = self.iso_path(id);
            let meta_path = self.meta_path(id);
            let _ = tokio::fs::remove_file(&iso).await;
            let _ = tokio::fs::remove_file(&meta_path).await;
        }
        self.inner.write().isos.remove(id);
        Ok(())
    }

    /// Register an externally-sourced ISO (e.g. NFS-mounted). Used by
    /// `NfsManager` after walking a freshly-mounted share. We do **not**
    /// persist a `meta.json` on disk for these — the source of truth is
    /// the share itself, and the NFS manager re-scans on startup.
    pub fn register_external(
        &self,
        id: String,
        filename: String,
        size_bytes: u64,
        introspection: IntrospectionReport,
        boot_entries: Vec<BootEntry>,
        source: IsoSource,
    ) {
        let meta = IsoMeta {
            id: id.clone(),
            filename,
            size_bytes,
            sha256_hex: None,
            uploaded_at: OffsetDateTime::now_utc(),
            introspection,
            boot_entries,
            source,
        };
        self.inner.write().isos.insert(id, meta);
    }

    /// Drop every entry that belongs to `mount_id`. Used by the NFS
    /// manager when an operator removes a share, or before re-scanning
    /// to clean out stale entries.
    pub fn drop_external_source(&self, mount_id: &str) {
        let mut g = self.inner.write();
        g.isos.retain(|_, m| {
            !matches!(&m.source, IsoSource::Nfs { mount_id: mid, .. } if mid == mount_id)
        });
    }
}

fn slugify(filename: &str) -> String {
    let stem = Path::new(filename)
        .file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or("iso");
    slugify_str(stem)
}

/// Slugify an arbitrary string to lowercase ASCII alphanumerics, hyphens,
/// and underscores. Public so the NFS manager can mint ids that follow the
/// same rules as upload-time ISO ids.
#[must_use]
pub fn slugify_str(input: &str) -> String {
    input
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
                c.to_ascii_lowercase()
            } else {
                '-'
            }
        })
        .collect::<String>()
        .trim_matches('-')
        .to_string()
}

/// Public wrapper around [`generate_boot_entries`] so the NFS manager can
/// build entries for shares it just scanned, using the same logic as the
/// upload pipeline. Re-exported via the crate root.
#[must_use]
pub fn generate_boot_entries_for(
    id: &str,
    filename: &str,
    r: &IntrospectionReport,
) -> Vec<BootEntry> {
    generate_boot_entries(id, filename, r)
}

/// Build `BootEntry`s from the introspection report. URLs are relative —
/// the HTTP layer rewrites them with the public base URL per request.
fn generate_boot_entries(id: &str, filename: &str, r: &IntrospectionReport) -> Vec<BootEntry> {
    let title = r.volume_label.clone().unwrap_or_else(|| filename.to_string());
    match r.family {
        DistroFamily::WindowsPe if r.has_boot_wim => {
            // Standard wimboot chain. Paths are in-ISO; the HTTP layer maps
            // `iso/<id>/<path>` to on-disk extraction via ISO9660 lookup.
            let base = format!("iso/{id}");
            vec![BootEntry {
                id: format!("{id}-winpe"),
                title: format!("{title} (Windows / wimboot)"),
                kind: BootKind::Wimboot {
                    wimboot_url: "ipxe/wimboot".to_string(),
                    files: vec![
                        ("bootmgr".into(), format!("{base}/bootmgr")),
                        ("bootmgr.efi".into(), format!("{base}/bootmgr.efi")),
                        ("bcd".into(), format!("{base}/boot/bcd")),
                        ("boot.sdi".into(), format!("{base}/boot/boot.sdi")),
                        ("boot.wim".into(), format!("{base}/sources/boot.wim")),
                    ],
                },
            }]
        }
        fam if r.kernel_path.is_some() => {
            let base = format!("iso/{id}");
            let kernel_url = format!("{base}{}", r.kernel_path.as_deref().unwrap_or(""));
            let initrd_urls = r.initrd_paths.iter().map(|p| format!("{base}{p}")).collect();
            let args = KernelArgs { cmdline: linux_cmdline(fam, id) };
            vec![BootEntry {
                id: format!("{id}-linux"),
                title,
                kind: BootKind::LinuxKernel { kernel_url, initrd_urls, args },
            }]
        }
        _ => {
            // Last-resort SAN boot. Won't work for large modern ISOs, but
            // lets the ISO at least appear in the menu.
            vec![BootEntry {
                id: format!("{id}-sanboot"),
                title: format!("{title} (SAN boot — may fail for >1GiB ISOs)"),
                kind: BootKind::SanBootIso { iso_url: format!("iso/{id}.iso") },
            }]
        }
    }
}

fn linux_cmdline(family: DistroFamily, id: &str) -> String {
    // The HTTP layer resolves `${base-url}` at render time.
    let iso_url = format!("${{base-url}}/iso/{id}.iso");
    match family {
        DistroFamily::DebianUbuntu => format!(
            "boot=casper netboot=url url={iso_url} ip=dhcp ---"
        ),
        DistroFamily::RhelFedora => format!(
            "inst.repo={iso_url} inst.stage2={iso_url} ip=dhcp"
        ),
        DistroFamily::OpenSuse => format!(
            "install={iso_url} netsetup=dhcp"
        ),
        DistroFamily::Arch => {
            "archiso_http_srv=${base-url}/iso/ archisobasedir=arch ip=dhcp copytoram".to_string()
        }
        DistroFamily::Alpine => format!(
            "alpine_repo=${{base-url}}/iso/{id}/ modloop=${{base-url}}/iso/{id}/boot/modloop-lts ip=dhcp"
        ),
        _ => "ip=dhcp".into(),
    }
}

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

    #[test]
    fn slugify_basic() {
        // Upload filenames come from multipart parts (no path components);
        // file_stem drops the extension, then non-alphanumerics become `-`.
        assert_eq!(slugify("Ubuntu 24.04 Desktop.iso"), "ubuntu-24-04-desktop");
        assert_eq!(slugify("Rocky-9.4-x86_64-dvd.iso"), "rocky-9-4-x86_64-dvd");
        assert_eq!(slugify("arch.iso"), "arch");
        // If a path sneaks in, file_stem strips the directory — OK, not a hazard.
        assert_eq!(slugify("/etc/passwd"), "passwd");
    }
}