One-KVM/libs/ventoy-img-rs/src/resources.rs

//! Ventoy resources loader
//!
//! Loads Ventoy boot resources from external files in a resource directory.
//! Resource files (boot.img, core.img, ventoy.disk.img) should be pre-decompressed.

use crate::error::{Result, VentoyError};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::OnceLock;

/// Resource file names
const BOOT_IMG_NAME: &str = "boot.img";
const CORE_IMG_NAME: &str = "core.img";
const VENTOY_DISK_IMG_NAME: &str = "ventoy.disk.img";

/// Ventoy signature (16 bytes at MBR offset 0x190)
pub const VENTOY_SIGNATURE: [u8; 16] = [
    0x56, 0x54, 0x00, 0x47, 0x65, 0x00, 0x48, 0x44, 0x00, 0x52, 0x64, 0x00, 0x20, 0x45, 0x72, 0x0D,
];

/// Cached resources loaded from disk
struct ResourceCache {
    boot_img: Vec<u8>,
    core_img: Vec<u8>,
    ventoy_disk_img: Vec<u8>,
}

/// Global resource cache
static RESOURCE_CACHE: OnceLock<ResourceCache> = OnceLock::new();

/// Initialize resources from a directory
///
/// This function must be called before using `get_boot_img()`, `get_core_img()`,
/// or `get_ventoy_disk_img()`. It loads all resource files into memory.
///
/// # Arguments
/// * `resource_dir` - Path to directory containing boot.img, core.img, ventoy.disk.img
///
/// # Example
/// ```no_run
/// use ventoy_img::resources::init_resources;
/// use std::path::Path;
///
/// init_resources(Path::new("/var/lib/one-kvm/ventoy")).unwrap();
/// ```
pub fn init_resources(resource_dir: &Path) -> Result<()> {
    if RESOURCE_CACHE.get().is_some() {
        // Already initialized
        return Ok(());
    }

    let boot_path = resource_dir.join(BOOT_IMG_NAME);
    let core_path = resource_dir.join(CORE_IMG_NAME);
    let ventoy_disk_path = resource_dir.join(VENTOY_DISK_IMG_NAME);

    // Check all files exist
    if !boot_path.exists() {
        return Err(VentoyError::ResourceNotFound(format!(
            "boot.img not found at {}",
            boot_path.display()
        )));
    }
    if !core_path.exists() {
        return Err(VentoyError::ResourceNotFound(format!(
            "core.img not found at {}",
            core_path.display()
        )));
    }
    if !ventoy_disk_path.exists() {
        return Err(VentoyError::ResourceNotFound(format!(
            "ventoy.disk.img not found at {}",
            ventoy_disk_path.display()
        )));
    }

    // Load files
    let boot_img = fs::read(&boot_path).map_err(|e| {
        VentoyError::ResourceNotFound(format!("Failed to read {}: {}", boot_path.display(), e))
    })?;

    let core_img = fs::read(&core_path).map_err(|e| {
        VentoyError::ResourceNotFound(format!("Failed to read {}: {}", core_path.display(), e))
    })?;

    let ventoy_disk_img = fs::read(&ventoy_disk_path).map_err(|e| {
        VentoyError::ResourceNotFound(format!(
            "Failed to read {}: {}",
            ventoy_disk_path.display(),
            e
        ))
    })?;

    // Validate boot.img size
    if boot_img.len() != 512 {
        return Err(VentoyError::ResourceNotFound(format!(
            "boot.img has invalid size: {} bytes (expected 512)",
            boot_img.len()
        )));
    }

    let cache = ResourceCache {
        boot_img,
        core_img,
        ventoy_disk_img,
    };

    // Try to set the cache (ignore if already set by another thread)
    let _ = RESOURCE_CACHE.set(cache);

    Ok(())
}

/// Check if resources have been initialized
pub fn is_initialized() -> bool {
    RESOURCE_CACHE.get().is_some()
}

/// Get the boot.img data (512 bytes MBR boot code)
pub fn get_boot_img() -> Result<&'static [u8]> {
    RESOURCE_CACHE
        .get()
        .map(|c| c.boot_img.as_slice())
        .ok_or_else(|| {
            VentoyError::ResourceNotFound(
                "Resources not initialized. Call init_resources() first.".to_string(),
            )
        })
}

/// Get the core.img data (GRUB core image, ~1MB)
pub fn get_core_img() -> Result<&'static [u8]> {
    RESOURCE_CACHE
        .get()
        .map(|c| c.core_img.as_slice())
        .ok_or_else(|| {
            VentoyError::ResourceNotFound(
                "Resources not initialized. Call init_resources() first.".to_string(),
            )
        })
}

/// Get the ventoy.disk.img data (EFI partition, ~32MB)
pub fn get_ventoy_disk_img() -> Result<&'static [u8]> {
    RESOURCE_CACHE
        .get()
        .map(|c| c.ventoy_disk_img.as_slice())
        .ok_or_else(|| {
            VentoyError::ResourceNotFound(
                "Resources not initialized. Call init_resources() first.".to_string(),
            )
        })
}

/// Get the resource directory path for a given data directory
///
/// Returns `{data_dir}/ventoy`
pub fn get_resource_dir(data_dir: &Path) -> PathBuf {
    data_dir.join("ventoy")
}

/// List required resource files
pub fn required_files() -> &'static [&'static str] {
    &[BOOT_IMG_NAME, CORE_IMG_NAME, VENTOY_DISK_IMG_NAME]
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::TempDir;

    fn create_test_resources(dir: &Path) {
        // Create boot.img (512 bytes)
        let mut boot = std::fs::File::create(dir.join(BOOT_IMG_NAME)).unwrap();
        boot.write_all(&[0u8; 512]).unwrap();

        // Create core.img (fake, 1KB)
        let mut core = std::fs::File::create(dir.join(CORE_IMG_NAME)).unwrap();
        core.write_all(&[0u8; 1024]).unwrap();

        // Create ventoy.disk.img (fake, 1KB)
        let mut ventoy = std::fs::File::create(dir.join(VENTOY_DISK_IMG_NAME)).unwrap();
        ventoy.write_all(&[0u8; 1024]).unwrap();
    }

    #[test]
    fn test_required_files() {
        let files = required_files();
        assert_eq!(files.len(), 3);
        assert!(files.contains(&"boot.img"));
        assert!(files.contains(&"core.img"));
        assert!(files.contains(&"ventoy.disk.img"));
    }

    #[test]
    fn test_get_resource_dir() {
        let data_dir = Path::new("/var/lib/one-kvm");
        let resource_dir = get_resource_dir(data_dir);
        assert_eq!(resource_dir, PathBuf::from("/var/lib/one-kvm/ventoy"));
    }
}