Add support for none and unix URI schemes

Vfkit supports the schemes `unix` and `none` in addition to `tcp`.
Krunkit currently only implements `tcp`, so implement the remaining
schemes to meet vfkit parity.

Signed-off-by: Jake Correnti <jakecorrenti+github@proton.me>

Author: jakecorrenti

Cargo.lock

@@ -15,3 +15,4 @@ mac_address = "1.1.5"
 sysinfo = "0.31.4"
 log = "0.4.0"
 env_logger = "0.11.8"
+regex = "1.11.1"

Cargo.toml

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{status::RestfulUriAddr, virtio::VirtioDeviceConfig};
+use crate::{status::RestfulUri, virtio::VirtioDeviceConfig};
 
 use std::{collections::HashMap, path::PathBuf, str::FromStr};
 
@@ -29,7 +29,7 @@ pub struct Args {
 
     /// URI of the status/shutdown listener.
     #[arg(long = "restful-uri")]
-    pub restful_uri: Option<RestfulUriAddr>,
+    pub restful_uri: Option<RestfulUri>,
 
     /// GUI option for compatibility with vfkit (ignored).
     #[arg(long, default_value_t = false)]
@@ -514,8 +514,10 @@ mod tests {
 
         let restful_uri = args.restful_uri.expect("restful-uri argument not found");
 
-        assert_eq!(restful_uri.ip_addr, Ipv4Addr::new(127, 0, 0, 1));
-        assert_eq!(restful_uri.port, 49573);
+        assert_eq!(
+            restful_uri,
+            RestfulUri::Tcp(Ipv4Addr::new(127, 0, 0, 1), 49573)
+        );
 
         assert_eq!(args.gui, true);
         assert_eq!(args.krun_log_level, 5);

src/cmdline.rs

@@ -3,7 +3,7 @@
 use super::*;
 
 use crate::{
-    status::{get_shutdown_eventfd, status_listener},
+    status::{get_shutdown_eventfd, status_listener, RestfulUri},
     virtio::KrunContextSet,
 };
 
@@ -125,7 +125,9 @@ impl KrunContext {
         let shutdown_eventfd = unsafe { get_shutdown_eventfd(self.id) };
         let uri = self.args.restful_uri.clone();
 
-        thread::spawn(move || status_listener(shutdown_eventfd, uri).unwrap());
+        if uri != Some(RestfulUri::None) {
+            thread::spawn(move || status_listener(shutdown_eventfd, uri).unwrap());
+        }
 
         // Run the workload.
         if unsafe { krun_start_enter(self.id) } < 0 {

src/context.rs

@@ -2,14 +2,16 @@
 
 use std::{
     fs::File,
-    io::{Read, Write},
+    io::{ErrorKind, Read, Write},
     net::{Ipv4Addr, TcpListener},
-    os::fd::{FromRawFd, RawFd},
+    os::{
+        fd::{FromRawFd, RawFd},
+        unix::net::UnixListener,
+    },
     str::FromStr,
 };
 
 use anyhow::{anyhow, Context};
-use clap::Parser;
 
 #[link(name = "krun-efi")]
 extern "C" {
@@ -22,51 +24,84 @@ const HTTP_RUNNING: &str =
 const HTTP_STOPPING: &str =
     "HTTP/1.1 200 OK\r\nContent-type: application/json\r\n\r\n{\"state\": \"VirtualMachineStateStopping\"}\0";
 
-/// Socket address in which the restful URI socket should listen on. Identical to Rust's
-/// SocketAddrV4, but requires a modified FromStr implementation due to how the address is
-/// presented on the command line.
-#[derive(Clone, Debug, Parser)]
-pub struct RestfulUriAddr {
-    pub ip_addr: Ipv4Addr,
-    pub port: u16,
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub enum UriScheme {
+    #[default]
+    Tcp,
+    Unix,
+    None,
 }
 
-impl FromStr for RestfulUriAddr {
+impl FromStr for UriScheme {
     type Err = anyhow::Error;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        let mut string = String::from(s);
-
-        if let Some(removed) = string.strip_prefix("tcp://") {
-            string = String::from(removed);
+        match s {
+            "tcp" => Ok(Self::Tcp),
+            "unix" => Ok(Self::Unix),
+            "none" => Ok(Self::None),
+            _ => Err(anyhow!("invalid scheme")),
         }
+    }
+}
 
-        let mut parts: Vec<String> = string.split(':').map(|s| s.to_string()).collect();
-        if parts.len() != 2 {
-            return Err(anyhow!("restful URI formatted incorrectly"));
-        }
+/// Socket address in which the restful URI socket should listen on. Identical to Rust's
+/// SocketAddrV4, but requires a modified FromStr implementation due to how the address is
+/// presented on the command line.
+#[derive(Clone, Debug, PartialEq)]
+pub enum RestfulUri {
+    Tcp(Ipv4Addr, u16),
+    Unix(String),
+    None,
+}
 
-        // Ipv4Address's FromStr does not understand that the "localhost" IP address translates to
-        // 127.0.0.1, this must be manually translated.
-        if &parts[0][..] == "localhost" {
-            parts[0] = String::from("127.0.0.1");
+impl FromStr for RestfulUri {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let expression = regex::Regex::new(r"^(?P<scheme>none|tcp|unix)://(?P<value>.*)").unwrap();
+        let Some(cap) = expression.captures(s) else {
+            return Err(anyhow!("invalid scheme input"));
+        };
+        let scheme = &cap["scheme"];
+        let value = &cap["value"];
+        match UriScheme::from_str(scheme)? {
+            UriScheme::Tcp => {
+                let (ip_addr, port) = parse_tcp_input(value)?;
+                Ok(Self::Tcp(ip_addr, port))
+            }
+            UriScheme::Unix => {
+                if value.is_empty() {
+                    return Err(anyhow!("empty unix socket path"));
+                }
+                Ok(Self::Unix(value.to_string()))
+            }
+            UriScheme::None => Ok(Self::None),
         }
+    }
+}
 
-        let ip_addr = Ipv4Addr::from_str(&parts[0])
-            .context("restful URI IP address formatted incorrectly")?;
-        let port =
-            u16::from_str(&parts[1]).context("restful URI port number formatted incorrectly")?;
+fn parse_tcp_input(input: &str) -> Result<(Ipv4Addr, u16), anyhow::Error> {
+    let mut parts: Vec<String> = input.split(':').map(|s| s.to_string()).collect();
+    if parts.len() != 2 {
+        return Err(anyhow!("restful URI formatted incorrectly"));
+    }
 
-        Ok(Self { ip_addr, port })
+    // Ipv4Address's FromStr does not understand that the "localhost" IP address translates to
+    // 127.0.0.1, this must be manually translated.
+    if &parts[0][..] == "localhost" {
+        parts[0] = String::from("127.0.0.1");
     }
+
+    let ip_addr =
+        Ipv4Addr::from_str(&parts[0]).context("restful URI IP address formatted incorrectly")?;
+    let port = u16::from_str(&parts[1]).context("restful URI port number formatted incorrectly")?;
+    Ok((ip_addr, port))
 }
 
-impl Default for RestfulUriAddr {
+impl Default for RestfulUri {
     fn default() -> Self {
-        Self {
-            ip_addr: Ipv4Addr::new(127, 0, 0, 1),
-            port: 8081,
-        }
+        Self::Tcp(Ipv4Addr::new(127, 0, 0, 1), 8081)
     }
 }
 
@@ -79,42 +114,154 @@ pub unsafe fn get_shutdown_eventfd(ctx_id: u32) -> i32 {
     fd
 }
 
+fn handle_incoming_stream<T: Read + Write>(stream: &mut T, shutdown_fd: &mut File) {
+    let mut buf = [0u8; 4096];
+    match stream.read(&mut buf) {
+        Ok(_sz) => {
+            let request = String::from_utf8_lossy(&buf);
+            if request.contains("POST") {
+                // Send a VirtualMachineStateStopping message to the client.
+                if let Err(e) = stream.write_all(HTTP_STOPPING.as_bytes()) {
+                    println!("Error writing POST response: {e}");
+                }
+
+                // Shut down the VM.
+                if let Err(e) = shutdown_fd.write_all(&1u64.to_le_bytes()) {
+                    println!("Error writing to shutdown fd: {e}");
+                }
+            } else if let Err(e) = stream.write_all(HTTP_RUNNING.as_bytes()) {
+                println!("Error writing GET response: {e}");
+            }
+        }
+        Err(e) => println!("Error reading stream: {}", e),
+    }
+}
+
 /// Listen for status and shutdown requests from the client. Shut down the krun VM when prompted.
 pub fn status_listener(
     shutdown_eventfd: RawFd,
-    addr: Option<RestfulUriAddr>,
+    addr: Option<RestfulUri>,
 ) -> Result<(), anyhow::Error> {
     // VM is shut down by writing to the shutdown event file.
     let mut shutdown = unsafe { File::from_raw_fd(shutdown_eventfd) };
 
     let addr = addr.unwrap_or_default();
 
-    let listener = TcpListener::bind((addr.ip_addr, addr.port)).unwrap();
-
-    for stream in listener.incoming() {
-        let mut buf = [0u8; 4096];
-        let mut stream = stream.unwrap();
-
-        match stream.read(&mut buf) {
-            Ok(_sz) => {
-                let request = String::from_utf8_lossy(&buf);
-                if request.contains("POST") {
-                    // Send a VirtualMachineStateStopping message to the client.
-                    if let Err(e) = stream.write_all(HTTP_STOPPING.as_bytes()) {
-                        println!("Error writting POST response: {e}");
-                    }
-
-                    // Shut down the VM.
-                    if let Err(e) = shutdown.write_all(&1u64.to_le_bytes()) {
-                        println!("Error writting to shutdown fd: {e}");
-                    }
-                } else if let Err(e) = stream.write_all(HTTP_RUNNING.as_bytes()) {
-                    println!("Error writting GET response: {e}");
+    match addr {
+        RestfulUri::Tcp(addr, port) => {
+            let listener = TcpListener::bind((addr, port))
+                .map_err(|e| anyhow!("Unable to bind to TCP listener: {}", e))?;
+
+            for stream in listener.incoming() {
+                handle_incoming_stream(&mut stream.unwrap(), &mut shutdown)
+            }
+        }
+        RestfulUri::Unix(path) => {
+            if let Err(e) = std::fs::remove_file(&path) {
+                if e.kind() != ErrorKind::NotFound {
+                    return Err(anyhow!("failed to remove socket with error {e}"));
                 }
             }
-            Err(e) => println!("Error reading stream: {}", e),
+            let listener = UnixListener::bind(path)
+                .map_err(|e| anyhow!("Unable to bind to unix socket: {}", e))?;
+
+            for stream in listener.incoming() {
+                handle_incoming_stream(&mut stream.unwrap(), &mut shutdown)
+            }
         }
+        RestfulUri::None => unreachable!(),
     }
 
     Ok(())
 }
+
+#[allow(unused_imports)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parse_valid_unix_scheme() {
+        assert_eq!(
+            RestfulUri::Unix("/tmp/path".to_string()),
+            RestfulUri::from_str("unix:///tmp/path").unwrap()
+        );
+    }
+
+    #[test]
+    fn parse_unix_scheme_missing_path() {
+        assert_eq!(
+            anyhow!("empty unix socket path").to_string(),
+            RestfulUri::from_str("unix://").err().unwrap().to_string()
+        );
+    }
+
+    #[test]
+    fn parse_unix_scheme_missing_slashes() {
+        assert_eq!(
+            anyhow!("invalid scheme input").to_string(),
+            RestfulUri::from_str("unix:").err().unwrap().to_string()
+        );
+    }
+
+    #[test]
+    fn parse_unix_scheme_misspelling() {
+        assert_eq!(
+            anyhow!("invalid scheme input").to_string(),
+            RestfulUri::from_str("uni://path")
+                .err()
+                .unwrap()
+                .to_string()
+        );
+    }
+
+    #[test]
+    fn parse_valid_tcp_scheme() {
+        assert_eq!(
+            RestfulUri::Tcp(Ipv4Addr::new(127, 0, 0, 1), 8080),
+            RestfulUri::from_str("tcp://localhost:8080").unwrap(),
+        );
+    }
+
+    #[test]
+    fn parse_tcp_scheme_missing_port() {
+        assert_eq!(
+            anyhow!("restful URI formatted incorrectly").to_string(),
+            RestfulUri::from_str("tcp://localhost")
+                .err()
+                .unwrap()
+                .to_string()
+        );
+    }
+
+    #[test]
+    fn parse_tcp_scheme_with_unix_path() {
+        assert_eq!(
+            anyhow!("restful URI formatted incorrectly").to_string(),
+            RestfulUri::from_str("tcp:///tmp/path")
+                .err()
+                .unwrap()
+                .to_string(),
+        );
+    }
+
+    #[test]
+    fn parse_valid_none_scheme() {
+        assert_eq!(RestfulUri::None, RestfulUri::from_str("none://").unwrap());
+    }
+
+    #[test]
+    fn parse_none_scheme_missing_postfix() {
+        assert_eq!(
+            anyhow!("invalid scheme input").to_string(),
+            RestfulUri::from_str("none").err().unwrap().to_string(),
+        );
+    }
+
+    #[test]
+    fn parse_random_string_scheme() {
+        assert_eq!(
+            anyhow!("invalid scheme input").to_string(),
+            RestfulUri::from_str("foobar").err().unwrap().to_string(),
+        );
+    }
+}