lazaret/crates/lazaret-prepared-lower/src/store.rs

use std::fs::File;
use std::io;
use std::os::fd::AsFd;
use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::path::{Path, PathBuf};
use std::thread;
use std::time::{Duration, Instant};

use anyhow::{Context, Result, bail};
use lazaret_core::identity::{
    DeviceId, HostPath, IdMapDigest, InodeId, PolicyId, PrepareOperationId, PreparedArtifactId,
    SelectedManifestDigest, SourceIdentity,
};
use lazaret_core::records::{ArtifactState, PrepareOperation, PrepareStageState, PreparedArtifact};
use lazaret_linux::filesystem::{ExclusiveFileLock, try_exclusive_flock};
use lazaret_storage_catalog::{Catalog, CatalogError, CatalogRow};
use rustix::fs::{AtFlags, Mode, OFlags, ResolveFlags, mkdirat, openat2, renameat, unlinkat};
use sha2::{Digest, Sha256};

use crate::layer::{
    IdMap, POLICY_NAME, PreparationLimits, apply_layout, remove_children, verify_layout,
};

const PHYSICAL_LOCK_WAIT: Duration = Duration::from_secs(5);
const LOCK_RETRY: Duration = Duration::from_millis(2);

#[derive(Clone, Debug)]
pub struct PrepareRequest {
    pub operation_id: PrepareOperationId,
    pub layout_path: PathBuf,
    pub source_identity: Vec<u8>,
    pub selected_manifest_digest: [u8; 32],
    pub id_map: IdMap,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PrepareRecovery {
    ClearedUnmaterializedIntent,
    RemovedOwnedStage,
    PublishedCompletedStage,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ReclaimOutcome {
    Removed,
    ReplacementPreserved,
}

#[derive(Clone, Debug)]
pub struct PreparedLowerStore {
    root: PathBuf,
    catalog: Catalog,
    limits: PreparationLimits,
}

impl PreparedLowerStore {
    /// Creates the fixed prepared-lower directories beneath an existing Storage Root.
    ///
    /// # Errors
    ///
    /// Returns an error if the root cannot be opened or any fixed path is not a real directory.
    pub fn initialize(root: impl Into<PathBuf>, catalog: Catalog) -> Result<Self> {
        let store = Self {
            root: root.into(),
            catalog,
            limits: PreparationLimits::default(),
        };
        let root = File::open(&store.root).context("opening Storage Root")?;
        for name in ["artifacts", "stages", "quarantine"] {
            match mkdirat(&root, name, Mode::from_raw_mode(0o700)) {
                Ok(()) | Err(rustix::io::Errno::EXIST) => {}
                Err(error) => return Err(error.into()),
            }
            open_directory_at(&root, Path::new(name))?;
        }
        open_lock_at(&root)?;
        Ok(store)
    }

    #[must_use]
    pub fn with_limits(mut self, limits: PreparationLimits) -> Self {
        self.limits = limits;
        self
    }

    /// Materializes and atomically publishes one verified immutable prepared lower.
    ///
    /// # Errors
    ///
    /// Refuses corrupt or unsupported OCI input, forged physical ownership, catalog conflicts,
    /// resource limits, and publication failures.
    pub fn prepare(&self, request: &PrepareRequest) -> Result<PreparedArtifact> {
        if request.source_identity.is_empty() || request.source_identity.len() > 4096 {
            bail!("source identity is empty or exceeds 4096 bytes");
        }
        let verified = verify_layout(&request.layout_path, self.limits)?;
        let identities = Identities::new(request)?;
        if let Some(existing) = self.available_artifact(identities.artifact_id)? {
            identities.require_equivalent(&existing)?;
            self.require_artifact_identity(&existing)?;
            return Ok(existing);
        }

        let stage_relative = stage_path(request.operation_id);
        let intent = identities.intent(request.operation_id, &stage_relative);
        self.catalog.record_prepare(&intent)?;
        let stage = match self.create_stage(&stage_relative) {
            Ok(stage) => stage,
            Err(error) => {
                let _ = self.catalog.remove_prepare(&intent);
                return Err(error);
            }
        };
        let (device, inode) = file_identity(&stage)?;
        let owned = PrepareOperation {
            stage_state: PrepareStageState::Owned,
            device: Some(device),
            inode: Some(inode),
            ..intent.clone()
        };
        if let Err(error) = self.catalog.bind_prepare_stage(&intent, &owned) {
            if self
                .remove_exact_path(&stage_relative, device, inode)
                .is_ok()
            {
                let _ = self.catalog.remove_prepare(&intent);
            }
            return Err(error.into());
        }
        if let Err(error) = apply_layout(
            &request.layout_path,
            &stage,
            &verified,
            &request.id_map,
            self.limits,
        )
        .and_then(|()| sync_tree(&stage))
        {
            if self
                .remove_exact_path(&stage_relative, device, inode)
                .is_ok()
            {
                let _ = self.catalog.remove_prepare(&owned);
            }
            return Err(error);
        }
        self.publish_owned(&owned, &identities)
    }

    /// Recovers one exact durable prepare operation without adopting foreign paths.
    ///
    /// # Errors
    ///
    /// Refuses unexpected paths, missing ownership evidence, identity mismatch, and catalog or
    /// filesystem failures.
    pub fn recover_prepare(&self, operation: &PrepareOperation) -> Result<PrepareRecovery> {
        let expected_stage = stage_path(operation.operation_id);
        if operation.stage_path.as_bytes() != path_bytes(&expected_stage) {
            bail!("prepare operation stage path is not canonical");
        }
        match operation.stage_state {
            PrepareStageState::Intent => {
                if self.open_relative(&expected_stage)?.is_some() {
                    bail!("unowned stage path is present");
                }
                self.catalog.remove_prepare(operation)?;
                Ok(PrepareRecovery::ClearedUnmaterializedIntent)
            }
            PrepareStageState::Owned => {
                let device = operation.device.context("owned prepare lacks device")?;
                let inode = operation.inode.context("owned prepare lacks inode")?;
                if self.open_relative(&expected_stage)?.is_some() {
                    self.remove_exact_path(&expected_stage, device, inode)?;
                    self.catalog.remove_prepare(operation)?;
                    return Ok(PrepareRecovery::RemovedOwnedStage);
                }
                let artifact_relative = artifact_path(operation.expected_artifact_id);
                let Some(artifact_fd) = self.open_relative(&artifact_relative)? else {
                    bail!("owned preparation path is absent without exact recovery evidence");
                };
                if file_identity(&artifact_fd)? != (device, inode) {
                    bail!("published path does not match preparation ownership");
                }
                let artifact =
                    artifact_from_operation(operation, &artifact_relative, device, inode);
                self.publish_catalog_row(operation, &artifact)?;
                Ok(PrepareRecovery::PublishedCompletedStage)
            }
        }
    }

    /// Moves an exact zero-lease artifact to quarantine, deletes it, and closes reclaim intent.
    ///
    /// # Errors
    ///
    /// Refuses leases, forged catalog rows, physical identity mismatch, foreign paths, and
    /// filesystem or catalog failures.
    pub fn reclaim(&self, expected: &PreparedArtifact) -> Result<ReclaimOutcome> {
        let reclaiming = self.catalog.begin_reclaim(expected)?;
        self.resume_reclaim(&reclaiming)
    }

    /// Resumes an exact durable reclaim after interruption.
    ///
    /// # Errors
    ///
    /// Refuses non-reclaiming state, path/identity mismatch, missing physical evidence, and
    /// catalog or filesystem failures.
    pub fn resume_reclaim(&self, reclaiming: &PreparedArtifact) -> Result<ReclaimOutcome> {
        if reclaiming.state != ArtifactState::Reclaiming {
            bail!("artifact is not reclaiming");
        }
        let expected_artifact = artifact_path(reclaiming.prepared_artifact_id);
        if reclaiming.path.as_bytes() != path_bytes(&expected_artifact) {
            bail!("artifact path is not canonical");
        }
        let quarantine = quarantine_path(reclaiming);
        {
            let _lock = self.physical_lock()?;
            if let Some(canonical) = self.open_relative(&expected_artifact)? {
                require_identity(&canonical, reclaiming.device, reclaiming.inode)?;
                self.rename_relative(&expected_artifact, &quarantine)?;
                self.sync_fixed_directory("artifacts")?;
                self.sync_fixed_directory("quarantine")?;
            } else {
                let existing = self
                    .open_relative(&quarantine)?
                    .context("reclaiming artifact and quarantine are both absent")?;
                require_identity(&existing, reclaiming.device, reclaiming.inode)?;
            }
        }
        self.remove_exact_path(&quarantine, reclaiming.device, reclaiming.inode)?;
        match self.catalog.finish_reclaim(reclaiming) {
            Ok(()) => Ok(ReclaimOutcome::Removed),
            Err(CatalogError::OwnershipMismatch {
                row: CatalogRow::PreparedArtifact,
            }) => {
                let replacement = self
                    .available_artifact(reclaiming.prepared_artifact_id)?
                    .context("reclaim row changed without an available replacement")?;
                self.require_artifact_identity(&replacement)?;
                Ok(ReclaimOutcome::ReplacementPreserved)
            }
            Err(error) => Err(error.into()),
        }
    }

    fn publish_owned(
        &self,
        operation: &PrepareOperation,
        identities: &Identities,
    ) -> Result<PreparedArtifact> {
        let stage = stage_path(operation.operation_id);
        let artifact_path = artifact_path(operation.expected_artifact_id);
        let _lock = self.physical_lock()?;
        if let Some(existing) = self.open_relative(&artifact_path)? {
            if let Some(row) = self.available_artifact(operation.expected_artifact_id)? {
                identities.require_equivalent(&row)?;
                require_identity(&existing, row.device, row.inode)?;
                let device = operation.device.context("owned prepare lacks device")?;
                let inode = operation.inode.context("owned prepare lacks inode")?;
                self.remove_exact_path(&stage, device, inode)?;
                self.catalog.remove_prepare(operation)?;
                return Ok(row);
            }
            let reclaiming = self
                .catalog
                .snapshot()?
                .prepared_artifacts
                .into_iter()
                .find(|candidate| {
                    candidate.prepared_artifact_id == operation.expected_artifact_id
                        && candidate.state == ArtifactState::Reclaiming
                })
                .context("canonical artifact path exists without matching catalog ownership")?;
            require_identity(&existing, reclaiming.device, reclaiming.inode)?;
            let quarantine = quarantine_path(&reclaiming);
            self.rename_relative(&artifact_path, &quarantine)?;
            self.sync_fixed_directory("quarantine")?;
        }
        self.rename_relative(&stage, &artifact_path)?;
        self.sync_fixed_directory("artifacts")?;
        let artifact_fd = self
            .open_relative(&artifact_path)?
            .context("published artifact disappeared")?;
        let (device, inode) = file_identity(&artifact_fd)?;
        let artifact = identities.artifact(&artifact_path, device, inode);
        self.publish_catalog_row(operation, &artifact)?;
        Ok(artifact)
    }

    fn publish_catalog_row(
        &self,
        operation: &PrepareOperation,
        artifact: &PreparedArtifact,
    ) -> Result<()> {
        let snapshot = self.catalog.snapshot()?;
        if let Some(reclaiming) = snapshot.prepared_artifacts.into_iter().find(|candidate| {
            candidate.prepared_artifact_id == artifact.prepared_artifact_id
                && candidate.state == ArtifactState::Reclaiming
        }) {
            self.catalog
                .publish_replacement(operation, &reclaiming, artifact)?;
        } else {
            self.catalog.publish(operation, artifact)?;
        }
        Ok(())
    }

    fn available_artifact(&self, id: PreparedArtifactId) -> Result<Option<PreparedArtifact>> {
        Ok(self
            .catalog
            .snapshot()?
            .prepared_artifacts
            .into_iter()
            .find(|artifact| {
                artifact.prepared_artifact_id == id && artifact.state == ArtifactState::Available
            }))
    }

    fn require_artifact_identity(&self, artifact: &PreparedArtifact) -> Result<()> {
        let path = PathBuf::from(std::ffi::OsString::from_vec(
            artifact.path.as_bytes().to_vec(),
        ));
        let expected = artifact_path(artifact.prepared_artifact_id);
        if path != expected {
            bail!("artifact path is not canonical");
        }
        let file = self
            .open_relative(&path)?
            .context("artifact path is absent")?;
        require_identity(&file, artifact.device, artifact.inode)
    }

    fn create_stage(&self, relative: &Path) -> Result<File> {
        let stages = open_directory_at(&File::open(&self.root)?, Path::new("stages"))?;
        let name = relative.file_name().context("stage has no name")?;
        mkdirat(&stages, name, Mode::from_raw_mode(0o700))?;
        open_directory_at(&stages, Path::new(name))
    }

    fn open_relative(&self, relative: &Path) -> Result<Option<File>> {
        let root = File::open(&self.root)?;
        match openat2(
            &root,
            relative,
            OFlags::RDONLY | OFlags::DIRECTORY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
            Mode::empty(),
            secure_resolve(),
        ) {
            Ok(fd) => Ok(Some(File::from(fd))),
            Err(rustix::io::Errno::NOENT) => Ok(None),
            Err(error) => Err(error.into()),
        }
    }

    fn rename_relative(&self, from: &Path, to: &Path) -> Result<()> {
        let root = File::open(&self.root)?;
        renameat(&root, from, &root, to)?;
        Ok(())
    }

    fn remove_exact_path(&self, relative: &Path, device: DeviceId, inode: InodeId) -> Result<()> {
        let root = File::open(&self.root)?;
        let file = open_directory_at(&root, relative)?;
        require_identity(&file, device, inode)?;
        remove_children(&file)?;
        unlinkat(&root, relative, AtFlags::REMOVEDIR)?;
        Ok(())
    }

    fn physical_lock(&self) -> Result<ExclusiveFileLock> {
        let deadline = Instant::now() + PHYSICAL_LOCK_WAIT;
        loop {
            let file = open_lock_at(&File::open(&self.root)?)?;
            match try_exclusive_flock(file) {
                Ok(lock) => return Ok(lock),
                Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
                    let now = Instant::now();
                    if now >= deadline {
                        bail!("prepared-lower physical lock wait expired");
                    }
                    thread::sleep(LOCK_RETRY.min(deadline.saturating_duration_since(now)));
                }
                Err(error) => return Err(error.into()),
            }
        }
    }

    fn sync_fixed_directory(&self, name: &str) -> Result<()> {
        open_directory_at(&File::open(&self.root)?, Path::new(name))?.sync_all()?;
        Ok(())
    }
}

struct Identities {
    artifact_id: PreparedArtifactId,
    source: SourceIdentity,
    manifest: SelectedManifestDigest,
    policy: PolicyId,
    id_map: IdMapDigest,
}

impl Identities {
    fn new(request: &PrepareRequest) -> Result<Self> {
        let source = SourceIdentity::new(request.source_identity.clone());
        let manifest = SelectedManifestDigest::new(request.selected_manifest_digest);
        let policy = PolicyId::new(hash_parts(&[b"lazaret-policy-v1", POLICY_NAME.as_bytes()]));
        let id_map_bytes = postcard::to_stdvec(&request.id_map)?;
        let id_map = IdMapDigest::new(hash_parts(&[b"lazaret-id-map-v1", &id_map_bytes]));
        let artifact_id = PreparedArtifactId::new(hash_parts(&[
            b"lazaret-prepared-artifact-v1",
            source.as_bytes(),
            manifest.as_bytes(),
            policy.as_bytes(),
            id_map.as_bytes(),
        ]));
        Ok(Self {
            artifact_id,
            source,
            manifest,
            policy,
            id_map,
        })
    }

    fn intent(&self, operation_id: PrepareOperationId, stage: &Path) -> PrepareOperation {
        PrepareOperation {
            operation_id,
            expected_artifact_id: self.artifact_id,
            source_identity: self.source.clone(),
            selected_manifest_digest: self.manifest,
            policy_id: self.policy,
            id_map_digest: self.id_map,
            stage_path: HostPath::new(path_bytes(stage).to_vec()),
            stage_state: PrepareStageState::Intent,
            device: None,
            inode: None,
        }
    }

    fn artifact(&self, path: &Path, device: DeviceId, inode: InodeId) -> PreparedArtifact {
        PreparedArtifact {
            prepared_artifact_id: self.artifact_id,
            source_identity: self.source.clone(),
            selected_manifest_digest: self.manifest,
            policy_id: self.policy,
            id_map_digest: self.id_map,
            path: HostPath::new(path_bytes(path).to_vec()),
            device,
            inode,
            state: ArtifactState::Available,
        }
    }

    fn require_equivalent(&self, artifact: &PreparedArtifact) -> Result<()> {
        if artifact.prepared_artifact_id != self.artifact_id
            || artifact.source_identity != self.source
            || artifact.selected_manifest_digest != self.manifest
            || artifact.policy_id != self.policy
            || artifact.id_map_digest != self.id_map
        {
            bail!("prepared artifact identity inputs differ");
        }
        Ok(())
    }
}

fn artifact_from_operation(
    operation: &PrepareOperation,
    path: &Path,
    device: DeviceId,
    inode: InodeId,
) -> PreparedArtifact {
    PreparedArtifact {
        prepared_artifact_id: operation.expected_artifact_id,
        source_identity: operation.source_identity.clone(),
        selected_manifest_digest: operation.selected_manifest_digest,
        policy_id: operation.policy_id,
        id_map_digest: operation.id_map_digest,
        path: HostPath::new(path_bytes(path).to_vec()),
        device,
        inode,
        state: ArtifactState::Available,
    }
}

fn file_identity(file: &File) -> Result<(DeviceId, InodeId)> {
    let stat = rustix::fs::fstat(file)?;
    Ok((DeviceId::new(stat.st_dev), InodeId::new(stat.st_ino)))
}

fn require_identity(file: &File, device: DeviceId, inode: InodeId) -> Result<()> {
    if file_identity(file)? == (device, inode) {
        Ok(())
    } else {
        bail!("physical path identity mismatch")
    }
}

fn open_directory_at(parent: impl AsFd, path: &Path) -> Result<File> {
    Ok(File::from(openat2(
        parent,
        path,
        OFlags::RDONLY | OFlags::DIRECTORY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
        Mode::empty(),
        secure_resolve(),
    )?))
}

fn open_lock_at(root: impl AsFd) -> Result<File> {
    Ok(File::from(openat2(
        root,
        "prepared-lower.lock",
        OFlags::RDWR | OFlags::CREATE | OFlags::CLOEXEC | OFlags::NOFOLLOW,
        Mode::from_raw_mode(0o600),
        secure_resolve(),
    )?))
}

fn sync_tree(directory: &File) -> Result<()> {
    let mut buffer = [std::mem::MaybeUninit::uninit(); 8192];
    let names = {
        let mut names = Vec::new();
        let mut entries = rustix::fs::RawDir::new(directory.as_fd(), &mut buffer);
        while let Some(entry) = entries.next() {
            let entry = entry?;
            if entry.file_name().to_bytes() != b"." && entry.file_name().to_bytes() != b".." {
                names.push(entry.file_name().to_owned());
            }
        }
        names
    };
    for name in names {
        let path_fd = openat2(
            directory,
            &name,
            OFlags::PATH | OFlags::CLOEXEC | OFlags::NOFOLLOW,
            Mode::empty(),
            secure_resolve(),
        )?;
        let kind = rustix::fs::FileType::from_raw_mode(rustix::fs::fstat(&path_fd)?.st_mode);
        if kind == rustix::fs::FileType::Directory {
            let child = open_directory_at(
                directory,
                Path::new(std::ffi::OsStr::from_bytes(name.to_bytes())),
            )?;
            sync_tree(&child)?;
        } else if kind == rustix::fs::FileType::RegularFile {
            File::from(openat2(
                directory,
                &name,
                OFlags::RDONLY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
                Mode::empty(),
                secure_resolve(),
            )?)
            .sync_all()?;
        }
    }
    directory.sync_all()?;
    Ok(())
}

fn hash_parts(parts: &[&[u8]]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    for part in parts {
        hasher.update((part.len() as u64).to_be_bytes());
        hasher.update(part);
    }
    hasher.finalize().into()
}

fn stage_path(id: PrepareOperationId) -> PathBuf {
    Path::new("stages").join(hex(id.as_bytes()))
}

fn artifact_path(id: PreparedArtifactId) -> PathBuf {
    Path::new("artifacts").join(hex(id.as_bytes()))
}

fn quarantine_path(artifact: &PreparedArtifact) -> PathBuf {
    Path::new("quarantine").join(format!(
        "{}-{}-{}",
        hex(artifact.prepared_artifact_id.as_bytes()),
        artifact.device.get(),
        artifact.inode.get()
    ))
}

fn path_bytes(path: &Path) -> &[u8] {
    std::os::unix::ffi::OsStrExt::as_bytes(path.as_os_str())
}

fn hex(bytes: &[u8]) -> String {
    use std::fmt::Write as _;

    bytes.iter().fold(
        String::with_capacity(bytes.len() * 2),
        |mut output, byte| {
            write!(output, "{byte:02x}").expect("writing to String cannot fail");
            output
        },
    )
}

fn secure_resolve() -> ResolveFlags {
    ResolveFlags::BENEATH | ResolveFlags::NO_SYMLINKS | ResolveFlags::NO_MAGICLINKS
}