on cd76a4f help sign in
vibebox/src/instance.rs
521 lines 18 kB history raw
1use std::{
2 env,
3 fs,
4 io::{self, Write},
5 net::{SocketAddr, TcpStream},
6 os::unix::{
7 fs::PermissionsExt,
8 io::OwnedFd,
9 },
10 path::{Path, PathBuf},
11 process::{Command, Stdio},
12 sync::{
13 Arc, Mutex,
14 atomic::{AtomicBool, Ordering},
15 },
16 thread,
17};
19use serde::{Deserialize, Serialize};
20use uuid::Uuid;
21use std::sync::mpsc::Sender;
23use crate::{
24 session_manager::INSTANCE_DIR_NAME,
25 tui::{self, AppState},
26 vm::{self, DirectoryShare, LoginAction, VmInput},
27};
29const INSTANCE_TOML: &str = "instance.toml";
30const SSH_KEY_NAME: &str = "ssh_key";
31const SERIAL_LOG_NAME: &str = "serial.log";
32const SSH_GUEST_DIR: &str = "/root/.vibebox";
33const DEFAULT_SSH_USER: &str = "vibebox";
34const SSH_CONNECT_RETRIES: usize = 20;
35const SSH_CONNECT_DELAY_MS: u64 = 500;
37#[derive(Debug, Clone, Serialize, Deserialize)]
38struct InstanceConfig {
39 #[serde(default = "default_ssh_user")]
40 ssh_user: String,
41 #[serde(default)]
42 sudo_password: String,
43 #[serde(default)]
44 vm_ipv4: Option<String>,
47fn default_ssh_user() -> String {
48 DEFAULT_SSH_USER.to_string()
51pub fn run_with_ssh(
52 args: vm::CliArgs,
53 app: Arc<Mutex<AppState>>,
54) -> Result<(), Box<dyn std::error::Error>> {
55 let project_root = env::current_dir()?;
56 let project_name = project_root
57 .file_name()
58 .ok_or("Project directory has no name")?
59 .to_string_lossy()
60 .into_owned();
61 let instance_dir = ensure_instance_dir(&project_root)?;
62 let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
64 let config = load_or_create_instance_config(&instance_dir)?;
65 let config = Arc::new(Mutex::new(config));
67 let extra_shares = vec![DirectoryShare::new(
68 instance_dir.clone(),
69 SSH_GUEST_DIR.into(),
70 true,
71 )?];
73 let extra_login_actions = build_ssh_login_actions(
74 &config,
75 &project_name,
76 SSH_GUEST_DIR,
77 SSH_KEY_NAME,
78 );
80 vm::run_with_args_and_extras(
81 args,
82 |output_monitor, vm_output_fd, vm_input_fd| {
83 spawn_ssh_io(
84 app.clone(),
85 config.clone(),
86 instance_dir.clone(),
87 ssh_key.clone(),
88 output_monitor,
89 vm_output_fd,
90 vm_input_fd,
91 )
92 },
93 extra_login_actions,
94 extra_shares,
95 )
98fn ensure_instance_dir(project_root: &Path) -> Result<PathBuf, io::Error> {
99 let instance_dir = project_root.join(INSTANCE_DIR_NAME);
100 fs::create_dir_all(&instance_dir)?;
101 Ok(instance_dir)
104fn ensure_ssh_keypair(instance_dir: &Path) -> 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
127 .to_str()
128 .ok_or("ssh key path not utf-8")?,
129 "-C",
130 "vibebox",
131 ])
132 .status()?;
134 if !status.success() {
135 return Err("ssh-keygen failed".into());
136 }
138 fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?;
139 fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?;
141 Ok((private_key, public_key))
144fn load_or_create_instance_config(
145 instance_dir: &Path,
146) -> Result<InstanceConfig, Box<dyn std::error::Error>> {
147 let config_path = instance_dir.join(INSTANCE_TOML);
148 let mut config = if config_path.exists() {
149 let raw = fs::read_to_string(&config_path)?;
150 toml::from_str::<InstanceConfig>(&raw)?
151 } else {
152 InstanceConfig {
153 ssh_user: default_ssh_user(),
154 sudo_password: String::new(),
155 vm_ipv4: None,
156 }
157 };
159 let mut changed = false;
160 if config.ssh_user.trim().is_empty() {
161 config.ssh_user = default_ssh_user();
162 changed = true;
163 }
165 if config.sudo_password.trim().is_empty() {
166 config.sudo_password = generate_password();
167 changed = true;
168 }
170 if !config_path.exists() || changed {
171 write_instance_config(&config_path, &config)?;
172 }
174 Ok(config)
177fn write_instance_config(
178 path: &Path,
179 config: &InstanceConfig,
180) -> Result<(), Box<dyn std::error::Error>> {
181 let data = toml::to_string_pretty(config)?;
182 fs::write(path, data)?;
183 fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
184 Ok(())
187fn generate_password() -> String {
188 Uuid::now_v7().simple().to_string()
191fn extract_ipv4(line: &str) -> Option<String> {
192 let mut current = String::new();
193 let mut best: Option<String> = None;
195 for ch in line.chars().chain(std::iter::once(' ')) {
196 if ch.is_ascii_digit() || ch == '.' {
197 current.push(ch);
198 } else if !current.is_empty() {
199 if is_ipv4_candidate(&current) {
200 best = Some(current.clone());
201 break;
202 }
203 current.clear();
204 }
205 }
207 best
210fn is_ipv4_candidate(candidate: &str) -> bool {
211 let parts: Vec<&str> = candidate.split('.').collect();
212 if parts.len() != 4 {
213 return false;
214 }
215 for part in parts {
216 if part.is_empty() || part.len() > 3 {
217 return false;
218 }
219 if part.parse::<u8>().is_err() {
220 return false;
221 }
222 }
223 true
226fn ssh_port_open(ip: &str) -> bool {
227 let addr: SocketAddr = match format!("{ip}:22").parse() {
228 Ok(addr) => addr,
229 Err(_) => return false,
230 };
231 TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok()
234fn build_ssh_login_actions(
235 config: &Arc<Mutex<InstanceConfig>>,
236 project_name: &str,
237 guest_dir: &str,
238 key_name: &str,
239) -> Vec<LoginAction> {
240 let config_guard = config.lock().expect("config mutex poisoned");
241 let ssh_user = config_guard.ssh_user.clone();
242 let sudo_password = config_guard.sudo_password.clone();
243 drop(config_guard);
245 let key_path = format!("{guest_dir}/{key_name}.pub");
247 let mask_cache = format!(
248 "if [ -d /root/{project_name}/.vibebox ]; then mount -t tmpfs tmpfs /root/{project_name}/.vibebox; fi"
249 );
251 let setup_lines = [
252 "if ! command -v sshd >/dev/null 2>&1; then apt-get update && apt-get install -y openssh-server sudo; fi",
253 "systemctl enable ssh >/dev/null 2>&1 || true",
254 &format!("id -u {ssh_user} >/dev/null 2>&1 || useradd -m -s /bin/bash {ssh_user}"),
255 &format!("echo \"{ssh_user}:{sudo_password}\" | chpasswd"),
256 &format!("usermod -aG sudo {ssh_user}"),
257 &format!("install -d -m 700 /home/{ssh_user}/.ssh"),
258 &format!("install -m 600 {key_path} /home/{ssh_user}/.ssh/authorized_keys"),
259 &format!("chown -R {ssh_user}:{ssh_user} /home/{ssh_user}/.ssh"),
260 "mkdir -p /etc/ssh/sshd_config.d",
261 "cat >/etc/ssh/sshd_config.d/vibebox.conf <<'VIBEBOX_SSHD'",
262 "PasswordAuthentication no",
263 "KbdInteractiveAuthentication no",
264 "ChallengeResponseAuthentication no",
265 "PubkeyAuthentication yes",
266 "PermitRootLogin no",
267 &format!("AllowUsers {ssh_user}"),
268 "VIBEBOX_SSHD",
269 "systemctl restart ssh",
270 "echo VIBEBOX_SSH_READY",
271 ];
272 let setup = setup_lines.join("\n");
274 let ip_probe = "while true; do ip=$(ip -4 -o addr show scope global | awk '{print $4}' | cut -d/ -f1 | head -n 1); if [ -n \"$ip\" ]; then echo VIBEBOX_IPV4=$ip; break; fi; sleep 1; done";
276 vec![
277 LoginAction::Send(mask_cache),
278 LoginAction::Send(setup),
279 LoginAction::Send(ip_probe.to_string()),
280 ]
283fn spawn_ssh_io(
284 app: Arc<Mutex<AppState>>,
285 config: Arc<Mutex<InstanceConfig>>,
286 instance_dir: PathBuf,
287 ssh_key: PathBuf,
288 output_monitor: Arc<vm::OutputMonitor>,
289 vm_output_fd: OwnedFd,
290 vm_input_fd: OwnedFd,
291) -> vm::IoContext {
292 let io_control = vm::IoControl::new();
294 let log_path = instance_dir.join(SERIAL_LOG_NAME);
295 let log_file = fs::OpenOptions::new()
296 .create(true)
297 .append(true)
298 .open(&log_path)
299 .ok()
300 .map(|file| Arc::new(Mutex::new(file)));
302 let ssh_connected = Arc::new(AtomicBool::new(false));
303 let ssh_started = Arc::new(AtomicBool::new(false));
304 let ssh_ready = Arc::new(AtomicBool::new(false));
305 let input_tx_holder: Arc<Mutex<Option<Sender<VmInput>>>> =
306 Arc::new(Mutex::new(None));
308 let instance_path = instance_dir.join(INSTANCE_TOML);
309 let config_for_output = config.clone();
310 let log_for_output = log_file.clone();
311 let ssh_connected_for_output = ssh_connected.clone();
312 let ssh_ready_for_output = ssh_ready.clone();
314 let mut line_buf = String::new();
316 let on_line = {
317 let app = app.clone();
318 move |line: &str| {
319 if line == ":help" {
320 if let Ok(mut locked) = app.lock() {
321 let _ = tui::render_commands_component(&mut locked);
322 }
323 return true;
324 }
325 false
326 }
327 };
329 let on_output = move |bytes: &[u8]| {
330 if ssh_connected_for_output.load(Ordering::SeqCst) {
331 if let Some(log) = &log_for_output {
332 if let Ok(mut file) = log.lock() {
333 let _ = file.write_all(bytes);
334 }
335 }
336 }
338 let text = String::from_utf8_lossy(bytes);
339 line_buf.push_str(&text);
341 while let Some(pos) = line_buf.find('\n') {
342 let mut line = line_buf[..pos].to_string();
343 if line.ends_with('\r') {
344 line.pop();
345 }
346 line_buf.drain(..=pos);
348 let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' ');
350 if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") {
351 let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..];
352 let ip = extract_ipv4(ip_raw).unwrap_or_default();
353 if !ip.is_empty() {
354 if let Ok(mut cfg) = config_for_output.lock() {
355 if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) {
356 cfg.vm_ipv4 = Some(ip.clone());
357 let _ = write_instance_config(&instance_path, &cfg);
358 eprintln!("[vibebox] detected vm IPv4: {ip}");
359 }
360 }
361 }
362 }
364 if cleaned.contains("VIBEBOX_SSH_READY") {
365 ssh_ready_for_output.store(true, Ordering::SeqCst);
366 eprintln!("[vibebox] sshd ready");
367 }
368 }
369 };
371 let io_ctx = vm::spawn_vm_io_with_hooks(
372 output_monitor,
373 vm_output_fd,
374 vm_input_fd,
375 io_control.clone(),
376 on_line,
377 on_output,
378 );
380 *input_tx_holder.lock().unwrap() = Some(io_ctx.input_tx.clone());
382 let ssh_ready_for_thread = ssh_ready.clone();
383 let ssh_started_for_thread = ssh_started.clone();
384 let ssh_connected_for_thread = ssh_connected.clone();
385 let io_control_for_thread = io_control.clone();
386 let ssh_key_for_thread = ssh_key.clone();
387 let config_for_thread = config.clone();
388 let input_tx_holder_for_thread = input_tx_holder.clone();
390 thread::spawn(move || {
391 loop {
392 if ssh_started_for_thread.load(Ordering::SeqCst) {
393 break;
394 }
395 if !ssh_ready_for_thread.load(Ordering::SeqCst) {
396 thread::sleep(std::time::Duration::from_millis(200));
397 continue;
398 }
400 let ip = config_for_thread
401 .lock()
402 .ok()
403 .and_then(|cfg| cfg.vm_ipv4.clone());
404 let ssh_user = config_for_thread
405 .lock()
406 .map(|cfg| cfg.ssh_user.clone())
407 .unwrap_or_else(|_| DEFAULT_SSH_USER.to_string());
409 if let Some(ip) = ip {
410 if ssh_started_for_thread
411 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
412 .is_err()
413 {
414 break;
415 }
417 eprintln!("[vibebox] starting ssh to {ssh_user}@{ip}");
418 io_control_for_thread.request_terminal_restore();
419 io_control_for_thread.set_forward_output(false);
420 io_control_for_thread.set_forward_input(false);
421 ssh_connected_for_thread.store(true, Ordering::SeqCst);
423 let mut attempts = 0usize;
424 loop {
425 attempts += 1;
426 if !ssh_port_open(&ip) {
427 eprintln!(
428 "[vibebox] waiting for ssh port on {ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
429 );
430 if attempts >= SSH_CONNECT_RETRIES {
431 ssh_connected_for_thread.store(false, Ordering::SeqCst);
432 eprintln!(
433 "[vibebox] ssh port not ready after {SSH_CONNECT_RETRIES} attempts"
434 );
435 break;
436 }
437 thread::sleep(std::time::Duration::from_millis(
438 SSH_CONNECT_DELAY_MS,
439 ));
440 continue;
441 }
443 eprintln!(
444 "[vibebox] starting ssh to {ssh_user}@{ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})"
445 );
446 let status = Command::new("ssh")
447 .args([
448 "-i",
449 ssh_key_for_thread
450 .to_str()
451 .unwrap_or(".vibebox/ssh_key"),
452 "-o",
453 "IdentitiesOnly=yes",
454 "-o",
455 "StrictHostKeyChecking=no",
456 "-o",
457 "UserKnownHostsFile=/dev/null",
458 "-o",
459 "GlobalKnownHostsFile=/dev/null",
460 "-o",
461 "PasswordAuthentication=no",
462 "-o",
463 "BatchMode=yes",
464 "-o",
465 "LogLevel=ERROR",
466 "-o",
467 "ConnectTimeout=5",
468 ])
469 .arg(format!("{ssh_user}@{ip}"))
470 .stdin(Stdio::inherit())
471 .stdout(Stdio::inherit())
472 .stderr(Stdio::inherit())
473 .status();
475 match status {
476 Ok(status) if status.success() => {
477 ssh_connected_for_thread.store(false, Ordering::SeqCst);
478 eprintln!("[vibebox] ssh exited with status: {status}");
479 if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
480 let _ =
481 tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
482 }
483 break;
484 }
485 Ok(status) if status.code() == Some(255) => {
486 eprintln!("[vibebox] ssh connection failed: {status}");
487 if attempts >= SSH_CONNECT_RETRIES {
488 ssh_connected_for_thread.store(false, Ordering::SeqCst);
489 eprintln!(
490 "[vibebox] ssh failed after {SSH_CONNECT_RETRIES} attempts"
491 );
492 break;
493 }
494 thread::sleep(std::time::Duration::from_millis(500));
495 continue;
496 }
497 Ok(status) => {
498 ssh_connected_for_thread.store(false, Ordering::SeqCst);
499 eprintln!("[vibebox] ssh exited with status: {status}");
500 if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() {
501 let _ =
502 tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec()));
503 }
504 break;
505 }
506 Err(err) => {
507 ssh_connected_for_thread.store(false, Ordering::SeqCst);
508 eprintln!("[vibebox] failed to start ssh: {err}");
509 break;
510 }
511 }
512 }
513 break;
514 }
516 thread::sleep(std::time::Duration::from_millis(200));
517 }
518 });
520 io_ctx