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