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