git.iko.soy/vibebox
on main help sign in
vibebox/src/instance.rs
580 lines 19 kB history raw
1use anyhow::{Context, Result, bail};
2use serde::{Deserialize, Serialize};
3use std::os::unix::fs::FileTypeExt;
4use std::{
5 env, fs,
6 io::{self, IsTerminal, Read},
7 net::{SocketAddr, TcpStream},
8 os::unix::{fs::PermissionsExt, net::UnixStream},
9 path::{Path, PathBuf},
10 process::{Child, Command, Stdio},
11 sync::{
12 Arc, Mutex,
13 atomic::{AtomicBool, Ordering},
14 mpsc,
15 },
16 thread,
17 time::{Duration, Instant},
18};
19use time::{OffsetDateTime, format_description::well_known::Rfc3339};
20use uuid::Uuid;
22use crate::{
23 commands,
24 session_manager::INSTANCE_DIR_NAME,
25 session_manager::{VM_MANAGER_PID_NAME, VM_MANAGER_SOCKET_NAME},
26 vm::{self, LoginAction},
27};
29pub const STATUS_VM_ERROR_PREFIX: &str = "error:";
31const SSH_KEY_NAME: &str = "ssh_key";
32const INSTANCE_FILENAME: &str = "instance.toml";
33const DEFAULT_SSH_USER: &str = "vibecoder";
34const SSH_CONNECT_RETRIES: usize = 10;
35const SSH_CONNECT_DELAY_MS: u64 = 500;
36const SSH_SETUP_SCRIPT: &str = include_str!("ssh.sh");
37const STATUS_PREFIX: &str = "status:";
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum VmLiveness {
41 RunningWithSocket { pid: u32 },
42 RunningWithoutSocket { pid: u32 },
43 NotRunningOrMissing,
46pub fn vm_liveness(project_root: &Path) -> Result<VmLiveness> {
47 let instance_dir = ensure_instance_dir(project_root)?;
48 let pid_path = instance_dir.join(VM_MANAGER_PID_NAME);
49 let socket_path = instance_dir.join(VM_MANAGER_SOCKET_NAME);
50 fn pid_is_alive(pid: u32) -> bool {
51 let pid = pid as libc::pid_t;
52 let result = unsafe { libc::kill(pid, 0) };
53 if result == 0 {
54 return true;
55 }
56 match io::Error::last_os_error().raw_os_error() {
57 Some(code) if code == libc::EPERM => true,
58 Some(code) if code == libc::ESRCH => false,
59 _ => false,
60 }
61 }
62 let Ok(content) = fs::read_to_string(pid_path) else {
63 return Ok(VmLiveness::NotRunningOrMissing);
64 };
65 let Ok(pid) = content.trim().parse::<u32>() else {
66 return Ok(VmLiveness::NotRunningOrMissing);
67 };
68 if !pid_is_alive(pid) {
69 return Ok(VmLiveness::NotRunningOrMissing);
70 }
71 let has_socket = fs::metadata(socket_path)
72 .map(|meta| meta.file_type().is_socket())
73 .unwrap_or(false);
74 if has_socket {
75 Ok(VmLiveness::RunningWithSocket { pid })
76 } else {
77 Ok(VmLiveness::RunningWithoutSocket { pid })
78 }
81#[derive(Debug, thiserror::Error)]
82pub enum InstanceError {
83 #[error("unexpected disconnection from vm manager")]
84 UnexpectedDisconnection,
85 #[error("{0}")]
86 VMError(String),
89fn default_ssh_user() -> String {
90 DEFAULT_SSH_USER.to_string()
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct InstanceConfig {
95 #[serde(default)]
96 pub id: String,
97 #[serde(default = "default_ssh_user")]
98 pub ssh_user: String,
99 #[serde(default)]
100 sudo_password: String,
101 #[serde(default)]
102 pub last_active: Option<String>,
103 #[serde(default)]
104 pub vm_ipv4: Option<String>,
107impl Default for InstanceConfig {
108 fn default() -> Self {
109 Self {
110 id: Uuid::now_v7().to_string(),
111 ssh_user: DEFAULT_SSH_USER.to_string(),
112 sudo_password: Uuid::now_v7().simple().to_string(),
113 last_active: None,
114 vm_ipv4: None,
115 }
116 }
119pub fn run_with_ssh(manager_conn: UnixStream) -> Result<()> {
120 let project_root = env::current_dir()?;
121 tracing::info!(root = %project_root.display(), "starting ssh session");
122 let instance_dir = ensure_instance_dir(&project_root)?;
123 tracing::debug!(instance_dir = %instance_dir.display(), "instance dir ready");
124 let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
126 let config = load_or_create_instance_config(&project_root)?;
127 let ssh_user = config.ssh_user.clone();
128 tracing::debug!(ssh_user = %ssh_user, "loaded instance config");
130 wait_for_vm_ipv4(&project_root, Duration::from_secs(480), &manager_conn)?;
132 let ip = load_or_create_instance_config(&project_root)?
133 .vm_ipv4
134 .with_context(|| "failed to load instance IP address")?;
135 tracing::info!(ip = %ip, "vm ipv4 ready");
137 run_ssh_session(ssh_key, ssh_user, ip, manager_conn, project_root)
140pub fn ensure_instance_dir(project_root: &Path) -> Result<PathBuf, io::Error> {
141 let instance_dir = project_root.join(INSTANCE_DIR_NAME);
142 fs::create_dir_all(&instance_dir)?;
143 Ok(instance_dir)
146pub fn ensure_ssh_keypair(instance_dir: &Path) -> Result<(PathBuf, PathBuf)> {
147 let private_key = instance_dir.join(SSH_KEY_NAME);
148 let public_key = instance_dir.join(format!("{SSH_KEY_NAME}.pub"));
150 if private_key.exists() && public_key.exists() {
151 return Ok((private_key, public_key));
152 }
154 if private_key.exists() {
155 let _ = fs::remove_file(&private_key);
156 }
157 if public_key.exists() {
158 let _ = fs::remove_file(&public_key);
159 }
161 let status = Command::new("ssh-keygen")
162 .args([
163 "-t",
164 "ed25519",
165 "-N",
166 "",
167 "-f",
168 private_key
169 .to_str()
170 .with_context(|| "ssh key path not utf-8")?,
171 "-C",
172 "vibebox",
173 ])
174 .stdin(Stdio::null())
175 .stdout(Stdio::inherit())
176 .stderr(Stdio::inherit())
177 .status()?;
179 if !status.success() {
180 bail!("ssh-keygen failed");
181 }
183 fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?;
184 fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?;
186 Ok((private_key, public_key))
189pub fn load_or_create_instance_config(project_dir: &Path) -> Result<InstanceConfig> {
190 let mut exist = true;
191 let mut config = read_instance_config(project_dir).unwrap_or_else(|_| {
192 exist = false;
193 InstanceConfig::default()
194 });
196 let mut changed = false;
197 if config.ssh_user.trim().is_empty() {
198 config.ssh_user = InstanceConfig::default().ssh_user;
199 changed = true;
200 }
202 if config.id.trim().is_empty() {
203 config.id = InstanceConfig::default().id;
204 changed = true;
205 }
207 if config.sudo_password.trim().is_empty() {
208 config.sudo_password = InstanceConfig::default().sudo_password;
209 changed = true;
210 }
212 if !exist || changed {
213 write_instance_config(project_dir, &config)?;
214 }
216 Ok(config)
219pub fn read_instance_config(project_dir: &Path) -> Result<InstanceConfig> {
220 // todo maybe verify schema?
221 let config_path = project_dir.join(INSTANCE_DIR_NAME).join(INSTANCE_FILENAME);
222 if !config_path.exists() {
223 bail!("instance config file does not exist");
224 }
225 let raw = fs::read_to_string(&config_path)?;
226 let config = toml::from_str::<InstanceConfig>(&raw)?;
227 Ok(config)
230pub fn touch_last_active(project_dir: &Path) -> Result<()> {
231 let mut config = load_or_create_instance_config(project_dir)?;
232 let now = OffsetDateTime::now_utc().format(&Rfc3339)?;
233 config.last_active = Some(now);
234 write_instance_config(project_dir, &config)?;
235 Ok(())
238pub fn write_instance_config(project_dir: &Path, config: &InstanceConfig) -> Result<()> {
239 let path = project_dir.join(INSTANCE_DIR_NAME).join(INSTANCE_FILENAME);
240 let data = toml::to_string_pretty(config)?;
241 fs::create_dir_all(project_dir.join(INSTANCE_DIR_NAME))?;
242 fs::write(&path, data)?;
243 fs::set_permissions(&path, fs::Permissions::from_mode(0o600))?;
244 Ok(())
247#[cfg_attr(feature = "mock-vm", allow(dead_code))]
248pub fn extract_ipv4(line: &str) -> Option<String> {
249 let mut current = String::new();
250 let mut best: Option<String> = None;
252 for ch in line.chars().chain(std::iter::once(' ')) {
253 if ch.is_ascii_digit() || ch == '.' {
254 current.push(ch);
255 } else if !current.is_empty() {
256 if is_ipv4_candidate(&current) {
257 best = Some(current.clone());
258 break;
259 }
260 current.clear();
261 }
262 }
264 best
267fn handle_manager_line(line: &str, last_status: &mut Option<String>) -> Result<()> {
268 if let Some(status) = line.strip_prefix(STATUS_PREFIX) {
269 let status = status.trim();
270 if let Some(message) = status.strip_prefix(STATUS_VM_ERROR_PREFIX) {
271 let message = message.trim();
272 if message.is_empty() {
273 bail!("vm manager reported startup failure");
274 }
275 return Err(InstanceError::VMError(message.to_string()).into());
276 }
277 if !status.is_empty() && last_status.as_deref() != Some(status) {
278 tracing::info!("[background]: {}", status);
279 *last_status = Some(status.to_string());
280 }
281 }
282 Ok(())
285fn wait_for_vm_ipv4(
286 project_dir: &Path,
287 timeout: Duration,
288 manager_conn: &UnixStream,
289) -> Result<()> {
290 let start = Instant::now();
291 let mut next_log_at = start + Duration::from_secs(10);
292 let mut stream = manager_conn.try_clone()?;
293 let _ = stream.set_read_timeout(Some(Duration::from_millis(250)));
294 let mut read_buf = [0u8; 1024];
295 let mut pending = String::new();
296 tracing::info!("waiting for vm ipv4");
297 let mut last_status: Option<String> = None;
298 let mut once_hint = false;
299 loop {
300 match stream.read(&mut read_buf) {
301 Ok(0) => {
302 bail!("vm manager disconnected before VM became ready");
303 }
304 Ok(n) => {
305 pending.push_str(&String::from_utf8_lossy(&read_buf[..n]));
306 while let Some(pos) = pending.find('\n') {
307 let line = pending[..pos].trim().to_string();
308 pending.drain(..=pos);
309 handle_manager_line(&line, &mut last_status)?;
310 }
311 }
312 Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
313 Err(err) if err.kind() == io::ErrorKind::TimedOut => {}
314 Err(err) if err.kind() == io::ErrorKind::Interrupted => {}
315 Err(err) => {
316 tracing::warn!(error = %err, "failed to read vm manager status stream");
317 }
318 }
320 let config = load_or_create_instance_config(project_dir)?;
321 if config.vm_ipv4.is_some() {
322 return Ok(());
323 }
324 if start.elapsed() > timeout {
325 bail!("timed out waiting for VM IPv4");
326 }
328 let now = Instant::now();
329 if now >= next_log_at {
330 let waited = start.elapsed();
331 if waited.as_secs() > 15 && !once_hint {
332 tracing::info!(
333 "if vibebox is just initialized in this directory, it might take up to 1 minute depending on your machine, and then you can enjoy secure & speed vibecoding! go pack!"
334 );
335 once_hint = true;
336 }
337 tracing::info!("still waiting for vm ipv4, {}s elapsed", waited.as_secs(),);
338 next_log_at += Duration::from_secs(20);
339 }
340 }
343fn run_ssh_session(
344 ssh_key: PathBuf,
345 ssh_user: String,
346 ip: String,
347 manager_conn: UnixStream,
348 project_root: PathBuf,
349) -> Result<()> {
350 let mut attempts = 0usize;
351 loop {
352 if matches!(vm_liveness(&project_root)?, VmLiveness::NotRunningOrMissing) {
353 return Err(InstanceError::UnexpectedDisconnection.into());
354 }
355 attempts += 1;
356 if !ssh_port_open(&ip) {
357 tracing::debug!(attempts, "ssh port doesn't open yet");
358 tracing::info!(
359 attempts,
360 ip = %ip,
361 "waiting for ssh port ({}/{})",
362 attempts,
363 SSH_CONNECT_RETRIES
364 );
365 if attempts >= SSH_CONNECT_RETRIES {
366 bail!("ssh port not ready after {SSH_CONNECT_RETRIES} attempts");
367 }
368 thread::sleep(Duration::from_millis(SSH_CONNECT_DELAY_MS));
369 continue;
370 }
372 tracing::info!(
373 attempts,
374 user = %ssh_user,
375 ip = %ip,
376 "starting ssh ({}/{})",
377 attempts,
378 SSH_CONNECT_RETRIES
379 );
380 let child = Command::new("ssh")
381 .args([
382 "-i",
383 ssh_key.to_str().with_context(|| "invalid path")?,
384 "-o",
385 "IdentitiesOnly=yes",
386 "-o",
387 "StrictHostKeyChecking=no",
388 "-o",
389 "UserKnownHostsFile=/dev/null",
390 "-o",
391 "GlobalKnownHostsFile=/dev/null",
392 "-o",
393 "PasswordAuthentication=no",
394 "-o",
395 "BatchMode=yes",
396 "-o",
397 "LogLevel=ERROR",
398 "-o",
399 "ConnectTimeout=5",
400 ])
401 .env_remove("LC_CTYPE")
402 .env_remove("LC_ALL")
403 .env_remove("LANG")
404 .arg(format!("{ssh_user}@{ip}"))
405 .stdin(Stdio::inherit())
406 .stdout(Stdio::inherit())
407 .stderr(Stdio::inherit())
408 .spawn();
410 match child {
411 Ok(mut child) => {
412 let done = Arc::new(AtomicBool::new(false));
413 let done_for_monitor = done.clone();
414 let (disconnect_tx, disconnect_rx) = mpsc::channel::<()>();
415 let mut manager_stream = manager_conn.try_clone()?;
416 let _ = manager_stream.set_read_timeout(Some(Duration::from_millis(250)));
417 thread::spawn(move || {
418 let mut buf = [0u8; 1];
419 while !done_for_monitor.load(Ordering::Relaxed) {
420 match manager_stream.read(&mut buf) {
421 Ok(0) => {
422 let _ = disconnect_tx.send(());
423 return;
424 }
425 Ok(_) => {}
426 Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
427 Err(err) if err.kind() == io::ErrorKind::TimedOut => {}
428 Err(err) if err.kind() == io::ErrorKind::Interrupted => {}
429 Err(_) => {
430 let _ = disconnect_tx.send(());
431 return;
432 }
433 }
434 }
435 });
437 let status = loop {
438 if disconnect_rx.try_recv().is_ok() {
439 done.store(true, Ordering::Relaxed);
440 terminate_ssh_child(&mut child);
441 restore_terminal_after_disconnect();
442 return Err(InstanceError::UnexpectedDisconnection.into());
443 }
444 if let Some(status) = child.try_wait()? {
445 done.store(true, Ordering::Relaxed);
446 break status;
447 }
448 thread::sleep(Duration::from_millis(100));
449 };
451 if status.success() {
452 tracing::info!(status = %status, "ssh exited");
453 break;
454 }
455 if status.code() == Some(255) {
456 tracing::warn!(status = %status, "ssh connection failed");
457 if attempts >= SSH_CONNECT_RETRIES {
458 bail!("ssh failed after {SSH_CONNECT_RETRIES} attempts");
459 }
460 thread::sleep(Duration::from_millis(500));
461 continue;
462 }
463 tracing::info!(status = %status, "ssh exited");
464 break;
465 }
466 Err(err) => {
467 tracing::error!(error = %err, "failed to start ssh");
468 bail!("failed to start ssh: {err}");
469 }
470 }
471 }
473 Ok(())
476fn terminate_ssh_child(child: &mut Child) {
477 let pid = child.id() as i32;
478 unsafe {
479 libc::kill(pid, libc::SIGTERM);
480 }
481 let deadline = Instant::now() + Duration::from_millis(700);
482 while Instant::now() < deadline {
483 match child.try_wait() {
484 Ok(Some(_)) => return,
485 Ok(None) => thread::sleep(Duration::from_millis(50)),
486 Err(_) => break,
487 }
488 }
489 let _ = child.kill();
490 let _ = child.wait();
493fn restore_terminal_after_disconnect() {
494 if io::stdin().is_terminal() {
495 let _ = Command::new("stty").arg("sane").status();
496 }
497 eprintln!();
500#[cfg_attr(feature = "mock-vm", allow(dead_code))]
501fn is_ipv4_candidate(candidate: &str) -> bool {
502 let parts: Vec<&str> = candidate.split('.').collect();
503 if parts.len() != 4 {
504 return false;
505 }
506 for part in parts {
507 if part.is_empty() || part.len() > 3 {
508 return false;
509 }
510 if part.parse::<u8>().is_err() {
511 return false;
512 }
513 }
514 true
517fn ssh_port_open(ip: &str) -> bool {
518 let addr: SocketAddr = match format!("{ip}:22").parse() {
519 Ok(addr) => addr,
520 Err(_) => return false,
521 };
522 TcpStream::connect_timeout(&addr, Duration::from_millis(500)).is_ok()
525pub fn build_ssh_login_actions(
526 config: &Arc<Mutex<InstanceConfig>>,
527 project_name: &str,
528 project_guest_dir: &str,
529 guest_dir: &str,
530 key_name: &str,
531 home_links_script: &str,
532) -> Vec<LoginAction> {
533 let config_guard = config.lock().expect("config mutex poisoned");
534 let ssh_user = config_guard.ssh_user.clone();
535 let sudo_password = config_guard.sudo_password.clone();
536 drop(config_guard);
538 let key_path = format!("{guest_dir}/{key_name}.pub");
540 let ssh_script = SSH_SETUP_SCRIPT
541 .replace("__SSH_USER__", &ssh_user)
542 .replace("__SUDO_PASSWORD__", &sudo_password)
543 .replace("__PROJECT_NAME__", project_name)
544 .replace("__PROJECT_GUEST_DIR__", project_guest_dir)
545 .replace("__KEY_PATH__", &key_path)
546 .replace("__VIBEBOX_SHELL_SCRIPT__", &commands::render_shell_script())
547 .replace("__VIBEBOX_HOME_LINKS__", home_links_script);
548 let setup = vm::script_command_from_content("ssh.sh", &ssh_script)
549 .expect("ssh setup script contained invalid marker");
551 vec![LoginAction::Send(setup)]
554#[cfg(test)]
555mod tests {
556 use super::*;
558 #[test]
559 fn handle_manager_line_updates_status() {
560 let mut last_status = None;
561 handle_manager_line("status: preparing VM image...", &mut last_status)
562 .expect("status should be accepted");
563 assert_eq!(last_status.as_deref(), Some("preparing VM image..."));
564 }
566 #[test]
567 fn handle_manager_line_ignores_non_status_lines() {
568 let mut last_status = None;
569 handle_manager_line("pid=123", &mut last_status).expect("non-status lines are ignored");
570 assert!(last_status.is_none());
571 }
573 #[test]
574 fn handle_manager_line_surfaces_error_status() {
575 let mut last_status = None;
576 let err = handle_manager_line("status: error: vm failed to boot", &mut last_status)
577 .expect_err("error status should fail");
578 assert_eq!(err.to_string(), "vm failed to boot");
579 }