hyperhive/hive-priv/src/storage.rs

//! Agent state storage ops: btrfs subvolumes (`EnsureAgentSubvolume`,
//! `DeleteAgentSubvolume`, `UpgradeAgentSubvolume`), qgroup accounting and
//! quotas (`EnsureBtrfsQuota`, `ReadSubvolumeUsage`, `SetSubvolumeQuota`),
//! read-only snapshots and their export (`SnapshotAgentSubvolume`,
//! `DeleteAgentSnapshot`, `SendAgentSnapshotToFile`, `SendAgentSnapshotToFd`)
//! and the `du`-based `MeasureAgentDisk`.
//!
//! Every path is derived from a validated agent name (and, for snapshots, a
//! `hive-`-prefixed label) under `AGENT_STATE_ROOT`.

use std::os::fd::OwnedFd;
use std::path::{Path, PathBuf};

use anyhow::{Context as _, Result, anyhow, bail};
use hive_priv_sock::{AGENT_STATE_ROOT, MEASURE_AGENT_DISK_TIMEOUT, MIGRATE_STAGING_ROOT};
use tokio::process::Command;

use crate::shared::{container_system_name, validate_agent_name, validate_credential_name};

/// `EnsureAgentSubvolume`.
pub(crate) async fn exec_ensure_agent_subvolume(agent_name: &str) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    ensure_agent_subvolume(agent_name).await
}

/// `DeleteAgentSubvolume`.
pub(crate) async fn exec_delete_agent_subvolume(agent_name: &str) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    delete_agent_subvolume(agent_name).await
}

/// `ReadSubvolumeUsage`.
pub(crate) async fn exec_read_subvolume_usage(agent_name: &str) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    read_subvolume_usage(agent_name).await
}

/// `MeasureAgentDisk`.
pub(crate) async fn exec_measure_agent_disk(agent_name: &str) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    measure_agent_disk(agent_name).await
}

/// `SetSubvolumeQuota`.
pub(crate) async fn exec_set_subvolume_quota(
    agent_name: &str,
    limit_bytes: Option<u64>,
) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    set_subvolume_quota(agent_name, limit_bytes).await
}

/// `UpgradeAgentSubvolume`.
pub(crate) async fn exec_upgrade_agent_subvolume(agent_name: &str) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    upgrade_agent_subvolume(agent_name).await
}

/// `SnapshotAgentSubvolume`.
pub(crate) async fn exec_snapshot_agent_subvolume(
    agent_name: &str,
    snapshot_name: &str,
) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    validate_snapshot_name(snapshot_name)?;
    snapshot_agent_subvolume(agent_name, snapshot_name).await
}

/// `DeleteAgentSnapshot`.
pub(crate) async fn exec_delete_agent_snapshot(
    agent_name: &str,
    snapshot_name: &str,
) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    validate_snapshot_name(snapshot_name)?;
    delete_agent_snapshot(agent_name, snapshot_name).await
}

/// `SendAgentSnapshotToFile`.
pub(crate) async fn exec_send_agent_snapshot_to_file(
    agent_name: &str,
    snapshot_name: &str,
    parent_snapshot_name: Option<&str>,
    dest_file_name: &str,
) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    validate_snapshot_name(snapshot_name)?;
    if let Some(parent) = parent_snapshot_name {
        validate_snapshot_name(parent)?;
    }
    validate_credential_name(dest_file_name)?;
    send_agent_snapshot_to_file(
        agent_name,
        snapshot_name,
        parent_snapshot_name,
        dest_file_name,
    )
    .await
}

/// `SendAgentSnapshotToFd`.
pub(crate) async fn exec_send_agent_snapshot_to_fd(
    agent_name: &str,
    snapshot_name: &str,
    parent_snapshot_name: Option<&str>,
    fd: Option<OwnedFd>,
) -> Result<(String, String)> {
    validate_agent_name(agent_name)?;
    validate_snapshot_name(snapshot_name)?;
    if let Some(parent) = parent_snapshot_name {
        validate_snapshot_name(parent)?;
    }
    let dest = fd.context("no descriptor to stream into")?;
    send_agent_snapshot_to_fd(agent_name, snapshot_name, parent_snapshot_name, dest).await
}

/// A btrfs snapshot label must start with `hive-` — this doubles as an
/// allow-list: only names hivectl itself constructs (or an operator who
/// knows the convention) can reach the `btrfs subvolume snapshot`/`delete`
/// shellouts, so an arbitrary caller can't use the snapshot ops to probe or
/// churn unrelated paths under `AGENT_STATE_ROOT`. Beyond the prefix, the
/// same charset restriction as [`validate_credential_name`] applies (it's
/// interpolated straight into a filesystem path).
fn validate_snapshot_name(name: &str) -> Result<()> {
    if !name.starts_with("hive-") {
        bail!("invalid snapshot label {name:?}: must start with \"hive-\"");
    }
    validate_credential_name(name)
}

/// btrfs superblock magic, as reported by `statfs(2)`'s `f_type`.
const BTRFS_SUPER_MAGIC: i64 = 0x9123_683E;

/// Inode number of a btrfs subvolume root (`BTRFS_FIRST_FREE_OBJECTID`).
/// Every subvolume's top directory has this inode; plain directories do
/// not, so `statfs == btrfs && st_ino == 256` reliably identifies a
/// subvolume root.
const BTRFS_SUBVOL_ROOT_INO: u64 = 256;

/// Whether `path` lives on a btrfs filesystem (via `statfs(2)`).
fn is_on_btrfs(path: &Path) -> Result<bool> {
    use std::os::unix::ffi::OsStrExt as _;
    let c_path = std::ffi::CString::new(path.as_os_str().as_bytes())
        .with_context(|| format!("path {} has an interior null byte", path.display()))?;
    // SAFETY: `c_path` is a valid NUL-terminated C string that outlives the
    // call; `statfs` only writes into the zero-initialised `buf`.
    let mut buf: libc::statfs = unsafe { std::mem::zeroed() };
    let rc = unsafe { libc::statfs(c_path.as_ptr(), &raw mut buf) };
    if rc != 0 {
        return Err(std::io::Error::last_os_error())
            .with_context(|| format!("statfs {}", path.display()));
    }
    Ok(buf.f_type == BTRFS_SUPER_MAGIC)
}

/// Whether `path` is the root of a btrfs subvolume (on btrfs and inode 256).
fn is_btrfs_subvolume(path: &Path) -> bool {
    use std::os::unix::fs::MetadataExt as _;
    let on_btrfs = is_on_btrfs(path).unwrap_or(false);
    let ino_match = std::fs::metadata(path).is_ok_and(|m| m.ino() == BTRFS_SUBVOL_ROOT_INO);
    on_btrfs && ino_match
}

/// `EnsureAgentSubvolume` — make the agent's state root a btrfs subvolume
/// when the FS supports it. Idempotent + progressive: no-op when the root
/// already exists or the FS isn't btrfs. See the wire doc on the variant.
async fn ensure_agent_subvolume(agent_name: &str) -> Result<(String, String)> {
    use std::os::unix::fs::MetadataExt as _;

    let root = PathBuf::from(AGENT_STATE_ROOT);
    let agent_root = root.join(agent_name);

    // Progressive: existing agents (plain dir OR already a subvol) are left
    // untouched — never auto-migrated.
    if agent_root.exists() {
        return Ok((String::new(), String::new()));
    }

    // Only btrfs supports subvolumes; on anything else hive-c0re's normal
    // `create_dir_all` makes a plain dir (the pre-subvolume behaviour). The
    // parent must exist for both statfs and `btrfs subvolume create`.
    std::fs::create_dir_all(&root)
        .with_context(|| format!("create agents root {}", root.display()))?;
    if !is_on_btrfs(&root)? {
        return Ok((String::new(), String::new()));
    }

    let out = Command::new("btrfs")
        .args(["subvolume", "create"])
        .arg(&agent_root)
        .output()
        .await
        .with_context(|| format!("spawn btrfs subvolume create {}", agent_root.display()))?;
    if !out.status.success() {
        bail!(
            "btrfs subvolume create {} failed: {}",
            agent_root.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }

    // The subvol root is created root-owned, but hive-c0re (the `hive-core`
    // user) must be able to mkdir state/ claude/ harness/ inside it — exactly
    // as it would in a plain dir. Chown it to AGENT_STATE_ROOT's owner
    // (hive-core). This MUST succeed: a root-owned subvol would make the
    // downstream dir creation fail with a confusing permission error, and the
    // c0re-side exists-check would then skip re-running this op on retry,
    // wedging the agent. So on any failure roll the subvol back and bail — the
    // create path surfaces a clear error and a retry starts clean.
    let chown_result = std::fs::metadata(&root)
        .with_context(|| format!("stat agents root {} for ownership", root.display()))
        .and_then(|meta| {
            std::os::unix::fs::chown(&agent_root, Some(meta.uid()), Some(meta.gid())).with_context(
                || format!("chown subvol {} to agents-root owner", agent_root.display()),
            )
        });
    if let Err(e) = chown_result {
        // Best-effort rollback so we never leave a root-owned subvol behind.
        let _ = Command::new("btrfs")
            .args(["subvolume", "delete"])
            .arg(&agent_root)
            .output()
            .await;
        return Err(e.context(format!(
            "rolled back subvolume {} after chown failed",
            agent_root.display()
        )));
    }
    tracing::info!(agent = %agent_name, path = %agent_root.display(), "created agent state subvolume");
    Ok((String::new(), String::new()))
}

/// `DeleteAgentSubvolume` — delete the agent's state root iff it's a btrfs
/// subvolume (purge path only). No-op for plain dirs / missing paths;
/// hive-c0re's own `remove_dir_all` covers those. See the wire doc.
async fn delete_agent_subvolume(agent_name: &str) -> Result<(String, String)> {
    let agent_root = PathBuf::from(AGENT_STATE_ROOT).join(agent_name);
    if !is_btrfs_subvolume(&agent_root) {
        return Ok((String::new(), String::new()));
    }
    let out = Command::new("btrfs")
        .args(["subvolume", "delete"])
        .arg(&agent_root)
        .output()
        .await
        .with_context(|| format!("spawn btrfs subvolume delete {}", agent_root.display()))?;
    if !out.status.success() {
        bail!(
            "btrfs subvolume delete {} failed: {}",
            agent_root.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    tracing::info!(agent = %agent_name, path = %agent_root.display(), "deleted agent state subvolume");
    Ok((String::new(), String::new()))
}

/// `EnsureBtrfsQuota` — enable btrfs qgroup accounting on the filesystem
/// holding `AGENT_STATE_ROOT`. Idempotent + statfs-gated (no-op off btrfs).
/// Operator opt-in only; see the wire doc.
pub(crate) async fn ensure_btrfs_quota() -> Result<(String, String)> {
    let root = PathBuf::from(AGENT_STATE_ROOT);
    std::fs::create_dir_all(&root)
        .with_context(|| format!("create agents root {}", root.display()))?;
    if !is_on_btrfs(&root)? {
        // Non-btrfs host: quota/qgroups don't apply. No-op success so the
        // operator-facing verb degrades cleanly.
        return Ok((
            String::new(),
            "filesystem is not btrfs — quota not applicable".to_owned(),
        ));
    }
    let out = Command::new("btrfs")
        .args(["quota", "enable"])
        .arg(&root)
        .output()
        .await
        .with_context(|| format!("spawn btrfs quota enable {}", root.display()))?;
    if !out.status.success() {
        bail!(
            "btrfs quota enable {} failed: {}",
            root.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    tracing::info!(path = %root.display(), "enabled btrfs qgroup accounting");
    Ok((String::new(), String::new()))
}

/// `ReadSubvolumeUsage` — return an agent subvolume's qgroup rows
/// (`btrfs qgroup show -f --raw <…/agent_name>`) verbatim in stdout for
/// hive-c0re to parse. `-f` lists the qgroups impacting the given path,
/// excluding ancestral qgroups (per btrfs-qgroup-show(8)) — so it scopes
/// to this subvolume and never mixes in other agents'. hive-c0re then
/// selects the level-0 (`0/<subvolid>`) leaf row. See the wire doc.
async fn read_subvolume_usage(agent_name: &str) -> Result<(String, String)> {
    let agent_root = PathBuf::from(AGENT_STATE_ROOT).join(agent_name);
    let out = Command::new("btrfs")
        .args(["qgroup", "show", "-f", "--raw"])
        .arg(&agent_root)
        .output()
        .await
        .with_context(|| format!("spawn btrfs qgroup show {}", agent_root.display()))?;
    if !out.status.success() {
        // The common failure is "quota not enabled" — pass the stderr
        // through so hive-c0re can surface it gracefully.
        bail!(
            "btrfs qgroup show {} failed: {}",
            agent_root.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    Ok((
        String::from_utf8_lossy(&out.stdout).into_owned(),
        String::new(),
    ))
}

/// The host paths `MeasureAgentDisk` sums for `agent_name`: its `state/`
/// and `work/` dirs under `AGENT_STATE_ROOT` and its container rootfs.
/// These are the only paths that op ever walks.
// `state` / `work` must match hive-c0re's `agent_notes_dir` and
// `paths::WORK_SUBDIR`.
fn agent_disk_paths(agent_name: &str) -> [PathBuf; 3] {
    let agent_root = PathBuf::from(AGENT_STATE_ROOT).join(agent_name);
    [
        agent_root.join("state"),
        agent_root.join("work"),
        PathBuf::from("/var/lib/nixos-containers").join(container_system_name(agent_name)),
    ]
}

/// `MeasureAgentDisk` — one `du` over the [`agent_disk_paths`] that exist,
/// returning the grand total in bytes. An agent with none of them has no
/// measurable size, which is an error rather than a total of 0.
async fn measure_agent_disk(agent_name: &str) -> Result<(String, String)> {
    let mut paths = Vec::new();
    for path in agent_disk_paths(agent_name) {
        match std::fs::symlink_metadata(&path) {
            Ok(_) => paths.push(path),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
            Err(e) => return Err(e).with_context(|| format!("stat {}", path.display())),
        }
    }
    let total = du_total(agent_name, du_command(&paths)?, MEASURE_AGENT_DISK_TIMEOUT).await?;
    Ok((total.to_string(), String::new()))
}

/// `du -s -x -b -c -P -- <paths>`. `-x` keeps each walk on its operand's
/// filesystem (so the bind-mounted nix store is never counted), `-P` never
/// follows a symlink, `-b` counts apparent size, `-c` adds the `total` line
/// [`parse_du_total`] reads. Errors on no paths: `du` with no operand would
/// measure its working directory instead.
fn du_command(paths: &[PathBuf]) -> Result<Command> {
    if paths.is_empty() {
        bail!("none of the agent's state, work or rootfs paths exist");
    }
    let mut cmd = Command::new("du");
    cmd.args(["-s", "-x", "-b", "-c", "-P", "--"]).args(paths);
    Ok(cmd)
}

/// Run `cmd` (a [`du_command`]) for at most `limit` and parse its total.
/// The child is killed on timeout. A timeout, a spawn failure, a non-zero
/// exit or output without a total is an error, never a partial total.
async fn du_total(agent_name: &str, mut cmd: Command, limit: std::time::Duration) -> Result<u64> {
    cmd.kill_on_drop(true);
    let Ok(out) = tokio::time::timeout(limit, cmd.output()).await else {
        bail!(
            "du for {agent_name} timed out after {}s; killed it",
            limit.as_secs()
        );
    };
    let out = out.with_context(|| format!("spawn du for {agent_name}"))?;
    if !out.status.success() {
        bail!(
            "du for {agent_name} failed ({}): {}",
            out.status,
            String::from_utf8_lossy(&out.stderr)
                .lines()
                .last()
                .unwrap_or("")
                .trim()
        );
    }
    let stdout = String::from_utf8_lossy(&out.stdout);
    parse_du_total(&stdout)
        .ok_or_else(|| anyhow!("no `total` line in du output for {agent_name}: {stdout:?}"))
}

/// The byte count on `du -c`'s final `<bytes>\ttotal` line. `None` when the
/// last line isn't that shape, so truncated or unexpected output is never
/// read as a smaller total.
fn parse_du_total(stdout: &str) -> Option<u64> {
    let last = stdout.lines().rev().find(|l| !l.trim().is_empty())?;
    let (bytes, label) = last.split_once('\t')?;
    if label != "total" {
        return None;
    }
    bytes.parse().ok()
}

/// Best-effort removal of a leftover migration path from a prior aborted
/// upgrade: try `btrfs subvolume delete` (in case it's a half-created
/// subvolume) then a plain recursive remove. Both failures are ignored —
/// the path may simply not exist.
async fn cleanup_stale_path(path: &Path) {
    if path.exists() {
        let _ = Command::new("btrfs")
            .args(["subvolume", "delete"])
            .arg(path)
            .output()
            .await;
        let _ = std::fs::remove_dir_all(path);
    }
}

/// Stage a populated subvolume at `tmp` mirroring `agent_root`: create the
/// subvolume, copy `agent_root`'s contents into it preserving
/// ownership/permissions/xattrs, then match the subvolume root's owner + mode
/// to the original. On any failure the partially-staged `tmp` is cleaned up
/// (so the caller can bail with the original dir still untouched).
async fn stage_upgrade_subvolume(agent_root: &Path, tmp: &Path) -> Result<()> {
    use std::os::unix::fs::{MetadataExt as _, PermissionsExt as _};

    // Fresh subvolume to receive the copy.
    let out = Command::new("btrfs")
        .args(["subvolume", "create"])
        .arg(tmp)
        .output()
        .await
        .with_context(|| format!("spawn btrfs subvolume create {}", tmp.display()))?;
    if !out.status.success() {
        bail!(
            "btrfs subvolume create {} failed: {}",
            tmp.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }

    // Copy contents preserving everything (`-a` = --preserve=all → mode,
    // ownership, timestamps, links, xattrs); reflink for fast CoW clones on the
    // same btrfs. `<src>/.` copies the directory's contents (incl. dotfiles)
    // into the subvolume rather than nesting it.
    let copy = Command::new("cp")
        .arg("-a")
        .arg("--reflink=auto")
        .arg(format!("{}/.", agent_root.display()))
        .arg(tmp)
        .output()
        .await
        .with_context(|| format!("spawn cp into {}", tmp.display()))?;
    if !copy.status.success() {
        cleanup_stale_path(tmp).await;
        bail!(
            "copy {} -> {} failed (original left untouched): {}",
            agent_root.display(),
            tmp.display(),
            String::from_utf8_lossy(&copy.stderr).trim()
        );
    }

    // Match the new subvolume root's ownership + mode to the original dir.
    // `cp -a <src>/.` copies the *contents* but the subvolume root keeps its
    // create-time root ownership, so set it explicitly — the swapped-in
    // subvolume must be indistinguishable from the original to hive-c0re.
    let apply = std::fs::metadata(agent_root)
        .with_context(|| format!("stat {} for ownership", agent_root.display()))
        .and_then(|m| {
            std::os::unix::fs::chown(tmp, Some(m.uid()), Some(m.gid()))
                .with_context(|| format!("chown {} to match original", tmp.display()))?;
            std::fs::set_permissions(tmp, std::fs::Permissions::from_mode(m.mode()))
                .with_context(|| format!("chmod {} to match original", tmp.display()))?;
            Ok(())
        });
    if let Err(e) = apply {
        cleanup_stale_path(tmp).await;
        return Err(e.context("upgrade aborted before swap; original left untouched"));
    }
    Ok(())
}

/// `UpgradeAgentSubvolume` — convert an existing plain-dir agent state root
/// into a btrfs subvolume in place. Operator opt-in; the caller (hivectl)
/// stops the agent first and restarts it after. See the wire doc.
///
/// Migration: stage a sibling subvolume mirroring the dir
/// ([`stage_upgrade_subvolume`]), then rename the original aside and the
/// subvolume into place, then remove the original. Any failure before the
/// rename-swap leaves the original dir untouched.
async fn upgrade_agent_subvolume(agent_name: &str) -> Result<(String, String)> {
    let root = PathBuf::from(AGENT_STATE_ROOT);
    let agent_root = root.join(agent_name);

    if !agent_root.exists() {
        // A crash between the two swap renames (original → `.<name>.old`
        // succeeded, `.<name>.migrating` → agent_root did not) leaves the
        // agent root missing but the original data intact under `.<name>.old`.
        // Point at the recovery rather than a bare "nothing to do" so the
        // operator isn't left guessing where the data went.
        let old = root.join(format!(".{agent_name}.old"));
        if old.exists() {
            bail!(
                "no state dir at {agent} — but {old} holds the original data from an \
                 interrupted upgrade (host crashed mid-swap). Restore it with \
                 `mv {old} {agent}`, then re-run the upgrade.",
                agent = agent_root.display(),
                old = old.display(),
            );
        }
        bail!(
            "no state dir to upgrade at {} — nothing to do",
            agent_root.display()
        );
    }
    // Idempotent: already a subvolume → nothing to do.
    if is_btrfs_subvolume(&agent_root) {
        return Ok((String::new(), String::new()));
    }
    if !is_on_btrfs(&root)? {
        bail!(
            "{} is not on btrfs — subvolumes are unsupported, cannot upgrade",
            root.display()
        );
    }

    // Sibling temp paths on the same filesystem (so the copy can reflink and
    // the swap renames are atomic). Leading dots keep them out of the agent
    // namespace (`validate_agent_name` rejects dot-prefixed names).
    let tmp = root.join(format!(".{agent_name}.migrating"));
    let old = root.join(format!(".{agent_name}.old"));
    // Clear any debris from a previously interrupted run before starting.
    cleanup_stale_path(&tmp).await;
    cleanup_stale_path(&old).await;

    stage_upgrade_subvolume(&agent_root, &tmp).await?;

    // Swap. `rename` is atomic within a filesystem. The window between the
    //    two renames is the only unsafe point: a crash there leaves the agent
    //    root missing but both `.old` (original) and the new subvolume present
    //    — recoverable by hand, hence the loud logging.
    if let Err(e) = std::fs::rename(&agent_root, &old) {
        cleanup_stale_path(&tmp).await;
        return Err(anyhow::Error::new(e).context(format!(
            "rename {} -> {} failed; original left untouched",
            agent_root.display(),
            old.display()
        )));
    }
    if let Err(e) = std::fs::rename(&tmp, &agent_root) {
        // Restore the original from its renamed-aside copy.
        let restored = std::fs::rename(&old, &agent_root).is_ok();
        cleanup_stale_path(&tmp).await;
        return Err(anyhow::Error::new(e).context(format!(
            "rename {} -> {} failed; original {}",
            tmp.display(),
            agent_root.display(),
            if restored {
                "restored"
            } else {
                "COULD NOT BE RESTORED — manual recovery needed"
            }
        )));
    }

    // 5. Success: drop the original (a plain dir) and report.
    if let Err(e) = std::fs::remove_dir_all(&old) {
        // The migration succeeded; a leftover `.old` is cosmetic. Warn only.
        tracing::warn!(
            agent = %agent_name, path = %old.display(),
            "upgraded subvolume but failed to remove old dir: {e}"
        );
    }
    tracing::info!(
        agent = %agent_name, path = %agent_root.display(),
        "upgraded agent state dir to btrfs subvolume"
    );
    Ok((
        format!("upgraded {} to a btrfs subvolume", agent_root.display()),
        String::new(),
    ))
}

/// Derive a snapshot's path from the agent name + label: a dot-prefixed
/// sibling of the agent's state root so it can never collide with a real
/// agent directory (`validate_agent_name` rejects dot-prefixed names).
fn snapshot_path(agent_name: &str, snapshot_name: &str) -> PathBuf {
    PathBuf::from(AGENT_STATE_ROOT).join(format!(".{agent_name}.snapshot.{snapshot_name}"))
}

/// `SnapshotAgentSubvolume` — create a read-only btrfs snapshot of an
/// agent's state subvolume, for `btrfs send` to stream from during
/// inter-hive migration. See the wire doc.
async fn snapshot_agent_subvolume(
    agent_name: &str,
    snapshot_name: &str,
) -> Result<(String, String)> {
    let agent_root = PathBuf::from(AGENT_STATE_ROOT).join(agent_name);
    if !is_btrfs_subvolume(&agent_root) {
        bail!(
            "{} is not a btrfs subvolume — nothing to snapshot (run `hivectl agent <name> subvol upgrade` first)",
            agent_root.display()
        );
    }
    let snap = snapshot_path(agent_name, snapshot_name);
    if snap.exists() {
        bail!(
            "snapshot {} already exists — delete it first or pick a different name",
            snap.display()
        );
    }
    let out = Command::new("btrfs")
        .args(["subvolume", "snapshot", "-r"])
        .arg(&agent_root)
        .arg(&snap)
        .output()
        .await
        .with_context(|| {
            format!(
                "spawn btrfs subvolume snapshot -r {} {}",
                agent_root.display(),
                snap.display()
            )
        })?;
    if !out.status.success() {
        bail!(
            "btrfs subvolume snapshot -r {} {} failed: {}",
            agent_root.display(),
            snap.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    tracing::info!(
        agent = %agent_name, snapshot = %snap.display(),
        "created read-only agent state snapshot"
    );
    Ok((snap.display().to_string(), String::new()))
}

/// `DeleteAgentSnapshot` — delete a previously-created read-only snapshot.
/// No-op if the path doesn't exist. See the wire doc.
async fn delete_agent_snapshot(agent_name: &str, snapshot_name: &str) -> Result<(String, String)> {
    let snap = snapshot_path(agent_name, snapshot_name);
    if !snap.exists() {
        return Ok((String::new(), String::new()));
    }
    let out = Command::new("btrfs")
        .args(["subvolume", "delete"])
        .arg(&snap)
        .output()
        .await
        .with_context(|| format!("spawn btrfs subvolume delete {}", snap.display()))?;
    if !out.status.success() {
        bail!(
            "btrfs subvolume delete {} failed: {}",
            snap.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    tracing::info!(agent = %agent_name, snapshot = %snap.display(), "deleted agent state snapshot");
    Ok((String::new(), String::new()))
}

/// Both paths a `SendAgentSnapshotToFile` request reads from — the
/// snapshot itself, and its optional incremental parent — must exist
/// before we touch `dest`. Pulled out on its own so that invariant is
/// checkable without a `dest` path, a `MIGRATE_STAGING_ROOT`, or `btrfs`
/// at all: a caller cannot create the output file by accident while
/// still holding a reference to this check.
fn check_send_snapshot_preconditions(snap: &Path, parent_path: Option<&Path>) -> Result<()> {
    if !snap.exists() {
        bail!(
            "snapshot {} does not exist — create it with `subvol snapshot create` first",
            snap.display()
        );
    }
    if let Some(parent_path) = parent_path
        && !parent_path.exists()
    {
        bail!(
            "parent snapshot {} does not exist — pick an existing parent or omit it for a full send",
            parent_path.display()
        );
    }
    Ok(())
}

/// Validate preconditions, then atomically create `dest` for a snapshot
/// export. The two are combined in one function on purpose: keeping
/// `check_send_snapshot_preconditions` ahead of the `create_new` call is
/// the entire fix, and folding them together here means that ordering
/// can't drift back apart in a later edit to `send_agent_snapshot_to_file`.
/// No `btrfs` involved, so this is unit-testable against plain files.
fn open_export_dest(snap: &Path, parent_path: Option<&Path>, dest: &Path) -> Result<std::fs::File> {
    check_send_snapshot_preconditions(snap, parent_path)?;
    // `create_new` (O_CREAT|O_EXCL) makes the no-overwrite guarantee atomic
    // instead of a check-then-create race against a concurrent request.
    match std::fs::File::options()
        .write(true)
        .create_new(true)
        .open(dest)
    {
        Ok(f) => Ok(f),
        Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => bail!(
            "{} already exists — pick a different destination or remove it first \
             (send never overwrites an existing export)",
            dest.display()
        ),
        Err(e) => Err(e).with_context(|| format!("create {}", dest.display())),
    }
}

/// Deletes `path` when dropped, unless [`disarm`](Self::disarm)ed first.
/// Guards `dest` between the moment `send_agent_snapshot_to_file` creates
/// it and the moment the export fully succeeds, so any `?`/`bail!` on
/// that path — including ones a future edit adds — removes the partial
/// file instead of leaving a zero-byte export that then makes every
/// retry against the same `--dest` fail with "already exists".
struct PartialExportGuard<'a> {
    path: &'a Path,
    armed: bool,
}

impl<'a> PartialExportGuard<'a> {
    fn new(path: &'a Path) -> Self {
        Self { path, armed: true }
    }

    fn disarm(mut self) {
        self.armed = false;
    }
}

impl Drop for PartialExportGuard<'_> {
    fn drop(&mut self) {
        if !self.armed {
            return;
        }
        // Best-effort: warn (don't fail the whole call over it) if removal
        // itself fails, so a stuck partial file that later masquerades as a
        // completed export is at least visible in the log.
        if let Err(rm_err) = std::fs::remove_file(self.path) {
            tracing::warn!(
                dest = %self.path.display(), error = %rm_err,
                "failed to remove partial export after failure — \
                 next attempt at this dest will hit the already-exists guard"
            );
        }
    }
}

/// `SendAgentSnapshotToFile` — stream a read-only snapshot (optionally
/// incremental against `parent_name`) to a file under
/// `MIGRATE_STAGING_ROOT` via `btrfs send`. Local-file half of the
/// inter-hive migration transport; see `PrivRequest::SendAgentSnapshotToFile`
/// for the cross-hive follow-up.
async fn send_agent_snapshot_to_file(
    agent_name: &str,
    snapshot_name: &str,
    parent_name: Option<&str>,
    dest_file_name: &str,
) -> Result<(String, String)> {
    let snap = snapshot_path(agent_name, snapshot_name);
    let parent_path = parent_name.map(|parent| snapshot_path(agent_name, parent));

    std::fs::create_dir_all(MIGRATE_STAGING_ROOT)
        .with_context(|| format!("create {MIGRATE_STAGING_ROOT}"))?;
    let dest = Path::new(MIGRATE_STAGING_ROOT).join(dest_file_name);
    // Preconditions are checked *inside* this call, before `dest` is
    // created — see `open_export_dest`'s doc comment for why that ordering
    // matters and is kept in one place.
    let dest_file = open_export_dest(&snap, parent_path.as_deref(), &dest)?;
    let cleanup = PartialExportGuard::new(&dest);

    let mut cmd = Command::new("btrfs");
    cmd.arg("send");
    if let Some(parent_path) = &parent_path {
        cmd.arg("-p").arg(parent_path);
    }
    cmd.arg(&snap);
    cmd.stdout(std::process::Stdio::from(dest_file));
    cmd.stderr(std::process::Stdio::piped());

    let out = cmd
        .spawn()
        .with_context(|| format!("spawn btrfs send {}", snap.display()))?
        .wait_with_output()
        .await
        .with_context(|| format!("wait on btrfs send {}", snap.display()))?;
    if !out.status.success() {
        bail!(
            "btrfs send {} failed: {}",
            snap.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    cleanup.disarm();
    tracing::info!(
        agent = %agent_name, snapshot = %snap.display(), dest = %dest.display(),
        parent = ?parent_name, "exported agent snapshot to file"
    );
    Ok((dest.display().to_string(), String::new()))
}

/// `SendAgentSnapshotToFd` — stream a read-only snapshot (optionally
/// incremental against `parent_name`) straight into a descriptor the
/// caller passed us.
///
/// The network half of the inter-hive migration transport, arranged so
/// this helper never learns there *is* a network: hive-c0re connects to
/// the peer's snapshot store, writes the header itself, and hands the
/// connected socket over. We only ever see "a thing to write bytes into",
/// which keeps a root process out of any address, protocol or trust
/// decision — and keeps everyone out of the data path once `btrfs send`
/// starts, which matters at multi-gigabyte sizes.
async fn send_agent_snapshot_to_fd(
    agent_name: &str,
    snapshot_name: &str,
    parent_name: Option<&str>,
    dest: OwnedFd,
) -> Result<(String, String)> {
    let snap = snapshot_path(agent_name, snapshot_name);
    if !snap.exists() {
        bail!(
            "snapshot {} does not exist — create it with `subvol snapshot create` first",
            snap.display()
        );
    }

    let mut cmd = Command::new("btrfs");
    cmd.arg("send");
    if let Some(parent) = parent_name {
        let parent_path = snapshot_path(agent_name, parent);
        if !parent_path.exists() {
            bail!(
                "parent snapshot {} does not exist — pick an existing parent or omit it for a full send",
                parent_path.display()
            );
        }
        cmd.arg("-p").arg(&parent_path);
    }
    cmd.arg(&snap);
    cmd.stdout(std::process::Stdio::from(dest));
    cmd.stderr(std::process::Stdio::piped());

    let out = cmd
        .spawn()
        .with_context(|| format!("spawn btrfs send {}", snap.display()))?
        .wait_with_output()
        .await
        .with_context(|| format!("wait on btrfs send {}", snap.display()))?;
    if !out.status.success() {
        // Nothing to clean up: the destination isn't ours. A partial
        // stream is the receiving end's problem, and `btrfs receive`
        // refuses to commit an incomplete subvolume anyway.
        bail!(
            "btrfs send {} failed: {}",
            snap.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    tracing::info!(
        agent = %agent_name, snapshot = %snap.display(), parent = ?parent_name,
        "streamed agent snapshot into a passed descriptor"
    );
    Ok((String::new(), String::new()))
}

/// `SetSubvolumeQuota` — set or clear a qgroup size limit on an agent
/// subvolume (`btrfs qgroup limit <bytes|none> <…/agent_name>`). See the
/// wire doc.
async fn set_subvolume_quota(
    agent_name: &str,
    limit_bytes: Option<u64>,
) -> Result<(String, String)> {
    let agent_root = PathBuf::from(AGENT_STATE_ROOT).join(agent_name);
    let limit = limit_bytes.map_or_else(|| "none".to_owned(), |n| n.to_string());
    let out = Command::new("btrfs")
        .args(["qgroup", "limit", &limit])
        .arg(&agent_root)
        .output()
        .await
        .with_context(|| format!("spawn btrfs qgroup limit {}", agent_root.display()))?;
    if !out.status.success() {
        bail!(
            "btrfs qgroup limit {limit} {} failed: {}",
            agent_root.display(),
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    tracing::info!(agent = %agent_name, %limit, "set agent subvolume quota");
    Ok((String::new(), String::new()))
}
#[cfg(test)]
mod tests {
    use super::{
        Command, agent_disk_paths, du_command, du_total, exec_measure_agent_disk, open_export_dest,
        parse_du_total, validate_snapshot_name,
    };
    use std::path::PathBuf;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::time::Duration;

    /// A snapshot label must additionally carry the `hive-` prefix (it
    /// doubles as the allow-list gating the btrfs snapshot/delete
    /// shellouts) on top of [`validate_credential_name`]'s charset rule.
    #[test]
    fn validate_snapshot_name_requires_hive_prefix() {
        assert!(validate_snapshot_name("not-hive-prefixed").is_err());
        assert!(validate_snapshot_name("hive-valid").is_ok());
    }

    /// Unique scratch dir per test, no external tempfile dep. Same pattern
    /// as `shared::test_support::scratch()`, kept separate so a rename of one doesn't
    /// collide with the other's directory names under parallel test runs.
    fn snapshot_export_scratch() -> PathBuf {
        static CTR: AtomicU32 = AtomicU32::new(0);
        let n = CTR.fetch_add(1, Ordering::Relaxed);
        let dir = std::env::temp_dir().join(format!(
            "hive-priv-snapshot-export-test-{}-{n}",
            std::process::id()
        ));
        std::fs::create_dir_all(&dir).unwrap();
        dir
    }

    /// The fix: a missing parent must be caught before `dest` is created,
    /// not after. Fails against the pre-fix ordering (`dest` created via
    /// `create_new` first, parent checked second) — see the issue for the
    /// inverted run.
    #[test]
    fn missing_parent_leaves_no_file_at_dest() {
        let dir = snapshot_export_scratch();
        let snap = dir.join("snap");
        std::fs::write(&snap, b"").unwrap();
        let parent = dir.join("parent-that-does-not-exist");
        let dest = dir.join("dest");

        let err = open_export_dest(&snap, Some(&parent), &dest).unwrap_err();
        assert!(
            err.to_string().contains("does not exist"),
            "wrong error: {err}"
        );
        assert!(
            !dest.exists(),
            "a failed export must not leave a file at dest"
        );
    }

    /// A retry against the same `dest` after a precondition failure must
    /// not trip the no-overwrite guard — there is nothing there to
    /// overwrite, since the failed attempt never created `dest` (the test
    /// above). This is the behavior a "stale parent falls back to a full
    /// send, retried without -p" caller needs against a fixed `--dest`.
    #[test]
    fn retry_after_precondition_failure_succeeds() {
        let dir = snapshot_export_scratch();
        let snap = dir.join("snap");
        std::fs::write(&snap, b"").unwrap();
        let parent = dir.join("parent-that-does-not-exist");
        let dest = dir.join("dest");

        assert!(open_export_dest(&snap, Some(&parent), &dest).is_err());
        assert!(!dest.exists());

        // Retry without the missing parent, same dest.
        let file = open_export_dest(&snap, None, &dest);
        assert!(
            file.is_ok(),
            "retry against a dest no previous attempt touched must succeed: {:?}",
            file.err()
        );
        assert!(dest.exists());
    }

    /// Control for the two tests above: an export that actually landed at
    /// `dest` must still refuse a second one. Without this, a bug that
    /// simply stopped creating `dest` at all (rather than fixing the
    /// ordering) would also make the retry test above pass.
    #[test]
    fn existing_dest_still_refuses_overwrite() {
        let dir = snapshot_export_scratch();
        let snap = dir.join("snap");
        std::fs::write(&snap, b"").unwrap();
        let dest = dir.join("dest");
        std::fs::write(&dest, b"already here").unwrap();

        let err = open_export_dest(&snap, None, &dest).unwrap_err();
        assert!(
            err.to_string().contains("already exists"),
            "wrong error: {err}"
        );
        assert_eq!(
            std::fs::read(&dest).unwrap(),
            b"already here",
            "must not have touched the existing file"
        );
    }

    #[test]
    fn disk_paths_are_derived_from_the_name_only() {
        assert_eq!(
            agent_disk_paths("iris"),
            [
                PathBuf::from("/var/lib/hyperhive/agents/iris/state"),
                PathBuf::from("/var/lib/hyperhive/agents/iris/work"),
                PathBuf::from("/var/lib/nixos-containers/h-iris"),
            ]
        );
    }

    /// Every name that could steer the walk out of the agent's own paths is
    /// refused before any path is built, let alone `du` spawned.
    #[tokio::test]
    async fn measure_rejects_names_that_are_not_agent_names() {
        for bad in [
            "", "..", "../etc", "a/b", "/etc", "Iris", "a b", "a\0b", "foo_bar",
        ] {
            let err = exec_measure_agent_disk(bad)
                .await
                .expect_err(&format!("{bad:?} must be rejected"));
            assert!(
                err.to_string().contains("invalid name"),
                "{bad:?}: wrong error: {err}"
            );
        }
    }

    #[test]
    fn du_total_reads_the_total_line() {
        let out = "100\t/var/lib/hyperhive/agents/iris/state\n\
                   2000\t/var/lib/nixos-containers/h-iris\n\
                   2100\ttotal\n";
        assert_eq!(parse_du_total(out), Some(2100));
        assert_eq!(parse_du_total("0\ttotal"), Some(0));
    }

    #[test]
    fn du_total_rejects_output_without_a_total_line() {
        assert_eq!(parse_du_total(""), None);
        assert_eq!(parse_du_total("\n"), None);
        // A per-path line is not the total, even when it is the last one.
        assert_eq!(parse_du_total("100\t/var/lib/x\n"), None);
        assert_eq!(parse_du_total("abc\ttotal\n"), None);
        assert_eq!(parse_du_total("-1\ttotal\n"), None);
        assert_eq!(parse_du_total("100 total\n"), None);
    }

    #[test]
    fn no_existing_paths_is_an_error_not_zero() {
        let err = du_command(&[]).expect_err("no operands must not run du");
        assert!(err.to_string().contains("none of"), "{err}");
    }

    #[tokio::test]
    async fn non_zero_du_exit_is_an_error() {
        let err = du_total("iris", Command::new("false"), Duration::from_secs(5))
            .await
            .expect_err("a failed du must not yield a total");
        assert!(err.to_string().contains("failed"), "{err}");
    }

    #[tokio::test]
    async fn a_walk_past_the_limit_is_killed_and_errors() {
        let mut cmd = Command::new("sleep");
        cmd.arg("30");
        let started = std::time::Instant::now();
        let err = du_total("iris", cmd, Duration::from_millis(50))
            .await
            .expect_err("a walk past the limit must fail");
        assert!(err.to_string().contains("timed out"), "{err}");
        assert!(started.elapsed() < Duration::from_secs(10));
    }

    #[tokio::test]
    async fn du_total_sums_a_real_walk() {
        let dir = std::env::temp_dir().join(format!("hive-priv-du-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        std::fs::write(dir.join("f"), b"hello").unwrap();
        let cmd = du_command(std::slice::from_ref(&dir)).unwrap();
        let total = du_total("iris", cmd, Duration::from_secs(30)).await;
        std::fs::remove_dir_all(&dir).unwrap();
        assert!(total.unwrap() >= 5);
    }
}