git.iko.soy/vibebox
on main help sign in
feat: now added ssh
cd76a4frobcholz
src/bin/vibebox-cli.rs+3−5
compare with main
unifiedsplit
@@ -7,7 +7,7 @@ use std::{
7use color_eyre::Result;
7use color_eyre::Result;
8
8
9use vibebox::tui::{AppState, VmInfo};
9use vibebox::tui::{AppState, VmInfo};
10−use vibebox::{tui, vm};
10+use vibebox::{instance, tui, vm};
11
11
12fn main() -> Result<()> {
12fn main() -> Result<()> {
13 color_eyre::install()?;
13 color_eyre::install()?;
@@ -40,10 +40,8 @@ fn main() -> Result<()> {
40 stdout.flush()?;
40 stdout.flush()?;
41 }
41 }
42
42
43− vm::run_with_args(args, |output_monitor, vm_output_fd, vm_input_fd| {
43+ instance::run_with_ssh(args, app.clone())
44− tui::passthrough_vm_io(app.clone(), output_monitor, vm_output_fd, vm_input_fd)
45− })
46− .map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?;
44+ .map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?;
47
45
48 Ok(())
46 Ok(())
49}
47}
src/instance.rsnew file+521
@@ -0,0 +1,521 @@
1+use 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+};
18+
19+use serde::{Deserialize, Serialize};
20+use uuid::Uuid;
21+use std::sync::mpsc::Sender;
22+
23+use crate::{
24+ session_manager::INSTANCE_DIR_NAME,
25+ tui::{self, AppState},
26+ vm::{self, DirectoryShare, LoginAction, VmInput},
27+};
28+
29+const INSTANCE_TOML: &str = "instance.toml";
30+const SSH_KEY_NAME: &str = "ssh_key";
31+const SERIAL_LOG_NAME: &str = "serial.log";
32+const SSH_GUEST_DIR: &str = "/root/.vibebox";
33+const DEFAULT_SSH_USER: &str = "vibebox";
34+const SSH_CONNECT_RETRIES: usize = 20;
35+const SSH_CONNECT_DELAY_MS: u64 = 500;
36+
37+#[derive(Debug, Clone, Serialize, Deserialize)]
38+struct 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>,
45+}
46+
47+fn default_ssh_user() -> String {
48+ DEFAULT_SSH_USER.to_string()
49+}
50+
51+pub 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)?;
63+
64+ let config = load_or_create_instance_config(&instance_dir)?;
65+ let config = Arc::new(Mutex::new(config));
66+
67+ let extra_shares = vec![DirectoryShare::new(
68+ instance_dir.clone(),
69+ SSH_GUEST_DIR.into(),
70+ true,
71+ )?];
72+
73+ let extra_login_actions = build_ssh_login_actions(
74+ &config,
75+ &project_name,
76+ SSH_GUEST_DIR,
77+ SSH_KEY_NAME,
78+ );
79+
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+ )
96+}
97+
98+fn 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)
102+}
103+
104+fn 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"));
107+
108+ if private_key.exists() && public_key.exists() {
109+ return Ok((private_key, public_key));
110+ }
111+
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+ }
118+
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()?;
133+
134+ if !status.success() {
135+ return Err("ssh-keygen failed".into());
136+ }
137+
138+ fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?;
139+ fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?;
140+
141+ Ok((private_key, public_key))
142+}
143+
144+fn 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+ };
158+
159+ let mut changed = false;
160+ if config.ssh_user.trim().is_empty() {
161+ config.ssh_user = default_ssh_user();
162+ changed = true;
163+ }
164+
165+ if config.sudo_password.trim().is_empty() {
166+ config.sudo_password = generate_password();
167+ changed = true;
168+ }
169+
170+ if !config_path.exists() || changed {
171+ write_instance_config(&config_path, &config)?;
172+ }
173+
174+ Ok(config)
175+}
176+
177+fn 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(())
185+}
186+
187+fn generate_password() -> String {
188+ Uuid::now_v7().simple().to_string()
189+}
190+
191+fn extract_ipv4(line: &str) -> Option<String> {
192+ let mut current = String::new();
193+ let mut best: Option<String> = None;
194+
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+ }
206+
207+ best
208+}
209+
210+fn 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
224+}
225+
226+fn 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()
232+}
233+
234+fn 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);
244+
245+ let key_path = format!("{guest_dir}/{key_name}.pub");
246+
247+ let mask_cache = format!(
248+ "if [ -d /root/{project_name}/.vibebox ]; then mount -t tmpfs tmpfs /root/{project_name}/.vibebox; fi"
249+ );
250+
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");
273+
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";
275+
276+ vec![
277+ LoginAction::Send(mask_cache),
278+ LoginAction::Send(setup),
279+ LoginAction::Send(ip_probe.to_string()),
280+ ]
281+}
282+
283+fn 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();
293+
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)));
301+
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));
307+
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();
313+
314+ let mut line_buf = String::new();
315+
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+ };
328+
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+ }
337+
338+ let text = String::from_utf8_lossy(bytes);
339+ line_buf.push_str(&text);
340+
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);
347+
348+ let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' ');
349+
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+ }
363+
364+ if cleaned.contains("VIBEBOX_SSH_READY") {
365+ ssh_ready_for_output.store(true, Ordering::SeqCst);
366+ eprintln!("[vibebox] sshd ready");
367+ }
368+ }
369+ };
370+
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+ );
379+
380+ *input_tx_holder.lock().unwrap() = Some(io_ctx.input_tx.clone());
381+
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();
389+
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+ }
399+
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());
408+
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+ }
416+
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);
422+
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+ }
442+
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();
474+
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+ }
515+
516+ thread::sleep(std::time::Duration::from_millis(200));
517+ }
518+ });
519+
520+ io_ctx
521+}
src/lib.rs+1
@@ -1,5 +1,6 @@
1pub mod session_manager;
1pub mod session_manager;
2pub mod tui;
2pub mod tui;
3pub mod vm;
3pub mod vm;
4+pub mod instance;
4
5
5pub use session_manager::{SessionError, SessionManager, SessionRecord};
6pub use session_manager::{SessionError, SessionManager, SessionRecord};
src/provision.sh+6−1
@@ -13,11 +13,16 @@ apt-get install -y --no-install-recommends \
13 libssl-dev \
13 libssl-dev \
14 curl \
14 curl \
15 git \
15 git \
16− ripgrep
16+ ripgrep \
17+ openssh-server \
18+ sudo
17
19
18# Set hostname to "vibe" so it's clear that you're inside the VM.
20# Set hostname to "vibe" so it's clear that you're inside the VM.
19hostnamectl set-hostname vibe
21hostnamectl set-hostname vibe
20
22
23+# Enable SSH server so instances can use key-based auth.
24+systemctl enable ssh
25+
21# Set this env var so claude doesn't complain about running as root.'
26# Set this env var so claude doesn't complain about running as root.'
22echo "export IS_SANDBOX=1" >> .bashrc
27echo "export IS_SANDBOX=1" >> .bashrc
23
28
src/vm.rs+123−42
@@ -15,6 +15,7 @@ use std::{
15 process::{Command, Stdio},
15 process::{Command, Stdio},
16 sync::{
16 sync::{
17 Arc, Condvar, Mutex,
17 Arc, Condvar, Mutex,
18+ atomic::{AtomicBool, Ordering},
18 mpsc::{self, Receiver, Sender},
19 mpsc::{self, Receiver, Sender},
19 },
20 },
20 thread,
21 thread,
@@ -43,7 +44,7 @@ const LOGIN_EXPECT_TIMEOUT: Duration = Duration::from_secs(120);
43const PROVISION_SCRIPT: &str = include_str!("provision.sh");
44const PROVISION_SCRIPT: &str = include_str!("provision.sh");
44
45
45#[derive(Clone)]
46#[derive(Clone)]
46−enum LoginAction {
47+pub(crate) enum LoginAction {
47 Expect { text: String, timeout: Duration },
48 Expect { text: String, timeout: Duration },
48 Send(String),
49 Send(String),
49 Script { path: PathBuf, index: usize },
50 Script { path: PathBuf, index: usize },
@@ -51,14 +52,14 @@ enum LoginAction {
51use LoginAction::*;
52use LoginAction::*;
52
53
53#[derive(Clone)]
54#[derive(Clone)]
54−struct DirectoryShare {
55+pub(crate) struct DirectoryShare {
55 host: PathBuf,
56 host: PathBuf,
56 guest: PathBuf,
57 guest: PathBuf,
57 read_only: bool,
58 read_only: bool,
58}
59}
59
60
60impl DirectoryShare {
61impl DirectoryShare {
61− fn new(
62+ pub(crate) fn new(
62 host: PathBuf,
63 host: PathBuf,
63 mut guest: PathBuf,
64 mut guest: PathBuf,
64 read_only: bool,
65 read_only: bool,
@@ -76,7 +77,7 @@ impl DirectoryShare {
76 })
77 })
77 }
78 }
78
79
79− fn from_mount_spec(spec: &str) -> Result<Self, Box<dyn std::error::Error>> {
80+ pub(crate) fn from_mount_spec(spec: &str) -> Result<Self, Box<dyn std::error::Error>> {
80 let parts: Vec<&str> = spec.split(':').collect();
81 let parts: Vec<&str> = spec.split(':').collect();
81 if parts.len() < 2 || parts.len() > 3 {
82 if parts.len() < 2 || parts.len() > 3 {
82 return Err(format!("Invalid mount spec: {spec}").into());
83 return Err(format!("Invalid mount spec: {spec}").into());
@@ -115,23 +116,19 @@ impl DirectoryShare {
115 }
116 }
116}
117}
117
118
118−pub fn run_cli() -> Result<(), Box<dyn std::error::Error>> {
119+pub fn run_with_args<F>(args: CliArgs, io_handler: F) -> Result<(), Box<dyn std::error::Error>>
119− let args = parse_cli()?;
120+where
120−
121+ F: FnOnce(Arc<OutputMonitor>, OwnedFd, OwnedFd) -> IoContext,
121− if args.version() {
122+{
122− print_version();
123+ run_with_args_and_extras(args, io_handler, Vec::new(), Vec::new())
123− return Ok(());
124− }
125−
126− if args.help() {
127− print_help();
128− return Ok(());
129− }
130−
131− run_with_args(args, spawn_vm_io)
132}
124}
133
125
126+pub(crate) fn run_with_args_and_extras<F>(
127+ args: CliArgs,
128+ io_handler: F,
129+ extra_login_actions: Vec<LoginAction>,
130+ extra_directory_shares: Vec<DirectoryShare>,
134−pub fn run_with_args<F>(args: CliArgs, io_handler: F) -> Result<(), Box<dyn std::error::Error>>
131+) -> Result<(), Box<dyn std::error::Error>>
135where
132where
136 F: FnOnce(Arc<OutputMonitor>, OwnedFd, OwnedFd) -> IoContext,
133 F: FnOnce(Arc<OutputMonitor>, OwnedFd, OwnedFd) -> IoContext,
137{
134{
@@ -233,6 +230,8 @@ where
233 }
230 }
234 }
231 }
235
232
233+ directory_shares.extend(extra_directory_shares);
234+
236 for spec in &args.mounts {
235 for spec in &args.mounts {
237 directory_shares.push(DirectoryShare::from_mount_spec(spec)?);
236 directory_shares.push(DirectoryShare::from_mount_spec(spec)?);
238 }
237 }
@@ -241,6 +240,8 @@ where
241 login_actions.push(motd_action);
240 login_actions.push(motd_action);
242 }
241 }
243
242
243+ login_actions.extend(extra_login_actions);
244+
244 // Any user-provided login actions must come after our system ones
245 // Any user-provided login actions must come after our system ones
245 login_actions.extend(args.login_actions);
246 login_actions.extend(args.login_actions);
246
247
@@ -461,18 +462,6 @@ fn motd_login_action(directory_shares: &[DirectoryShare]) -> Option<LoginAction>
461 }
462 }
462
463
463 let mut output = String::new();
464 let mut output = String::new();
464− output.push_str(
465− "
466−░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░▒▓███████▓▒░░▒▓████████▓▒░
467−░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░
468− ░▒▓█▓▒▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░
469− ░▒▓█▓▒▒▓█▓▒░░▒▓█▓▒░▒▓███████▓▒░░▒▓██████▓▒░
470− ░▒▓█▓▓█▓▒░ ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░
471− ░▒▓█▓▓█▓▒░ ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░
472− ░▒▓██▓▒░ ░▒▓█▓▒░▒▓███████▓▒░░▒▓████████▓▒░
473−
474−",
475− );
476 output.push_str(&format!(
465 output.push_str(&format!(
477 "{host_header:<host_width$} {guest_header:<guest_width$} {mode_header}\n",
466 "{host_header:<host_width$} {guest_header:<guest_width$} {mode_header}\n",
478 host_width = host_width
467 host_width = host_width
@@ -548,6 +537,49 @@ impl OutputMonitor {
548 }
537 }
549}
538}
550
539
540+#[derive(Debug)]
541+pub struct IoControl {
542+ forward_input: AtomicBool,
543+ forward_output: AtomicBool,
544+ restore_terminal: AtomicBool,
545+}
546+
547+impl IoControl {
548+ pub fn new() -> Arc<Self> {
549+ Arc::new(Self {
550+ forward_input: AtomicBool::new(true),
551+ forward_output: AtomicBool::new(true),
552+ restore_terminal: AtomicBool::new(false),
553+ })
554+ }
555+
556+ pub fn set_forward_input(&self, enabled: bool) {
557+ self.forward_input.store(enabled, Ordering::SeqCst);
558+ }
559+
560+ pub fn set_forward_output(&self, enabled: bool) {
561+ self.forward_output.store(enabled, Ordering::SeqCst);
562+ }
563+
564+ pub fn request_terminal_restore(&self) {
565+ self.restore_terminal.store(true, Ordering::SeqCst);
566+ }
567+
568+ fn forward_input(&self) -> bool {
569+ self.forward_input.load(Ordering::SeqCst)
570+ }
571+
572+ fn forward_output(&self) -> bool {
573+ self.forward_output.load(Ordering::SeqCst)
574+ }
575+
576+ fn take_restore_terminal(&self) -> bool {
577+ self.restore_terminal
578+ .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst)
579+ .is_ok()
580+ }
581+}
582+
551fn ensure_base_image(
583fn ensure_base_image(
552 base_raw: &Path,
584 base_raw: &Path,
553 base_compressed: &Path,
585 base_compressed: &Path,
@@ -667,14 +699,17 @@ pub fn create_pipe() -> (OwnedFd, OwnedFd) {
667 (read_stream.into(), write_stream.into())
699 (read_stream.into(), write_stream.into())
668}
700}
669
701
670−pub fn spawn_vm_io_with_line_handler<F>(
702+pub fn spawn_vm_io_with_hooks<F, G>(
671 output_monitor: Arc<OutputMonitor>,
703 output_monitor: Arc<OutputMonitor>,
672 vm_output_fd: OwnedFd,
704 vm_output_fd: OwnedFd,
673 vm_input_fd: OwnedFd,
705 vm_input_fd: OwnedFd,
706+ io_control: Arc<IoControl>,
674 mut on_line: F,
707 mut on_line: F,
708+ mut on_output: G,
675) -> IoContext
709) -> IoContext
676where
710where
677 F: FnMut(&str) -> bool + ::std::marker::Send + 'static,
711 F: FnMut(&str) -> bool + ::std::marker::Send + 'static,
712+ G: FnMut(&[u8]) + ::std::marker::Send + 'static,
678{
713{
679 let (input_tx, input_rx): (Sender<VmInput>, Receiver<VmInput>) = mpsc::channel();
714 let (input_tx, input_rx): (Sender<VmInput>, Receiver<VmInput>) = mpsc::channel();
680
715
@@ -721,17 +756,36 @@ where
721 }
756 }
722 }
757 }
723
758
759+ fn poll_wakeup_only(wakeup_fd: RawFd, timeout_ms: i32) -> bool {
760+ let mut fds = [libc::pollfd {
761+ fd: wakeup_fd,
762+ events: libc::POLLIN,
763+ revents: 0,
764+ }];
765+
766+ let ret = unsafe { libc::poll(fds.as_mut_ptr(), 1, timeout_ms) };
767+ ret > 0 && fds[0].revents & libc::POLLIN != 0
768+ }
769+
724 // Copies from stdin to the VM; also polls wakeup_read to exit the thread when it's time to shutdown.
770 // Copies from stdin to the VM; also polls wakeup_read to exit the thread when it's time to shutdown.
725 let stdin_thread = thread::spawn({
771 let stdin_thread = thread::spawn({
726 let input_tx = input_tx.clone();
772 let input_tx = input_tx.clone();
727 let raw_guard = raw_guard.clone();
773 let raw_guard = raw_guard.clone();
728 let wakeup_read = wakeup_read.try_clone().unwrap();
774 let wakeup_read = wakeup_read.try_clone().unwrap();
775+ let io_control = io_control.clone();
729
776
730 move || {
777 move || {
731 let mut buf = [0u8; 1024];
778 let mut buf = [0u8; 1024];
732 let mut pending_command: Vec<u8> = Vec::new();
779 let mut pending_command: Vec<u8> = Vec::new();
733 let mut command_mode = false;
780 let mut command_mode = false;
734 loop {
781 loop {
782+ if !io_control.forward_input() {
783+ if poll_wakeup_only(wakeup_read.as_raw_fd(), 100) {
784+ break;
785+ }
786+ continue;
787+ }
788+
735 match poll_with_wakeup(libc::STDIN_FILENO, wakeup_read.as_raw_fd(), &mut buf) {
789 match poll_with_wakeup(libc::STDIN_FILENO, wakeup_read.as_raw_fd(), &mut buf) {
736 PollResult::Shutdown | PollResult::Error => break,
790 PollResult::Shutdown | PollResult::Error => break,
737 PollResult::Spurious => continue,
791 PollResult::Spurious => continue,
@@ -776,28 +830,36 @@ where
776 let stdout_thread = thread::spawn({
830 let stdout_thread = thread::spawn({
777 let raw_guard = raw_guard.clone();
831 let raw_guard = raw_guard.clone();
778 let wakeup_read = wakeup_read.try_clone().unwrap();
832 let wakeup_read = wakeup_read.try_clone().unwrap();
833+ let io_control = io_control.clone();
779
834
780 move || {
835 move || {
781− let mut stdout = std::io::stdout().lock();
782 let mut buf = [0u8; 1024];
836 let mut buf = [0u8; 1024];
783 loop {
837 loop {
838+ if io_control.take_restore_terminal() {
839+ let mut guard = raw_guard.lock().unwrap();
840+ *guard = None;
841+ }
784 match poll_with_wakeup(vm_output_fd.as_raw_fd(), wakeup_read.as_raw_fd(), &mut buf)
842 match poll_with_wakeup(vm_output_fd.as_raw_fd(), wakeup_read.as_raw_fd(), &mut buf)
785 {
843 {
786 PollResult::Shutdown | PollResult::Error => break,
844 PollResult::Shutdown | PollResult::Error => break,
787 PollResult::Spurious => continue,
845 PollResult::Spurious => continue,
788 PollResult::Ready(bytes) => {
846 PollResult::Ready(bytes) => {
847+ if io_control.forward_output() {
789− // enable raw mode, if we haven't already
848+ // enable raw mode, if we haven't already
790− if raw_guard.lock().unwrap().is_none()
849+ if raw_guard.lock().unwrap().is_none()
791− && let Ok(guard) = enable_raw_mode(libc::STDIN_FILENO)
850+ && let Ok(guard) = enable_raw_mode(libc::STDIN_FILENO)
792− {
851+ {
793− *raw_guard.lock().unwrap() = Some(guard);
852+ *raw_guard.lock().unwrap() = Some(guard);
794− }
853+ }
795
854
855+ let mut stdout = std::io::stdout().lock();
796− if stdout.write_all(bytes).is_err() {
856+ if stdout.write_all(bytes).is_err() {
797− break;
857+ break;
858+ }
859+ let _ = stdout.flush();
798 }
860 }
799− let _ = stdout.flush();
800 output_monitor.push(bytes);
861 output_monitor.push(bytes);
862+ on_output(bytes);
801 }
863 }
802 }
864 }
803 }
865 }
@@ -829,6 +891,25 @@ where
829 }
891 }
830}
892}
831
893
894+pub fn spawn_vm_io_with_line_handler<F>(
895+ output_monitor: Arc<OutputMonitor>,
896+ vm_output_fd: OwnedFd,
897+ vm_input_fd: OwnedFd,
898+ on_line: F,
899+) -> IoContext
900+where
901+ F: FnMut(&str) -> bool + ::std::marker::Send + 'static,
902+{
903+ spawn_vm_io_with_hooks(
904+ output_monitor,
905+ vm_output_fd,
906+ vm_input_fd,
907+ IoControl::new(),
908+ on_line,
909+ |_| {},
910+ )
911+}
912+
832pub fn spawn_vm_io(
913pub fn spawn_vm_io(
833 output_monitor: Arc<OutputMonitor>,
914 output_monitor: Arc<OutputMonitor>,
834 vm_output_fd: OwnedFd,
915 vm_output_fd: OwnedFd,