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