229 lines
7.2 KiB
Rust
229 lines
7.2 KiB
Rust
//! 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");
|
|
}
|