lazaret/crates/lazaret-linux/tests/acquisition.rs

use std::ffi::OsString;
use std::fs::File;
use std::io::{Seek, Write};
use std::os::fd::AsFd;
use std::process::{ExitCode, Stdio};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Barrier};

use lazaret_core::identity::{
    CoordinatorId, DomainId, GenerationId, HostPidLocator, PidfsProcessIdentityV1,
};
use lazaret_core::records::{
    HostProcessOwnershipSlot, ProcessOwnershipCommitAck, StagedProcessOwnership,
};
use lazaret_linux::process::{
    ChildStandardIo, ProcessOwnershipContext, TargetCommand, append_parent_identity, child_broker,
    current_process_identity, dispatch_private_process_role, verify_direct_child_pidfd,
};
use rustix::pipe::{PipeFlags, pipe_with};
use rustix::process::{PidfdFlags, pidfd_open};

static SEQUENCE: AtomicU64 = AtomicU64::new(0);

fn main() -> ExitCode {
    let arguments = std::env::args_os().skip(1).collect::<Vec<_>>();
    if let Some(result) = dispatch_private_process_role(&arguments) {
        return exit(result);
    }
    if arguments.first() == Some(&OsString::from("--run-acquisition")) {
        return exit(run_acquisition(&arguments[1..]));
    }
    exit(run_all())
}

fn exit(result: std::io::Result<()>) -> ExitCode {
    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("acquisition fixture failed: {error}");
            ExitCode::FAILURE
        }
    }
}

fn run_all() -> std::io::Result<()> {
    successful_acquisition()?;
    lifetime_pipe_descriptor_contract()?;
    failed_stages_and_acknowledgments()?;
    production_wrapper_refuses_former_harness_path()?;
    malformed_barriers_refuse()?;
    sa_nocldwait_refuses()?;
    foreign_non_child_refuses()?;
    concurrent_acquisitions_serialize()
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn lifetime_pipe_descriptor_contract() -> std::io::Result<()> {
    let held_pipe_marker = marker("lifetime-pipe");
    let parent = current_process_identity()?;
    let target = TargetCommand::new("/usr/bin/touch")
        .argument(&held_pipe_marker)
        .with_lifetime_pipe_parent_edge();
    let (lifetime_reader, lifetime_writer) = pipe_with(PipeFlags::CLOEXEC)?;
    let process = child_broker().acquire_with_stdio(
        parent,
        context(),
        &target,
        ChildStandardIo::private_control_with_lifetime_pipe(
            Stdio::null(),
            Stdio::from(lifetime_reader),
        ),
        |_| Ok(()),
        |staged| Ok(ProcessOwnershipCommitAck::v1(staged)),
    )?;
    child_broker().reap(process)?;
    drop(lifetime_writer);
    if !held_pipe_marker.is_file() {
        return Err(std::io::Error::other(
            "lifetime-pipe acquisition did not execute its target",
        ));
    }
    std::fs::remove_file(held_pipe_marker)?;

    let wrong_descriptor_marker = marker("lifetime-pipe-wrong-descriptor");
    run_lifetime_pipe_target(&wrong_descriptor_marker, Stdio::null())?;
    if wrong_descriptor_marker.exists() {
        return Err(std::io::Error::other(
            "non-pipe lifetime descriptor released its target",
        ));
    }

    let closed_pipe_marker = marker("lifetime-pipe-closed");
    let (closed_reader, closed_writer) = pipe_with(PipeFlags::CLOEXEC)?;
    drop(closed_writer);
    run_lifetime_pipe_target(&closed_pipe_marker, Stdio::from(closed_reader))?;
    if closed_pipe_marker.exists() {
        return Err(std::io::Error::other(
            "closed lifetime pipe released its target",
        ));
    }
    Ok(())
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn run_lifetime_pipe_target(marker: &std::path::Path, stderr: Stdio) -> std::io::Result<()> {
    let parent = current_process_identity()?;
    let target = TargetCommand::new("/usr/bin/touch")
        .argument(marker)
        .with_lifetime_pipe_parent_edge();
    let process = child_broker().acquire_with_stdio(
        parent,
        context(),
        &target,
        ChildStandardIo::private_control_with_lifetime_pipe(Stdio::null(), stderr),
        |_| Ok(()),
        |staged| Ok(ProcessOwnershipCommitAck::v1(staged)),
    )?;
    child_broker().reap(process)?;
    Ok(())
}

fn marker(name: &str) -> std::path::PathBuf {
    std::env::temp_dir().join(format!(
        "lazaret-acquisition-{name}-{}-{}",
        std::process::id(),
        SEQUENCE.fetch_add(1, Ordering::Relaxed)
    ))
}

fn production_wrapper() -> &'static std::path::Path {
    std::path::Path::new(env!("CARGO_BIN_EXE_lazaret-launch-wrapper"))
}

fn context() -> ProcessOwnershipContext {
    ProcessOwnershipContext {
        domain_id: DomainId::new([1; 32]),
        generation_id: GenerationId::new(2),
        coordinator_id: CoordinatorId::new([3; 32]),
        slot: HostProcessOwnershipSlot::SandboxRuntimeSupervisor,
    }
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn run_acquisition(arguments: &[OsString]) -> std::io::Result<()> {
    let scenario = arguments
        .first()
        .and_then(|argument| argument.to_str())
        .ok_or_else(|| std::io::Error::other("scenario is missing"))?;
    let marker = arguments
        .get(1)
        .ok_or_else(|| std::io::Error::other("marker is missing"))?;
    let parent = current_process_identity()?;
    let target = TargetCommand::new("/usr/bin/touch").argument(marker);
    let expected = std::cell::Cell::new(None);
    let stage_count = std::cell::Cell::new(0_u8);
    let acknowledgment_count = std::cell::Cell::new(0_u8);
    let process = child_broker().acquire(
        parent,
        context(),
        &target,
        |staged| {
            stage_count.set(stage_count.get() + 1);
            if scenario == "stage-fail" {
                return Err(std::io::Error::other("injected stage failure"));
            }
            expected.set(Some(staged));
            Ok(())
        },
        |staged| {
            acknowledgment_count.set(acknowledgment_count.get() + 1);
            if scenario == "ack-loss" {
                return Err(std::io::Error::other("injected acknowledgment loss"));
            }
            let committed = mismatched_acknowledgment(scenario, staged);
            if expected.get() != Some(staged) {
                return Err(std::io::Error::other("staged ownership changed"));
            }
            Ok(ProcessOwnershipCommitAck::v1(committed))
        },
    )?;
    if stage_count.get() != 1 || acknowledgment_count.get() != 1 {
        return Err(std::io::Error::other(
            "stage and acknowledgment were not consumed exactly once",
        ));
    }
    child_broker().reap(process)
}

fn mismatched_acknowledgment(
    scenario: &str,
    staged: StagedProcessOwnership,
) -> StagedProcessOwnership {
    match scenario {
        "ack-domain-mismatch" => StagedProcessOwnership::v1(
            DomainId::new([9; 32]),
            staged.generation_id(),
            staged.coordinator_id(),
            staged.slot(),
            staged.identity(),
        ),
        "ack-generation-mismatch" => StagedProcessOwnership::v1(
            staged.domain_id(),
            GenerationId::new(staged.generation_id().get() + 1),
            staged.coordinator_id(),
            staged.slot(),
            staged.identity(),
        ),
        "ack-coordinator-mismatch" => StagedProcessOwnership::v1(
            staged.domain_id(),
            staged.generation_id(),
            CoordinatorId::new([9; 32]),
            staged.slot(),
            staged.identity(),
        ),
        "ack-slot-mismatch" => StagedProcessOwnership::v1(
            staged.domain_id(),
            staged.generation_id(),
            staged.coordinator_id(),
            HostProcessOwnershipSlot::SandboxAttachedRunsc,
            staged.identity(),
        ),
        "ack-identity-mismatch" => {
            let identity = staged.identity();
            StagedProcessOwnership::v1(
                staged.domain_id(),
                staged.generation_id(),
                staged.coordinator_id(),
                staged.slot(),
                PidfsProcessIdentityV1::new(
                    identity.authority(),
                    HostPidLocator::new(identity.host_pid().get() + 1),
                    identity.start_time_ticks(),
                ),
            )
        }
        _ => staged,
    }
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn successful_acquisition() -> std::io::Result<()> {
    let marker = marker("success");
    let status = std::process::Command::new(std::env::current_exe()?)
        .arg("--run-acquisition")
        .arg("success")
        .arg(&marker)
        .status()?;
    if !status.success() || !std::fs::read(&marker)?.is_empty() {
        return Err(std::io::Error::other(
            "successful target did not run exactly once",
        ));
    }
    std::fs::remove_file(marker)
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn failed_stages_and_acknowledgments() -> std::io::Result<()> {
    for scenario in [
        "stage-fail",
        "ack-loss",
        "ack-domain-mismatch",
        "ack-generation-mismatch",
        "ack-coordinator-mismatch",
        "ack-slot-mismatch",
        "ack-identity-mismatch",
    ] {
        let marker = marker(scenario);
        let status = std::process::Command::new(std::env::current_exe()?)
            .arg("--run-acquisition")
            .arg(scenario)
            .arg(&marker)
            .status()?;
        if status.success() || marker.exists() {
            return Err(std::io::Error::other(format!(
                "{scenario} released target execution"
            )));
        }
    }
    Ok(())
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn production_wrapper_refuses_former_harness_path() -> std::io::Result<()> {
    let marker = marker("former-harness");
    let status = std::process::Command::new(production_wrapper())
        .arg("--acquisition-harness")
        .arg("success")
        .arg(&marker)
        .status()?;
    if status.success() || marker.exists() {
        Err(std::io::Error::other(
            "production wrapper accepted former test harness path",
        ))
    } else {
        Ok(())
    }
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn malformed_barriers_refuse() -> std::io::Result<()> {
    for bytes in [
        Vec::new(),
        b"partial".to_vec(),
        vec![b'x'; 32],
        b"LAZARET-CAUSAL-BARRIER-RELEASE!!x".to_vec(),
    ] {
        let marker = marker("barrier-refusal");
        let parent = current_process_identity()?;
        let mut command = std::process::Command::new(production_wrapper());
        command
            .arg("--lazaret-private-process-role")
            .arg("launch-wrapper");
        append_parent_identity(&mut command, parent);
        command
            .args(["--parent-action", "kill", "--", "/usr/bin/touch"])
            .arg(&marker);
        command.stdin(Stdio::from(temporary_file(&bytes)?));
        if command.status()?.success() || marker.exists() {
            return Err(std::io::Error::other("malformed barrier released target"));
        }
    }
    Ok(())
}

fn temporary_file(bytes: &[u8]) -> std::io::Result<File> {
    let path = marker("barrier-input");
    let mut file = File::options()
        .create_new(true)
        .read(true)
        .write(true)
        .open(&path)?;
    file.write_all(bytes)?;
    file.rewind()?;
    std::fs::remove_file(path)?;
    Ok(file)
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn sa_nocldwait_refuses() -> std::io::Result<()> {
    let marker = marker("sa-nocldwait");
    let script = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("sa_nocldwait.py");
    let output = std::process::Command::new("uv")
        .args([
            "--cache-dir",
            "/tmp/lazaret-uv-cache",
            "run",
            "--no-project",
            "python",
        ])
        .arg(script)
        .arg(std::env::current_exe()?)
        .arg(&marker)
        .output()?;
    let stderr = String::from_utf8_lossy(&output.stderr);
    if output.status.success()
        || marker.exists()
        || !stderr.contains("SIGCHLD semantic probe was not reapable")
    {
        Err(std::io::Error::other("SA_NOCLDWAIT was accepted"))
    } else {
        Ok(())
    }
}

#[allow(clippy::disallowed_methods)]
fn foreign_non_child_refuses() -> std::io::Result<()> {
    let pidfd = pidfd_open(rustix::process::getpid(), PidfdFlags::empty())?;
    if verify_direct_child_pidfd(pidfd.as_fd()).is_err() {
        Ok(())
    } else {
        Err(std::io::Error::other(
            "foreign non-child pidfd was accepted",
        ))
    }
}

#[allow(clippy::disallowed_methods)]
fn concurrent_acquisitions_serialize() -> std::io::Result<()> {
    let parent = current_process_identity()?;
    let start = Arc::new(Barrier::new(9));
    let inside = Arc::new(AtomicUsize::new(0));
    let maximum = Arc::new(AtomicUsize::new(0));
    let workers = (0..8)
        .map(|_| {
            let start = Arc::clone(&start);
            let inside = Arc::clone(&inside);
            let maximum = Arc::clone(&maximum);
            std::thread::spawn(move || -> std::io::Result<()> {
                start.wait();
                let process = child_broker().acquire(
                    parent,
                    context(),
                    &TargetCommand::new("/bin/true"),
                    |staged| {
                        let active = inside.fetch_add(1, Ordering::SeqCst) + 1;
                        maximum.fetch_max(active, Ordering::SeqCst);
                        std::thread::yield_now();
                        inside.fetch_sub(1, Ordering::SeqCst);
                        let _ = staged;
                        Ok(())
                    },
                    |staged| Ok(ProcessOwnershipCommitAck::v1(staged)),
                )?;
                child_broker().reap(process)
            })
        })
        .collect::<Vec<_>>();
    start.wait();
    for worker in workers {
        worker
            .join()
            .map_err(|_| std::io::Error::other("worker panicked"))??;
    }
    if maximum.load(Ordering::SeqCst) == 1 {
        Ok(())
    } else {
        Err(std::io::Error::other("broker acquisition overlapped"))
    }
}