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