diff --git a/src/bin/vibebox-cli.rs b/src/bin/vibebox-cli.rs index 5c27167..6b14c60 100644 --- a/src/bin/vibebox-cli.rs +++ b/src/bin/vibebox-cli.rs @@ -7,7 +7,7 @@ use std::{ use color_eyre::Result; use vibebox::tui::{AppState, VmInfo}; -use vibebox::{tui, vm}; +use vibebox::{instance, tui, vm}; fn main() -> Result<()> { color_eyre::install()?; @@ -40,10 +40,8 @@ fn main() -> Result<()> { stdout.flush()?; } - vm::run_with_args(args, |output_monitor, vm_output_fd, vm_input_fd| { - tui::passthrough_vm_io(app.clone(), output_monitor, vm_output_fd, vm_input_fd) - }) - .map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + instance::run_with_ssh(args, app.clone()) + .map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; Ok(()) } diff --git a/src/instance.rs b/src/instance.rs new file mode 100644 index 0000000..e10b0ad --- /dev/null +++ b/src/instance.rs @@ -0,0 +1,521 @@ +use std::{ + env, + fs, + io::{self, Write}, + net::{SocketAddr, TcpStream}, + os::unix::{ + fs::PermissionsExt, + io::OwnedFd, + }, + path::{Path, PathBuf}, + process::{Command, Stdio}, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, + thread, +}; + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; +use std::sync::mpsc::Sender; + +use crate::{ + session_manager::INSTANCE_DIR_NAME, + tui::{self, AppState}, + vm::{self, DirectoryShare, LoginAction, VmInput}, +}; + +const INSTANCE_TOML: &str = "instance.toml"; +const SSH_KEY_NAME: &str = "ssh_key"; +const SERIAL_LOG_NAME: &str = "serial.log"; +const SSH_GUEST_DIR: &str = "/root/.vibebox"; +const DEFAULT_SSH_USER: &str = "vibebox"; +const SSH_CONNECT_RETRIES: usize = 20; +const SSH_CONNECT_DELAY_MS: u64 = 500; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct InstanceConfig { + #[serde(default = "default_ssh_user")] + ssh_user: String, + #[serde(default)] + sudo_password: String, + #[serde(default)] + vm_ipv4: Option, +} + +fn default_ssh_user() -> String { + DEFAULT_SSH_USER.to_string() +} + +pub fn run_with_ssh( + args: vm::CliArgs, + app: Arc>, +) -> Result<(), Box> { + let project_root = env::current_dir()?; + let project_name = project_root + .file_name() + .ok_or("Project directory has no name")? + .to_string_lossy() + .into_owned(); + let instance_dir = ensure_instance_dir(&project_root)?; + let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?; + + let config = load_or_create_instance_config(&instance_dir)?; + let config = Arc::new(Mutex::new(config)); + + let extra_shares = vec![DirectoryShare::new( + instance_dir.clone(), + SSH_GUEST_DIR.into(), + true, + )?]; + + let extra_login_actions = build_ssh_login_actions( + &config, + &project_name, + SSH_GUEST_DIR, + SSH_KEY_NAME, + ); + + vm::run_with_args_and_extras( + args, + |output_monitor, vm_output_fd, vm_input_fd| { + spawn_ssh_io( + app.clone(), + config.clone(), + instance_dir.clone(), + ssh_key.clone(), + output_monitor, + vm_output_fd, + vm_input_fd, + ) + }, + extra_login_actions, + extra_shares, + ) +} + +fn ensure_instance_dir(project_root: &Path) -> Result { + let instance_dir = project_root.join(INSTANCE_DIR_NAME); + fs::create_dir_all(&instance_dir)?; + Ok(instance_dir) +} + +fn ensure_ssh_keypair(instance_dir: &Path) -> Result<(PathBuf, PathBuf), Box> { + let private_key = instance_dir.join(SSH_KEY_NAME); + let public_key = instance_dir.join(format!("{SSH_KEY_NAME}.pub")); + + if private_key.exists() && public_key.exists() { + return Ok((private_key, public_key)); + } + + if private_key.exists() { + let _ = fs::remove_file(&private_key); + } + if public_key.exists() { + let _ = fs::remove_file(&public_key); + } + + let status = Command::new("ssh-keygen") + .args([ + "-t", + "ed25519", + "-N", + "", + "-f", + private_key + .to_str() + .ok_or("ssh key path not utf-8")?, + "-C", + "vibebox", + ]) + .status()?; + + if !status.success() { + return Err("ssh-keygen failed".into()); + } + + fs::set_permissions(&private_key, fs::Permissions::from_mode(0o600))?; + fs::set_permissions(&public_key, fs::Permissions::from_mode(0o644))?; + + Ok((private_key, public_key)) +} + +fn load_or_create_instance_config( + instance_dir: &Path, +) -> Result> { + let config_path = instance_dir.join(INSTANCE_TOML); + let mut config = if config_path.exists() { + let raw = fs::read_to_string(&config_path)?; + toml::from_str::(&raw)? + } else { + InstanceConfig { + ssh_user: default_ssh_user(), + sudo_password: String::new(), + vm_ipv4: None, + } + }; + + let mut changed = false; + if config.ssh_user.trim().is_empty() { + config.ssh_user = default_ssh_user(); + changed = true; + } + + if config.sudo_password.trim().is_empty() { + config.sudo_password = generate_password(); + changed = true; + } + + if !config_path.exists() || changed { + write_instance_config(&config_path, &config)?; + } + + Ok(config) +} + +fn write_instance_config( + path: &Path, + config: &InstanceConfig, +) -> Result<(), Box> { + let data = toml::to_string_pretty(config)?; + fs::write(path, data)?; + fs::set_permissions(path, fs::Permissions::from_mode(0o600))?; + Ok(()) +} + +fn generate_password() -> String { + Uuid::now_v7().simple().to_string() +} + +fn extract_ipv4(line: &str) -> Option { + let mut current = String::new(); + let mut best: Option = None; + + for ch in line.chars().chain(std::iter::once(' ')) { + if ch.is_ascii_digit() || ch == '.' { + current.push(ch); + } else if !current.is_empty() { + if is_ipv4_candidate(¤t) { + best = Some(current.clone()); + break; + } + current.clear(); + } + } + + best +} + +fn is_ipv4_candidate(candidate: &str) -> bool { + let parts: Vec<&str> = candidate.split('.').collect(); + if parts.len() != 4 { + return false; + } + for part in parts { + if part.is_empty() || part.len() > 3 { + return false; + } + if part.parse::().is_err() { + return false; + } + } + true +} + +fn ssh_port_open(ip: &str) -> bool { + let addr: SocketAddr = match format!("{ip}:22").parse() { + Ok(addr) => addr, + Err(_) => return false, + }; + TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok() +} + +fn build_ssh_login_actions( + config: &Arc>, + project_name: &str, + guest_dir: &str, + key_name: &str, +) -> Vec { + let config_guard = config.lock().expect("config mutex poisoned"); + let ssh_user = config_guard.ssh_user.clone(); + let sudo_password = config_guard.sudo_password.clone(); + drop(config_guard); + + let key_path = format!("{guest_dir}/{key_name}.pub"); + + let mask_cache = format!( + "if [ -d /root/{project_name}/.vibebox ]; then mount -t tmpfs tmpfs /root/{project_name}/.vibebox; fi" + ); + + let setup_lines = [ + "if ! command -v sshd >/dev/null 2>&1; then apt-get update && apt-get install -y openssh-server sudo; fi", + "systemctl enable ssh >/dev/null 2>&1 || true", + &format!("id -u {ssh_user} >/dev/null 2>&1 || useradd -m -s /bin/bash {ssh_user}"), + &format!("echo \"{ssh_user}:{sudo_password}\" | chpasswd"), + &format!("usermod -aG sudo {ssh_user}"), + &format!("install -d -m 700 /home/{ssh_user}/.ssh"), + &format!("install -m 600 {key_path} /home/{ssh_user}/.ssh/authorized_keys"), + &format!("chown -R {ssh_user}:{ssh_user} /home/{ssh_user}/.ssh"), + "mkdir -p /etc/ssh/sshd_config.d", + "cat >/etc/ssh/sshd_config.d/vibebox.conf <<'VIBEBOX_SSHD'", + "PasswordAuthentication no", + "KbdInteractiveAuthentication no", + "ChallengeResponseAuthentication no", + "PubkeyAuthentication yes", + "PermitRootLogin no", + &format!("AllowUsers {ssh_user}"), + "VIBEBOX_SSHD", + "systemctl restart ssh", + "echo VIBEBOX_SSH_READY", + ]; + let setup = setup_lines.join("\n"); + + 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"; + + vec![ + LoginAction::Send(mask_cache), + LoginAction::Send(setup), + LoginAction::Send(ip_probe.to_string()), + ] +} + +fn spawn_ssh_io( + app: Arc>, + config: Arc>, + instance_dir: PathBuf, + ssh_key: PathBuf, + output_monitor: Arc, + vm_output_fd: OwnedFd, + vm_input_fd: OwnedFd, +) -> vm::IoContext { + let io_control = vm::IoControl::new(); + + let log_path = instance_dir.join(SERIAL_LOG_NAME); + let log_file = fs::OpenOptions::new() + .create(true) + .append(true) + .open(&log_path) + .ok() + .map(|file| Arc::new(Mutex::new(file))); + + let ssh_connected = Arc::new(AtomicBool::new(false)); + let ssh_started = Arc::new(AtomicBool::new(false)); + let ssh_ready = Arc::new(AtomicBool::new(false)); + let input_tx_holder: Arc>>> = + Arc::new(Mutex::new(None)); + + let instance_path = instance_dir.join(INSTANCE_TOML); + let config_for_output = config.clone(); + let log_for_output = log_file.clone(); + let ssh_connected_for_output = ssh_connected.clone(); + let ssh_ready_for_output = ssh_ready.clone(); + + let mut line_buf = String::new(); + + let on_line = { + let app = app.clone(); + move |line: &str| { + if line == ":help" { + if let Ok(mut locked) = app.lock() { + let _ = tui::render_commands_component(&mut locked); + } + return true; + } + false + } + }; + + let on_output = move |bytes: &[u8]| { + if ssh_connected_for_output.load(Ordering::SeqCst) { + if let Some(log) = &log_for_output { + if let Ok(mut file) = log.lock() { + let _ = file.write_all(bytes); + } + } + } + + let text = String::from_utf8_lossy(bytes); + line_buf.push_str(&text); + + while let Some(pos) = line_buf.find('\n') { + let mut line = line_buf[..pos].to_string(); + if line.ends_with('\r') { + line.pop(); + } + line_buf.drain(..=pos); + + let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' '); + + if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") { + let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..]; + let ip = extract_ipv4(ip_raw).unwrap_or_default(); + if !ip.is_empty() { + if let Ok(mut cfg) = config_for_output.lock() { + if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) { + cfg.vm_ipv4 = Some(ip.clone()); + let _ = write_instance_config(&instance_path, &cfg); + eprintln!("[vibebox] detected vm IPv4: {ip}"); + } + } + } + } + + if cleaned.contains("VIBEBOX_SSH_READY") { + ssh_ready_for_output.store(true, Ordering::SeqCst); + eprintln!("[vibebox] sshd ready"); + } + } + }; + + let io_ctx = vm::spawn_vm_io_with_hooks( + output_monitor, + vm_output_fd, + vm_input_fd, + io_control.clone(), + on_line, + on_output, + ); + + *input_tx_holder.lock().unwrap() = Some(io_ctx.input_tx.clone()); + + let ssh_ready_for_thread = ssh_ready.clone(); + let ssh_started_for_thread = ssh_started.clone(); + let ssh_connected_for_thread = ssh_connected.clone(); + let io_control_for_thread = io_control.clone(); + let ssh_key_for_thread = ssh_key.clone(); + let config_for_thread = config.clone(); + let input_tx_holder_for_thread = input_tx_holder.clone(); + + thread::spawn(move || { + loop { + if ssh_started_for_thread.load(Ordering::SeqCst) { + break; + } + if !ssh_ready_for_thread.load(Ordering::SeqCst) { + thread::sleep(std::time::Duration::from_millis(200)); + continue; + } + + let ip = config_for_thread + .lock() + .ok() + .and_then(|cfg| cfg.vm_ipv4.clone()); + let ssh_user = config_for_thread + .lock() + .map(|cfg| cfg.ssh_user.clone()) + .unwrap_or_else(|_| DEFAULT_SSH_USER.to_string()); + + if let Some(ip) = ip { + if ssh_started_for_thread + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_err() + { + break; + } + + eprintln!("[vibebox] starting ssh to {ssh_user}@{ip}"); + io_control_for_thread.request_terminal_restore(); + io_control_for_thread.set_forward_output(false); + io_control_for_thread.set_forward_input(false); + ssh_connected_for_thread.store(true, Ordering::SeqCst); + + let mut attempts = 0usize; + loop { + attempts += 1; + if !ssh_port_open(&ip) { + eprintln!( + "[vibebox] waiting for ssh port on {ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})" + ); + if attempts >= SSH_CONNECT_RETRIES { + ssh_connected_for_thread.store(false, Ordering::SeqCst); + eprintln!( + "[vibebox] ssh port not ready after {SSH_CONNECT_RETRIES} attempts" + ); + break; + } + thread::sleep(std::time::Duration::from_millis( + SSH_CONNECT_DELAY_MS, + )); + continue; + } + + eprintln!( + "[vibebox] starting ssh to {ssh_user}@{ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})" + ); + let status = Command::new("ssh") + .args([ + "-i", + ssh_key_for_thread + .to_str() + .unwrap_or(".vibebox/ssh_key"), + "-o", + "IdentitiesOnly=yes", + "-o", + "StrictHostKeyChecking=no", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "GlobalKnownHostsFile=/dev/null", + "-o", + "PasswordAuthentication=no", + "-o", + "BatchMode=yes", + "-o", + "LogLevel=ERROR", + "-o", + "ConnectTimeout=5", + ]) + .arg(format!("{ssh_user}@{ip}")) + .stdin(Stdio::inherit()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status(); + + match status { + Ok(status) if status.success() => { + ssh_connected_for_thread.store(false, Ordering::SeqCst); + eprintln!("[vibebox] ssh exited with status: {status}"); + if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() { + let _ = + tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec())); + } + break; + } + Ok(status) if status.code() == Some(255) => { + eprintln!("[vibebox] ssh connection failed: {status}"); + if attempts >= SSH_CONNECT_RETRIES { + ssh_connected_for_thread.store(false, Ordering::SeqCst); + eprintln!( + "[vibebox] ssh failed after {SSH_CONNECT_RETRIES} attempts" + ); + break; + } + thread::sleep(std::time::Duration::from_millis(500)); + continue; + } + Ok(status) => { + ssh_connected_for_thread.store(false, Ordering::SeqCst); + eprintln!("[vibebox] ssh exited with status: {status}"); + if let Some(tx) = input_tx_holder_for_thread.lock().unwrap().clone() { + let _ = + tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec())); + } + break; + } + Err(err) => { + ssh_connected_for_thread.store(false, Ordering::SeqCst); + eprintln!("[vibebox] failed to start ssh: {err}"); + break; + } + } + } + break; + } + + thread::sleep(std::time::Duration::from_millis(200)); + } + }); + + io_ctx +} diff --git a/src/lib.rs b/src/lib.rs index 4155675..f0c19a1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ pub mod session_manager; pub mod tui; pub mod vm; +pub mod instance; pub use session_manager::{SessionError, SessionManager, SessionRecord}; diff --git a/src/provision.sh b/src/provision.sh index fbad94d..22db145 100644 --- a/src/provision.sh +++ b/src/provision.sh @@ -13,11 +13,16 @@ apt-get install -y --no-install-recommends \ libssl-dev \ curl \ git \ - ripgrep + ripgrep \ + openssh-server \ + sudo # Set hostname to "vibe" so it's clear that you're inside the VM. hostnamectl set-hostname vibe +# Enable SSH server so instances can use key-based auth. +systemctl enable ssh + # Set this env var so claude doesn't complain about running as root.' echo "export IS_SANDBOX=1" >> .bashrc diff --git a/src/vm.rs b/src/vm.rs index 0e021fa..85eb16b 100644 --- a/src/vm.rs +++ b/src/vm.rs @@ -15,6 +15,7 @@ use std::{ process::{Command, Stdio}, sync::{ Arc, Condvar, Mutex, + atomic::{AtomicBool, Ordering}, mpsc::{self, Receiver, Sender}, }, thread, @@ -43,7 +44,7 @@ const LOGIN_EXPECT_TIMEOUT: Duration = Duration::from_secs(120); const PROVISION_SCRIPT: &str = include_str!("provision.sh"); #[derive(Clone)] -enum LoginAction { +pub(crate) enum LoginAction { Expect { text: String, timeout: Duration }, Send(String), Script { path: PathBuf, index: usize }, @@ -51,14 +52,14 @@ enum LoginAction { use LoginAction::*; #[derive(Clone)] -struct DirectoryShare { +pub(crate) struct DirectoryShare { host: PathBuf, guest: PathBuf, read_only: bool, } impl DirectoryShare { - fn new( + pub(crate) fn new( host: PathBuf, mut guest: PathBuf, read_only: bool, @@ -76,7 +77,7 @@ impl DirectoryShare { }) } - fn from_mount_spec(spec: &str) -> Result> { + pub(crate) fn from_mount_spec(spec: &str) -> Result> { let parts: Vec<&str> = spec.split(':').collect(); if parts.len() < 2 || parts.len() > 3 { return Err(format!("Invalid mount spec: {spec}").into()); @@ -115,23 +116,19 @@ impl DirectoryShare { } } -pub fn run_cli() -> Result<(), Box> { - let args = parse_cli()?; - - if args.version() { - print_version(); - return Ok(()); - } - - if args.help() { - print_help(); - return Ok(()); - } - - run_with_args(args, spawn_vm_io) +pub fn run_with_args(args: CliArgs, io_handler: F) -> Result<(), Box> +where + F: FnOnce(Arc, OwnedFd, OwnedFd) -> IoContext, +{ + run_with_args_and_extras(args, io_handler, Vec::new(), Vec::new()) } -pub fn run_with_args(args: CliArgs, io_handler: F) -> Result<(), Box> +pub(crate) fn run_with_args_and_extras( + args: CliArgs, + io_handler: F, + extra_login_actions: Vec, + extra_directory_shares: Vec, +) -> Result<(), Box> where F: FnOnce(Arc, OwnedFd, OwnedFd) -> IoContext, { @@ -233,6 +230,8 @@ where } } + directory_shares.extend(extra_directory_shares); + for spec in &args.mounts { directory_shares.push(DirectoryShare::from_mount_spec(spec)?); } @@ -241,6 +240,8 @@ where login_actions.push(motd_action); } + login_actions.extend(extra_login_actions); + // Any user-provided login actions must come after our system ones login_actions.extend(args.login_actions); @@ -461,18 +462,6 @@ fn motd_login_action(directory_shares: &[DirectoryShare]) -> Option } let mut output = String::new(); - output.push_str( - " -░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░▒▓███████▓▒░░▒▓████████▓▒░ -░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░ - ░▒▓█▓▒▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░ - ░▒▓█▓▒▒▓█▓▒░░▒▓█▓▒░▒▓███████▓▒░░▒▓██████▓▒░ - ░▒▓█▓▓█▓▒░ ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░ - ░▒▓█▓▓█▓▒░ ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░ - ░▒▓██▓▒░ ░▒▓█▓▒░▒▓███████▓▒░░▒▓████████▓▒░ - -", - ); output.push_str(&format!( "{host_header: Arc { + Arc::new(Self { + forward_input: AtomicBool::new(true), + forward_output: AtomicBool::new(true), + restore_terminal: AtomicBool::new(false), + }) + } + + pub fn set_forward_input(&self, enabled: bool) { + self.forward_input.store(enabled, Ordering::SeqCst); + } + + pub fn set_forward_output(&self, enabled: bool) { + self.forward_output.store(enabled, Ordering::SeqCst); + } + + pub fn request_terminal_restore(&self) { + self.restore_terminal.store(true, Ordering::SeqCst); + } + + fn forward_input(&self) -> bool { + self.forward_input.load(Ordering::SeqCst) + } + + fn forward_output(&self) -> bool { + self.forward_output.load(Ordering::SeqCst) + } + + fn take_restore_terminal(&self) -> bool { + self.restore_terminal + .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + } +} + fn ensure_base_image( base_raw: &Path, base_compressed: &Path, @@ -667,14 +699,17 @@ pub fn create_pipe() -> (OwnedFd, OwnedFd) { (read_stream.into(), write_stream.into()) } -pub fn spawn_vm_io_with_line_handler( +pub fn spawn_vm_io_with_hooks( output_monitor: Arc, vm_output_fd: OwnedFd, vm_input_fd: OwnedFd, + io_control: Arc, mut on_line: F, + mut on_output: G, ) -> IoContext where F: FnMut(&str) -> bool + ::std::marker::Send + 'static, + G: FnMut(&[u8]) + ::std::marker::Send + 'static, { let (input_tx, input_rx): (Sender, Receiver) = mpsc::channel(); @@ -721,17 +756,36 @@ where } } + fn poll_wakeup_only(wakeup_fd: RawFd, timeout_ms: i32) -> bool { + let mut fds = [libc::pollfd { + fd: wakeup_fd, + events: libc::POLLIN, + revents: 0, + }]; + + let ret = unsafe { libc::poll(fds.as_mut_ptr(), 1, timeout_ms) }; + ret > 0 && fds[0].revents & libc::POLLIN != 0 + } + // Copies from stdin to the VM; also polls wakeup_read to exit the thread when it's time to shutdown. let stdin_thread = thread::spawn({ let input_tx = input_tx.clone(); let raw_guard = raw_guard.clone(); let wakeup_read = wakeup_read.try_clone().unwrap(); + let io_control = io_control.clone(); move || { let mut buf = [0u8; 1024]; let mut pending_command: Vec = Vec::new(); let mut command_mode = false; loop { + if !io_control.forward_input() { + if poll_wakeup_only(wakeup_read.as_raw_fd(), 100) { + break; + } + continue; + } + match poll_with_wakeup(libc::STDIN_FILENO, wakeup_read.as_raw_fd(), &mut buf) { PollResult::Shutdown | PollResult::Error => break, PollResult::Spurious => continue, @@ -776,28 +830,36 @@ where let stdout_thread = thread::spawn({ let raw_guard = raw_guard.clone(); let wakeup_read = wakeup_read.try_clone().unwrap(); + let io_control = io_control.clone(); move || { - let mut stdout = std::io::stdout().lock(); let mut buf = [0u8; 1024]; loop { + if io_control.take_restore_terminal() { + let mut guard = raw_guard.lock().unwrap(); + *guard = None; + } match poll_with_wakeup(vm_output_fd.as_raw_fd(), wakeup_read.as_raw_fd(), &mut buf) { PollResult::Shutdown | PollResult::Error => break, PollResult::Spurious => continue, PollResult::Ready(bytes) => { - // enable raw mode, if we haven't already - if raw_guard.lock().unwrap().is_none() - && let Ok(guard) = enable_raw_mode(libc::STDIN_FILENO) - { - *raw_guard.lock().unwrap() = Some(guard); - } + if io_control.forward_output() { + // enable raw mode, if we haven't already + if raw_guard.lock().unwrap().is_none() + && let Ok(guard) = enable_raw_mode(libc::STDIN_FILENO) + { + *raw_guard.lock().unwrap() = Some(guard); + } - if stdout.write_all(bytes).is_err() { - break; + let mut stdout = std::io::stdout().lock(); + if stdout.write_all(bytes).is_err() { + break; + } + let _ = stdout.flush(); } - let _ = stdout.flush(); output_monitor.push(bytes); + on_output(bytes); } } } @@ -829,6 +891,25 @@ where } } +pub fn spawn_vm_io_with_line_handler( + output_monitor: Arc, + vm_output_fd: OwnedFd, + vm_input_fd: OwnedFd, + on_line: F, +) -> IoContext +where + F: FnMut(&str) -> bool + ::std::marker::Send + 'static, +{ + spawn_vm_io_with_hooks( + output_monitor, + vm_output_fd, + vm_input_fd, + IoControl::new(), + on_line, + |_| {}, + ) +} + pub fn spawn_vm_io( output_monitor: Arc, vm_output_fd: OwnedFd,