git.iko.soy/vibebox
on 518ca14 help sign in
vibebox/src/instance.rs
609 lines 21 kB history raw
1use std::{
2 env, fs,
3 io::{self, Write},
4 net::{SocketAddr, TcpStream},
5 os::unix::{fs::PermissionsExt, io::OwnedFd, net::UnixStream},
6 path::{Path, PathBuf},
7 process::{Command, Stdio},
8 sync::{
9 Arc, Mutex,
10 atomic::{AtomicBool, Ordering},
11 },
12 thread,
13 time::{Duration, Instant},
14};
16use serde::{Deserialize, Serialize};
17use std::sync::mpsc::Sender;
18use time::{OffsetDateTime, format_description::well_known::Rfc3339};
19use uuid::Uuid;
21use crate::{
22 commands,
23 session_manager::{INSTANCE_DIR_NAME, INSTANCE_FILENAME},
24 tui::AppState,
25 vm::{self, LoginAction, VmInput},
26};
28const SSH_KEY_NAME: &str = "ssh_key";
29pub(crate) const SERIAL_LOG_NAME: &str = "serial.log";
30pub(crate) const DEFAULT_SSH_USER: &str = "vibecoder";
31const SSH_CONNECT_RETRIES: usize = 30;
32const SSH_CONNECT_DELAY_MS: u64 = 500;
33const SSH_SETUP_SCRIPT: &str = include_str!("ssh.sh");
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub(crate) struct InstanceConfig {
37 #[serde(default)]
38 id: String,
39 #[serde(default = "default_ssh_user")]
40 ssh_user: String,
41 #[serde(default)]
42 sudo_password: String,
43 #[serde(default)]
44 last_active: Option<String>,
45 #[serde(default)]
46 pub(crate) vm_ipv4: Option<String>,
49fn default_ssh_user() -> String {
50 DEFAULT_SSH_USER.to_string()
53pub fn run_with_ssh(manager_conn: UnixStream) -> Result<(), Box<dyn std::error::Error>> {
54 let project_root = env::current_dir()?;
55 tracing::info!(root = %project_root.display(), "starting ssh session");
56 let instance_dir = ensure_instance_dir(&project_root)?;
57 tracing::debug!(instance_dir = %instance_dir.display(), "instance dir ready");
58 let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
60 let config = load_or_create_instance_config(&instance_dir)?;
61 let ssh_user = config.ssh_user.clone();
62 tracing::debug!(ssh_user = %ssh_user, "loaded instance config");
64 let _manager_conn = manager_conn;
65 tracing::debug!("waiting for vm ipv4");
66 wait_for_vm_ipv4(&instance_dir, Duration::from_secs(120))?;
68 let ip = load_or_create_instance_config(&instance_dir)?
69 .vm_ipv4
70 .ok_or("VM IPv4 not available")?;
71 tracing::info!(ip = %ip, "vm ipv4 ready");
73 run_ssh_session(ssh_key, ssh_user, ip)
76pub fn ensure_instance_dir(project_root: &Path) -> Result<PathBuf, io::Error> {
77 let instance_dir = project_root.join(INSTANCE_DIR_NAME);
78 fs::create_dir_all(&instance_dir)?;
79 Ok(instance_dir)
82pub(crate) fn ensure_ssh_keypair(
83 instance_dir: &Path,
84) -> Result<(PathBuf, PathBuf), Box<dyn std::error::Error>> {
85 let private_key = instance_dir.join(SSH_KEY_NAME);
86 let public_key = instance_dir.join(format!("{SSH_KEY_NAME}.pub"));
88 if private_key.exists() && public_key.exists() {
89 return Ok((private_key, public_key));
90 }
92 if private_key.exists() {
93 let _ = fs::remove_file(&private_key);
94 }
95 if public_key.exists() {
96 let _ = fs::remove_file(&public_key);
97 }
99 let status = Command::new("ssh-keygen")
100 .args([
101 "-t",
102 "ed25519",
103 "-N",
104 "",
105 "-f",
106 private_key.to_str().ok_or("ssh key path not utf-8")?,
107 "-C",
108 "vibebox",
109 ])
110 .stdin(Stdio::null())
111 .stdout(Stdio::null())
112 .stderr(Stdio::inherit())
113 .status()?;
115 if !status.success() {
116 return Err("ssh-keygen failed".into());
117 }
119 fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?;
120 fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?;
122 Ok((private_key, public_key))
125pub(crate) fn load_or_create_instance_config(
126 instance_dir: &Path,
127) -> Result<InstanceConfig, Box<dyn std::error::Error>> {
128 let config_path = instance_dir.join(INSTANCE_FILENAME);
129 let mut config = if config_path.exists() {
130 let raw = fs::read_to_string(&config_path)?;
131 toml::from_str::<InstanceConfig>(&raw)?
132 } else {
133 InstanceConfig {
134 id: String::new(),
135 ssh_user: default_ssh_user(),
136 sudo_password: String::new(),
137 last_active: None,
138 vm_ipv4: None,
139 }
140 };
142 let mut changed = false;
143 if config.ssh_user.trim().is_empty() {
144 config.ssh_user = default_ssh_user();
145 changed = true;
146 }
148 if config.id.trim().is_empty() {
149 config.id = Uuid::now_v7().to_string();
150 changed = true;
151 }
153 if config.sudo_password.trim().is_empty() {
154 config.sudo_password = generate_password();
155 changed = true;
156 }
158 if !config_path.exists() || changed {
159 write_instance_config(&config_path, &config)?;
160 }
162 Ok(config)
165fn read_instance_config(
166 instance_dir: &Path,
167) -> Result<Option<InstanceConfig>, Box<dyn std::error::Error>> {
168 let config_path = instance_dir.join(INSTANCE_FILENAME);
169 if !config_path.exists() {
170 return Ok(None);
171 }
172 let raw = fs::read_to_string(&config_path)?;
173 let config = toml::from_str::<InstanceConfig>(&raw)?;
174 Ok(Some(config))
177pub fn read_instance_vm_ip(
178 instance_dir: &Path,
179) -> Result<Option<String>, Box<dyn std::error::Error>> {
180 let config = read_instance_config(instance_dir)?;
181 Ok(config.and_then(|cfg| cfg.vm_ipv4))
184pub fn read_instance_ssh_user(
185 instance_dir: &Path,
186) -> Result<Option<String>, Box<dyn std::error::Error>> {
187 let config = read_instance_config(instance_dir)?;
188 Ok(config
189 .map(|cfg| cfg.ssh_user)
190 .filter(|user| !user.trim().is_empty()))
193pub fn touch_last_active(instance_dir: &Path) -> Result<(), Box<dyn std::error::Error>> {
194 let mut config = load_or_create_instance_config(instance_dir)?;
195 let now = OffsetDateTime::now_utc().format(&Rfc3339)?;
196 config.last_active = Some(now);
197 write_instance_config(&instance_dir.join(INSTANCE_FILENAME), &config)?;
198 Ok(())
201pub(crate) fn write_instance_config(
202 path: &Path,
203 config: &InstanceConfig,
204) -> Result<(), Box<dyn std::error::Error>> {
205 let data = toml::to_string_pretty(config)?;
206 fs::write(path, data)?;
207 fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
208 Ok(())
211fn generate_password() -> String {
212 Uuid::now_v7().simple().to_string()
215pub(crate) fn extract_ipv4(line: &str) -> Option<String> {
216 let mut current = String::new();
217 let mut best: Option<String> = None;
219 for ch in line.chars().chain(std::iter::once(' ')) {
220 if ch.is_ascii_digit() || ch == '.' {
221 current.push(ch);
222 } else if !current.is_empty() {
223 if is_ipv4_candidate(&current) {
224 best = Some(current.clone());
225 break;
226 }
227 current.clear();
228 }
229 }
231 best
234fn wait_for_vm_ipv4(
235 instance_dir: &Path,
236 timeout: Duration,
237) -> Result<(), Box<dyn std::error::Error>> {
238 let start = Instant::now();
239 loop {
240 let config = load_or_create_instance_config(instance_dir)?;
241 if config.vm_ipv4.is_some() {
242 return Ok(());
243 }
244 if start.elapsed() > timeout {
245 return Err("Timed out waiting for VM IPv4".into());
246 }
247 if start.elapsed().as_secs() % 5 == 0 {
248 tracing::debug!("still waiting for vm ipv4");
249 }
250 thread::sleep(Duration::from_millis(200));
251 }
254fn run_ssh_session(
255 ssh_key: PathBuf,
256 ssh_user: String,
257 ip: String,
258) -> Result<(), Box<dyn std::error::Error>> {
259 let mut attempts = 0usize;
260 loop {
261 attempts += 1;
262 if !ssh_port_open(&ip) {
263 tracing::debug!(attempts, "ssh port not open yet");
264 eprintln!(
265 "[vibebox] waiting for ssh port on {ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
266 );
267 if attempts >= SSH_CONNECT_RETRIES {
268 return Err(
269 format!("ssh port not ready after {SSH_CONNECT_RETRIES} attempts").into(),
270 );
271 }
272 thread::sleep(Duration::from_millis(SSH_CONNECT_DELAY_MS));
273 continue;
274 }
276 eprintln!(
277 "[vibebox] starting ssh to {ssh_user}@{ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
278 );
279 tracing::info!(attempts, user = %ssh_user, ip = %ip, "starting ssh");
280 let status = Command::new("ssh")
281 .args([
282 "-i",
283 ssh_key.to_str().unwrap_or(".vibebox/ssh_key"),
284 "-o",
285 "IdentitiesOnly=yes",
286 "-o",
287 "StrictHostKeyChecking=no",
288 "-o",
289 "UserKnownHostsFile=/dev/null",
290 "-o",
291 "GlobalKnownHostsFile=/dev/null",
292 "-o",
293 "PasswordAuthentication=no",
294 "-o",
295 "BatchMode=yes",
296 "-o",
297 "LogLevel=ERROR",
298 "-o",
299 "ConnectTimeout=5",
300 ])
301 .env_remove("LC_CTYPE")
302 .env_remove("LC_ALL")
303 .env_remove("LANG")
304 .arg(format!("{ssh_user}@{ip}"))
305 .stdin(Stdio::inherit())
306 .stdout(Stdio::inherit())
307 .stderr(Stdio::inherit())
308 .status();
310 match status {
311 Ok(status) if status.success() => {
312 eprintln!("[vibebox] ssh exited with status: {status}");
313 break;
314 }
315 Ok(status) if status.code() == Some(255) => {
316 eprintln!("[vibebox] ssh connection failed: {status}");
317 if attempts >= SSH_CONNECT_RETRIES {
318 return Err(format!("ssh failed after {SSH_CONNECT_RETRIES} attempts").into());
319 }
320 thread::sleep(Duration::from_millis(500));
321 }
322 Ok(status) => {
323 eprintln!("[vibebox] ssh exited with status: {status}");
324 break;
325 }
326 Err(err) => {
327 return Err(format!("failed to start ssh: {err}").into());
328 }
329 }
330 }
332 Ok(())
335fn is_ipv4_candidate(candidate: &str) -> bool {
336 let parts: Vec<&str> = candidate.split('.').collect();
337 if parts.len() != 4 {
338 return false;
339 }
340 for part in parts {
341 if part.is_empty() || part.len() > 3 {
342 return false;
343 }
344 if part.parse::<u8>().is_err() {
345 return false;
346 }
347 }
348 true
351fn ssh_port_open(ip: &str) -> bool {
352 let addr: SocketAddr = match format!("{ip}:22").parse() {
353 Ok(addr) => addr,
354 Err(_) => return false,
355 };
356 TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok()
359pub(crate) fn build_ssh_login_actions(
360 config: &Arc<Mutex<InstanceConfig>>,
361 project_name: &str,
362 guest_dir: &str,
363 key_name: &str,
364) -> Vec<LoginAction> {
365 let config_guard = config.lock().expect("config mutex poisoned");
366 let ssh_user = config_guard.ssh_user.clone();
367 let sudo_password = config_guard.sudo_password.clone();
368 drop(config_guard);
370 let key_path = format!("{guest_dir}/{key_name}.pub");
372 let setup_script = SSH_SETUP_SCRIPT
373 .replace("__SSH_USER__", &ssh_user)
374 .replace("__SUDO_PASSWORD__", &sudo_password)
375 .replace("__PROJECT_NAME__", project_name)
376 .replace("__KEY_PATH__", &key_path)
377 .replace("__VIBEBOX_SHELL_SCRIPT__", &commands::render_shell_script());
378 let setup = vm::script_command_from_content("ssh_setup", &setup_script)
379 .expect("ssh setup script contained invalid marker");
381 vec![LoginAction::Send(setup)]
384fn spawn_ssh_io(
385 app: Arc<Mutex<AppState>>,
386 config: Arc<Mutex<InstanceConfig>>,
387 instance_dir: PathBuf,
388 ssh_key: PathBuf,
389 output_monitor: Arc<vm::OutputMonitor>,
390 vm_output_fd: OwnedFd,
391 vm_input_fd: OwnedFd,
392) -> vm::IoContext {
393 let io_control = vm::IoControl::new();
395 let log_path = instance_dir.join(SERIAL_LOG_NAME);
396 let log_file = fs::OpenOptions::new()
397 .create(true)
398 .write(true)
399 .truncate(true)
400 .open(&log_path)
401 .ok()
402 .map(|file| Arc::new(Mutex::new(file)));
404 let ssh_connected = Arc::new(AtomicBool::new(false));
405 let ssh_started = Arc::new(AtomicBool::new(false));
406 let ssh_ready = Arc::new(AtomicBool::new(false));
407 let input_tx_holder: Arc<Mutex<Option<Sender<VmInput>>>> = Arc::new(Mutex::new(None));
409 let instance_path = instance_dir.join(INSTANCE_FILENAME);
410 let config_for_output = config.clone();
411 let log_for_output = log_file.clone();
412 let ssh_connected_for_output = ssh_connected.clone();
413 let ssh_ready_for_output = ssh_ready.clone();
415 let mut line_buf = String::new();
417 let handlers = commands::build_handlers(app.clone(), io_control.clone());
418 let on_line = move |line: &str| handlers.handle(line);
420 let on_output = move |bytes: &[u8]| {
421 if ssh_connected_for_output.load(Ordering::SeqCst) {
422 if let Some(log) = &log_for_output {
423 if let Ok(mut file) = log.lock() {
424 let _ = file.write_all(bytes);
425 }
426 }
427 }
429 let text = String::from_utf8_lossy(bytes);
430 line_buf.push_str(&text);
432 while let Some(pos) = line_buf.find('\n') {
433 let mut line = line_buf[..pos].to_string();
434 if line.ends_with('\r') {
435 line.pop();
436 }
437 line_buf.drain(..=pos);
439 let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' ');
441 if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") {
442 let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..];
443 let ip = extract_ipv4(ip_raw).unwrap_or_default();
444 if !ip.is_empty() {
445 if let Ok(mut cfg) = config_for_output.lock() {
446 if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) {
447 cfg.vm_ipv4 = Some(ip.clone());
448 let _ = write_instance_config(&instance_path, &cfg);
449 eprintln!("[vibebox] detected vm IPv4: {ip}");
450 }
451 }
452 }
453 }
455 if cleaned.contains("VIBEBOX_SSH_READY") {
456 ssh_ready_for_output.store(true, Ordering::SeqCst);
457 eprintln!("[vibebox] sshd ready");
458 }
459 }
460 };
462 let io_ctx = vm::spawn_vm_io_with_hooks(
463 output_monitor,
464 vm_output_fd,
465 vm_input_fd,
466 io_control.clone(),
467 on_line,
468 on_output,
469 );
471 *input_tx_holder.lock().unwrap() = Some(io_ctx.input_tx.clone());
473 let ssh_ready_for_thread = ssh_ready.clone();
474 let ssh_started_for_thread = ssh_started.clone();
475 let ssh_connected_for_thread = ssh_connected.clone();
476 let io_control_for_thread = io_control.clone();
477 let ssh_key_for_thread = ssh_key.clone();
478 let config_for_thread = config.clone();
479 let input_tx_holder_for_thread = input_tx_holder.clone();
481 thread::spawn(move || {
482 loop {
483 if ssh_started_for_thread.load(Ordering::SeqCst) {
484 break;
485 }
486 if !ssh_ready_for_thread.load(Ordering::SeqCst) {
487 thread::sleep(std::time::Duration::from_millis(200));
488 continue;
489 }
491 let ip = config_for_thread
492 .lock()
493 .ok()
494 .and_then(|cfg| cfg.vm_ipv4.clone());
495 let ssh_user = config_for_thread
496 .lock()
497 .map(|cfg| cfg.ssh_user.clone())
498 .unwrap_or_else(|_| DEFAULT_SSH_USER.to_string());
500 if let Some(ip) = ip {
501 if ssh_started_for_thread
502 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
503 .is_err()
504 {
505 break;
506 }
508 eprintln!("[vibebox] starting ssh to {ssh_user}@{ip}");
509 io_control_for_thread.request_terminal_restore();
510 io_control_for_thread.set_forward_output(false);
511 io_control_for_thread.set_forward_input(false);
512 ssh_connected_for_thread.store(true, Ordering::SeqCst);
514 let mut attempts = 0usize;
515 loop {
516 attempts += 1;
517 if !ssh_port_open(&ip) {
518 eprintln!(
519 "[vibebox] waiting for ssh port on {ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
520 );
521 if attempts >= SSH_CONNECT_RETRIES {
522 ssh_connected_for_thread.store(false, Ordering::SeqCst);
523 eprintln!(
524 "[vibebox] ssh port not ready after {SSH_CONNECT_RETRIES} attempts"
525 );
526 break;
527 }
528 thread::sleep(std::time::Duration::from_millis(SSH_CONNECT_DELAY_MS));
529 continue;
530 }
532 eprintln!(
533 "[vibebox] starting ssh to {ssh_user}@{ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
534 );
535 let status = Command::new("ssh")
536 .args([
537 "-i",
538 ssh_key_for_thread.to_str().unwrap_or(".vibebox/ssh_key"),
539 "-o",
540 "IdentitiesOnly=yes",
541 "-o",
542 "StrictHostKeyChecking=no",
543 "-o",
544 "UserKnownHostsFile=/dev/null",
545 "-o",
546 "GlobalKnownHostsFile=/dev/null",
547 "-o",
548 "PasswordAuthentication=no",
549 "-o",
550 "BatchMode=yes",
551 "-o",
552 "LogLevel=ERROR",
553 "-o",
554 "ConnectTimeout=5",
555 ])
556 .env_remove("LC_CTYPE")
557 .env_remove("LC_ALL")
558 // .env_remove("LANG")
559 .arg(format!("{ssh_user}@{ip}"))
560 .stdin(Stdio::inherit())
561 .stdout(Stdio::inherit())
562 .stderr(Stdio::inherit())
563 .status();
565 match status {
566 Ok(status) if status.success() => {
567 ssh_connected_for_thread.store(false, Ordering::SeqCst);
568 eprintln!("[vibebox] ssh exited with status: {status}");
569 if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
570 let _ = tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
571 }
572 break;
573 }
574 Ok(status) if status.code() == Some(255) => {
575 eprintln!("[vibebox] ssh connection failed: {status}");
576 if attempts >= SSH_CONNECT_RETRIES {
577 ssh_connected_for_thread.store(false, Ordering::SeqCst);
578 eprintln!(
579 "[vibebox] ssh failed after {SSH_CONNECT_RETRIES} attempts"
580 );
581 break;
582 }
583 thread::sleep(std::time::Duration::from_millis(500));
584 continue;
585 }
586 Ok(status) => {
587 ssh_connected_for_thread.store(false, Ordering::SeqCst);
588 eprintln!("[vibebox] ssh exited with status: {status}");
589 if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
590 let _ = tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
591 }
592 break;
593 }
594 Err(err) => {
595 ssh_connected_for_thread.store(false, Ordering::SeqCst);
596 eprintln!("[vibebox] failed to start ssh: {err}");
597 break;
598 }
599 }
600 }
601 break;
602 }
604 thread::sleep(std::time::Duration::from_millis(200));
605 }
606 });
608 io_ctx