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

use std::fs::File;
use std::io::{Seek, Write};
use std::os::fd::AsFd;
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Barrier};
use std::thread;
use std::time::Duration;

use lazaret_linux::process::{
    TargetCommand, append_parent_identity, child_broker, current_process_identity,
    verify_direct_child_pidfd,
};
use rustix::process::{PidfdFlags, pidfd_open};

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

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 target(path: &std::path::Path) -> TargetCommand {
    TargetCommand::new("/bin/sh")
        .argument("-c")
        .argument(format!("printf target > '{}'", path.display()))
}

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

#[test]
fn target_runs_only_after_exact_stage_ack_and_release() {
    let marker = marker("released");
    let parent = current_process_identity().unwrap();
    let process = child_broker()
        .acquire(
            wrapper(),
            parent,
            &target(&marker),
            |identity| {
                assert!(!marker.exists());
                Ok(identity)
            },
            |staged, identity| {
                assert_eq!(*staged, identity);
                assert!(!marker.exists());
                Ok(())
            },
        )
        .unwrap();
    for _ in 0..100 {
        if marker.exists() {
            child_broker().reap(process).unwrap();
            assert_eq!(std::fs::read(&marker).unwrap(), b"target");
            std::fs::remove_file(marker).unwrap();
            return;
        }
        thread::sleep(Duration::from_millis(10));
    }
    panic!("released target did not run");
}

#[test]
fn stage_and_ack_failures_leave_target_absent() {
    for failure in ["stage", "ack"] {
        let marker = marker(failure);
        let parent = current_process_identity().unwrap();
        let result = child_broker().acquire(
            wrapper(),
            parent,
            &target(&marker),
            |_| {
                if failure == "stage" {
                    Err(std::io::Error::other("injected stage failure"))
                } else {
                    Ok(())
                }
            },
            |(), _| {
                if failure == "ack" {
                    Err(std::io::Error::other("injected acknowledgment failure"))
                } else {
                    Ok(())
                }
            },
        );
        assert!(result.is_err());
        assert!(!marker.exists());
    }
}

#[test]
#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn inherited_barrier_refuses_eof_partial_and_malformed_release() {
    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().unwrap();
        let mut command = std::process::Command::new(wrapper());
        append_parent_identity(&mut command, parent);
        command
            .arg("--")
            .arg("/bin/sh")
            .arg("-c")
            .arg(format!("printf target > '{}'", marker.display()));
        let mut input = tempfile_file(&bytes);
        input.flush().unwrap();
        command.stdin(Stdio::from(input));
        assert!(!command.status().unwrap().success());
        assert!(!marker.exists());
    }
}

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

#[test]
#[allow(clippy::disallowed_methods)]
fn foreign_non_child_pidfd_is_refused_without_action() {
    let pidfd = pidfd_open(rustix::process::getpid(), PidfdFlags::empty()).unwrap();
    assert!(verify_direct_child_pidfd(pidfd.as_fd()).is_err());
}

#[test]
fn process_global_broker_serializes_concurrent_acquisitions() {
    let start = Arc::new(Barrier::new(9));
    let inside = Arc::new(AtomicU64::new(0));
    let maximum = Arc::new(AtomicU64::new(0));
    let workers = (0..8)
        .map(|_| {
            let start = Arc::clone(&start);
            let inside = Arc::clone(&inside);
            let maximum = Arc::clone(&maximum);
            thread::spawn(move || {
                start.wait();
                let parent = current_process_identity().unwrap();
                let process = child_broker()
                    .acquire(
                        wrapper(),
                        parent,
                        &TargetCommand::new("/bin/true"),
                        |_| {
                            let active = inside.fetch_add(1, Ordering::SeqCst) + 1;
                            maximum.fetch_max(active, Ordering::SeqCst);
                            thread::sleep(Duration::from_millis(5));
                            inside.fetch_sub(1, Ordering::SeqCst);
                            Ok(())
                        },
                        |(), _| Ok(()),
                    )
                    .unwrap();
                child_broker().reap(process).unwrap();
            })
        })
        .collect::<Vec<_>>();
    start.wait();
    for worker in workers {
        worker.join().unwrap();
    }
    assert_eq!(maximum.load(Ordering::SeqCst), 1);
}