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

use std::ffi::OsString;
use std::fs::File;
use std::io::{Read, 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::{
    AcknowledgedChildAcquisition, ChildStandardIo, ProcessOwnershipContext, TargetCommand,
    append_parent_identity, child_broker, current_process_identity, dispatch_private_process_role,
    verify_direct_child_pidfd,
};
use rustix::fs::{OFlags, fcntl_getfl, fcntl_setfl};
use rustix::pipe::{PipeFlags, pipe_with};
use rustix::process::{PidfdFlags, pidfd_open};
use rustix::pty::{OpenptFlags, grantpt, ioctl_tiocgptpeer, openpt, unlockpt};
use rustix::termios::{Winsize, tcgetwinsize, tcsetwinsize};

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()?;
    transient_captured_child_preserves_streams_and_exit()?;
    terminal_child_restores_slave_stdin_and_merges_output()?;
    acknowledged_transition_precedes_exact_wrapper_release()?;
    failed_acknowledged_transition_refuses_exact_wrapper_release()?;
    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()
}

fn terminal_child_restores_slave_stdin_and_merges_output() -> std::io::Result<()> {
    let master = openpt(OpenptFlags::RDWR | OpenptFlags::NOCTTY | OpenptFlags::CLOEXEC)?;
    grantpt(&master)?;
    unlockpt(&master)?;
    let slave = ioctl_tiocgptpeer(
        &master,
        OpenptFlags::RDWR | OpenptFlags::NOCTTY | OpenptFlags::CLOEXEC,
    )?;
    tcsetwinsize(
        &master,
        Winsize {
            ws_row: 44,
            ws_col: 132,
            ws_xpixel: 0,
            ws_ypixel: 0,
        },
    )?;
    if tcgetwinsize(&master)?.ws_row != 44 || tcgetwinsize(&master)?.ws_col != 132 {
        return Err(std::io::Error::other(
            "PTY resize did not update the master",
        ));
    }
    let flags = fcntl_getfl(&master)?;
    fcntl_setfl(&master, flags | OFlags::NONBLOCK)?;
    let slave = File::from(slave);
    let target = TargetCommand::new("/bin/sh")
        .argument("-c")
        .argument(
            "read value; printf 'stdin-fd\\n' >&0; printf 'out:%s\\n' \"$value\"; printf 'err\\n' >&2",
        )
        .with_target_stdin(slave.try_clone()?)?;
    let expected_parent = current_process_identity()?;
    let ownership = ProcessOwnershipContext {
        domain_id: DomainId::new([41; 32]),
        generation_id: GenerationId::new(42),
        coordinator_id: CoordinatorId::new([43; 32]),
        slot: HostProcessOwnershipSlot::SessionTerminalExecClient(
            lazaret_core::identity::SessionId::new([44; 32]),
        ),
    };
    let process = child_broker().acquire_with_stdio(
        expected_parent,
        ownership,
        &target,
        ChildStandardIo::terminal(slave)?,
        |_| Ok(()),
        |staged| Ok(ProcessOwnershipCommitAck::v1(staged)),
    )?;
    let mut master = File::from(master);
    master.write_all(b"hello\n")?;
    let mut merged = Vec::new();
    while !process.has_exited()? {
        let _ = read_available(&mut master, &mut merged);
        std::thread::yield_now();
    }
    let _ = read_available(&mut master, &mut merged);
    let exit = child_broker().reap_status(process)?;
    let text = String::from_utf8_lossy(&merged);
    let stdin = text.find("stdin-fd");
    let out = text.find("out:hello");
    let err = text.find("err");
    if !matches!(
        (stdin, out, err),
        (Some(stdin), Some(out), Some(err)) if stdin < out && out < err
    ) || exit.exit_code != Some(0)
        || exit.terminating_signal.is_some()
    {
        return Err(std::io::Error::other(format!(
            "terminal stdio or merged ordering failed: {text:?}, {exit:?}"
        )));
    }
    Ok(())
}

fn transient_captured_child_preserves_streams_and_exit() -> std::io::Result<()> {
    let target = TargetCommand::new("/bin/sh")
        .argument("-c")
        .argument("printf transient-out; printf transient-err >&2; exit 7");
    let mut child = child_broker().run_transient_captured(&target)?;
    let mut stdout = Vec::new();
    let mut stderr = Vec::new();
    loop {
        let (out, err) = child.streams();
        let out_done = read_available(out, &mut stdout)?;
        let err_done = read_available(err, &mut stderr)?;
        if out_done && err_done && child.process().has_exited()? {
            break;
        }
        std::thread::yield_now();
    }
    let exit = child_broker().reap_status(child.into_process())?;
    if stdout != b"transient-out"
        || stderr != b"transient-err"
        || exit.exit_code != Some(7)
        || exit.terminating_signal.is_some()
    {
        return Err(std::io::Error::other(
            "transient captured child changed streams or exit",
        ));
    }
    Ok(())
}

fn read_available(file: &mut File, output: &mut Vec<u8>) -> std::io::Result<bool> {
    let mut buffer = [0; 4096];
    loop {
        match file.read(&mut buffer) {
            Ok(0) => return Ok(true),
            Ok(length) => output.extend_from_slice(&buffer[..length]),
            Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => return Ok(false),
            Err(error) => return Err(error),
        }
    }
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn failed_acknowledged_transition_refuses_exact_wrapper_release() -> std::io::Result<()> {
    let target_marker = marker("failed-transition-target");
    let target = TargetCommand::new("/usr/bin/touch").argument(&target_marker);
    let result = child_broker().acquire_with_stdio_after_ack(
        AcknowledgedChildAcquisition {
            expected_parent: current_process_identity()?,
            ownership: context(),
            target: &target,
            standard_io: ChildStandardIo::inherit(),
        },
        |_| Ok(()),
        |staged| Ok(ProcessOwnershipCommitAck::v1(staged)),
        || Err(std::io::Error::other("handoff authority changed")),
    );
    if result.is_ok() || target_marker.exists() {
        return Err(std::io::Error::other(
            "failed transition released the exact wrapper",
        ));
    }
    Ok(())
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn acknowledged_transition_precedes_exact_wrapper_release() -> std::io::Result<()> {
    let target_marker = marker("transition-target");
    let parent = current_process_identity()?;
    let staged = std::cell::Cell::new(None);
    let acknowledged = std::cell::Cell::new(false);
    let target = TargetCommand::new("/usr/bin/touch").argument(&target_marker);
    let process = child_broker().acquire_with_stdio_after_ack(
        AcknowledgedChildAcquisition {
            expected_parent: parent,
            ownership: context(),
            target: &target,
            standard_io: ChildStandardIo::inherit(),
        },
        |ownership| {
            staged.set(Some(ownership));
            Ok(())
        },
        |ownership| {
            if staged.get() != Some(ownership) {
                return Err(std::io::Error::other("transition staged ownership changed"));
            }
            acknowledged.set(true);
            Ok(ProcessOwnershipCommitAck::v1(ownership))
        },
        || {
            if !acknowledged.get() || target_marker.exists() {
                return Err(std::io::Error::other(
                    "wrapper was released outside the acknowledged transition boundary",
                ));
            }
            Ok(())
        },
    )?;
    child_broker().reap(process)?;
    if !target_marker.is_file() {
        return Err(std::io::Error::other(
            "acknowledged transition did not release the exact wrapper",
        ));
    }
    std::fs::remove_file(target_marker)
}

#[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"))
    }
}