zlayer-agent 0.14.1

//! Runtime-agnostic storage-volume preparation.
//!
//! Every runtime that backs a container with a real filesystem needs the same
//! up-front step: turn the declarative [`StorageSpec`] list into concrete host
//! directories (a persistent dir per *named* volume, a scratch dir per
//! *anonymous* volume, a FUSE mount per *S3* volume) before the container's
//! mounts can point at them. The youki runtime grew this logic first; it is
//! hoisted here verbatim so the macOS (`macos_sandbox`, `macos_vz_linux`) and
//! Windows (`wsl2_delegate`) runtimes prepare volumes *identically* rather than
//! each re-deriving a slightly different version.
//!
//! [`prepare_storage_volumes`] returns a `target/name -> host_path` map keyed
//! exactly the way [`crate::bundle`] consumes it (volume name for `Named`,
//! `_anon_<target>` for `Anonymous`, `_s3_<bucket>_<prefix>` for `S3`).
//! [`resolve_host_source`] inverts that map for runtimes that materialise their
//! own mounts (virtiofs shares, Seatbelt rootfs rebasing) rather than going
//! through `bundle::build_storage_mounts`.

use std::collections::HashMap;
use std::path::PathBuf;

use zlayer_spec::StorageSpec;

use crate::error::{AgentError, Result};
use crate::storage_manager::StorageManager;

/// Materialise the host-side backing directories for a container's storage.
///
/// Mirrors the original youki `prepare_storage_volumes`: `Named` volumes are
/// ensured (and layer-synced) + attached to the container, `Anonymous` volumes
/// get a per-container scratch dir, `S3` buckets are FUSE-mounted, and
/// `Bind`/`Tmpfs` need no preparation. The returned map is keyed so it can be
/// fed straight into [`crate::bundle`]'s `build_storage_mounts`.
// Used by the youki (Linux) and macOS storage-prep paths; the Windows WSL2
// delegate shells `youki` inside the WSL2 distro and never calls this Rust-side
// prep, so it is legitimately unused on `all(windows, wsl)`.
#[cfg_attr(all(target_os = "windows", feature = "wsl"), allow(dead_code))]
pub(crate) async fn prepare_storage_volumes(
    storage_manager: &mut StorageManager,
    container_id: &str,
    storage: &[StorageSpec],
) -> Result<HashMap<String, PathBuf>> {
    let mut volume_paths = HashMap::new();

    for spec in storage {
        match spec {
            StorageSpec::Named { name, .. } => {
                let path = storage_manager
                    .ensure_volume_with_sync(name)
                    .await
                    .map_err(|e| AgentError::CreateFailed {
                        id: container_id.to_string(),
                        reason: format!("failed to ensure volume '{name}': {e}"),
                    })?;
                storage_manager
                    .attach_volume(name, container_id)
                    .map_err(|e| AgentError::CreateFailed {
                        id: container_id.to_string(),
                        reason: format!("failed to attach volume '{name}': {e}"),
                    })?;
                volume_paths.insert(name.clone(), path);
            }

            StorageSpec::Anonymous { target, .. } => {
                let path = storage_manager
                    .create_anonymous(container_id, target)
                    .map_err(|e| AgentError::CreateFailed {
                        id: container_id.to_string(),
                        reason: format!("failed to create anonymous volume for '{target}': {e}"),
                    })?;
                volume_paths.insert(anon_key(target), path);
            }

            // Bind and tmpfs mounts don't need a prepared host directory.
            StorageSpec::Bind { .. } | StorageSpec::Tmpfs { .. } => {}

            StorageSpec::S3 {
                bucket,
                prefix,
                endpoint,
                ..
            } => {
                let path = storage_manager
                    .mount_s3(bucket, prefix.as_deref(), endpoint.as_deref(), container_id)
                    .map_err(|e| AgentError::CreateFailed {
                        id: container_id.to_string(),
                        reason: format!("failed to mount S3 bucket '{bucket}': {e}"),
                    })?;
                volume_paths.insert(s3_key(bucket, prefix.as_deref()), path);
            }
        }
    }

    Ok(volume_paths)
}

/// Release the host-side space a container held: detach its named volumes
/// (dropping the refcount so an unreferenced volume can later be reclaimed),
/// unmount its S3 mounts, and delete its anonymous scratch dirs.
///
/// Keyed by `container_id` so cleanup reclaims exactly the space *this* container
/// attached — the inverse of [`prepare_storage_volumes`]. Best-effort: a failure
/// on one volume is logged and the rest still run, so a container never leaks all
/// its volumes because one detach failed. Persistent named volumes are detached,
/// not deleted (another container may still use them).
///
/// macOS-only for now: the sandbox and VZ-Linux runtimes both retain the full
/// `ServiceSpec` at removal time and so can detach exactly what they attached;
/// youki removal only has the container id and keeps its narrower anonymous-only
/// reclaim.
#[cfg(target_os = "macos")]
pub(crate) fn cleanup_storage_volumes(
    storage_manager: &mut StorageManager,
    container_id: &str,
    storage: &[StorageSpec],
) {
    for spec in storage {
        match spec {
            StorageSpec::Named { name, .. } => {
                if let Err(e) = storage_manager.detach_volume(name, container_id) {
                    tracing::warn!(
                        volume = %name,
                        container = %container_id,
                        error = %e,
                        "failed to detach volume"
                    );
                }
            }
            StorageSpec::S3 { bucket, prefix, .. } => {
                if let Err(e) = storage_manager.unmount_s3(bucket, prefix.as_deref(), container_id)
                {
                    tracing::warn!(
                        bucket = %bucket,
                        container = %container_id,
                        error = %e,
                        "failed to unmount S3 bucket"
                    );
                }
            }
            StorageSpec::Anonymous { .. }
            | StorageSpec::Bind { .. }
            | StorageSpec::Tmpfs { .. } => {}
        }
    }

    // Anonymous volumes are this container's alone — delete their backing dirs.
    if let Err(e) = storage_manager.cleanup_anonymous(container_id) {
        tracing::warn!(
            container = %container_id,
            error = %e,
            "failed to cleanup anonymous volumes"
        );
    }
}

/// The `volume_paths` key [`prepare_storage_volumes`] uses for an anonymous
/// volume at `target` — kept in lockstep with `bundle::build_storage_mounts`.
// Consumed only by `prepare_storage_volumes` and the macOS `resolve_host_source`;
// on `all(windows, wsl)` neither is reachable (the WSL2 delegate uses in-distro
// youki), so this key helper goes dead there. Used on macOS/Linux storage prep.
#[cfg_attr(all(target_os = "windows", feature = "wsl"), allow(dead_code))]
#[must_use]
pub(crate) fn anon_key(target: &str) -> String {
    format!("_anon_{}", target.trim_start_matches('/').replace('/', "_"))
}

/// The `volume_paths` key [`prepare_storage_volumes`] uses for an S3 mount.
// Consumed only by `prepare_storage_volumes` and the macOS `resolve_host_source`;
// on `all(windows, wsl)` neither is reachable (the WSL2 delegate uses in-distro
// youki), so this key helper goes dead there. Used on macOS/Linux storage prep.
#[cfg_attr(all(target_os = "windows", feature = "wsl"), allow(dead_code))]
#[must_use]
pub(crate) fn s3_key(bucket: &str, prefix: Option<&str>) -> String {
    format!("_s3_{}_{}", bucket, prefix.unwrap_or(""))
}

/// A single resolved mount: where it lives on the host, where it should appear
/// inside the container, and whether it is read-only.
///
/// Only the macOS runtimes build their own mounts from this (virtiofs shares for
/// VZ-Linux, rootfs rebasing for Seatbelt); youki and the WSL2 delegate consume
/// the `volume_paths` map through `bundle::build_storage_mounts` instead.
#[cfg(target_os = "macos")]
pub(crate) struct ResolvedMount {
    pub host: PathBuf,
    pub target: String,
    pub readonly: bool,
}

/// Resolve a single [`StorageSpec`] to its host source directory using the map
/// returned by [`prepare_storage_volumes`].
///
/// Returns `None` for storage kinds that have no host-directory source the
/// caller can bind/rebase (`Tmpfs`, or a volume missing from `volume_paths`).
/// Used by the virtiofs (VZ-Linux) and Seatbelt-rootfs runtimes, which build
/// their own mounts instead of OCI `bundle` mounts.
#[cfg(target_os = "macos")]
#[must_use]
pub(crate) fn resolve_host_source(
    spec: &StorageSpec,
    volume_paths: &HashMap<String, PathBuf>,
) -> Option<ResolvedMount> {
    match spec {
        StorageSpec::Bind {
            source,
            target,
            readonly,
        } => Some(ResolvedMount {
            host: PathBuf::from(source),
            target: target.clone(),
            readonly: *readonly,
        }),
        StorageSpec::Named {
            name,
            target,
            readonly,
            ..
        } => volume_paths.get(name).map(|host| ResolvedMount {
            host: host.clone(),
            target: target.clone(),
            readonly: *readonly,
        }),
        StorageSpec::Anonymous { target, .. } => {
            volume_paths
                .get(&anon_key(target))
                .map(|host| ResolvedMount {
                    host: host.clone(),
                    target: target.clone(),
                    readonly: false,
                })
        }
        StorageSpec::S3 {
            bucket,
            prefix,
            target,
            readonly,
            ..
        } => volume_paths
            .get(&s3_key(bucket, prefix.as_deref()))
            .map(|host| ResolvedMount {
                host: host.clone(),
                target: target.clone(),
                readonly: *readonly,
            }),
        StorageSpec::Tmpfs { .. } => None,
    }
}