@@ -0,0 +1,538 @@
1+use std::{
2+ env, fs,
3+ io::{Read, Write},
4+ os::unix::{
5+ fs::PermissionsExt,
6+ io::AsRawFd,
7+ net::{UnixListener, UnixStream},
8+ process::CommandExt,
9+ },
10+ path::{Path, PathBuf},
11+ process::{Command, Stdio},
12+ sync::{Arc, Mutex, mpsc},
13+ thread,
14+ time::{Duration, Instant},
15+};
16+
17+use crate::{
18+ instance::{
19+ InstanceConfig, build_ssh_login_actions, ensure_instance_dir, ensure_ssh_keypair,
20+ extract_ipv4, load_or_create_instance_config, write_instance_config,
21+ },
22+ session_manager::GLOBAL_DIR_NAME,
23+ vm::{self, DirectoryShare, LoginAction, VmInput},
24+};
25+
26+const VM_MANAGER_SOCKET_NAME: &str = "vm.sock";
27+
28+pub fn ensure_manager(
29+ raw_args: &[std::ffi::OsString],
30+ auto_shutdown_ms: u64,
31+) -> Result<UnixStream, Box<dyn std::error::Error>> {
32+ let project_root = env::current_dir()?;
33+ tracing::debug!(root = %project_root.display(), "ensure vm manager");
34+ let instance_dir = ensure_instance_dir(&project_root)?;
35+ let socket_path = instance_dir.join(VM_MANAGER_SOCKET_NAME);
36+
37+ if let Ok(stream) = UnixStream::connect(&socket_path) {
38+ send_client_pid(&stream);
39+ tracing::info!(path = %socket_path.display(), "connected to existing vm manager");
40+ return Ok(stream);
41+ }
42+
43+ tracing::info!(path = %socket_path.display(), "spawning vm manager");
44+ spawn_manager_process(raw_args, auto_shutdown_ms, &instance_dir)?;
45+
46+ let start = Instant::now();
47+ let timeout = Duration::from_secs(10);
48+ loop {
49+ match UnixStream::connect(&socket_path) {
50+ Ok(stream) => {
51+ send_client_pid(&stream);
52+ tracing::info!(path = %socket_path.display(), "connected to vm manager");
53+ return Ok(stream);
54+ }
55+ Err(err) => {
56+ tracing::debug!(error = %err, "waiting for vm manager socket");
57+ if start.elapsed() > timeout {
58+ return Err(format!(
59+ "Timed out waiting for vm manager socket: {} ({})",
60+ socket_path.display(),
61+ err
62+ )
63+ .into());
64+ }
65+ thread::sleep(Duration::from_millis(100));
66+ }
67+ }
68+ }
69+}
70+
71+pub fn run_manager(
72+ args: vm::CliArgs,
73+ auto_shutdown_ms: u64,
74+) -> Result<(), Box<dyn std::error::Error>> {
75+ let project_root = env::current_dir()?;
76+ tracing::info!(root = %project_root.display(), "vm manager starting");
77+ run_manager_with(
78+ &project_root,
79+ args,
80+ auto_shutdown_ms,
81+ &RealVmExecutor,
82+ ManagerOptions {
83+ ensure_signed: true,
84+ detach: true,
85+ prepare_vm: true,
86+ },
87+ )
88+}
89+
90+fn spawn_manager_process(
91+ raw_args: &[std::ffi::OsString],
92+ auto_shutdown_ms: u64,
93+ instance_dir: &Path,
94+) -> Result<(), Box<dyn std::error::Error>> {
95+ let exe = env::current_exe()?;
96+ let mut supervisor_exe = exe.clone();
97+ supervisor_exe.set_file_name("vibebox-supervisor");
98+ let use_supervisor = supervisor_exe.exists();
99+ let mut cmd = if use_supervisor {
100+ Command::new(supervisor_exe)
101+ } else {
102+ let mut cmd = Command::new(exe);
103+ cmd.arg0("vibebox-supervisor");
104+ cmd
105+ };
106+ if raw_args.len() > 1 {
107+ cmd.args(&raw_args[1..]);
108+ }
109+ if !use_supervisor {
110+ cmd.env("VIBEBOX_VM_MANAGER", "1");
111+ }
112+ cmd.env("VIBEBOX_LOG_NO_COLOR", "1");
113+ cmd.env("VIBEBOX_AUTO_SHUTDOWN_MS", auto_shutdown_ms.to_string());
114+ tracing::info!(auto_shutdown_ms, "vm manager process spawn requested");
115+ let log_path = instance_dir.join("vm_manager.log");
116+ let log_file = fs::OpenOptions::new()
117+ .create(true)
118+ .append(true)
119+ .open(&log_path)
120+ .ok();
121+ if let Some(file) = log_file {
122+ let stderr = Stdio::from(file);
123+ cmd.stdin(Stdio::null())
124+ .stdout(Stdio::null())
125+ .stderr(stderr);
126+ } else {
127+ cmd.stdin(Stdio::null())
128+ .stdout(Stdio::null())
129+ .stderr(Stdio::null());
130+ }
131+ let _child = cmd.spawn()?;
132+ Ok(())
133+}
134+
135+fn detach_from_terminal() {
136+ unsafe {
137+ libc::setsid();
138+ }
139+
140+ if let Ok(devnull) = fs::OpenOptions::new()
141+ .read(true)
142+ .write(true)
143+ .open("/dev/null")
144+ {
145+ let _ = unsafe { libc::dup2(devnull.as_raw_fd(), libc::STDIN_FILENO) };
146+ let _ = unsafe { libc::dup2(devnull.as_raw_fd(), libc::STDOUT_FILENO) };
147+ }
148+}
149+
150+fn wait_for_disconnect(mut stream: UnixStream) {
151+ let mut buf = [0u8; 64];
152+ loop {
153+ match stream.read(&mut buf) {
154+ Ok(0) => break,
155+ Ok(_) => continue,
156+ Err(_) => break,
157+ }
158+ }
159+}
160+
161+fn send_client_pid(stream: &UnixStream) {
162+ let pid = std::process::id();
163+ let payload = format!("pid={pid}\n");
164+ if let Ok(mut stream) = stream.try_clone() {
165+ let _ = stream.write_all(payload.as_bytes());
166+ let _ = stream.flush();
167+ }
168+}
169+
170+fn read_client_pid(stream: &UnixStream) -> Option<u32> {
171+ let mut stream = stream.try_clone().ok()?;
172+ let _ = stream.set_read_timeout(Some(Duration::from_millis(200)));
173+ let mut buf = [0u8; 64];
174+ let mut len = 0usize;
175+ loop {
176+ match stream.read(&mut buf[len..]) {
177+ Ok(0) => break,
178+ Ok(n) => {
179+ len += n;
180+ if buf[..len].iter().any(|b| *b == b'\n') || len == buf.len() {
181+ break;
182+ }
183+ }
184+ Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => break,
185+ Err(err) if err.kind() == std::io::ErrorKind::TimedOut => break,
186+ Err(_) => break,
187+ }
188+ }
189+ let _ = stream.set_read_timeout(None);
190+ if len == 0 {
191+ return None;
192+ }
193+ let line = String::from_utf8_lossy(&buf[..len]);
194+ let trimmed = line.trim();
195+ if let Some(value) = trimmed.strip_prefix("pid=") {
196+ value.parse::<u32>().ok()
197+ } else {
198+ None
199+ }
200+}
201+
202+fn spawn_manager_io(
203+ config: Arc<Mutex<InstanceConfig>>,
204+ instance_dir: PathBuf,
205+ output_monitor: Arc<vm::OutputMonitor>,
206+ vm_output_fd: std::os::unix::io::OwnedFd,
207+ vm_input_fd: std::os::unix::io::OwnedFd,
208+) -> vm::IoContext {
209+ let log_path = instance_dir.join("serial.log");
210+ let log_file = fs::OpenOptions::new()
211+ .create(true)
212+ .append(true)
213+ .open(&log_path)
214+ .ok()
215+ .map(|file| Arc::new(Mutex::new(file)));
216+
217+ let instance_path = instance_dir.join("instance.toml");
218+ let config_for_output = config.clone();
219+ let log_for_output = log_file.clone();
220+ let mut line_buf = String::new();
221+
222+ let on_output = move |bytes: &[u8]| {
223+ if let Some(log) = &log_for_output {
224+ if let Ok(mut file) = log.lock() {
225+ let _ = file.write_all(bytes);
226+ }
227+ }
228+
229+ let text = String::from_utf8_lossy(bytes);
230+ line_buf.push_str(&text);
231+
232+ while let Some(pos) = line_buf.find('\n') {
233+ let mut line = line_buf[..pos].to_string();
234+ if line.ends_with('\r') {
235+ line.pop();
236+ }
237+ line_buf.drain(..=pos);
238+
239+ let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' ');
240+ if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") {
241+ let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..];
242+ let ip = extract_ipv4(ip_raw).unwrap_or_default();
243+ if !ip.is_empty() {
244+ if let Ok(mut cfg) = config_for_output.lock() {
245+ if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) {
246+ cfg.vm_ipv4 = Some(ip.clone());
247+ let _ = write_instance_config(&instance_path, &cfg);
248+ }
249+ }
250+ }
251+ }
252+ }
253+ };
254+
255+ vm::spawn_vm_io_with_hooks(
256+ output_monitor,
257+ vm_output_fd,
258+ vm_input_fd,
259+ vm::IoControl::new(),
260+ |_| false,
261+ on_output,
262+ )
263+}
264+
265+enum ManagerEvent {
266+ Inc(Option<u32>),
267+ Dec(Option<u32>),
268+ VmExited(Option<String>),
269+}
270+
271+#[cfg(test)]
272+mod tests {
273+ use super::*;
274+ use std::{sync::mpsc, time::Duration};
275+
276+ #[test]
277+ fn manager_powers_off_after_grace_when_no_refs() {
278+ let _temp = tempfile::Builder::new()
279+ .prefix("vb")
280+ .tempdir_in("/tmp")
281+ .expect("tempdir");
282+
283+ let (event_tx, event_rx) = mpsc::channel::<ManagerEvent>();
284+ let (vm_tx, vm_rx) = mpsc::channel::<VmInput>();
285+ let vm_input_tx = Arc::new(Mutex::new(Some(vm_tx)));
286+
287+ let manager_thread = thread::spawn(move || {
288+ manager_event_loop(event_rx, vm_input_tx, 50).expect("event loop");
289+ });
290+
291+ event_tx.send(ManagerEvent::Inc(None)).unwrap();
292+ assert!(vm_rx.recv_timeout(Duration::from_millis(100)).is_err());
293+
294+ event_tx.send(ManagerEvent::Dec(None)).unwrap();
295+ let msg = vm_rx
296+ .recv_timeout(Duration::from_secs(2))
297+ .expect("poweroff");
298+ match msg {
299+ VmInput::Bytes(data) => {
300+ assert_eq!(data, b"systemctl poweroff\n");
301+ }
302+ _ => panic!("unexpected vm input"),
303+ }
304+ let _ = event_tx.send(ManagerEvent::VmExited(None));
305+ let _ = manager_thread.join();
306+ }
307+}
308+
309+struct ManagerOptions {
310+ ensure_signed: bool,
311+ detach: bool,
312+ prepare_vm: bool,
313+}
314+
315+trait VmExecutor {
316+ fn run_vm(
317+ &self,
318+ args: vm::CliArgs,
319+ extra_login_actions: Vec<LoginAction>,
320+ extra_shares: Vec<DirectoryShare>,
321+ config: Arc<Mutex<InstanceConfig>>,
322+ instance_dir: PathBuf,
323+ vm_input_tx: Arc<Mutex<Option<mpsc::Sender<VmInput>>>>,
324+ ) -> Result<(), Box<dyn std::error::Error>>;
325+}
326+
327+struct RealVmExecutor;
328+
329+impl VmExecutor for RealVmExecutor {
330+ fn run_vm(
331+ &self,
332+ args: vm::CliArgs,
333+ extra_login_actions: Vec<LoginAction>,
334+ extra_shares: Vec<DirectoryShare>,
335+ config: Arc<Mutex<InstanceConfig>>,
336+ instance_dir: PathBuf,
337+ vm_input_tx: Arc<Mutex<Option<mpsc::Sender<VmInput>>>>,
338+ ) -> Result<(), Box<dyn std::error::Error>> {
339+ let result = vm::run_with_args_and_extras(
340+ args,
341+ |output_monitor, vm_output_fd, vm_input_fd| {
342+ let io_ctx = spawn_manager_io(
343+ config.clone(),
344+ instance_dir.clone(),
345+ output_monitor,
346+ vm_output_fd,
347+ vm_input_fd,
348+ );
349+ *vm_input_tx.lock().unwrap() = Some(io_ctx.input_tx.clone());
350+ io_ctx
351+ },
352+ extra_login_actions,
353+ extra_shares,
354+ );
355+ result
356+ }
357+}
358+
359+fn run_manager_with(
360+ project_root: &Path,
361+ args: vm::CliArgs,
362+ auto_shutdown_ms: u64,
363+ executor: &dyn VmExecutor,
364+ options: ManagerOptions,
365+) -> Result<(), Box<dyn std::error::Error>> {
366+ if options.ensure_signed {
367+ let _had_skip = env::var("VIBEBOX_SKIP_CODESIGN").ok();
368+ unsafe {
369+ env::remove_var("VIBEBOX_SKIP_CODESIGN");
370+ }
371+ vm::ensure_signed();
372+ unsafe {
373+ env::set_var("VIBEBOX_SKIP_CODESIGN", "1");
374+ }
375+ }
376+ if options.detach {
377+ detach_from_terminal();
378+ }
379+
380+ let project_name = project_root
381+ .file_name()
382+ .ok_or("Project directory has no name")?
383+ .to_string_lossy()
384+ .into_owned();
385+ let instance_dir = ensure_instance_dir(project_root)?;
386+ if options.prepare_vm {
387+ let _ = ensure_ssh_keypair(&instance_dir)?;
388+ }
389+
390+ let mut config = load_or_create_instance_config(&instance_dir)?;
391+ if config.vm_ipv4.is_some() {
392+ config.vm_ipv4 = None;
393+ write_instance_config(&instance_dir.join("instance.toml"), &config)?;
394+ }
395+ let config = Arc::new(Mutex::new(config));
396+
397+ let ssh_guest_dir = format!("/root/{}", GLOBAL_DIR_NAME);
398+ let extra_shares = vec![DirectoryShare::new(
399+ instance_dir.clone(),
400+ ssh_guest_dir.clone().into(),
401+ true,
402+ )?];
403+ let extra_login_actions =
404+ build_ssh_login_actions(&config, &project_name, ssh_guest_dir.as_str(), "ssh_key");
405+
406+ let socket_path = instance_dir.join(VM_MANAGER_SOCKET_NAME);
407+ if let Ok(stream) = UnixStream::connect(&socket_path) {
408+ drop(stream);
409+ return Ok(());
410+ }
411+ if socket_path.exists() {
412+ let _ = fs::remove_file(&socket_path);
413+ }
414+
415+ let listener = UnixListener::bind(&socket_path)?;
416+ let _ = fs::set_permissions(&socket_path, fs::Permissions::from_mode(0o600));
417+ tracing::info!(path = %socket_path.display(), "vm manager socket bound");
418+
419+ let (event_tx, event_rx) = mpsc::channel::<ManagerEvent>();
420+ let event_tx_accept = event_tx.clone();
421+ thread::spawn(move || {
422+ for stream in listener.incoming() {
423+ match stream {
424+ Ok(stream) => {
425+ let event_tx_conn = event_tx_accept.clone();
426+ thread::spawn(move || {
427+ let pid = read_client_pid(&stream);
428+ let _ = event_tx_conn.send(ManagerEvent::Inc(pid));
429+ wait_for_disconnect(stream);
430+ let _ = event_tx_conn.send(ManagerEvent::Dec(pid));
431+ });
432+ }
433+ Err(_) => break,
434+ }
435+ }
436+ });
437+
438+ let vm_input_tx: Arc<Mutex<Option<mpsc::Sender<VmInput>>>> = Arc::new(Mutex::new(None));
439+ let vm_input_for_loop = vm_input_tx.clone();
440+ let event_loop_handle =
441+ thread::spawn(move || manager_event_loop(event_rx, vm_input_for_loop, auto_shutdown_ms));
442+
443+ tracing::info!("vm manager launching vm");
444+ let vm_result = executor.run_vm(
445+ args,
446+ extra_login_actions,
447+ extra_shares,
448+ config.clone(),
449+ instance_dir.clone(),
450+ vm_input_tx.clone(),
451+ );
452+ tracing::info!("vm manager vm run completed");
453+ let vm_err = vm_result.err().map(|e| e.to_string());
454+ let _ = event_tx.send(ManagerEvent::VmExited(vm_err.clone()));
455+ let event_loop_result: Result<(), String> = event_loop_handle
456+ .join()
457+ .unwrap_or_else(|_| Err("vm manager event loop panicked".into()))
458+ .map_err(|err| err.to_string());
459+ let _ = fs::remove_file(&socket_path);
460+ if let Err(err) = &event_loop_result {
461+ tracing::error!(error = %err, "vm manager exiting due to event loop error");
462+ return Err(err.to_string().into());
463+ }
464+ if let Some(err) = vm_err {
465+ tracing::error!(error = %err, "vm manager exiting due to vm error");
466+ return Err(err.into());
467+ }
468+ tracing::info!("vm manager exiting");
469+ Ok(event_loop_result?)
470+}
471+
472+fn manager_event_loop(
473+ event_rx: mpsc::Receiver<ManagerEvent>,
474+ vm_input_tx: Arc<Mutex<Option<mpsc::Sender<VmInput>>>>,
475+ auto_shutdown_ms: u64,
476+) -> Result<(), String> {
477+ let mut ref_count: usize = 0;
478+ let mut shutdown_deadline: Option<Instant> = None;
479+ let mut shutdown_sent = false;
480+ let grace = Duration::from_millis(auto_shutdown_ms.max(1));
481+
482+ loop {
483+ let timeout = match shutdown_deadline {
484+ Some(deadline) => deadline.saturating_duration_since(Instant::now()),
485+ None => Duration::from_secs(1),
486+ };
487+
488+ match event_rx.recv_timeout(timeout) {
489+ Ok(ManagerEvent::Inc(pid)) => {
490+ ref_count = ref_count.saturating_add(1);
491+ tracing::info!(
492+ ref_count,
493+ pid = pid.unwrap_or(0),
494+ pid_known = pid.is_some(),
495+ "vm manager refcount increment"
496+ );
497+ shutdown_deadline = None;
498+ shutdown_sent = false;
499+ }
500+ Ok(ManagerEvent::Dec(pid)) => {
501+ if ref_count > 0 {
502+ ref_count -= 1;
503+ }
504+ tracing::info!(
505+ ref_count,
506+ pid = pid.unwrap_or(0),
507+ pid_known = pid.is_some(),
508+ "vm manager refcount decrement"
509+ );
510+ if ref_count == 0 {
511+ shutdown_deadline = Some(Instant::now() + grace);
512+ tracing::info!(grace_ms = auto_shutdown_ms, "shutdown scheduled");
513+ }
514+ }
515+ Ok(ManagerEvent::VmExited(err)) => {
516+ if let Some(err) = err {
517+ tracing::error!(error = %err, "vm exited with error");
518+ }
519+ break;
520+ }
521+ Err(mpsc::RecvTimeoutError::Timeout) => {
522+ if let Some(deadline) = shutdown_deadline {
523+ if Instant::now() >= deadline && !shutdown_sent {
524+ if let Some(tx) = vm_input_tx.lock().unwrap().clone() {
525+ let _ = tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
526+ }
527+ tracing::info!("shutdown command sent");
528+ shutdown_sent = true;
529+ shutdown_deadline = None;
530+ }
531+ }
532+ }
533+ Err(mpsc::RecvTimeoutError::Disconnected) => break,
534+ }
535+ }
536+
537+ Ok(())
538+}