git.iko.soy/vibebox
on 2ef4011 help sign in
vibebox/src/instance.rs
653 lines 22 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 VM_ROOT_LOG_NAME: &str = "vm_root.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>,
49impl InstanceConfig {
50 pub(crate) fn ssh_user_display(&self) -> String {
51 if self.ssh_user.trim().is_empty() {
52 DEFAULT_SSH_USER.to_string()
53 } else {
54 self.ssh_user.clone()
55 }
56 }
59fn default_ssh_user() -> String {
60 DEFAULT_SSH_USER.to_string()
63pub fn run_with_ssh(manager_conn: UnixStream) -> Result<(), Box<dyn std::error::Error>> {
64 let project_root = env::current_dir()?;
65 tracing::info!(root = %project_root.display(), "starting ssh session");
66 let instance_dir = ensure_instance_dir(&project_root)?;
67 tracing::debug!(instance_dir = %instance_dir.display(), "instance dir ready");
68 let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
70 let config = load_or_create_instance_config(&instance_dir)?;
71 let ssh_user = config.ssh_user.clone();
72 tracing::debug!(ssh_user = %ssh_user, "loaded instance config");
74 let _manager_conn = manager_conn;
75 wait_for_vm_ipv4(&instance_dir, Duration::from_secs(120))?;
77 let ip = load_or_create_instance_config(&instance_dir)?
78 .vm_ipv4
79 .ok_or("VM IPv4 not available")?;
80 tracing::info!(ip = %ip, "vm ipv4 ready");
82 run_ssh_session(ssh_key, ssh_user, ip)
85pub fn ensure_instance_dir(project_root: &Path) -> Result<PathBuf, io::Error> {
86 let instance_dir = project_root.join(INSTANCE_DIR_NAME);
87 fs::create_dir_all(&instance_dir)?;
88 Ok(instance_dir)
91pub(crate) fn ensure_ssh_keypair(
92 instance_dir: &Path,
93) -> Result<(PathBuf, PathBuf), Box<dyn std::error::Error>> {
94 let private_key = instance_dir.join(SSH_KEY_NAME);
95 let public_key = instance_dir.join(format!("{SSH_KEY_NAME}.pub"));
97 if private_key.exists() && public_key.exists() {
98 return Ok((private_key, public_key));
99 }
101 if private_key.exists() {
102 let _ = fs::remove_file(&private_key);
103 }
104 if public_key.exists() {
105 let _ = fs::remove_file(&public_key);
106 }
108 let status = Command::new("ssh-keygen")
109 .args([
110 "-t",
111 "ed25519",
112 "-N",
113 "",
114 "-f",
115 private_key.to_str().ok_or("ssh key path not utf-8")?,
116 "-C",
117 "vibebox",
118 ])
119 .stdin(Stdio::null())
120 .stdout(Stdio::inherit())
121 .stderr(Stdio::inherit())
122 .status()?;
124 if !status.success() {
125 return Err("ssh-keygen failed".into());
126 }
128 fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?;
129 fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?;
131 Ok((private_key, public_key))
134pub(crate) fn load_or_create_instance_config(
135 instance_dir: &Path,
136) -> Result<InstanceConfig, Box<dyn std::error::Error>> {
137 let config_path = instance_dir.join(INSTANCE_FILENAME);
138 let mut config = if config_path.exists() {
139 let raw = fs::read_to_string(&config_path)?;
140 toml::from_str::<InstanceConfig>(&raw)?
141 } else {
142 InstanceConfig {
143 id: String::new(),
144 ssh_user: default_ssh_user(),
145 sudo_password: String::new(),
146 last_active: None,
147 vm_ipv4: None,
148 }
149 };
151 let mut changed = false;
152 if config.ssh_user.trim().is_empty() {
153 config.ssh_user = default_ssh_user();
154 changed = true;
155 }
157 if config.id.trim().is_empty() {
158 config.id = Uuid::now_v7().to_string();
159 changed = true;
160 }
162 if config.sudo_password.trim().is_empty() {
163 config.sudo_password = generate_password();
164 changed = true;
165 }
167 if !config_path.exists() || changed {
168 write_instance_config(&config_path, &config)?;
169 }
171 Ok(config)
174fn read_instance_config(
175 instance_dir: &Path,
176) -> Result<Option<InstanceConfig>, Box<dyn std::error::Error>> {
177 let config_path = instance_dir.join(INSTANCE_FILENAME);
178 if !config_path.exists() {
179 return Ok(None);
180 }
181 let raw = fs::read_to_string(&config_path)?;
182 let config = toml::from_str::<InstanceConfig>(&raw)?;
183 Ok(Some(config))
186pub fn read_instance_vm_ip(
187 instance_dir: &Path,
188) -> Result<Option<String>, Box<dyn std::error::Error>> {
189 let config = read_instance_config(instance_dir)?;
190 Ok(config.and_then(|cfg| cfg.vm_ipv4))
193pub fn read_instance_ssh_user(
194 instance_dir: &Path,
195) -> Result<Option<String>, Box<dyn std::error::Error>> {
196 let config = read_instance_config(instance_dir)?;
197 Ok(config
198 .map(|cfg| cfg.ssh_user)
199 .filter(|user| !user.trim().is_empty()))
202pub fn touch_last_active(instance_dir: &Path) -> Result<(), Box<dyn std::error::Error>> {
203 let mut config = load_or_create_instance_config(instance_dir)?;
204 let now = OffsetDateTime::now_utc().format(&Rfc3339)?;
205 config.last_active = Some(now);
206 write_instance_config(&instance_dir.join(INSTANCE_FILENAME), &config)?;
207 Ok(())
210pub(crate) fn write_instance_config(
211 path: &Path,
212 config: &InstanceConfig,
213) -> Result<(), Box<dyn std::error::Error>> {
214 let data = toml::to_string_pretty(config)?;
215 fs::write(path, data)?;
216 fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
217 Ok(())
220fn generate_password() -> String {
221 Uuid::now_v7().simple().to_string()
224pub(crate) fn extract_ipv4(line: &str) -> Option<String> {
225 let mut current = String::new();
226 let mut best: Option<String> = None;
228 for ch in line.chars().chain(std::iter::once(' ')) {
229 if ch.is_ascii_digit() || ch == '.' {
230 current.push(ch);
231 } else if !current.is_empty() {
232 if is_ipv4_candidate(&current) {
233 best = Some(current.clone());
234 break;
235 }
236 current.clear();
237 }
238 }
240 best
243fn wait_for_vm_ipv4(
244 instance_dir: &Path,
245 timeout: Duration,
246) -> Result<(), Box<dyn std::error::Error>> {
247 let start = Instant::now();
248 let mut next_log_at = start + Duration::from_secs(10);
249 tracing::info!("waiting for vm ipv4");
250 let mut once_hint = false;
251 loop {
252 let config = load_or_create_instance_config(instance_dir)?;
253 if config.vm_ipv4.is_some() {
254 return Ok(());
255 }
256 if start.elapsed() > timeout {
257 return Err("Timed out waiting for VM IPv4".into());
258 }
259 if Instant::now() >= next_log_at {
260 let waited = start.elapsed();
261 if waited.as_secs() > 15 && !once_hint {
262 tracing::info!(
263 "if vibebox is just initialized in this directory, it might take up to 1 minutes depending on your machine, and then you can enjoy the speed vibecoding! go pack!"
264 );
265 once_hint = true;
266 }
267 tracing::info!("still waiting for vm ipv4, {}s elapsed", waited.as_secs(),);
268 next_log_at += Duration::from_secs(10);
269 }
270 thread::sleep(Duration::from_millis(200));
271 }
274fn run_ssh_session(
275 ssh_key: PathBuf,
276 ssh_user: String,
277 ip: String,
278) -> Result<(), Box<dyn std::error::Error>> {
279 let mut attempts = 0usize;
280 loop {
281 attempts += 1;
282 if !ssh_port_open(&ip) {
283 tracing::debug!(attempts, "ssh port doesn't open yet");
284 tracing::info!(
285 attempts,
286 ip = %ip,
287 "waiting for ssh port ({}/{})",
288 attempts,
289 SSH_CONNECT_RETRIES
290 );
291 if attempts >= SSH_CONNECT_RETRIES {
292 return Err(
293 format!("ssh port not ready after {SSH_CONNECT_RETRIES} attempts").into(),
294 );
295 }
296 thread::sleep(Duration::from_millis(SSH_CONNECT_DELAY_MS));
297 continue;
298 }
300 tracing::info!(
301 attempts,
302 user = %ssh_user,
303 ip = %ip,
304 "starting ssh ({}/{})",
305 attempts,
306 SSH_CONNECT_RETRIES
307 );
308 let status = Command::new("ssh")
309 .args([
310 "-i",
311 ssh_key.to_str().unwrap_or(".vibebox/ssh_key"),
312 "-o",
313 "IdentitiesOnly=yes",
314 "-o",
315 "StrictHostKeyChecking=no",
316 "-o",
317 "UserKnownHostsFile=/dev/null",
318 "-o",
319 "GlobalKnownHostsFile=/dev/null",
320 "-o",
321 "PasswordAuthentication=no",
322 "-o",
323 "BatchMode=yes",
324 "-o",
325 "LogLevel=ERROR",
326 "-o",
327 "ConnectTimeout=5",
328 ])
329 .env_remove("LC_CTYPE")
330 .env_remove("LC_ALL")
331 .env_remove("LANG")
332 .arg(format!("{ssh_user}@{ip}"))
333 .stdin(Stdio::inherit())
334 .stdout(Stdio::inherit())
335 .stderr(Stdio::inherit())
336 .status();
338 match status {
339 Ok(status) if status.success() => {
340 tracing::info!(status = %status, "ssh exited");
341 break;
342 }
343 Ok(status) if status.code() == Some(255) => {
344 tracing::warn!(status = %status, "ssh connection failed");
345 if attempts >= SSH_CONNECT_RETRIES {
346 return Err(format!("ssh failed after {SSH_CONNECT_RETRIES} attempts").into());
347 }
348 thread::sleep(Duration::from_millis(500));
349 }
350 Ok(status) => {
351 tracing::info!(status = %status, "ssh exited");
352 break;
353 }
354 Err(err) => {
355 tracing::error!(error = %err, "failed to start ssh");
356 return Err(format!("failed to start ssh: {err}").into());
357 }
358 }
359 }
361 Ok(())
364fn is_ipv4_candidate(candidate: &str) -> bool {
365 let parts: Vec<&str> = candidate.split('.').collect();
366 if parts.len() != 4 {
367 return false;
368 }
369 for part in parts {
370 if part.is_empty() || part.len() > 3 {
371 return false;
372 }
373 if part.parse::<u8>().is_err() {
374 return false;
375 }
376 }
377 true
380fn ssh_port_open(ip: &str) -> bool {
381 let addr: SocketAddr = match format!("{ip}:22").parse() {
382 Ok(addr) => addr,
383 Err(_) => return false,
384 };
385 TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok()
388pub(crate) fn build_ssh_login_actions(
389 config: &Arc<Mutex<InstanceConfig>>,
390 project_name: &str,
391 guest_dir: &str,
392 key_name: &str,
393 home_links_script: &str,
394) -> Vec<LoginAction> {
395 let config_guard = config.lock().expect("config mutex poisoned");
396 let ssh_user = config_guard.ssh_user.clone();
397 let sudo_password = config_guard.sudo_password.clone();
398 drop(config_guard);
400 let key_path = format!("{guest_dir}/{key_name}.pub");
402 let setup_script = SSH_SETUP_SCRIPT
403 .replace("__SSH_USER__", &ssh_user)
404 .replace("__SUDO_PASSWORD__", &sudo_password)
405 .replace("__PROJECT_NAME__", project_name)
406 .replace("__KEY_PATH__", &key_path)
407 .replace("__VIBEBOX_SHELL_SCRIPT__", &commands::render_shell_script())
408 .replace("__VIBEBOX_HOME_LINKS__", home_links_script);
409 let setup = vm::script_command_from_content("ssh_setup", &setup_script)
410 .expect("ssh setup script contained invalid marker");
412 vec![LoginAction::Send(setup)]
415fn spawn_ssh_io(
416 app: Arc<Mutex<AppState>>,
417 config: Arc<Mutex<InstanceConfig>>,
418 instance_dir: PathBuf,
419 ssh_key: PathBuf,
420 output_monitor: Arc<vm::OutputMonitor>,
421 vm_output_fd: OwnedFd,
422 vm_input_fd: OwnedFd,
423) -> vm::IoContext {
424 let io_control = vm::IoControl::new();
426 let log_path = instance_dir.join(VM_ROOT_LOG_NAME);
427 let log_file = fs::OpenOptions::new()
428 .create(true)
429 .write(true)
430 .truncate(true)
431 .open(&log_path)
432 .ok()
433 .map(|file| Arc::new(Mutex::new(file)));
435 let ssh_connected = Arc::new(AtomicBool::new(false));
436 let ssh_started = Arc::new(AtomicBool::new(false));
437 let ssh_ready = Arc::new(AtomicBool::new(false));
438 let input_tx_holder: Arc<Mutex<Option<Sender<VmInput>>>> = Arc::new(Mutex::new(None));
440 let instance_path = instance_dir.join(INSTANCE_FILENAME);
441 let config_for_output = config.clone();
442 let log_for_output = log_file.clone();
443 let ssh_connected_for_output = ssh_connected.clone();
444 let ssh_ready_for_output = ssh_ready.clone();
446 let mut line_buf = String::new();
448 let handlers = commands::build_handlers(app.clone(), io_control.clone());
449 let on_line = move |line: &str| handlers.handle(line);
451 let on_output = move |bytes: &[u8]| {
452 if ssh_connected_for_output.load(Ordering::SeqCst) {
453 if let Some(log) = &log_for_output {
454 if let Ok(mut file) = log.lock() {
455 let _ = file.write_all(bytes);
456 }
457 }
458 }
460 let text = String::from_utf8_lossy(bytes);
461 line_buf.push_str(&text);
463 while let Some(pos) = line_buf.find('\n') {
464 let mut line = line_buf[..pos].to_string();
465 if line.ends_with('\r') {
466 line.pop();
467 }
468 line_buf.drain(..=pos);
470 let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' ');
472 if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") {
473 let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..];
474 let ip = extract_ipv4(ip_raw).unwrap_or_default();
475 if !ip.is_empty() {
476 if let Ok(mut cfg) = config_for_output.lock() {
477 if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) {
478 cfg.vm_ipv4 = Some(ip.clone());
479 let _ = write_instance_config(&instance_path, &cfg);
480 tracing::info!(ip = %ip, "detected vm ipv4");
481 }
482 }
483 }
484 }
486 if cleaned.contains("VIBEBOX_SSH_READY") {
487 ssh_ready_for_output.store(true, Ordering::SeqCst);
488 tracing::info!("sshd ready");
489 }
490 }
491 };
493 let io_ctx = vm::spawn_vm_io_with_hooks(
494 output_monitor,
495 vm_output_fd,
496 vm_input_fd,
497 io_control.clone(),
498 on_line,
499 on_output,
500 );
502 *input_tx_holder.lock().unwrap() = Some(io_ctx.input_tx.clone());
504 let ssh_ready_for_thread = ssh_ready.clone();
505 let ssh_started_for_thread = ssh_started.clone();
506 let ssh_connected_for_thread = ssh_connected.clone();
507 let io_control_for_thread = io_control.clone();
508 let ssh_key_for_thread = ssh_key.clone();
509 let config_for_thread = config.clone();
510 let input_tx_holder_for_thread = input_tx_holder.clone();
512 thread::spawn(move || {
513 loop {
514 if ssh_started_for_thread.load(Ordering::SeqCst) {
515 break;
516 }
517 if !ssh_ready_for_thread.load(Ordering::SeqCst) {
518 thread::sleep(std::time::Duration::from_millis(200));
519 continue;
520 }
522 let ip = config_for_thread
523 .lock()
524 .ok()
525 .and_then(|cfg| cfg.vm_ipv4.clone());
526 let ssh_user = config_for_thread
527 .lock()
528 .map(|cfg| cfg.ssh_user.clone())
529 .unwrap_or_else(|_| DEFAULT_SSH_USER.to_string());
531 if let Some(ip) = ip {
532 if ssh_started_for_thread
533 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
534 .is_err()
535 {
536 break;
537 }
539 tracing::info!(user = %ssh_user, ip = %ip, "starting ssh");
540 io_control_for_thread.request_terminal_restore();
541 io_control_for_thread.set_forward_output(false);
542 io_control_for_thread.set_forward_input(false);
543 ssh_connected_for_thread.store(true, Ordering::SeqCst);
545 let mut attempts = 0usize;
546 loop {
547 attempts += 1;
548 if !ssh_port_open(&ip) {
549 tracing::info!(
550 attempts,
551 ip = %ip,
552 "waiting for ssh port ({}/{})",
553 attempts,
554 SSH_CONNECT_RETRIES
555 );
556 if attempts >= SSH_CONNECT_RETRIES {
557 ssh_connected_for_thread.store(false, Ordering::SeqCst);
558 tracing::warn!(
559 attempts,
560 "ssh port not ready after {} attempts",
561 SSH_CONNECT_RETRIES
562 );
563 break;
564 }
565 thread::sleep(std::time::Duration::from_millis(SSH_CONNECT_DELAY_MS));
566 continue;
567 }
569 tracing::info!(
570 attempts,
571 user = %ssh_user,
572 ip = %ip,
573 "starting ssh ({}/{})",
574 attempts,
575 SSH_CONNECT_RETRIES
576 );
577 let status = Command::new("ssh")
578 .args([
579 "-i",
580 ssh_key_for_thread.to_str().unwrap_or(".vibebox/ssh_key"),
581 "-o",
582 "IdentitiesOnly=yes",
583 "-o",
584 "StrictHostKeyChecking=no",
585 "-o",
586 "UserKnownHostsFile=/dev/null",
587 "-o",
588 "GlobalKnownHostsFile=/dev/null",
589 "-o",
590 "PasswordAuthentication=no",
591 "-o",
592 "BatchMode=yes",
593 "-o",
594 "LogLevel=ERROR",
595 "-o",
596 "ConnectTimeout=5",
597 ])
598 .env_remove("LC_CTYPE")
599 .env_remove("LC_ALL")
600 // .env_remove("LANG")
601 .arg(format!("{ssh_user}@{ip}"))
602 .stdin(Stdio::inherit())
603 .stdout(Stdio::inherit())
604 .stderr(Stdio::inherit())
605 .status();
607 match status {
608 Ok(status) if status.success() => {
609 ssh_connected_for_thread.store(false, Ordering::SeqCst);
610 tracing::info!(status = %status, "ssh exited");
611 if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
612 let _ = tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
613 }
614 break;
615 }
616 Ok(status) if status.code() == Some(255) => {
617 tracing::warn!(status = %status, "ssh connection failed");
618 if attempts >= SSH_CONNECT_RETRIES {
619 ssh_connected_for_thread.store(false, Ordering::SeqCst);
620 tracing::warn!(
621 attempts,
622 "ssh failed after {} attempts",
623 SSH_CONNECT_RETRIES
624 );
625 break;
626 }
627 thread::sleep(std::time::Duration::from_millis(500));
628 continue;
629 }
630 Ok(status) => {
631 ssh_connected_for_thread.store(false, Ordering::SeqCst);
632 tracing::info!(status = %status, "ssh exited");
633 if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
634 let _ = tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
635 }
636 break;
637 }
638 Err(err) => {
639 ssh_connected_for_thread.store(false, Ordering::SeqCst);
640 tracing::error!(error = %err, "failed to start ssh");
641 break;
642 }
643 }
644 }
645 break;
646 }
648 thread::sleep(std::time::Duration::from_millis(200));
649 }
650 });
652 io_ctx