git.iko.soy/vibebox
on 3e1b069 help sign in
vibebox/src/instance.rs
419 lines 13 kB history raw
1use std::{
2 env, fs,
3 io::{self},
4 net::{SocketAddr, TcpStream},
5 os::unix::{fs::PermissionsExt, net::UnixStream},
6 path::{Path, PathBuf},
7 process::{Command, Stdio},
8 sync::{Arc, Mutex},
9 thread,
10 time::{Duration, Instant},
11};
13use serde::{Deserialize, Serialize};
14use time::{OffsetDateTime, format_description::well_known::Rfc3339};
15use uuid::Uuid;
17use crate::{
18 commands,
19 session_manager::{INSTANCE_DIR_NAME, INSTANCE_FILENAME},
20 vm::{self, LoginAction},
21};
23const SSH_KEY_NAME: &str = "ssh_key";
24pub(crate) const VM_ROOT_LOG_NAME: &str = "vm_root.log";
25pub(crate) const STATUS_FILE_NAME: &str = "status.txt";
26pub(crate) const DEFAULT_SSH_USER: &str = "vibecoder";
27const SSH_CONNECT_RETRIES: usize = 30;
28const SSH_CONNECT_DELAY_MS: u64 = 500;
29const SSH_SETUP_SCRIPT: &str = include_str!("ssh.sh");
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub(crate) struct InstanceConfig {
33 #[serde(default)]
34 id: String,
35 #[serde(default = "default_ssh_user")]
36 ssh_user: String,
37 #[serde(default)]
38 sudo_password: String,
39 #[serde(default)]
40 last_active: Option<String>,
41 #[serde(default)]
42 pub(crate) vm_ipv4: Option<String>,
45impl InstanceConfig {
46 pub(crate) fn ssh_user_display(&self) -> String {
47 if self.ssh_user.trim().is_empty() {
48 DEFAULT_SSH_USER.to_string()
49 } else {
50 self.ssh_user.clone()
51 }
52 }
55fn default_ssh_user() -> String {
56 DEFAULT_SSH_USER.to_string()
59pub fn run_with_ssh(manager_conn: UnixStream) -> Result<(), Box<dyn std::error::Error>> {
60 let project_root = env::current_dir()?;
61 tracing::info!(root = %project_root.display(), "starting ssh session");
62 let instance_dir = ensure_instance_dir(&project_root)?;
63 tracing::debug!(instance_dir = %instance_dir.display(), "instance dir ready");
64 let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
66 let config = load_or_create_instance_config(&instance_dir)?;
67 let ssh_user = config.ssh_user.clone();
68 tracing::debug!(ssh_user = %ssh_user, "loaded instance config");
70 let _manager_conn = manager_conn;
71 wait_for_vm_ipv4(&instance_dir, Duration::from_secs(120))?;
73 let ip = load_or_create_instance_config(&instance_dir)?
74 .vm_ipv4
75 .ok_or("VM IPv4 not available")?;
76 tracing::info!(ip = %ip, "vm ipv4 ready");
78 run_ssh_session(ssh_key, ssh_user, ip)
81pub fn ensure_instance_dir(project_root: &Path) -> Result<PathBuf, io::Error> {
82 let instance_dir = project_root.join(INSTANCE_DIR_NAME);
83 fs::create_dir_all(&instance_dir)?;
84 Ok(instance_dir)
87pub(crate) fn ensure_ssh_keypair(
88 instance_dir: &Path,
89) -> Result<(PathBuf, PathBuf), Box<dyn std::error::Error>> {
90 let private_key = instance_dir.join(SSH_KEY_NAME);
91 let public_key = instance_dir.join(format!("{SSH_KEY_NAME}.pub"));
93 if private_key.exists() && public_key.exists() {
94 return Ok((private_key, public_key));
95 }
97 if private_key.exists() {
98 let _ = fs::remove_file(&private_key);
99 }
100 if public_key.exists() {
101 let _ = fs::remove_file(&public_key);
102 }
104 let status = Command::new("ssh-keygen")
105 .args([
106 "-t",
107 "ed25519",
108 "-N",
109 "",
110 "-f",
111 private_key.to_str().ok_or("ssh key path not utf-8")?,
112 "-C",
113 "vibebox",
114 ])
115 .stdin(Stdio::null())
116 .stdout(Stdio::inherit())
117 .stderr(Stdio::inherit())
118 .status()?;
120 if !status.success() {
121 return Err("ssh-keygen failed".into());
122 }
124 fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?;
125 fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?;
127 Ok((private_key, public_key))
130pub(crate) fn load_or_create_instance_config(
131 instance_dir: &Path,
132) -> Result<InstanceConfig, Box<dyn std::error::Error>> {
133 let config_path = instance_dir.join(INSTANCE_FILENAME);
134 let mut config = if config_path.exists() {
135 let raw = fs::read_to_string(&config_path)?;
136 toml::from_str::<InstanceConfig>(&raw)?
137 } else {
138 InstanceConfig {
139 id: String::new(),
140 ssh_user: default_ssh_user(),
141 sudo_password: String::new(),
142 last_active: None,
143 vm_ipv4: None,
144 }
145 };
147 let mut changed = false;
148 if config.ssh_user.trim().is_empty() {
149 config.ssh_user = default_ssh_user();
150 changed = true;
151 }
153 if config.id.trim().is_empty() {
154 config.id = Uuid::now_v7().to_string();
155 changed = true;
156 }
158 if config.sudo_password.trim().is_empty() {
159 config.sudo_password = generate_password();
160 changed = true;
161 }
163 if !config_path.exists() || changed {
164 write_instance_config(&config_path, &config)?;
165 }
167 Ok(config)
170fn read_instance_config(
171 instance_dir: &Path,
172) -> Result<Option<InstanceConfig>, Box<dyn std::error::Error>> {
173 let config_path = instance_dir.join(INSTANCE_FILENAME);
174 if !config_path.exists() {
175 return Ok(None);
176 }
177 let raw = fs::read_to_string(&config_path)?;
178 let config = toml::from_str::<InstanceConfig>(&raw)?;
179 Ok(Some(config))
182pub fn read_instance_vm_ip(
183 instance_dir: &Path,
184) -> Result<Option<String>, Box<dyn std::error::Error>> {
185 let config = read_instance_config(instance_dir)?;
186 Ok(config.and_then(|cfg| cfg.vm_ipv4))
189pub fn read_instance_ssh_user(
190 instance_dir: &Path,
191) -> Result<Option<String>, Box<dyn std::error::Error>> {
192 let config = read_instance_config(instance_dir)?;
193 Ok(config
194 .map(|cfg| cfg.ssh_user)
195 .filter(|user| !user.trim().is_empty()))
198pub fn touch_last_active(instance_dir: &Path) -> Result<(), Box<dyn std::error::Error>> {
199 let mut config = load_or_create_instance_config(instance_dir)?;
200 let now = OffsetDateTime::now_utc().format(&Rfc3339)?;
201 config.last_active = Some(now);
202 write_instance_config(&instance_dir.join(INSTANCE_FILENAME), &config)?;
203 Ok(())
206pub(crate) fn write_instance_config(
207 path: &Path,
208 config: &InstanceConfig,
209) -> Result<(), Box<dyn std::error::Error>> {
210 let data = toml::to_string_pretty(config)?;
211 fs::write(path, data)?;
212 fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
213 Ok(())
216fn generate_password() -> String {
217 Uuid::now_v7().simple().to_string()
220pub(crate) fn extract_ipv4(line: &str) -> Option<String> {
221 let mut current = String::new();
222 let mut best: Option<String> = None;
224 for ch in line.chars().chain(std::iter::once(' ')) {
225 if ch.is_ascii_digit() || ch == '.' {
226 current.push(ch);
227 } else if !current.is_empty() {
228 if is_ipv4_candidate(&current) {
229 best = Some(current.clone());
230 break;
231 }
232 current.clear();
233 }
234 }
236 best
239fn wait_for_vm_ipv4(
240 instance_dir: &Path,
241 timeout: Duration,
242) -> Result<(), Box<dyn std::error::Error>> {
243 let start = Instant::now();
244 let mut next_log_at = start + Duration::from_secs(10);
245 tracing::info!("waiting for vm ipv4");
246 let status_path = instance_dir.join(STATUS_FILE_NAME);
247 let mut last_status: Option<String> = None;
248 let mut once_hint = false;
249 loop {
250 let config = load_or_create_instance_config(instance_dir)?;
251 if config.vm_ipv4.is_some() {
252 return Ok(());
253 }
254 if start.elapsed() > timeout {
255 return Err("Timed out waiting for VM IPv4".into());
256 }
257 if Instant::now() >= next_log_at {
258 let waited = start.elapsed();
259 if waited.as_secs() > 15 && !once_hint {
260 tracing::info!(
261 "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!"
262 );
263 once_hint = true;
264 }
265 if let Ok(status) = fs::read_to_string(&status_path) {
266 let status = status.trim().to_string();
267 if !status.is_empty() && last_status.as_deref() != Some(status.as_str()) {
268 tracing::info!("[background]: {}", status);
269 last_status = Some(status);
270 }
271 } else {
272 tracing::info!("still waiting for vm ipv4, {}s elapsed", waited.as_secs(),);
273 }
274 next_log_at += Duration::from_secs(10);
275 }
276 thread::sleep(Duration::from_millis(200));
277 }
280fn run_ssh_session(
281 ssh_key: PathBuf,
282 ssh_user: String,
283 ip: String,
284) -> Result<(), Box<dyn std::error::Error>> {
285 let mut attempts = 0usize;
286 loop {
287 attempts += 1;
288 if !ssh_port_open(&ip) {
289 tracing::debug!(attempts, "ssh port doesn't open yet");
290 tracing::info!(
291 attempts,
292 ip = %ip,
293 "waiting for ssh port ({}/{})",
294 attempts,
295 SSH_CONNECT_RETRIES
296 );
297 if attempts >= SSH_CONNECT_RETRIES {
298 return Err(
299 format!("ssh port not ready after {SSH_CONNECT_RETRIES} attempts").into(),
300 );
301 }
302 thread::sleep(Duration::from_millis(SSH_CONNECT_DELAY_MS));
303 continue;
304 }
306 tracing::info!(
307 attempts,
308 user = %ssh_user,
309 ip = %ip,
310 "starting ssh ({}/{})",
311 attempts,
312 SSH_CONNECT_RETRIES
313 );
314 let status = Command::new("ssh")
315 .args([
316 "-i",
317 ssh_key.to_str().unwrap_or(".vibebox/ssh_key"),
318 "-o",
319 "IdentitiesOnly=yes",
320 "-o",
321 "StrictHostKeyChecking=no",
322 "-o",
323 "UserKnownHostsFile=/dev/null",
324 "-o",
325 "GlobalKnownHostsFile=/dev/null",
326 "-o",
327 "PasswordAuthentication=no",
328 "-o",
329 "BatchMode=yes",
330 "-o",
331 "LogLevel=ERROR",
332 "-o",
333 "ConnectTimeout=5",
334 ])
335 .env_remove("LC_CTYPE")
336 .env_remove("LC_ALL")
337 .env_remove("LANG")
338 .arg(format!("{ssh_user}@{ip}"))
339 .stdin(Stdio::inherit())
340 .stdout(Stdio::inherit())
341 .stderr(Stdio::inherit())
342 .status();
344 match status {
345 Ok(status) if status.success() => {
346 tracing::info!(status = %status, "ssh exited");
347 break;
348 }
349 Ok(status) if status.code() == Some(255) => {
350 tracing::warn!(status = %status, "ssh connection failed");
351 if attempts >= SSH_CONNECT_RETRIES {
352 return Err(format!("ssh failed after {SSH_CONNECT_RETRIES} attempts").into());
353 }
354 thread::sleep(Duration::from_millis(500));
355 }
356 Ok(status) => {
357 tracing::info!(status = %status, "ssh exited");
358 break;
359 }
360 Err(err) => {
361 tracing::error!(error = %err, "failed to start ssh");
362 return Err(format!("failed to start ssh: {err}").into());
363 }
364 }
365 }
367 Ok(())
370fn is_ipv4_candidate(candidate: &str) -> bool {
371 let parts: Vec<&str> = candidate.split('.').collect();
372 if parts.len() != 4 {
373 return false;
374 }
375 for part in parts {
376 if part.is_empty() || part.len() > 3 {
377 return false;
378 }
379 if part.parse::<u8>().is_err() {
380 return false;
381 }
382 }
383 true
386fn ssh_port_open(ip: &str) -> bool {
387 let addr: SocketAddr = match format!("{ip}:22").parse() {
388 Ok(addr) => addr,
389 Err(_) => return false,
390 };
391 TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok()
394pub(crate) fn build_ssh_login_actions(
395 config: &Arc<Mutex<InstanceConfig>>,
396 project_name: &str,
397 guest_dir: &str,
398 key_name: &str,
399 home_links_script: &str,
400) -> Vec<LoginAction> {
401 let config_guard = config.lock().expect("config mutex poisoned");
402 let ssh_user = config_guard.ssh_user.clone();
403 let sudo_password = config_guard.sudo_password.clone();
404 drop(config_guard);
406 let key_path = format!("{guest_dir}/{key_name}.pub");
408 let setup_script = SSH_SETUP_SCRIPT
409 .replace("__SSH_USER__", &ssh_user)
410 .replace("__SUDO_PASSWORD__", &sudo_password)
411 .replace("__PROJECT_NAME__", project_name)
412 .replace("__KEY_PATH__", &key_path)
413 .replace("__VIBEBOX_SHELL_SCRIPT__", &commands::render_shell_script())
414 .replace("__VIBEBOX_HOME_LINKS__", home_links_script);
415 let setup = vm::script_command_from_content("ssh_setup", &setup_script)
416 .expect("ssh setup script contained invalid marker");
418 vec![LoginAction::Send(setup)]