axbuild 0.4.15

An OS build lib toolkit used by arceos
Documentation
use std::{
    fs,
    path::{Path, PathBuf},
};

use anyhow::anyhow;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

pub const DEFAULT_REGISTRY_URL: &str =
    "https://raw.githubusercontent.com/rcore-os/tgosimages/refs/heads/main/registry/default.toml";
pub const DEFAULT_FALLBACK_REGISTRY_URL: &str =
    "https://raw.githubusercontent.com/rcore-os/tgosimages/refs/heads/main/registry/v0.0.6.toml";
pub const IMAGE_CONFIG_FILENAME: &str = ".image.toml";
const DEFAULT_AUTO_SYNC_THRESHOLD: u64 = 60 * 60 * 24 * 7;
const LOCAL_STORAGE_ENV: &str = "TGOS_IMAGE_LOCAL_STORAGE";
const FALLBACK_REGISTRY_ENV: &str = "TGOS_IMAGE_REGISTRY_FALLBACK_URL";

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
pub struct ImageConfig {
    pub local_storage: PathBuf,
    pub registry: String,
    pub auto_sync: bool,
    pub auto_sync_threshold: u64,
}

impl ImageConfig {
    pub fn new_default(base_dir: &Path) -> Self {
        Self {
            local_storage: crate::context::axbuild_tmp_dir(base_dir).join("rootfs"),
            registry: DEFAULT_REGISTRY_URL.to_string(),
            auto_sync: true,
            auto_sync_threshold: DEFAULT_AUTO_SYNC_THRESHOLD,
        }
    }

    pub fn get_config_file_path(base_dir: &Path) -> PathBuf {
        crate::context::axbuild_tmp_dir(base_dir).join(IMAGE_CONFIG_FILENAME)
    }

    pub fn read_config(base_dir: &Path) -> anyhow::Result<Self> {
        let path = Self::get_config_file_path(base_dir);

        let mut config = if !path.exists() {
            let config = Self::new_default(base_dir);
            Self::write_config(base_dir, &config)?;
            config
        } else {
            let s = fs::read_to_string(&path)?;
            toml::from_str(&s).map_err(|e| anyhow!("Invalid image config file: {e}"))?
        };

        if let Some(local_storage) = non_empty_env(LOCAL_STORAGE_ENV) {
            config.local_storage = PathBuf::from(local_storage);
        }

        Ok(config)
    }

    pub fn write_config(base_dir: &Path, config: &Self) -> anyhow::Result<()> {
        let path = Self::get_config_file_path(base_dir);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)
                .map_err(|e| anyhow!("Failed to create image config directory: {e}"))?;
        }
        fs::write(path, toml::to_string(config)?)
            .map_err(|e| anyhow!("Failed to write image config file: {e}"))
    }
}

pub(crate) fn fallback_registry_url() -> String {
    non_empty_env(FALLBACK_REGISTRY_ENV)
        .unwrap_or_else(|| DEFAULT_FALLBACK_REGISTRY_URL.to_string())
}

fn non_empty_env(key: &str) -> Option<String> {
    std::env::var(key)
        .ok()
        .filter(|value| !value.trim().is_empty())
}

#[cfg(test)]
mod tests {
    use std::{
        ffi::{OsStr, OsString},
        sync::{LazyLock, Mutex},
    };

    use tempfile::tempdir;

    use super::*;

    static ENV_LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));

    struct TempEnvVar {
        key: &'static str,
        original: Option<OsString>,
    }

    impl TempEnvVar {
        fn set(key: &'static str, value: impl AsRef<OsStr>) -> Self {
            let original = std::env::var_os(key);
            unsafe {
                std::env::set_var(key, value);
            }
            Self { key, original }
        }

        fn unset(key: &'static str) -> Self {
            let original = std::env::var_os(key);
            unsafe {
                std::env::remove_var(key);
            }
            Self { key, original }
        }
    }

    impl Drop for TempEnvVar {
        fn drop(&mut self) {
            match self.original.as_ref() {
                Some(value) => unsafe {
                    std::env::set_var(self.key, value);
                },
                None => unsafe {
                    std::env::remove_var(self.key);
                },
            }
        }
    }

    #[test]
    fn read_config_creates_default_when_missing() {
        let _lock = ENV_LOCK.lock().unwrap();
        let _env = TempEnvVar::unset(LOCAL_STORAGE_ENV);
        let dir = tempdir().unwrap();

        let config = ImageConfig::read_config(dir.path()).unwrap();

        assert_eq!(config, ImageConfig::new_default(dir.path()));
        assert_eq!(config.local_storage, dir.path().join("tmp/axbuild/rootfs"));
        assert_eq!(
            ImageConfig::get_config_file_path(dir.path()),
            dir.path().join("tmp/axbuild/.image.toml")
        );
        assert!(ImageConfig::get_config_file_path(dir.path()).exists());
    }

    #[test]
    fn read_config_prefers_local_storage_env_override() {
        let _lock = ENV_LOCK.lock().unwrap();
        let dir = tempdir().unwrap();
        let override_path = dir.path().join("persistent-cache");
        let _env = TempEnvVar::set(LOCAL_STORAGE_ENV, override_path.as_os_str());

        let config = ImageConfig::read_config(dir.path()).unwrap();

        assert_eq!(config.local_storage, override_path);
    }

    #[test]
    fn fallback_registry_url_prefers_new_env() {
        let _lock = ENV_LOCK.lock().unwrap();
        let _env = TempEnvVar::set(FALLBACK_REGISTRY_ENV, "https://example.com/new.toml");

        assert_eq!(fallback_registry_url(), "https://example.com/new.toml");
    }
}