endorphin-aster-vm/crates/host/src/metadata.rs

use std::fs::{self, OpenOptions};
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};

use serde::{Deserialize, Serialize};

static NEXT_TEMP: AtomicU64 = AtomicU64::new(0);
const FORMAT_VERSION: u32 = 1;

#[derive(Deserialize)]
struct LegacyRecord {
    key: String,
    value: String,
}

#[derive(Serialize, Deserialize)]
struct Record {
    format_version: u32,
    key: String,
    value: String,
}

/// One independently replaceable JSON file per ASCII metadata key. Clones
/// share a directory, not a cached map; reads always observe the file system.
#[derive(Debug, Clone)]
pub struct MetadataStore {
    directory: PathBuf,
}

impl MetadataStore {
    /// Does not create the directory until the first write.
    pub fn new(directory: impl Into<PathBuf>) -> io::Result<Self> {
        let directory = directory.into();
        if directory.as_os_str().is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "metadata directory is empty",
            ));
        }
        Ok(Self { directory })
    }

    /// Preserve OS paths without requiring the environment value to be UTF-8.
    pub fn from_env() -> io::Result<Self> {
        Self::new(
            std::env::var_os("ASTER_STATE_DIR")
                .map(PathBuf::from)
                .unwrap_or_else(|| PathBuf::from("data/host")),
        )
    }

    pub fn directory(&self) -> &Path {
        &self.directory
    }

    pub fn get(&self, key: &str) -> io::Result<Option<String>> {
        let path = self.path(key)?;
        let bytes = match fs::read(path) {
            Ok(bytes) => bytes,
            Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(None),
            Err(error) => return Err(error),
        };
        let record: Record = serde_json::from_slice(&bytes)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        if record.format_version != FORMAT_VERSION || record.key != key {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "metadata key or format version mismatch",
            ));
        }
        Ok(Some(record.value))
    }

    /// Writes a new sibling file, syncs it, then publishes it with rename.
    /// Different keys do not share a read/modify/write map. For the same key,
    /// the last successful rename wins. This is not a multi-key transaction.
    pub fn put(&self, key: &str, value: &str) -> io::Result<()> {
        self.store(key, value, true).map(|_| ())
    }

    /// One-time migration input: a JSON array of key/value string records.
    /// Validate all keys first, and never overwrite an existing file record.
    pub fn import_missing(&self, reader: impl Read) -> io::Result<usize> {
        let records: Vec<LegacyRecord> = serde_json::from_reader(reader)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        for record in &records {
            self.path(&record.key)?;
        }
        let mut imported = 0;
        for record in records {
            if self.store(&record.key, &record.value, false)? {
                imported += 1;
            } else if self.get(&record.key)?.is_none() {
                return Err(io::Error::new(
                    io::ErrorKind::NotFound,
                    "existing metadata disappeared during import",
                ));
            }
        }
        Ok(imported)
    }

    fn store(&self, key: &str, value: &str, replace: bool) -> io::Result<bool> {
        let destination = self.path(key)?;
        let mut contents = serde_json::to_vec_pretty(&Record {
            format_version: FORMAT_VERSION,
            key: key.into(),
            value: value.into(),
        })
        .map_err(io::Error::other)?;
        contents.push(b'\n');
        fs::create_dir_all(&self.directory)?;
        for _ in 0..32 {
            let sequence = NEXT_TEMP.fetch_add(1, Ordering::Relaxed);
            let temporary = self
                .directory
                .join(format!(".{key}.{}.{sequence}.tmp", std::process::id()));
            let mut options = OpenOptions::new();
            options.write(true).create_new(true);
            #[cfg(unix)]
            {
                use std::os::unix::fs::OpenOptionsExt;
                options.mode(0o600);
            }
            let mut file = match options.open(&temporary) {
                Ok(file) => file,
                Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
                Err(error) => return Err(error),
            };
            let result = (|| {
                file.write_all(&contents)?;
                file.sync_all()?;
                drop(file);
                if replace {
                    fs::rename(&temporary, &destination)?;
                    Ok(true)
                } else {
                    // Atomic create-if-absent: a concurrent writer's record
                    // cannot be replaced by the legacy import.
                    match fs::hard_link(&temporary, &destination) {
                        Ok(()) => Ok(true),
                        Err(error) if error.kind() == io::ErrorKind::AlreadyExists => Ok(false),
                        Err(error) => Err(error),
                    }
                }
            })();
            let _ = fs::remove_file(&temporary);
            return result;
        }
        Err(io::Error::new(
            io::ErrorKind::AlreadyExists,
            "cannot reserve a metadata temporary file",
        ))
    }

    fn path(&self, key: &str) -> io::Result<PathBuf> {
        if key.is_empty()
            || key.len() > 128
            || !key.as_bytes()[0].is_ascii_alphanumeric()
            || !key
                .bytes()
                .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-'))
        {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "metadata keys must start with an ASCII letter/digit and contain at most 128 ASCII letters, digits, dots, underscores or hyphens",
            ));
        }
        Ok(self.directory.join(format!("{key}.json")))
    }
}