ahri_tre_lake/

scratch.rs

use std::fs::{self, File};
use std::os::unix::fs::{MetadataExt, PermissionsExt};
use std::path::{Component, Path, PathBuf};
use std::sync::Arc;

use ahri_tre_security::{REDACTED_VALUE, redact_restricted_path};
use rustix::fs::{FlockOperation, Mode, OFlags, flock, openat};
use thiserror::Error;

/// Opaque identifier for one trusted staging attempt.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ScratchAttemptId(Box<str>);

impl ScratchAttemptId {
    pub fn new(value: &str) -> Result<Self, ScratchError> {
        if value.len() != 32
            || !value
                .bytes()
                .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
        {
            return Err(ScratchError::InvalidAttemptId);
        }
        Ok(Self(value.into()))
    }

    fn as_str(&self) -> &str {
        &self.0
    }
}

/// Validated trusted-service scratch root. It owns an open descriptor so the
/// validated root remains the authority for attempt creation.
pub struct TrustedScratch {
    pub(super) root: PathBuf,
    pub(super) root_directory: Arc<File>,
}

impl std::fmt::Debug for TrustedScratch {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("TrustedScratch")
            .field("root", &REDACTED_VALUE)
            .finish()
    }
}

impl TrustedScratch {
    pub fn open<'a>(
        root: &Path,
        filesystem_lake_locations: impl IntoIterator<Item = &'a Path>,
    ) -> Result<Self, ScratchError> {
        validate_absolute_normalized(root)?;
        if root == Path::new("/") || protected_roots().iter().any(|path| overlaps(root, path)) {
            return Err(ScratchError::ProtectedRoot);
        }
        if filesystem_lake_locations
            .into_iter()
            .any(|lake| overlaps(root, lake))
        {
            return Err(ScratchError::LakeOverlap);
        }

        let slash = File::open("/").map_err(|_| ScratchError::Unavailable)?;
        let mut root_directory = slash;
        for component in root.components() {
            let Component::Normal(name) = component else {
                continue;
            };
            let descriptor = openat(
                &root_directory,
                Path::new(name),
                OFlags::RDONLY | OFlags::DIRECTORY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
                Mode::empty(),
            )
            .map_err(|_| ScratchError::Unavailable)?;
            root_directory = File::from(descriptor);
        }
        let metadata = root_directory
            .metadata()
            .map_err(|_| ScratchError::Unavailable)?;
        if !metadata.is_dir()
            || metadata.uid() != rustix::process::geteuid().as_raw()
            || metadata.permissions().mode() & 0o777 != 0o700
        {
            return Err(ScratchError::UnsafeRoot);
        }
        Ok(Self {
            root: root.to_path_buf(),
            root_directory: Arc::new(root_directory),
        })
    }

    /// Session-local DuckDB spill is distinct from durable Dataset attempt scratch.
    /// A retained OS lease lets a restarted runtime reap only stopped Sessions.
    pub fn create_session_attempt(
        &self,
        identifier: ScratchAttemptId,
    ) -> Result<ScratchAttempt, ScratchError> {
        let mut attempt = self.create_attempt(identifier)?;
        let directory = File::from(
            openat(
                &*self.root_directory,
                attempt
                    .path
                    .file_name()
                    .ok_or(ScratchError::UnsafeAttempt)?,
                OFlags::RDONLY | OFlags::DIRECTORY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
                Mode::empty(),
            )
            .map_err(|_| ScratchError::Unavailable)?,
        );
        // Lock the directory before publishing the Session marker.
        flock(&directory, FlockOperation::NonBlockingLockExclusive)
            .map_err(|_| ScratchError::Unavailable)?;
        let _marker = File::from(
            openat(
                &directory,
                ".session-lease",
                OFlags::RDWR | OFlags::CREATE | OFlags::EXCL | OFlags::CLOEXEC | OFlags::NOFOLLOW,
                Mode::from_bits_retain(0o600),
            )
            .map_err(|_| ScratchError::Unavailable)?,
        );
        attempt._session_lease = Some(directory);
        Ok(attempt)
    }

    /// Reaps abandoned Session spill only. Dataset attempt directories have no
    /// Session marker and must use shared durable attempt recovery, even when no
    /// public operation exists. A live lease always prevents deletion.
    pub fn reap_abandoned_sessions(&self) -> Result<(), ScratchError> {
        let mut cleanup_failed = false;
        for entry in fs::read_dir(&self.root).map_err(|_| ScratchError::Unavailable)? {
            let entry = entry.map_err(|_| ScratchError::Unavailable)?;
            let name = entry.file_name();
            if name
                .to_str()
                .is_none_or(|name| ScratchAttemptId::new(name).is_err())
            {
                continue;
            }
            let Ok(directory) = openat(
                &*self.root_directory,
                name.as_os_str(),
                OFlags::RDONLY | OFlags::DIRECTORY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
                Mode::empty(),
            ) else {
                continue;
            };
            let directory = File::from(directory);
            let metadata = directory
                .metadata()
                .map_err(|_| ScratchError::Unavailable)?;
            if metadata.uid() != rustix::process::geteuid().as_raw()
                || metadata.permissions().mode() & 0o777 != 0o700
            {
                continue;
            }
            let Ok(lease) = openat(
                &directory,
                ".session-lease",
                OFlags::RDWR | OFlags::CLOEXEC | OFlags::NOFOLLOW,
                Mode::empty(),
            ) else {
                continue;
            };
            let lease = File::from(lease);
            let metadata = lease.metadata().map_err(|_| ScratchError::Unavailable)?;
            if !metadata.is_file()
                || metadata.nlink() != 1
                || metadata.uid() != rustix::process::geteuid().as_raw()
                || metadata.permissions().mode() & 0o777 != 0o600
            {
                continue;
            }
            if flock(&directory, FlockOperation::NonBlockingLockExclusive).is_err() {
                continue;
            }
            cleanup_failed |= cleanup_session_tree(&entry.path()).is_err();
        }
        if cleanup_failed {
            Err(ScratchError::CleanupFailed)
        } else {
            Ok(())
        }
    }

    pub fn create_attempt(
        &self,
        identifier: ScratchAttemptId,
    ) -> Result<ScratchAttempt, ScratchError> {
        rustix::fs::mkdirat(
            &*self.root_directory,
            identifier.as_str(),
            Mode::from_bits_retain(0o700),
        )
        .map_err(|error| {
            if error == rustix::io::Errno::EXIST {
                ScratchError::AttemptExists
            } else {
                ScratchError::Unavailable
            }
        })?;
        let path = self.root.join(identifier.as_str());
        let metadata = fs::symlink_metadata(&path).map_err(|_| ScratchError::Unavailable)?;
        if !metadata.is_dir()
            || metadata.file_type().is_symlink()
            || metadata.uid() != rustix::process::geteuid().as_raw()
            || metadata.permissions().mode() & 0o777 != 0o700
        {
            return Err(ScratchError::UnsafeAttempt);
        }
        Ok(ScratchAttempt {
            root: self.root.clone(),
            path,
            _root_directory: Arc::clone(&self.root_directory),
            _session_lease: None,
        })
    }
}

/// One exclusively-created private attempt directory.
pub struct ScratchAttempt {
    root: PathBuf,
    path: PathBuf,
    _root_directory: Arc<File>,
    _session_lease: Option<File>,
}

impl std::fmt::Debug for ScratchAttempt {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("ScratchAttempt")
            .field(
                "path",
                &redact_restricted_path(self.path.display().to_string()),
            )
            .finish()
    }
}

impl ScratchAttempt {
    fn configured_root(&self) -> TrustedScratch {
        TrustedScratch {
            root: self.root.clone(),
            root_directory: Arc::clone(&self._root_directory),
        }
    }

    /// Creates output-owned scratch using the durable private attempt identity.
    pub fn create_dataset_attempt(
        &self,
        authority: &mut dyn ahri_tre_core::DatasetOutputAuthority,
    ) -> Result<Self, crate::LakeError> {
        authority.ensure_current()?;
        use sha2::Digest;
        let identity = crate::file_staging::hex_digest(sha2::Sha256::digest(
            self.root.as_os_str().as_encoded_bytes(),
        ));
        authority.authorize_scratch_root(&identity)?;
        let identifier = ScratchAttemptId::new(&authority.attempt_id().simple().to_string())?;
        self.configured_root()
            .create_attempt(identifier)
            .map_err(Into::into)
    }

    /// Reconciles only the recorded attempt under this retained configured root.
    /// Executor exclusion and metadata outcome must have been claimed first.
    pub fn cleanup_dataset_attempt(
        &self,
        authority: &mut dyn ahri_tre_core::DatasetOutputAuthority,
    ) -> Result<(), crate::LakeError> {
        authority.ensure_current()?;
        use sha2::Digest;
        let identity = crate::file_staging::hex_digest(sha2::Sha256::digest(
            self.root.as_os_str().as_encoded_bytes(),
        ));
        authority.authorize_scratch_root(&identity)?;
        let attempt = Self {
            root: self.root.clone(),
            path: self.root.join(authority.attempt_id().simple().to_string()),
            _root_directory: Arc::clone(&self._root_directory),
            _session_lease: None,
        };
        attempt.cleanup().map_err(Into::into)
    }

    /// Creates an exclusive child attempt under this retained scratch capability.
    pub fn create_child(&self, identifier: ScratchAttemptId) -> Result<Self, ScratchError> {
        TrustedScratch::open(self.path(), std::iter::empty::<&Path>())?.create_attempt(identifier)
    }

    /// Restricted trusted-adapter path; callers must not serialize or project it.
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Deletes the complete attempt tree without following child symlinks.
    /// Missing attempts are already clean and therefore succeed.
    pub fn cleanup(&self) -> Result<(), ScratchError> {
        if self.path.parent() != Some(self.root.as_path()) {
            return Err(ScratchError::UnsafeAttempt);
        }
        match fs::symlink_metadata(&self.path) {
            Ok(metadata) if metadata.is_dir() && !metadata.file_type().is_symlink() => {
                if self._session_lease.is_some() {
                    cleanup_session_tree(&self.path)
                } else {
                    fs::remove_dir_all(&self.path).map_err(|_| ScratchError::CleanupFailed)
                }
            }
            Ok(_) => Err(ScratchError::UnsafeAttempt),
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
            Err(_) => Err(ScratchError::CleanupFailed),
        }
    }
}

// Keep the retry marker until every spill entry is removed. Partial cleanup
// must not turn an abandoned Session tree into unidentifiable Dataset scratch.
fn cleanup_session_tree(path: &Path) -> Result<(), ScratchError> {
    for entry in fs::read_dir(path).map_err(|_| ScratchError::CleanupFailed)? {
        let entry = entry.map_err(|_| ScratchError::CleanupFailed)?;
        if entry.file_name() == ".session-lease" {
            continue;
        }
        if entry
            .file_type()
            .map_err(|_| ScratchError::CleanupFailed)?
            .is_dir()
        {
            fs::remove_dir_all(entry.path()).map_err(|_| ScratchError::CleanupFailed)?;
        } else {
            fs::remove_file(entry.path()).map_err(|_| ScratchError::CleanupFailed)?;
        }
    }
    fs::remove_file(path.join(".session-lease")).map_err(|_| ScratchError::CleanupFailed)?;
    if fs::remove_dir(path).is_err() {
        // No spill remains, but retain retry eligibility if directory removal fails.
        use std::os::unix::fs::OpenOptionsExt;
        let _ = fs::OpenOptions::new()
            .write(true)
            .create_new(true)
            .mode(0o600)
            .open(path.join(".session-lease"));
        return Err(ScratchError::CleanupFailed);
    }
    Ok(())
}

impl Drop for ScratchAttempt {
    fn drop(&mut self) {
        let _ = self.cleanup();
    }
}

fn validate_absolute_normalized(path: &Path) -> Result<(), ScratchError> {
    if !path.is_absolute()
        || path
            .components()
            .any(|component| !matches!(component, Component::RootDir | Component::Normal(_)))
    {
        return Err(ScratchError::InvalidRoot);
    }
    Ok(())
}

fn protected_roots() -> [PathBuf; 5] {
    [
        "/etc/ahri-tre",
        "/run/secrets",
        "/var/lib/ahri-tre/secrets",
        "/run/ahri-tre",
        "/etc/ahri-tre/config.toml",
    ]
    .map(PathBuf::from)
}

fn overlaps(left: &Path, right: &Path) -> bool {
    left.starts_with(right) || right.starts_with(left)
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum ScratchError {
    #[error("trusted scratch root must be an absolute normalized path")]
    InvalidRoot,
    #[error("trusted scratch root overlaps a protected path")]
    ProtectedRoot,
    #[error("trusted scratch root overlaps a filesystem Lake")]
    LakeOverlap,
    #[error("trusted scratch root is unavailable without following links")]
    Unavailable,
    #[error("trusted scratch root has unsafe ownership, type, or permissions")]
    UnsafeRoot,
    #[error("scratch attempt identifier is not canonical")]
    InvalidAttemptId,
    #[error("scratch attempt already exists")]
    AttemptExists,
    #[error("scratch attempt has unsafe ownership, type, or permissions")]
    UnsafeAttempt,
    #[error("scratch attempt cleanup failed")]
    CleanupFailed,
}