boxmaker/docs/plans/M5b/files/crates/inferproxy/tests/instance.rs

//! One `inferproxy` per socket (M5b spec, section 9a): the lock is `<socket>.lock`, beside the
//! `--listen` socket. A second `inferproxy` on the same socket is refused, and the first keeps
//! forwarding; once the first is killed, a new one removes the stale socket and forwards. Every
//! wait has a deadline, and every process is killed and every path removed by a guard (tip T44).
//! Do not edit.

use std::io::{Read, Write};
use std::net::TcpListener;
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::UnixStream;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicU32, Ordering};
use std::thread;
use std::time::{Duration, Instant};

const POINTER: &str = "see docs/runbook.md#already-running";

static NEXT: AtomicU32 = AtomicU32::new(0);

/// A temporary directory for the socket, removed when dropped, first with its owner's permissions
/// put back.
struct Dir(PathBuf);

impl Dir {
    fn new() -> Dir {
        let n = NEXT.fetch_add(1, Ordering::SeqCst);
        let dir = std::env::temp_dir().join(format!("ip-lock-{}-{n}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        Dir(dir)
    }

    fn socket(&self) -> PathBuf {
        self.0.join("infer.sock")
    }
}

impl Drop for Dir {
    fn drop(&mut self) {
        let _ = std::fs::set_permissions(&self.0, std::fs::Permissions::from_mode(0o700));
        let _ = std::fs::remove_dir_all(&self.0);
    }
}

/// An upstream that reads four bytes from each connection, answers `pong` and closes.
fn upstream() -> String {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap().to_string();
    thread::spawn(move || {
        for stream in listener.incoming() {
            let Ok(mut stream) = stream else { break };
            thread::spawn(move || {
                let mut ping = [0u8; 4];
                if stream.read_exact(&mut ping).is_ok() {
                    let _ = stream.write_all(b"pong");
                }
            });
        }
    });
    addr
}

/// A running `inferproxy`, killed and waited for when dropped. Declared after the `Dir`, so it
/// drops first.
struct Proxy(Child);

impl Drop for Proxy {
    fn drop(&mut self) {
        let _ = self.0.kill();
        let _ = self.0.wait();
    }
}

impl Proxy {
    fn running(&mut self) -> bool {
        matches!(self.0.try_wait(), Ok(None))
    }
}

fn command(socket: &Path, upstream: &str) -> Command {
    let mut command = Command::new(env!("CARGO_BIN_EXE_inferproxy"));
    command
        .arg("--listen")
        .arg(socket)
        .args(["--upstream", upstream])
        .stdin(Stdio::null())
        .stdout(Stdio::null());
    command
}

fn start(socket: &Path, upstream: &str) -> Proxy {
    Proxy(
        command(socket, upstream)
            .stderr(Stdio::null())
            .spawn()
            .unwrap(),
    )
}

/// Runs `inferproxy` to its end, at most 10 s (then it is killed), and gives its exit code
/// (`None` if it had to be killed) and its stderr.
fn run_to_end(socket: &Path, upstream: &str) -> (Option<i32>, String) {
    let mut child = Proxy(
        command(socket, upstream)
            .stderr(Stdio::piped())
            .spawn()
            .unwrap(),
    );
    let deadline = Instant::now() + Duration::from_secs(10);
    let mut code = None;
    while Instant::now() < deadline {
        if let Ok(Some(status)) = child.0.try_wait() {
            code = status.code();
            break;
        }
        thread::sleep(Duration::from_millis(20));
    }
    let _ = child.0.kill();
    let _ = child.0.wait();
    let mut stderr = String::new();
    if let Some(mut pipe) = child.0.stderr.take() {
        let _ = pipe.read_to_string(&mut stderr);
    }
    (code, stderr)
}

/// True when a `ping` sent to `socket` comes back from the upstream as `pong` within 5 s.
fn forwards(socket: &Path) -> bool {
    let Ok(mut stream) = UnixStream::connect(socket) else {
        return false;
    };
    let limit = Some(Duration::from_secs(5));
    if stream.set_read_timeout(limit).is_err() || stream.set_write_timeout(limit).is_err() {
        return false;
    }
    if stream.write_all(b"ping").is_err() {
        return false;
    }
    let mut answer = Vec::new();
    let _ = stream.read_to_end(&mut answer);
    answer == b"pong"
}

/// Waits up to 10 s for an `inferproxy` to forward on `socket`.
fn wait_until_it_forwards(socket: &Path) -> bool {
    let deadline = Instant::now() + Duration::from_secs(10);
    while Instant::now() < deadline {
        if forwards(socket) {
            return true;
        }
        thread::sleep(Duration::from_millis(50));
    }
    false
}

fn inode(path: &Path) -> u64 {
    use std::os::unix::fs::MetadataExt;
    std::fs::symlink_metadata(path)
        .map(|m| m.ino())
        .unwrap_or(0)
}

#[test]
fn a_second_inferproxy_on_the_same_socket_is_refused_and_the_first_keeps_forwarding() {
    let dir = Dir::new();
    let socket = dir.socket();
    let upstream = upstream();
    let mut first = start(&socket, &upstream);
    assert!(wait_until_it_forwards(&socket), "the first never forwarded");
    let before = inode(&socket);

    let (code, stderr) = run_to_end(&socket, &upstream);
    let line = format!(
        "inferproxy: another inferproxy is running (it holds {}.lock); {POINTER}",
        socket.display()
    );
    assert_eq!(code, Some(1), "the second inferproxy: {stderr}");
    assert!(
        stderr.lines().any(|l| l == line),
        "wanted the line\n{line}\nstderr: {stderr}"
    );

    assert!(first.running(), "the first inferproxy is still running");
    assert_eq!(inode(&socket), before, "the first's socket is still there");
    assert!(forwards(&socket), "the first still forwards");
}

#[test]
fn after_the_first_is_killed_a_new_inferproxy_removes_the_stale_socket_and_forwards() {
    let dir = Dir::new();
    let socket = dir.socket();
    let upstream = upstream();
    let mut first = start(&socket, &upstream);
    assert!(wait_until_it_forwards(&socket), "the first never forwarded");
    let _ = first.0.kill();
    let _ = first.0.wait();
    assert!(socket.exists(), "a killed inferproxy leaves its socket");
    assert!(!forwards(&socket), "nobody serves the stale socket");

    let _next = start(&socket, &upstream);
    assert!(
        wait_until_it_forwards(&socket),
        "the new one never forwarded"
    );
    let lock = dir.0.join("infer.sock.lock");
    let mode = std::fs::metadata(&lock).map(|m| m.permissions().mode() & 0o777);
    assert_eq!(mode.ok(), Some(0o600), "the lock file, beside the socket");
}

#[test]
fn a_socket_directory_that_cannot_be_written_stops_inferproxy_naming_the_error() {
    let dir = Dir::new();
    let socket = dir.socket();
    std::fs::set_permissions(&dir.0, std::fs::Permissions::from_mode(0o500)).unwrap();

    let (code, stderr) = run_to_end(&socket, &upstream());
    let line = format!(
        "inferproxy: cannot take the lock {}.lock: Permission denied (os error 13); {POINTER}",
        socket.display()
    );
    assert_eq!(code, Some(1), "{stderr}");
    assert!(
        stderr.lines().any(|l| l == line),
        "wanted the line\n{line}\nstderr: {stderr}"
    );
    assert!(!socket.exists(), "no socket");
}