muvm: Replace socat for interactive mode.

Socat's tty mode has issues with window resizing as it does not pass through SIGWINCH. Also add a tty-less interactive mode, for usage from scripts and similar. Signed-off-by: Sasha Finkelstein <[email protected]>
WhatAmISupposedToPutHere · Dec 7, 2024 · 38720e2 · 38720e2
1 parent 59b4934
commit 38720e2
Show file tree

Hide file tree

Showing 13 changed files with 467 additions and 73 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1 +1,2 @@
 /target
+.idea
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/muvm/Cargo.toml b/crates/muvm/Cargo.toml
@@ -19,7 +19,7 @@ krun-sys = { path = "../krun-sys", version = "1.9.1", default-features = false,
 log = { version = "0.4.21", default-features = false, features = ["kv"] }
 nix = { version = "0.29.0", default-features = false, features = ["event", "fs", "ioctl", "mman", "ptrace", "signal", "socket", "uio", "user"] }
 procfs = { version = "0.17.0", default-features = false, features = [] }
-rustix = { version = "0.38.34", default-features = false, features = ["fs", "mount", "process", "std", "stdio", "system", "use-libc-auxv"] }
+rustix = { version = "0.38.34", default-features = false, features = ["fs", "mount", "process", "pty", "std", "stdio", "system", "termios", "use-libc-auxv"] }
 serde = { version = "1.0.203", default-features = false, features = ["derive"] }
 serde_json = { version = "1.0.117", default-features = false, features = ["std"] }
 tempfile = { version = "3.10.1", default-features = false, features = [] }

diff --git a/crates/muvm/src/bin/muvm.rs b/crates/muvm/src/bin/muvm.rs
@@ -3,6 +3,7 @@ use std::ffi::{c_char, CString};
 use std::io::Write;
 use std::os::fd::{IntoRawFd, OwnedFd};
 use std::path::Path;
+use std::process::ExitCode;
 
 use anyhow::{anyhow, Context, Result};
 use krun_sys::{
@@ -59,7 +60,7 @@ fn add_ro_disk(ctx_id: u32, label: &str, path: &str) -> Result<()> {
     }
 }
 
-fn main() -> Result<()> {
+fn main() -> Result<ExitCode> {
     env_logger::init();
 
     if getuid().as_raw() == 0 || geteuid().as_raw() == 0 {
@@ -75,11 +76,12 @@ fn main() -> Result<()> {
         options.command_args,
         options.env,
         options.interactive,
+        options.tty,
     )? {
-        LaunchResult::LaunchRequested => {
+        LaunchResult::LaunchRequested(code) => {
             // There was a muvm instance already running and we've requested it
             // to launch the command successfully, so all the work is done.
-            return Ok(());
+            return Ok(code);
         },
         LaunchResult::LockAcquired {
             cookie,

diff --git a/crates/muvm/src/cli_options.rs b/crates/muvm/src/cli_options.rs
@@ -18,6 +18,7 @@ pub struct Options {
     pub fex_images: Vec<String>,
     pub sommelier: bool,
     pub interactive: bool,
+    pub tty: bool,
     pub command: PathBuf,
     pub command_args: Vec<String>,
 }
@@ -110,7 +111,11 @@ pub fn options() -> OptionParser<Options> {
         .switch();
     let interactive = long("interactive")
         .short('i')
-        .help("Allocate a tty guest-side and connect it to the current stdin/out")
+        .help("Attach to the command's stdin/out after starting it")
+        .switch();
+    let tty = long("tty")
+        .short('t')
+        .help("Allocate a tty for the command")
         .switch();
     let command = positional("COMMAND").help("the command you want to execute in the vm");
     let command_args = any::<String, _, _>("COMMAND_ARGS", |arg| {
@@ -130,6 +135,7 @@ pub fn options() -> OptionParser<Options> {
         fex_images,
         sommelier,
         interactive,
+        tty,
         // positionals
         command,
         command_args,

diff --git a/crates/muvm/src/launch.rs b/crates/muvm/src/launch.rs
@@ -11,17 +11,18 @@ use anyhow::{anyhow, Context, Result};
 use rustix::fs::{flock, FlockOperation};
 use uuid::Uuid;
 
-use super::utils::env::find_in_path;
 use crate::env::prepare_env_vars;
+use crate::tty::{run_io_host, RawTerminal};
 use crate::utils::launch::Launch;
-use rustix::path::Arg;
+use nix::unistd::unlink;
 use std::ops::Range;
-use std::process::{Child, Command};
+use std::os::unix::net::UnixListener;
+use std::process::ExitCode;
 
 pub const DYNAMIC_PORT_RANGE: Range<u32> = 50000..50200;
 
 pub enum LaunchResult {
-    LaunchRequested,
+    LaunchRequested(ExitCode),
     LockAcquired {
         cookie: Uuid,
         lock_file: File,
@@ -56,64 +57,75 @@ impl Display for LaunchError {
     }
 }
 
-fn start_socat() -> Result<(Child, u32)> {
+fn acquire_socket_lock() -> Result<(File, u32)> {
     let run_path = env::var("XDG_RUNTIME_DIR")
         .map_err(|e| anyhow!("unable to get XDG_RUNTIME_DIR: {:?}", e))?;
     let socket_dir = Path::new(&run_path).join("krun/socket");
-    let socat_path =
-        find_in_path("socat")?.ok_or_else(|| anyhow!("Unable to find socat in PATH"))?;
     for port in DYNAMIC_PORT_RANGE {
-        let path = socket_dir.join(format!("port-{}", port));
-        if path.exists() {
-            continue;
-        }
-        let child = Command::new(&socat_path)
-            .arg(format!("unix-l:{}", path.as_os_str().to_string_lossy()))
-            .arg("-,raw,echo=0")
-            .spawn()?;
-        return Ok((child, port));
-    }
-    Err(anyhow!("Ran out of ports."))
-}
-
-fn escape_for_socat(s: String) -> String {
-    let mut ret = String::with_capacity(s.len());
-    for c in s.chars() {
-        match c {
-            ':' | ',' | '!' | '"' | '\'' | '\\' | '(' | '[' | '{' => {
-                ret.push('\\');
+        let path = socket_dir.join(format!("port-{port}.lock"));
+        return Ok((
+            if !path.exists() {
+                let lock_file = File::create(path).context("Failed to create socket lock")?;
+                flock(&lock_file, FlockOperation::NonBlockingLockExclusive)
+                    .context("Failed to acquire socket lock")?;
+                lock_file
+            } else {
+                let lock_file = File::options()
+                    .write(true)
+                    .read(true)
+                    .open(path)
+                    .context("Failed to open lock file")?;
+                if flock(&lock_file, FlockOperation::NonBlockingLockExclusive).is_err() {
+                    continue;
+                }
+                lock_file
             },
-            _ => {},
-        }
-        ret.push(c);
+            port,
+        ));
     }
-    ret
+    Err(anyhow!("Ran out of ports."))
 }
 
 fn wrapped_launch(
     server_port: u32,
     cookie: Uuid,
-    mut command: PathBuf,
-    mut command_args: Vec<String>,
+    command: PathBuf,
+    command_args: Vec<String>,
     env: HashMap<String, String>,
     interactive: bool,
-) -> Result<()> {
+    tty: bool,
+) -> Result<ExitCode> {
     if !interactive {
-        return request_launch(server_port, cookie, command, command_args, env);
+        request_launch(server_port, cookie, command, command_args, env, 0, false)?;
+        return Ok(ExitCode::from(0));
     }
-    let (mut socat, vsock_port) = start_socat()?;
-    command_args.insert(0, command.to_string_lossy().into_owned());
-    command_args = vec![
-        format!("vsock:2:{}", vsock_port),
-        format!(
-            "exec:{},pty,setsid,stderr",
-            escape_for_socat(command_args.join(" "))
-        ),
-    ];
-    command = "socat".into();
-    request_launch(server_port, cookie, command, command_args, env)?;
-    socat.wait()?;
-    Ok(())
+    let run_path = env::var("XDG_RUNTIME_DIR")
+        .map_err(|e| anyhow!("unable to get XDG_RUNTIME_DIR: {:?}", e))?;
+    let socket_dir = Path::new(&run_path).join("krun/socket");
+    let (_lock, vsock_port) = acquire_socket_lock()?;
+    let path = socket_dir.join(format!("port-{vsock_port}"));
+    _ = unlink(&path);
+    let listener = UnixListener::bind(path).context("Failed to listen on vm socket")?;
+    let raw_tty = if tty {
+        Some(
+            RawTerminal::set()
+                .context("Asked to allocate a tty for the command, but stdin is not a tty")?,
+        )
+    } else {
+        None
+    };
+    request_launch(
+        server_port,
+        cookie,
+        command,
+        command_args,
+        env,
+        vsock_port,
+        tty,
+    )?;
+    let code = run_io_host(listener, tty)?;
+    drop(raw_tty);
+    Ok(ExitCode::from(code))
 }
 
 pub fn launch_or_lock(
@@ -122,16 +134,17 @@ pub fn launch_or_lock(
     command_args: Vec<String>,
     env: Vec<(String, Option<String>)>,
     interactive: bool,
+    tty: bool,
 ) -> Result<LaunchResult> {
     let running_server_port = env::var("MUVM_SERVER_PORT").ok();
     if let Some(port) = running_server_port {
         let port: u32 = port.parse()?;
         let env = prepare_env_vars(env)?;
         let cookie = read_cookie()?;
-        if let Err(err) = wrapped_launch(port, cookie, command, command_args, env, interactive) {
-            return Err(anyhow!("could not request launch to server: {err}"));
-        }
-        return Ok(LaunchResult::LaunchRequested);
+        return match wrapped_launch(port, cookie, command, command_args, env, interactive, tty) {
+            Err(err) => Err(anyhow!("could not request launch to server: {err}")),
+            Ok(code) => Ok(LaunchResult::LaunchRequested(code)),
+        };
     }
 
     let (lock_file, cookie) = lock_file()?;
@@ -154,6 +167,7 @@ pub fn launch_or_lock(
                     command_args.clone(),
                     env.clone(),
                     interactive,
+                    tty,
                 ) {
                     Err(err) => match err.downcast_ref::<LaunchError>() {
                         Some(&LaunchError::Connection(_)) => {
@@ -167,7 +181,7 @@ pub fn launch_or_lock(
                             return Err(anyhow!("could not request launch to server: {err}"));
                         },
                     },
-                    Ok(_) => return Ok(LaunchResult::LaunchRequested),
+                    Ok(code) => return Ok(LaunchResult::LaunchRequested(code)),
                 }
             }
         },
@@ -222,6 +236,8 @@ pub fn request_launch(
     command: PathBuf,
     command_args: Vec<String>,
     env: HashMap<String, String>,
+    interactive_vsock: u32,
+    tty: bool,
 ) -> Result<()> {
     let mut stream =
         TcpStream::connect(format!("127.0.0.1:{server_port}")).map_err(LaunchError::Connection)?;
@@ -231,6 +247,8 @@ pub fn request_launch(
         command,
         command_args,
         env,
+        interactive_vsock,
+        tty,
     };
 
     stream

diff --git a/crates/muvm/src/lib.rs b/crates/muvm/src/lib.rs
@@ -10,4 +10,5 @@ pub mod types;
 
 pub mod guest;
 pub mod server;
+pub mod tty;
 pub mod utils;
diff --git a/crates/muvm/src/monitor.rs b/crates/muvm/src/monitor.rs
@@ -44,7 +44,7 @@ fn set_guest_pressure(server_port: u32, cookie: Uuid, pressure: GuestPressure) -
         let command = PathBuf::from("/muvmdropcaches");
         let command_args = vec![];
         let env = HashMap::new();
-        request_launch(server_port, cookie, command, command_args, env)?;
+        request_launch(server_port, cookie, command, command_args, env, 0, false)?;
     }
 
     let wsf: u32 = pressure.into();
@@ -53,7 +53,7 @@ fn set_guest_pressure(server_port: u32, cookie: Uuid, pressure: GuestPressure) -
     let command = PathBuf::from("/sbin/sysctl");
     let command_args = vec![format!("vm.watermark_scale_factor={}", wsf)];
     let env = HashMap::new();
-    request_launch(server_port, cookie, command, command_args, env)
+    request_launch(server_port, cookie, command, command_args, env, 0, false)
 }
 
 fn run(server_port: u32, cookie: Uuid) {