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