aion_server/worker/auto_provision/

documents.rs

//! Where the server keeps the documents its OWN managed workers read.
//!
//! # Why a file at all
//!
//! `aion worker agent` serves the queue an `.awl` document declares, and the
//! whole launch — harness kind, concurrency, reconnect budget, the agent's own
//! command and environment — is IN that document. A managed-worker record
//! replays its argv verbatim, so the argv has to name a path that still holds
//! the same bytes on the next restart, a week later, across a server upgrade.
//! The server has those bytes: they ride in the deployed archive's `awl/`
//! provenance tree. So it writes them down.
//!
//! # Content-addressed, and addressed by the DOCUMENT
//!
//! The directory name carries a digest of the document text, NOT the package's
//! content hash. Those are different facts and the difference is load-bearing:
//! adding or editing a `harness` section changes no lowering (the MIR ratchet
//! pins that), so two packages that differ only in their harness section share
//! one content hash. Keying the snapshot on the package hash would therefore
//! hand a re-deploy the OLD document, and the re-mint that is supposed to move
//! the worker onto the new launch would be a no-op.
//!
//! Content-addressing is also what makes a re-mint honest in the other
//! direction: a redeploy of unchanged bytes resolves to the same path, so the
//! record's argv is unchanged and the running worker is left alone.
//!
//! # Not the studio's tree
//!
//! `<AION_HOME>/authoring` is the AWL studio's root: an operator-writable
//! document tree served over HTTP. These snapshots are neither — they are
//! immutable server-owned state that a running process reads by path, and an
//! edit through the studio would silently change what a worker serves on its
//! next restart while the record still names the old digest. They live under
//! `<AION_HOME>/workers/documents/` instead.

use std::collections::BTreeSet;
use std::path::{Component, Path, PathBuf};

use aion_package::AwlSource;
use sha2::{Digest, Sha256};

use crate::filesystem::ConfinedDir;
use crate::worker::lowercase_hex;

/// The home-relative root every snapshot is written under.
const DOCUMENTS_ROOT: &str = "workers/documents";

/// Why a deployed document could not be staged for a managed worker.
#[derive(Debug, thiserror::Error)]
pub enum DocumentStoreError {
    /// The Aion home could not be resolved, so there is nowhere server-owned
    /// to put the document.
    #[error(
        "the Aion home could not be resolved, so the deployed document has nowhere to live: {message}"
    )]
    HomeUnresolved {
        /// The resolution failure, verbatim.
        message: String,
    },
    /// A name from the archive is not a single ordinary path component.
    #[error(
        "the deployed package names `{name}` inside its archived AWL tree, which is not a \
         relative path of ordinary components; nothing was written"
    )]
    UnsafeName {
        /// The rejected name, verbatim.
        name: String,
    },
    /// The filesystem refused.
    #[error("the deployed document could not be written under `{path}`: {message}")]
    Write {
        /// The path being written, for the operator to inspect.
        path: String,
        /// The operating system's own diagnosis.
        message: String,
    },
}

/// One deployed document, staged on disk with the schema files it imports.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StagedDocument {
    /// Absolute path of the document itself — the path a record's argv names.
    pub path: PathBuf,
    /// Lowercase hexadecimal SHA-256 of the document text. Two deploys of the
    /// same bytes produce the same digest and therefore the same argv.
    pub digest: String,
}

/// Write `source`'s document and every schema it imports under the server's own
/// managed-worker document root, and return the document's absolute path.
///
/// Idempotent by construction: the directory is named by the document's digest,
/// so re-staging identical bytes rewrites identical bytes at the same path.
///
/// # Errors
///
/// Returns [`DocumentStoreError::HomeUnresolved`] when the Aion home cannot be
/// resolved, [`DocumentStoreError::UnsafeName`] for an archive entry whose name
/// is not a relative path of ordinary components, and
/// [`DocumentStoreError::Write`] when the filesystem refuses.
pub fn stage(
    root: &Path,
    source: &AwlSource,
    workflow_type: &str,
) -> Result<StagedDocument, DocumentStoreError> {
    let digest = lowercase_hex(&Sha256::digest(source.document().as_bytes()));
    let version = version_directory(workflow_type, &digest);
    let absolute = root.join(&version);
    // ATOMIC AS A DIRECTORY, not merely file by file. Each write below is
    // already atomic on its own, but a crash between the document and the last
    // schema would leave a tree that LOOKS complete — the directory name is a
    // digest, so nothing downstream re-checks it — and the next boot's
    // reconcile would hand a worker a document whose `schema(…)` import is
    // missing, crash-looping it. So the tree is built beside its final name and
    // renamed into place, which is one atomic step.
    let staging = root.join(format!("{version}.partial"));
    remove_tree(&staging)?;
    // Every missing component is created owner-only by the capability itself,
    // and every file it writes is owner-only too, so there is no separate
    // hardening pass to forget: a document that decides how an agent is
    // launched is never world-readable.
    let dir =
        ConfinedDir::open_or_create(&staging).map_err(|error| write_failure(&staging, &error))?;

    let document_name = safe_component(source.document_name())?;
    dir.atomic_write(&document_name, source.document().as_bytes())
        .map_err(|error| write_failure(&staging.join(&document_name), &error))?;

    for (schema_path, bytes) in source.schemas() {
        let relative = safe_relative(schema_path)?;
        if let Some(parent) = relative
            .parent()
            .filter(|parent| !parent.as_os_str().is_empty())
        {
            dir.create_dir_all(parent)
                .map_err(|error| write_failure(&staging.join(parent), &error))?;
        }
        dir.atomic_write(&relative, bytes)
            .map_err(|error| write_failure(&staging.join(&relative), &error))?;
    }
    drop(dir);

    // The rename cannot land on top of an existing directory, and an existing
    // one is already the SAME bytes (the name is their digest), so the staged
    // copy is simply discarded rather than swapped in.
    if absolute.exists() {
        remove_tree(&staging)?;
    } else {
        std::fs::rename(&staging, &absolute).map_err(|error| write_failure(&absolute, &error))?;
    }

    Ok(StagedDocument {
        path: absolute.join(document_name),
        digest,
    })
}

/// The server-owned root every snapshot lives under, resolved from the Aion
/// home.
///
/// # Errors
///
/// Returns [`DocumentStoreError::HomeUnresolved`] when the Aion home cannot be
/// resolved. Nothing is defaulted: a snapshot written somewhere nobody chose is
/// a path a record would then replay forever.
pub fn root() -> Result<PathBuf, DocumentStoreError> {
    crate::config::aion_home()
        .map(|home| home.path.join(DOCUMENTS_ROOT))
        .map_err(|error| DocumentStoreError::HomeUnresolved {
            message: error.to_string(),
        })
}

/// Remove every snapshot directory under `root` that no live record names.
///
/// A redeploy of a changed document writes a NEW digest directory and the old
/// one is not automatically anyone's to delete — a record minted against it may
/// still be replaying that argv until its own re-mint converges. But once no
/// record's argv names it, nothing will ever read it again, and leaving it is
/// an unbounded leak in the server's own state directory: a CI loop redeploying
/// a document a hundred times a day accumulates a hundred trees a day.
///
/// `live` is every path a current worker-deployment argv names. A path this
/// cannot classify is KEPT: deleting a snapshot on an unreadable answer is how
/// a running worker loses its document.
///
/// # Errors
///
/// Returns [`DocumentStoreError::Write`] when the root cannot be listed or a
/// directory cannot be removed.
pub fn prune(root: &Path, live: &BTreeSet<PathBuf>) -> Result<Vec<PathBuf>, DocumentStoreError> {
    if !root.exists() {
        return Ok(Vec::new());
    }
    let entries = std::fs::read_dir(root).map_err(|error| write_failure(root, &error))?;
    let mut removed = Vec::new();
    for entry in entries {
        let entry = entry.map_err(|error| write_failure(root, &error))?;
        let path = entry.path();
        if !path.is_dir() {
            continue;
        }
        // A snapshot is named by any live path whose ancestor it is: the argv
        // names the DOCUMENT inside the directory, not the directory.
        if live.iter().any(|document| document.starts_with(&path)) {
            continue;
        }
        remove_tree(&path)?;
        removed.push(path);
    }
    Ok(removed)
}

/// Remove one directory tree, treating "it was not there" as success.
fn remove_tree(path: &Path) -> Result<(), DocumentStoreError> {
    match std::fs::remove_dir_all(path) {
        Ok(()) => Ok(()),
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
        Err(error) => Err(write_failure(path, &error)),
    }
}

/// The per-version directory name: the workflow type, then its document digest.
///
/// The type is sanitised rather than trusted. It is an AWL identifier by the
/// time a package exists, but this reads an archive a stranger may have built,
/// and a directory name is not the place to find out.
fn version_directory(workflow_type: &str, digest: &str) -> String {
    format!("{}@{digest}", sanitise_workflow_type(workflow_type))
}

/// The workflow type as it is spelled in a staged directory's name.
///
/// One spelling, used by the deploy that names the directory and by every
/// reader that has to compare a record's staged type against a package's
/// workflow type — a type that contains a character the directory cannot is
/// compared in the form the directory carries, never the raw one.
pub(super) fn sanitise_workflow_type(workflow_type: &str) -> String {
    let sanitised: String = workflow_type
        .chars()
        .map(|character| {
            if character.is_ascii_alphanumeric() || character == '_' || character == '-' {
                character
            } else {
                '_'
            }
        })
        .collect();
    if sanitised.is_empty() {
        "workflow".to_owned()
    } else {
        sanitised
    }
}

/// One archive entry name reduced to a single ordinary path component.
fn safe_component(name: &str) -> Result<PathBuf, DocumentStoreError> {
    let path = Path::new(name);
    let mut components = path.components();
    match (components.next(), components.next()) {
        (Some(Component::Normal(single)), None) => Ok(PathBuf::from(single)),
        _ => Err(DocumentStoreError::UnsafeName {
            name: name.to_owned(),
        }),
    }
}

/// One archive entry name reduced to a relative path of ordinary components.
///
/// `ConfinedDir` is capability-rooted and would refuse an escape anyway; this
/// refuses it BY NAME first, so the log says which entry was rejected rather
/// than reporting a bare permission error against a path nobody wrote.
fn safe_relative(name: &str) -> Result<PathBuf, DocumentStoreError> {
    let path = Path::new(name);
    let mut safe = PathBuf::new();
    let mut any = false;
    for component in path.components() {
        match component {
            Component::Normal(part) => {
                safe.push(part);
                any = true;
            }
            _ => {
                return Err(DocumentStoreError::UnsafeName {
                    name: name.to_owned(),
                });
            }
        }
    }
    if any {
        Ok(safe)
    } else {
        Err(DocumentStoreError::UnsafeName {
            name: name.to_owned(),
        })
    }
}

fn write_failure(path: &Path, error: &std::io::Error) -> DocumentStoreError {
    DocumentStoreError::Write {
        path: path.display().to_string(),
        message: error.to_string(),
    }
}

#[cfg(test)]
#[path = "documents_tests.rs"]
mod tests;