#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn main() {
server::run();
}
#[cfg(not(all(target_os = "linux", target_arch = "x86_64")))]
fn main() {
eprintln!("_tree_world_server requires linux/x86_64 KVM");
std::process::exit(1);
}
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
mod server {
use serde_json::{json, Value};
use std::collections::HashMap;
use std::fs::OpenOptions;
use std::io::{self, BufRead, Write};
use std::os::unix::fs::OpenOptionsExt;
use std::os::unix::io::AsRawFd;
use std::path::{Path, PathBuf};
use std::sync::{
atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
Arc,
};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use supermachine::{Image, TcpForwardEvent, TcpForwardObserver, TcpForwarder, Vm, VmConfig};
const APP_PORT: u16 = 3000;
const CDP_PORT: u16 = 9222;
pub fn run() {
redirect_stderr_to_workdir_file();
let mut state = match State::from_env() {
Ok(state) => state,
Err(reason) => {
reply(json!({ "ok": false, "reason": reason }));
return;
}
};
let stdin = io::stdin();
for line in stdin.lock().lines() {
let line = match line {
Ok(line) => line,
Err(e) => {
reply(json!({ "ok": false, "reason": format!("read stdin: {e}") }));
continue;
}
};
if line.trim().is_empty() {
continue;
}
let request: Value = match serde_json::from_str(&line) {
Ok(v) => v,
Err(e) => {
reply(json!({ "ok": false, "reason": format!("malformed json: {e}") }));
continue;
}
};
let request_id = request.get("id").cloned();
let response = state.handle(request);
let should_exit = matches!(response.get("released"), Some(Value::Bool(true)));
let mut public_response = response;
if let Some(obj) = public_response.as_object_mut() {
obj.remove("released");
if let Some(id) = request_id {
obj.insert("id".to_owned(), id);
}
}
reply(public_response);
if should_exit {
return;
}
}
}
struct State {
snapshot: String,
base_snap: PathBuf,
workdir: PathBuf,
chrome: String,
cfg: VmConfig,
cdp_host_port: u16,
parent: Option<RunningWorld>,
forks: Vec<Option<RunningWorld>>,
max_slots: usize,
mem_mib: u32,
mem_budget_mib: u32,
captures: HashMap<String, Image>,
debug_relay: bool,
}
struct RunningWorld {
vm: Vm,
cdp_forwarder: Option<TcpForwarder>,
relay: Arc<RelayStats>,
}
struct RelayStats {
label: String,
cdp_port: u16,
debug: bool,
listening: AtomicBool,
active: AtomicUsize,
guest_connected: AtomicBool,
next_conn: AtomicU64,
}
struct RestoreTimings {
start: Instant,
}
impl RestoreTimings {
fn start() -> Self {
Self {
start: Instant::now(),
}
}
fn mark_ms(&self) -> f64 {
self.start.elapsed().as_secs_f64() * 1e3
}
}
impl State {
fn from_env() -> Result<Self, String> {
let snapshot = env("SUPERMACHINE_WORLD_SNAPSHOT")?;
let workdir = PathBuf::from(env("SUPERMACHINE_WORLD_WORKDIR")?);
let chrome = env("SUPERMACHINE_WORLD_CHROME")?;
let mem: u32 = std::env::var("SUPERMACHINE_WORLD_MEM_MIB")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(8192);
let cdp_host_port = std::env::var("SUPERMACHINE_WORLD_CDP_HOST_PORT")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(CDP_PORT);
let max_slots = std::env::var("SUPERMACHINE_WORLD_MAX_SLOTS")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(2);
if max_slots == 0 {
return Err("SUPERMACHINE_WORLD_MAX_SLOTS must be at least 1".to_owned());
}
if max_slots > u16::MAX as usize - cdp_host_port as usize {
return Err("SUPERMACHINE_WORLD_MAX_SLOTS overflows CDP host ports".to_owned());
}
let default_budget = mem.saturating_mul(max_slots as u32);
let mem_budget_mib = std::env::var("SUPERMACHINE_WORLD_MEM_BUDGET_MIB")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default_budget);
std::fs::create_dir_all(&workdir).map_err(|e| format!("create workdir: {e}"))?;
let forks = (0..max_slots).map(|_| None).collect();
Ok(Self {
base_snap: Path::new(&snapshot).join("restore.snap"),
snapshot,
workdir,
chrome,
cfg: VmConfig::new().with_memory_mib(mem),
cdp_host_port,
parent: None,
forks,
max_slots,
mem_mib: mem,
mem_budget_mib,
captures: HashMap::new(),
debug_relay: std::env::var("SUPERMACHINE_WORLD_DEBUG_RELAY")
.ok()
.as_deref()
== Some("1"),
})
}
fn handle(&mut self, request: Value) -> Value {
match request
.get("op")
.or_else(|| request.get("method"))
.and_then(Value::as_str)
{
Some("ensure") => self.ensure(),
Some("capture") => self.capture(ref_id(&request), slot(&request)),
Some("restore") => {
self.restore(ref_id(&request), slot(&request).map(|s| s.unwrap_or(0)))
}
Some("relayState") => self.relay_state(),
Some("release") => self.release(
slot(&request),
request.get("all").and_then(Value::as_bool).unwrap_or(false),
),
Some(op) => json!({ "ok": false, "reason": format!("unknown op: {op}") }),
None => json!({ "ok": false, "reason": "missing string op" }),
}
}
fn ensure(&mut self) -> Value {
if self.parent.is_none() {
if let Err(reason) = self.check_mem_budget(1) {
return json!({ "ok": false, "reason": reason, "reasonCode": "world_mem_budget_exceeded" });
}
let result = (|| -> Result<RunningWorld, String> {
let vm = Image::from_snapshot(&self.snapshot)
.map_err(|e| format!("from_snapshot: {e}"))?
.start(&self.cfg)
.map_err(|e| format!("start: {e}"))?;
wait_exec(&vm)?;
start_world(&vm, &self.chrome)?;
let relay = RelayStats::new("parent", self.cdp_host_port, self.debug_relay);
relay.log("expose_tcp bind/listen start");
let cdp_forwarder = vm
.expose_tcp_observed(self.cdp_host_port, CDP_PORT, Some(relay.clone()))
.map_err(|e| format!("expose cdp: {e}"))?;
wait_http(
&format!("127.0.0.1:{}", self.cdp_host_port),
"/json/version",
30,
)?;
relay.listening.store(true, Ordering::SeqCst);
relay.log("expose_tcp listening");
Ok(RunningWorld {
vm,
cdp_forwarder: Some(cdp_forwarder),
relay,
})
})();
match result {
Ok(world) => self.parent = Some(world),
Err(reason) => return json!({ "ok": false, "reason": reason }),
}
}
json!({ "ok": true, "cdpEndpoint": self.cdp_endpoint(0) })
}
fn capture(
&mut self,
ref_id: Result<String, String>,
slot: Result<Option<usize>, String>,
) -> Value {
let ref_id = match ref_id {
Ok(r) => r,
Err(reason) => return json!({ "ok": false, "reason": reason }),
};
if !safe_ref(&ref_id) {
return json!({ "ok": false, "reason": "ref must contain only [A-Za-z0-9._-]" });
}
let slot = match slot {
Ok(slot) => slot,
Err(reason) => return json!({ "ok": false, "reason": reason }),
};
if self.parent.is_none() && self.forks.iter().all(Option::is_none) && slot.is_none() {
let ensured = self.ensure();
if ensured.get("ok") != Some(&Value::Bool(true)) {
return ensured;
}
}
let world = match self.capture_world(slot) {
Ok(world) => world,
Err(reason) => return json!({ "ok": false, "reason": reason }),
};
let dir = self.workdir.join(&ref_id);
let _ = std::fs::remove_dir_all(&dir);
let t = Instant::now();
let image = match world.vm.snapshot_diff_live_dirty(&dir, &self.base_snap) {
Ok(image) => image,
Err(e) => {
return json!({ "ok": false, "reason": format!("capture dirty diff: {e}") })
}
};
let capture_ms = t.elapsed().as_secs_f64() * 1e3;
let bytes = dir_bytes(&dir);
self.captures.insert(ref_id.clone(), image);
json!({ "ok": true, "ref": ref_id, "captureMs": capture_ms, "bytes": bytes })
}
fn capture_world(&self, slot: Option<usize>) -> Result<&RunningWorld, String> {
let fork_live: Vec<bool> = self.forks.iter().map(Option::is_some).collect();
match select_capture_world(self.parent.is_some(), &fork_live, self.max_slots, slot)? {
CaptureWorld::Parent => self
.parent
.as_ref()
.ok_or_else(|| "capture requires a live world".to_owned()),
CaptureWorld::Fork(slot) => self.forks[slot]
.as_ref()
.ok_or_else(|| format!("slot {slot} is empty")),
}
}
fn restore(
&mut self,
ref_id: Result<String, String>,
slot: Result<usize, String>,
) -> Value {
let ref_id = match ref_id {
Ok(r) => r,
Err(reason) => return json!({ "ok": false, "reason": reason }),
};
let slot = match slot {
Ok(slot) => slot,
Err(reason) => return json!({ "ok": false, "reason": reason }),
};
if slot >= self.max_slots {
return json!({ "ok": false, "reason": format!("slot {slot} out of range 0..{}", self.max_slots) });
}
let replacing = self.forks[slot].is_some();
let parent_port_conflict =
self.parent.is_some() && self.slot_port(slot) == self.cdp_host_port;
let additional_worlds = usize::from(!replacing && !parent_port_conflict);
if let Err(reason) = self.check_mem_budget(additional_worlds) {
return json!({ "ok": false, "reason": reason, "reasonCode": "world_mem_budget_exceeded" });
}
stop_world(self.forks[slot].take(), "restore replacing world");
if parent_port_conflict {
stop_world(self.parent.take(), "restore replacing world");
}
if !self.captures.contains_key(&ref_id) {
if !safe_ref(&ref_id) {
return json!({ "ok": false, "reason": "ref must contain only [A-Za-z0-9._-]" });
}
let dir = self.workdir.join(&ref_id);
if dir.join("metadata.json").is_file() {
match Image::from_snapshot(&dir) {
Ok(image) => {
self.captures.insert(ref_id.clone(), image);
}
Err(e) => {
return json!({
"ok": false,
"reason": format!("rehydrate {}: {e}", dir.display()),
})
}
}
}
}
let image = match self.captures.get(&ref_id) {
Some(image) => image,
None => return json!({ "ok": false, "reason": format!("unknown ref: {ref_id}") }),
};
let timings = RestoreTimings::start();
let vm = match image.start(&self.cfg) {
Ok(vm) => vm,
Err(e) => return json!({ "ok": false, "reason": format!("restore start: {e}") }),
};
let fork_start_ms = timings.mark_ms();
if let Err(reason) = wait_exec(&vm) {
return json!({ "ok": false, "reason": reason });
}
let first_exec_success_ms = timings.mark_ms();
let cdp_host_port = self.slot_port(slot);
let relay = RelayStats::new(&format!("slot {slot}"), cdp_host_port, self.debug_relay);
relay.log("expose_tcp bind/listen start");
let cdp_forwarder =
match vm.expose_tcp_observed(cdp_host_port, CDP_PORT, Some(relay.clone())) {
Ok(f) => f,
Err(e) => {
return json!({ "ok": false, "reason": format!("expose restored cdp: {e}") })
}
};
let relay_listening_ms = timings.mark_ms();
if let Err(reason) =
wait_http(&format!("127.0.0.1:{}", cdp_host_port), "/json/version", 30)
{
return json!({ "ok": false, "reason": reason });
}
let host_cdp_200_ms = timings.mark_ms();
let restore_ms = host_cdp_200_ms;
relay.listening.store(true, Ordering::SeqCst);
relay.log("expose_tcp listening");
self.forks[slot] = Some(RunningWorld {
vm,
cdp_forwarder: Some(cdp_forwarder),
relay,
});
json!({
"ok": true,
"ref": ref_id,
"slot": slot,
"restoreMs": restore_ms,
"timings": {
"forkStartMs": fork_start_ms,
"firstExecSuccessMs": first_exec_success_ms,
"relayListeningMs": relay_listening_ms,
"hostCdp200Ms": host_cdp_200_ms
},
"cdpEndpoint": self.cdp_endpoint(slot)
})
}
fn release(&mut self, slot: Result<Option<usize>, String>, all: bool) -> Value {
let slot = match slot {
Ok(slot) => slot,
Err(reason) => return json!({ "ok": false, "reason": reason }),
};
if let Some(slot) = slot {
if slot >= self.max_slots {
return json!({ "ok": false, "reason": format!("slot {slot} out of range 0..{}", self.max_slots) });
}
stop_world(self.forks[slot].take(), "restore replacing world");
return json!({ "ok": true, "slot": slot });
}
for fork in &mut self.forks {
stop_world(fork.take(), "world released");
}
stop_world(self.parent.take(), "restore replacing world");
json!({ "ok": true, "released": all || slot.is_none() })
}
fn relay_state(&self) -> Value {
let mut worlds = Vec::new();
if let Some(parent) = &self.parent {
worlds.push(parent.relay.json(json!({ "parent": true })));
}
for (slot, world) in self.forks.iter().enumerate() {
if let Some(world) = world {
worlds.push(world.relay.json(json!({ "slot": slot })));
}
}
json!({ "ok": true, "worlds": worlds })
}
fn live_worlds(&self) -> usize {
usize::from(self.parent.is_some()) + self.forks.iter().filter(|f| f.is_some()).count()
}
fn check_mem_budget(&self, additional_worlds: usize) -> Result<(), String> {
let worlds = self.live_worlds().saturating_add(additional_worlds);
let required = worlds.saturating_mul(self.mem_mib as usize);
if required > self.mem_budget_mib as usize {
Err(format!(
"world_mem_budget_exceeded: {worlds} live worlds require {required} MiB, budget is {} MiB",
self.mem_budget_mib
))
} else {
Ok(())
}
}
fn slot_port(&self, slot: usize) -> u16 {
self.cdp_host_port + slot as u16
}
fn cdp_endpoint(&self, slot: usize) -> String {
format!("http://127.0.0.1:{}", self.slot_port(slot))
}
}
impl RelayStats {
fn new(label: &str, cdp_port: u16, debug: bool) -> Arc<Self> {
Arc::new(Self {
label: label.to_owned(),
cdp_port,
debug,
listening: AtomicBool::new(false),
active: AtomicUsize::new(0),
guest_connected: AtomicBool::new(false),
next_conn: AtomicU64::new(1),
})
}
fn log(&self, msg: &str) {
if self.debug {
let ts = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
nonblocking_stderr_line(&format!(
"[{ts}] relay {} port={} {msg}",
self.label, self.cdp_port
));
}
}
fn json(&self, mut base: Value) -> Value {
let obj = base.as_object_mut().unwrap();
obj.insert("cdpPort".to_owned(), json!(self.cdp_port));
obj.insert(
"listening".to_owned(),
json!(self.listening.load(Ordering::SeqCst)),
);
obj.insert(
"activeConnections".to_owned(),
json!(self.active.load(Ordering::SeqCst)),
);
obj.insert(
"guestConnected".to_owned(),
json!(self.guest_connected.load(Ordering::SeqCst)),
);
base
}
}
fn redirect_stderr_to_workdir_file() {
let path = if let Ok(path) = std::env::var("SUPERMACHINE_TREE_WORLD_SERVER_STDERR") {
PathBuf::from(path)
} else if let Ok(workdir) = std::env::var("SUPERMACHINE_WORLD_WORKDIR") {
Path::new(&workdir).join("tree-world-server.stderr")
} else {
return;
};
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let Ok(file) = OpenOptions::new().create(true).append(true).open(path) else {
return;
};
let _ = unsafe { libc::dup2(file.as_raw_fd(), libc::STDERR_FILENO) };
}
fn nonblocking_stderr_line(line: &str) {
let Ok(mut stderr) = OpenOptions::new()
.write(true)
.custom_flags(libc::O_NONBLOCK | libc::O_CLOEXEC)
.open("/proc/self/fd/2")
else {
return;
};
let _ = stderr.write(line.as_bytes());
let _ = stderr.write(b"\n");
}
#[cfg(test)]
mod relay_log_tests {
use super::*;
use std::io::Read;
use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
use std::time::Instant;
#[test]
fn relay_debug_log_drops_instead_of_blocking_on_full_stderr_pipe() {
let mut pipe_fds = [-1; 2];
assert_eq!(
unsafe { libc::pipe2(pipe_fds.as_mut_ptr(), libc::O_CLOEXEC) },
0
);
let read_fd = unsafe { OwnedFd::from_raw_fd(pipe_fds[0]) };
let write_fd = unsafe { OwnedFd::from_raw_fd(pipe_fds[1]) };
let saved_stderr = unsafe { libc::dup(libc::STDERR_FILENO) };
assert!(saved_stderr >= 0);
let saved_stderr = unsafe { OwnedFd::from_raw_fd(saved_stderr) };
assert_eq!(
unsafe { libc::dup2(write_fd.as_raw_fd(), libc::STDERR_FILENO) },
libc::STDERR_FILENO
);
let fill = vec![b'x'; 1024 * 1024];
let start = Instant::now();
for _ in 0..2048 {
nonblocking_stderr_line(std::str::from_utf8(&fill).unwrap());
if start.elapsed() > Duration::from_millis(250) {
break;
}
}
assert!(
start.elapsed() < Duration::from_secs(1),
"relay debug logging blocked behind an undrained stderr pipe"
);
assert_eq!(
unsafe { libc::dup2(saved_stderr.as_raw_fd(), libc::STDERR_FILENO) },
libc::STDERR_FILENO
);
drop(write_fd);
let mut drained = Vec::new();
let mut reader = std::fs::File::from(read_fd);
reader.read_to_end(&mut drained).unwrap();
assert!(
drained.len() < fill.len() * 2048,
"full stderr pipe should drop relay debug bytes rather than buffer unboundedly"
);
}
}
impl TcpForwardObserver for RelayStats {
fn on_tcp_forward_event(&self, event: TcpForwardEvent) {
match event {
TcpForwardEvent::Accept { peer } => {
let id = self.next_conn.fetch_add(1, Ordering::SeqCst);
self.active.fetch_add(1, Ordering::SeqCst);
self.log(&format!("conn#{id} inbound accept peer={peer:?}"));
}
TcpForwardEvent::WebSocketUpgrade => self.log("CDP websocket upgrade seen"),
TcpForwardEvent::GuestOpen => {
self.guest_connected.store(true, Ordering::SeqCst);
self.log("guest-side connection open");
}
TcpForwardEvent::Close {
host_to_guest_bytes,
guest_to_host_bytes,
} => {
self.active
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |v| {
Some(v.saturating_sub(1))
})
.ok();
if self.active.load(Ordering::SeqCst) == 0 {
self.guest_connected.store(false, Ordering::SeqCst);
}
self.log(&format!("guest-side connection close hostToGuestBytes={host_to_guest_bytes} guestToHostBytes={guest_to_host_bytes}"));
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CaptureWorld {
Parent,
Fork(usize),
}
fn select_capture_world(
parent_live: bool,
fork_live: &[bool],
max_slots: usize,
slot: Option<usize>,
) -> Result<CaptureWorld, String> {
if let Some(slot) = slot {
if slot >= max_slots {
return Err(format!("slot {slot} out of range 0..{max_slots}"));
}
return match fork_live.get(slot).copied().unwrap_or(false) {
true => Ok(CaptureWorld::Fork(slot)),
false => Err(format!("slot {slot} is empty")),
};
}
let mut live = Vec::new();
for (slot, is_live) in fork_live.iter().copied().enumerate() {
if is_live {
live.push(CaptureWorld::Fork(slot));
}
}
if parent_live {
live.push(CaptureWorld::Parent);
}
match live.as_slice() {
[world] => Ok(*world),
[] => Err("capture requires a live world".to_owned()),
_ => Err("capture is ambiguous with multiple live worlds; pass slot".to_owned()),
}
}
fn stop_world(world: Option<RunningWorld>, why: &str) {
if let Some(mut world) = world {
world.relay.log(&format!("relay teardown: {why}"));
world.relay.listening.store(false, Ordering::SeqCst);
drop(world.cdp_forwarder.take());
let _ = world.vm.stop();
}
}
fn slot(request: &Value) -> Result<Option<usize>, String> {
match request.get("slot") {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(n)) => n
.as_u64()
.and_then(|v| usize::try_from(v).ok())
.map(Some)
.ok_or_else(|| "slot must be a non-negative integer".to_owned()),
_ => Err("slot must be a non-negative integer".to_owned()),
}
}
fn ref_id(request: &Value) -> Result<String, String> {
request
.get("ref")
.or_else(|| request.get("prefixHash"))
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.ok_or_else(|| "missing string ref".to_owned())
}
fn safe_ref(ref_id: &str) -> bool {
!ref_id.is_empty()
&& ref_id
.bytes()
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-'))
}
fn start_world(vm: &Vm, chrome: &str) -> Result<(), String> {
let app = r#"const http=require('node:http');
http.createServer((q,res)=>{console.log('REQ',q.method,q.url);res.setHeader('content-type','text/html');res.end(`<!doctype html><button id="shared">shared</button><script>window.__worldReady=true</script>`)}).listen(3000,'127.0.0.1',()=>console.log('app ready'));"#;
let b64 = base64(app.as_bytes());
exec_ok(vm, &format!("echo {b64} | base64 -d > /root/app.cjs"))?;
let relay = format!(
"const net=require('node:net');net.createServer(c=>{{const u=net.connect({backend},'127.0.0.1');c.pipe(u);u.pipe(c);c.on('error',()=>u.destroy());u.on('error',()=>c.destroy())}}).listen({front},'0.0.0.0',()=>console.log('relay ready'));",
backend = CDP_PORT + 2,
front = CDP_PORT
);
let relay_b64 = base64(relay.as_bytes());
exec_ok(
vm,
&format!("echo {relay_b64} | base64 -d > /root/cdp-relay.cjs"),
)?;
exec_ok(
vm,
"nohup node /root/app.cjs > /root/app.log 2>&1 & sleep 0.2; true",
)?;
exec_ok(
vm,
&format!(
"nohup {chrome} --headless --no-sandbox --disable-gpu --disable-dev-shm-usage \
--disable-quic --disable-background-networking --disable-component-update \
--disable-sync --disable-default-apps --no-first-run \
--remote-debugging-port={} --remote-debugging-address=127.0.0.1 \
--user-data-dir=/root/cdp-profile about:blank \
> /root/chrome.log 2>&1 & sleep 0.2; true",
CDP_PORT + 2
),
)?;
exec_ok(
vm,
"nohup node /root/cdp-relay.cjs > /root/relay.log 2>&1 & sleep 0.2; true",
)?;
exec_retry(vm, &format!("node -e \"require('http').get('http://127.0.0.1:{APP_PORT}',r=>process.exit(0)).on('error',()=>process.exit(1))\""), 30)?;
exec_retry(vm, &format!("node -e \"require('http').get('http://127.0.0.1:{}/json/version',r=>process.exit(0)).on('error',()=>process.exit(1))\"", CDP_PORT + 2), 30)?;
exec_retry(vm, &format!("node -e \"require('http').get('http://127.0.0.1:{CDP_PORT}/json/version',r=>process.exit(0)).on('error',()=>process.exit(1))\""), 30)
}
fn wait_exec(vm: &Vm) -> Result<(), String> {
for _ in 0..120 {
if let Ok(o) = vm
.exec_builder()
.argv(["/bin/sh", "-c", "echo up"])
.output()
{
if o.success() {
return Ok(());
}
}
std::thread::sleep(Duration::from_millis(1000));
}
Err("guest exec never came up".to_owned())
}
fn exec_ok(vm: &Vm, cmd: &str) -> Result<(), String> {
let o = vm
.exec_builder()
.argv(["/bin/sh", "-c", cmd])
.output()
.map_err(|e| format!("exec: {e}"))?;
if o.success() {
Ok(())
} else {
Err(format!(
"guest cmd failed: {cmd}\n{}",
String::from_utf8_lossy(&o.stderr)
))
}
}
fn exec_retry(vm: &Vm, cmd: &str, secs: u32) -> Result<(), String> {
for _ in 0..secs {
if let Ok(o) = vm.exec_builder().argv(["/bin/sh", "-c", cmd]).output() {
if o.success() {
return Ok(());
}
}
std::thread::sleep(Duration::from_millis(1000));
}
Err(format!("guest readiness timed out: {cmd}"))
}
fn wait_http(addr: &str, path: &str, secs: u32) -> Result<(), String> {
for _ in 0..secs {
if let Ok(mut s) = std::net::TcpStream::connect(addr) {
let _ = s.set_read_timeout(Some(Duration::from_secs(2)));
let req =
format!("GET {path} HTTP/1.1\r\nHost: {addr}\r\nConnection: close\r\n\r\n");
if s.write_all(req.as_bytes()).is_ok() {
let mut buf = String::new();
let _ = std::io::Read::read_to_string(&mut s, &mut buf);
if buf.starts_with("HTTP/1.1 200") {
return Ok(());
}
}
}
std::thread::sleep(Duration::from_millis(1000));
}
Err(format!(
"host->guest http readiness timed out: {addr}{path}"
))
}
fn dir_bytes(path: &Path) -> u64 {
let mut sum = 0;
if let Ok(rd) = std::fs::read_dir(path) {
for e in rd.flatten() {
let p = e.path();
if let Ok(md) = e.metadata() {
if md.is_dir() {
sum += dir_bytes(&p);
} else {
sum += md.len();
}
}
}
}
sum
}
fn base64(data: &[u8]) -> String {
const T: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut out = String::new();
for chunk in data.chunks(3) {
let b = [
chunk[0],
*chunk.get(1).unwrap_or(&0),
*chunk.get(2).unwrap_or(&0),
];
let n = ((b[0] as u32) << 16) | ((b[1] as u32) << 8) | b[2] as u32;
out.push(T[(n >> 18) as usize & 63] as char);
out.push(T[(n >> 12) as usize & 63] as char);
out.push(if chunk.len() > 1 {
T[(n >> 6) as usize & 63] as char
} else {
'='
});
out.push(if chunk.len() > 2 {
T[(n & 63) as usize] as char
} else {
'='
});
}
out
}
fn env(name: &str) -> Result<String, String> {
std::env::var(name).map_err(|_| format!("missing env {name}"))
}
fn reply(value: Value) {
println!("{}", value);
let _ = io::stdout().flush();
}
#[cfg(test)]
mod tests {
use super::{select_capture_world, CaptureWorld};
#[test]
fn capture_with_slot_selects_that_live_fork() {
assert_eq!(
select_capture_world(false, &[true, true], 2, Some(1)).unwrap(),
CaptureWorld::Fork(1)
);
}
#[test]
fn capture_with_slot_refuses_empty_or_out_of_range_slot() {
assert_eq!(
select_capture_world(false, &[true, false], 2, Some(1)).unwrap_err(),
"slot 1 is empty"
);
assert_eq!(
select_capture_world(false, &[true, false], 2, Some(2)).unwrap_err(),
"slot 2 out of range 0..2"
);
}
#[test]
fn capture_without_slot_allows_exactly_one_live_world() {
assert_eq!(
select_capture_world(true, &[false, false], 2, None).unwrap(),
CaptureWorld::Parent
);
assert_eq!(
select_capture_world(false, &[false, true], 2, None).unwrap(),
CaptureWorld::Fork(1)
);
}
#[test]
fn capture_without_slot_refuses_ambiguous_live_worlds() {
assert_eq!(
select_capture_world(true, &[false, true], 2, None).unwrap_err(),
"capture is ambiguous with multiple live worlds; pass slot"
);
assert_eq!(
select_capture_world(false, &[true, true], 2, None).unwrap_err(),
"capture is ambiguous with multiple live worlds; pass slot"
);
}
}
}