Protect muvm-server, spawn one as root and monitor available host memory

crates/muvm/Cargo.toml

@@ -16,13 +16,14 @@ env_logger = { version = "0.11.3", default-features = false, features = ["auto-c
 krun-sys = { path = "../krun-sys", version = "1.9.1", default-features = false, features = [] }
 log = { version = "0.4.21", default-features = false, features = ["kv"] }
 nix = { version = "0.28.0", default-features = false, features = ["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"] }
 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 = [] }
 tokio = { version = "1.38.0", default-features = false, features = ["io-util", "macros", "net", "process", "rt-multi-thread", "sync"] }
 tokio-stream = { version = "0.1.15", default-features = false, features = ["net", "sync"] }
-uuid = { version = "1.10.0", default-features = false, features = ["std", "v7"] }
+uuid = { version = "1.10.0", default-features = false, features = ["serde", "std", "v7"] }
 
 [features]
 default = []

crates/muvm/src/bin/muvm.rs

@@ -16,6 +16,7 @@ use muvm::cli_options::options;
 use muvm::cpu::{get_fallback_cores, get_performance_cores};
 use muvm::env::{find_muvm_exec, prepare_env_vars};
 use muvm::launch::{launch_or_lock, LaunchResult};
+use muvm::monitor::spawn_monitor;
 use muvm::net::{connect_to_passt, start_passt};
 use muvm::types::MiB;
 use nix::sys::sysinfo::sysinfo;
@@ -67,7 +68,7 @@ fn main() -> Result<()> {
 
     let options = options().fallback_to_usage().run();
 
-    let (_lock_file, command, command_args, env) = match launch_or_lock(
+    let (cookie, _lock_file, command, command_args, env) = match launch_or_lock(
         options.server_port,
         options.command,
         options.command_args,
@@ -79,11 +80,12 @@ fn main() -> Result<()> {
             return Ok(());
         },
         LaunchResult::LockAcquired {
+            cookie,
             lock_file,
             command,
             command_args,
             env,
-        } => (lock_file, command, command_args, env),
+        } => (cookie, lock_file, command, command_args, env),
     };
 
     {
@@ -253,7 +255,7 @@ fn main() -> Result<()> {
                 .context("Failed to connect to `passt`")?
                 .into()
         } else {
-            start_passt(options.server_port)
+            start_passt(options.server_port, options.root_server_port)
                 .context("Failed to start `passt`")?
                 .into()
         };
@@ -374,6 +376,11 @@ fn main() -> Result<()> {
         "MUVM_SERVER_PORT".to_owned(),
         options.server_port.to_string(),
     );
+    env.insert(
+        "MUVM_ROOT_SERVER_PORT".to_owned(),
+        options.root_server_port.to_string(),
+    );
+    env.insert("MUVM_SERVER_COOKIE".to_owned(), cookie.to_string());
 
     if options.direct_x11 {
         let display =
@@ -429,6 +436,8 @@ fn main() -> Result<()> {
         }
     }
 
+    spawn_monitor(options.root_server_port, cookie);
+
     {
         // Start and enter the microVM. Unless there is some error while creating the
         // microVM this function never returns.

crates/muvm/src/cli_options.rs

@@ -13,6 +13,7 @@ pub struct Options {
     pub mem: Option<MiB>,
     pub vram: Option<MiB>,
     pub passt_socket: Option<PathBuf>,
+    pub root_server_port: u32,
     pub server_port: u32,
     pub fex_images: Vec<String>,
     pub direct_x11: bool,
@@ -93,6 +94,12 @@ pub fn options() -> OptionParser<Options> {
         .help("Instead of starting passt, connect to passt socket at PATH")
         .argument("PATH")
         .optional();
+    let root_server_port = long("root-server-port")
+        .short('r')
+        .help("Set the port to be used in root server mode")
+        .argument("ROOT_SERVER_PORT")
+        .fallback(3335)
+        .display_fallback();
     let server_port = long("server-port")
         .short('p')
         .help("Set the port to be used in server mode")
@@ -117,6 +124,7 @@ pub fn options() -> OptionParser<Options> {
         mem,
         vram,
         passt_socket,
+        root_server_port,
         server_port,
         fex_images,
         direct_x11,

crates/muvm/src/guest/bin/muvm-guest.rs

@@ -62,6 +62,12 @@ fn main() -> Result<()> {
             .spawn()?;
     }
 
+    // Before switching to the user, start another instance of muvm-server to serve
+    // launch requests as root.
+    Command::new("muvm-server")
+        .spawn()
+        .context("Failed to execute `muvm-server` as child process")?;
+
     let run_path = match setup_user(options.username, options.uid, options.gid) {
         Ok(p) => p,
         Err(err) => return Err(err).context("Failed to set up user, bailing out"),

crates/muvm/src/launch.rs

@@ -2,21 +2,22 @@ use std::collections::HashMap;
 use std::env;
 use std::error::Error;
 use std::fmt::{Display, Formatter};
-use std::fs::File;
+use std::fs::{self, File};
 use std::io::{BufRead, BufReader, Read, Write};
 use std::net::TcpStream;
 use std::path::{Path, PathBuf};
 
 use anyhow::{anyhow, Context, Result};
 use rustix::fs::{flock, FlockOperation};
-use rustix::path::Arg;
+use uuid::Uuid;
 
 use crate::env::prepare_env_vars;
 use crate::utils::launch::Launch;
 
 pub enum LaunchResult {
     LaunchRequested,
     LockAcquired {
+        cookie: Uuid,
         lock_file: File,
         command: PathBuf,
         command_args: Vec<String>,
@@ -59,53 +60,63 @@ pub fn launch_or_lock(
     if let Some(port) = running_server_port {
         let port: u32 = port.parse()?;
         let env = prepare_env_vars(env)?;
-        if let Err(err) = request_launch(port, command, command_args, env) {
+        let cookie = read_cookie()?;
+        if let Err(err) = request_launch(port, cookie, command, command_args, env) {
             return Err(anyhow!("could not request launch to server: {err}"));
         }
         return Ok(LaunchResult::LaunchRequested);
     }
 
-    let (lock_file, running_server_port) = lock_file(server_port)?;
+    let (lock_file, cookie) = lock_file()?;
     match lock_file {
         Some(lock_file) => Ok(LaunchResult::LockAcquired {
+            cookie,
             lock_file,
             command,
             command_args,
             env,
         }),
         None => {
-            if let Some(port) = running_server_port {
-                let env = prepare_env_vars(env)?;
-                let mut tries = 0;
-                loop {
-                    match request_launch(port, command.clone(), command_args.clone(), env.clone()) {
-                        Err(err) => match err.downcast_ref::<LaunchError>() {
-                            Some(&LaunchError::Connection(_)) => {
-                                if tries == 3 {
-                                    return Err(anyhow!(
-                                        "could not request launch to server: {err}"
-                                    ));
-                                } else {
-                                    tries += 1;
-                                }
-                            },
-                            _ => {
+            let env = prepare_env_vars(env)?;
+            let mut tries = 0;
+            loop {
+                match request_launch(
+                    server_port,
+                    cookie,
+                    command.clone(),
+                    command_args.clone(),
+                    env.clone(),
+                ) {
+                    Err(err) => match err.downcast_ref::<LaunchError>() {
+                        Some(&LaunchError::Connection(_)) => {
+                            if tries == 3 {
                                 return Err(anyhow!("could not request launch to server: {err}"));
-                            },
+                            } else {
+                                tries += 1;
+                            }
                         },
-                        Ok(_) => return Ok(LaunchResult::LaunchRequested),
-                    }
+                        _ => {
+                            return Err(anyhow!("could not request launch to server: {err}"));
+                        },
+                    },
+                    Ok(_) => return Ok(LaunchResult::LaunchRequested),
                 }
-            } else {
-                Err(anyhow!(
-                    "muvm is already running but couldn't find its server port, bailing out"
-                ))
             }
         },
     }
 }
 
-fn lock_file(server_port: u32) -> Result<(Option<File>, Option<u32>)> {
+fn read_cookie() -> Result<Uuid> {
+    let run_path = env::var("XDG_RUNTIME_DIR")
+        .context("Failed to read XDG_RUNTIME_DIR environment variable")?;
+    let lock_path = Path::new(&run_path).join("muvm.lock");
+    let data: Vec<u8> = fs::read(lock_path).context("Failed to read lock file")?;
+    assert!(data.len() == 16);
+
+    Uuid::from_slice(&data).context("Failed to read cookie from lock file")
+}
+
+fn lock_file() -> Result<(Option<File>, Uuid)> {
     let run_path = env::var("XDG_RUNTIME_DIR")
         .context("Failed to read XDG_RUNTIME_DIR environment variable")?;
     let lock_path = Path::new(&run_path).join("muvm.lock");
@@ -123,30 +134,23 @@ fn lock_file(server_port: u32) -> Result<(Option<File>, Option<u32>)> {
             .context("Failed to create lock file")?;
         let ret = flock(&lock_file, FlockOperation::NonBlockingLockExclusive);
         if ret.is_err() {
-            let mut data: Vec<u8> = Vec::with_capacity(5);
+            let mut data: Vec<u8> = Vec::with_capacity(16);
             lock_file.read_to_end(&mut data)?;
-            let port = match data.to_string_lossy().parse::<u32>() {
-                Ok(port) => {
-                    if port > 1024 {
-                        Some(port)
-                    } else {
-                        None
-                    }
-                },
-                Err(_) => None,
-            };
-            return Ok((None, port));
+            let cookie = Uuid::from_slice(&data).context("Failed to read cookie from lock file")?;
+            return Ok((None, cookie));
         }
         lock_file
     };
 
+    let cookie = Uuid::now_v7();
     lock_file.set_len(0)?;
-    lock_file.write_all(format!("{server_port}").as_bytes())?;
-    Ok((Some(lock_file), None))
+    lock_file.write_all(cookie.as_bytes())?;
+    Ok((Some(lock_file), cookie))
 }
 
-fn request_launch(
+pub fn request_launch(
     server_port: u32,
+    cookie: Uuid,
     command: PathBuf,
     command_args: Vec<String>,
     env: HashMap<String, String>,
@@ -155,6 +159,7 @@ fn request_launch(
         TcpStream::connect(format!("127.0.0.1:{server_port}")).map_err(LaunchError::Connection)?;
 
     let launch = Launch {
+        cookie,
         command,
         command_args,
         env,

crates/muvm/src/lib.rs

@@ -2,6 +2,7 @@ pub mod cli_options;
 pub mod cpu;
 pub mod env;
 pub mod launch;
+pub mod monitor;
 pub mod net;
 pub mod types;