git.iko.soy/vibebox
on a49606e help sign in
vibebox/src/instance.rs
493 lines 17 kB history raw
1use std::{
2 env, fs,
3 io::{self, Write},
4 net::{SocketAddr, TcpStream},
5 os::unix::{fs::PermissionsExt, io::OwnedFd},
6 path::{Path, PathBuf},
7 process::{Command, Stdio},
8 sync::{
9 Arc, Mutex,
10 atomic::{AtomicBool, Ordering},
11 },
12 thread,
13};
15use serde::{Deserialize, Serialize};
16use std::sync::mpsc::Sender;
17use uuid::Uuid;
19use crate::{
20 session_manager::INSTANCE_DIR_NAME,
21 tui::{self, AppState},
22 vm::{self, DirectoryShare, LoginAction, VmInput},
23};
25const INSTANCE_TOML: &str = "instance.toml";
26const SSH_KEY_NAME: &str = "ssh_key";
27const SERIAL_LOG_NAME: &str = "serial.log";
28const SSH_GUEST_DIR: &str = "/root/.vibebox";
29const DEFAULT_SSH_USER: &str = "vibecoder";
30const SSH_CONNECT_RETRIES: usize = 30;
31const SSH_CONNECT_DELAY_MS: u64 = 500;
32const SSH_SETUP_SCRIPT: &str = include_str!("ssh.sh");
34#[derive(Debug, Clone, Serialize, Deserialize)]
35struct InstanceConfig {
36 #[serde(default = "default_ssh_user")]
37 ssh_user: String,
38 #[serde(default)]
39 sudo_password: String,
40 #[serde(default)]
41 vm_ipv4: Option<String>,
44fn default_ssh_user() -> String {
45 DEFAULT_SSH_USER.to_string()
48pub fn run_with_ssh(
49 args: vm::CliArgs,
50 app: Arc<Mutex<AppState>>,
51) -> Result<(), Box<dyn std::error::Error>> {
52 let project_root = env::current_dir()?;
53 let project_name = project_root
54 .file_name()
55 .ok_or("Project directory has no name")?
56 .to_string_lossy()
57 .into_owned();
58 let instance_dir = ensure_instance_dir(&project_root)?;
59 let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
61 let mut config = load_or_create_instance_config(&instance_dir)?;
62 // Clear cached IP to avoid reusing a stale address on startup.
63 if config.vm_ipv4.is_some() {
64 config.vm_ipv4 = None;
65 write_instance_config(&instance_dir.join(INSTANCE_TOML), &config)?;
66 }
67 let config = Arc::new(Mutex::new(config));
69 let extra_shares = vec![DirectoryShare::new(
70 instance_dir.clone(),
71 SSH_GUEST_DIR.into(),
72 true,
73 )?];
75 let extra_login_actions =
76 build_ssh_login_actions(&config, &project_name, SSH_GUEST_DIR, SSH_KEY_NAME);
78 vm::run_with_args_and_extras(
79 args,
80 |output_monitor, vm_output_fd, vm_input_fd| {
81 spawn_ssh_io(
82 app.clone(),
83 config.clone(),
84 instance_dir.clone(),
85 ssh_key.clone(),
86 output_monitor,
87 vm_output_fd,
88 vm_input_fd,
89 )
90 },
91 extra_login_actions,
92 extra_shares,
93 )
96fn ensure_instance_dir(project_root: &Path) -> Result<PathBuf, io::Error> {
97 let instance_dir = project_root.join(INSTANCE_DIR_NAME);
98 fs::create_dir_all(&instance_dir)?;
99 Ok(instance_dir)
102fn ensure_ssh_keypair(
103 instance_dir: &Path,
104) -> Result<(PathBuf, PathBuf), Box<dyn std::error::Error>> {
105 let private_key = instance_dir.join(SSH_KEY_NAME);
106 let public_key = instance_dir.join(format!("{SSH_KEY_NAME}.pub"));
108 if private_key.exists() && public_key.exists() {
109 return Ok((private_key, public_key));
110 }
112 if private_key.exists() {
113 let _ = fs::remove_file(&private_key);
114 }
115 if public_key.exists() {
116 let _ = fs::remove_file(&public_key);
117 }
119 let status = Command::new("ssh-keygen")
120 .args([
121 "-t",
122 "ed25519",
123 "-N",
124 "",
125 "-f",
126 private_key.to_str().ok_or("ssh key path not utf-8")?,
127 "-C",
128 "vibebox",
129 ])
130 .stdin(Stdio::null())
131 .stdout(Stdio::null())
132 .stderr(Stdio::inherit())
133 .status()?;
135 if !status.success() {
136 return Err("ssh-keygen failed".into());
137 }
139 fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?;
140 fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?;
142 Ok((private_key, public_key))
145fn load_or_create_instance_config(
146 instance_dir: &Path,
147) -> Result<InstanceConfig, Box<dyn std::error::Error>> {
148 let config_path = instance_dir.join(INSTANCE_TOML);
149 let mut config = if config_path.exists() {
150 let raw = fs::read_to_string(&config_path)?;
151 toml::from_str::<InstanceConfig>(&raw)?
152 } else {
153 InstanceConfig {
154 ssh_user: default_ssh_user(),
155 sudo_password: String::new(),
156 vm_ipv4: None,
157 }
158 };
160 let mut changed = false;
161 if config.ssh_user.trim().is_empty() {
162 config.ssh_user = default_ssh_user();
163 changed = true;
164 }
166 if config.sudo_password.trim().is_empty() {
167 config.sudo_password = generate_password();
168 changed = true;
169 }
171 if !config_path.exists() || changed {
172 write_instance_config(&config_path, &config)?;
173 }
175 Ok(config)
178fn write_instance_config(
179 path: &Path,
180 config: &InstanceConfig,
181) -> Result<(), Box<dyn std::error::Error>> {
182 let data = toml::to_string_pretty(config)?;
183 fs::write(path, data)?;
184 fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
185 Ok(())
188fn generate_password() -> String {
189 Uuid::now_v7().simple().to_string()
192fn extract_ipv4(line: &str) -> Option<String> {
193 let mut current = String::new();
194 let mut best: Option<String> = None;
196 for ch in line.chars().chain(std::iter::once(' ')) {
197 if ch.is_ascii_digit() || ch == '.' {
198 current.push(ch);
199 } else if !current.is_empty() {
200 if is_ipv4_candidate(&current) {
201 best = Some(current.clone());
202 break;
203 }
204 current.clear();
205 }
206 }
208 best
211fn is_ipv4_candidate(candidate: &str) -> bool {
212 let parts: Vec<&str> = candidate.split('.').collect();
213 if parts.len() != 4 {
214 return false;
215 }
216 for part in parts {
217 if part.is_empty() || part.len() > 3 {
218 return false;
219 }
220 if part.parse::<u8>().is_err() {
221 return false;
222 }
223 }
224 true
227fn ssh_port_open(ip: &str) -> bool {
228 let addr: SocketAddr = match format!("{ip}:22").parse() {
229 Ok(addr) => addr,
230 Err(_) => return false,
231 };
232 TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok()
235fn build_ssh_login_actions(
236 config: &Arc<Mutex<InstanceConfig>>,
237 project_name: &str,
238 guest_dir: &str,
239 key_name: &str,
240) -> Vec<LoginAction> {
241 let config_guard = config.lock().expect("config mutex poisoned");
242 let ssh_user = config_guard.ssh_user.clone();
243 let sudo_password = config_guard.sudo_password.clone();
244 drop(config_guard);
246 let key_path = format!("{guest_dir}/{key_name}.pub");
248 let setup_script = SSH_SETUP_SCRIPT
249 .replace("__SSH_USER__", &ssh_user)
250 .replace("__SUDO_PASSWORD__", &sudo_password)
251 .replace("__PROJECT_NAME__", project_name)
252 .replace("__KEY_PATH__", &key_path);
253 let setup = vm::script_command_from_content("ssh_setup", &setup_script)
254 .expect("ssh setup script contained invalid marker");
256 vec![LoginAction::Send(setup)]
259fn spawn_ssh_io(
260 app: Arc<Mutex<AppState>>,
261 config: Arc<Mutex<InstanceConfig>>,
262 instance_dir: PathBuf,
263 ssh_key: PathBuf,
264 output_monitor: Arc<vm::OutputMonitor>,
265 vm_output_fd: OwnedFd,
266 vm_input_fd: OwnedFd,
267) -> vm::IoContext {
268 let io_control = vm::IoControl::new();
270 let log_path = instance_dir.join(SERIAL_LOG_NAME);
271 let log_file = fs::OpenOptions::new()
272 .create(true)
273 .append(true)
274 .open(&log_path)
275 .ok()
276 .map(|file| Arc::new(Mutex::new(file)));
278 let ssh_connected = Arc::new(AtomicBool::new(false));
279 let ssh_started = Arc::new(AtomicBool::new(false));
280 let ssh_ready = Arc::new(AtomicBool::new(false));
281 let input_tx_holder: Arc<Mutex<Option<Sender<VmInput>>>> = Arc::new(Mutex::new(None));
283 let instance_path = instance_dir.join(INSTANCE_TOML);
284 let config_for_output = config.clone();
285 let log_for_output = log_file.clone();
286 let ssh_connected_for_output = ssh_connected.clone();
287 let ssh_ready_for_output = ssh_ready.clone();
289 let mut line_buf = String::new();
291 let on_line = {
292 let app = app.clone();
293 move |line: &str| {
294 if line == ":help" {
295 if let Ok(mut locked) = app.lock() {
296 let _ = tui::render_commands_component(&mut locked);
297 }
298 return true;
299 }
300 false
301 }
302 };
304 let on_output = move |bytes: &[u8]| {
305 if ssh_connected_for_output.load(Ordering::SeqCst) {
306 if let Some(log) = &log_for_output {
307 if let Ok(mut file) = log.lock() {
308 let _ = file.write_all(bytes);
309 }
310 }
311 }
313 let text = String::from_utf8_lossy(bytes);
314 line_buf.push_str(&text);
316 while let Some(pos) = line_buf.find('\n') {
317 let mut line = line_buf[..pos].to_string();
318 if line.ends_with('\r') {
319 line.pop();
320 }
321 line_buf.drain(..=pos);
323 let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' ');
325 if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") {
326 let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..];
327 let ip = extract_ipv4(ip_raw).unwrap_or_default();
328 if !ip.is_empty() {
329 if let Ok(mut cfg) = config_for_output.lock() {
330 if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) {
331 cfg.vm_ipv4 = Some(ip.clone());
332 let _ = write_instance_config(&instance_path, &cfg);
333 eprintln!("[vibebox] detected vm IPv4: {ip}");
334 }
335 }
336 }
337 }
339 if cleaned.contains("VIBEBOX_SSH_READY") {
340 ssh_ready_for_output.store(true, Ordering::SeqCst);
341 eprintln!("[vibebox] sshd ready");
342 }
343 }
344 };
346 let io_ctx = vm::spawn_vm_io_with_hooks(
347 output_monitor,
348 vm_output_fd,
349 vm_input_fd,
350 io_control.clone(),
351 on_line,
352 on_output,
353 );
355 *input_tx_holder.lock().unwrap() = Some(io_ctx.input_tx.clone());
357 let ssh_ready_for_thread = ssh_ready.clone();
358 let ssh_started_for_thread = ssh_started.clone();
359 let ssh_connected_for_thread = ssh_connected.clone();
360 let io_control_for_thread = io_control.clone();
361 let ssh_key_for_thread = ssh_key.clone();
362 let config_for_thread = config.clone();
363 let input_tx_holder_for_thread = input_tx_holder.clone();
365 thread::spawn(move || {
366 loop {
367 if ssh_started_for_thread.load(Ordering::SeqCst) {
368 break;
369 }
370 if !ssh_ready_for_thread.load(Ordering::SeqCst) {
371 thread::sleep(std::time::Duration::from_millis(200));
372 continue;
373 }
375 let ip = config_for_thread
376 .lock()
377 .ok()
378 .and_then(|cfg| cfg.vm_ipv4.clone());
379 let ssh_user = config_for_thread
380 .lock()
381 .map(|cfg| cfg.ssh_user.clone())
382 .unwrap_or_else(|_| DEFAULT_SSH_USER.to_string());
384 if let Some(ip) = ip {
385 if ssh_started_for_thread
386 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
387 .is_err()
388 {
389 break;
390 }
392 eprintln!("[vibebox] starting ssh to {ssh_user}@{ip}");
393 io_control_for_thread.request_terminal_restore();
394 io_control_for_thread.set_forward_output(false);
395 io_control_for_thread.set_forward_input(false);
396 ssh_connected_for_thread.store(true, Ordering::SeqCst);
398 let mut attempts = 0usize;
399 loop {
400 attempts += 1;
401 if !ssh_port_open(&ip) {
402 eprintln!(
403 "[vibebox] waiting for ssh port on {ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
404 );
405 if attempts >= SSH_CONNECT_RETRIES {
406 ssh_connected_for_thread.store(false, Ordering::SeqCst);
407 eprintln!(
408 "[vibebox] ssh port not ready after {SSH_CONNECT_RETRIES} attempts"
409 );
410 break;
411 }
412 thread::sleep(std::time::Duration::from_millis(SSH_CONNECT_DELAY_MS));
413 continue;
414 }
416 eprintln!(
417 "[vibebox] starting ssh to {ssh_user}@{ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
418 );
419 let status = Command::new("ssh")
420 .args([
421 "-i",
422 ssh_key_for_thread.to_str().unwrap_or(".vibebox/ssh_key"),
423 "-o",
424 "IdentitiesOnly=yes",
425 "-o",
426 "StrictHostKeyChecking=no",
427 "-o",
428 "UserKnownHostsFile=/dev/null",
429 "-o",
430 "GlobalKnownHostsFile=/dev/null",
431 "-o",
432 "PasswordAuthentication=no",
433 "-o",
434 "BatchMode=yes",
435 "-o",
436 "LogLevel=ERROR",
437 "-o",
438 "ConnectTimeout=5",
439 ])
440 .env_remove("LC_CTYPE")
441 .env_remove("LC_ALL")
442 // .env_remove("LANG")
443 .arg(format!("{ssh_user}@{ip}"))
444 .stdin(Stdio::inherit())
445 .stdout(Stdio::inherit())
446 .stderr(Stdio::inherit())
447 .status();
449 match status {
450 Ok(status) if status.success() => {
451 ssh_connected_for_thread.store(false, Ordering::SeqCst);
452 eprintln!("[vibebox] ssh exited with status: {status}");
453 if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
454 let _ = tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
455 }
456 break;
457 }
458 Ok(status) if status.code() == Some(255) => {
459 eprintln!("[vibebox] ssh connection failed: {status}");
460 if attempts >= SSH_CONNECT_RETRIES {
461 ssh_connected_for_thread.store(false, Ordering::SeqCst);
462 eprintln!(
463 "[vibebox] ssh failed after {SSH_CONNECT_RETRIES} attempts"
464 );
465 break;
466 }
467 thread::sleep(std::time::Duration::from_millis(500));
468 continue;
469 }
470 Ok(status) => {
471 ssh_connected_for_thread.store(false, Ordering::SeqCst);
472 eprintln!("[vibebox] ssh exited with status: {status}");
473 if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
474 let _ = tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
475 }
476 break;
477 }
478 Err(err) => {
479 ssh_connected_for_thread.store(false, Ordering::SeqCst);
480 eprintln!("[vibebox] failed to start ssh: {err}");
481 break;
482 }
483 }
484 }
485 break;
486 }
488 thread::sleep(std::time::Duration::from_millis(200));
489 }
490 });
492 io_ctx