Clean vsock sockets after krunkit exit

Cargo.toml

@@ -16,5 +16,6 @@ sysinfo = "0.31.4"
 log = "0.4.0"
 env_logger = "0.11.8"
 regex = "1.11.1"
+ctrlc = { version = "3.5", features = ["termination"] }
 core-foundation = "0.10.1"
 objc2-io-kit = "0.3.2"

src/cleanup.rs

@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: Apache-2.0
+
+use std::path::PathBuf;
+use std::sync::Arc;
+
+/// Remove each socket file, logging the outcome. Missing files are silently ignored.
+fn remove_sockets(paths: &[PathBuf]) {
+    for path in paths {
+        match std::fs::remove_file(path) {
+            Ok(()) => log::info!("removed vsock socket: {}", path.display()),
+            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
+            Err(e) => log::warn!("failed to remove vsock socket {}: {}", path.display(), e),
+        }
+    }
+}
+
+/// Guard that cleans up vsock socket files when dropped (normal exit) and
+/// also installs a Ctrl-C handler so the sockets are removed on SIGINT.
+pub struct VsockCleanupGuard {
+    paths: Arc<Vec<PathBuf>>,
+}
+
+impl VsockCleanupGuard {
+    /// Create the guard **and** register the SIGINT handler.
+    pub fn new(paths: Vec<PathBuf>) -> Result<Self, anyhow::Error> {
+        let paths = Arc::new(paths);
+        let handler_paths = Arc::clone(&paths);
+        ctrlc::set_handler(move || {
+            remove_sockets(&handler_paths);
+            std::process::exit(0);
+        })?;
+        Ok(Self { paths })
+    }
+}
+
+impl Drop for VsockCleanupGuard {
+    fn drop(&mut self) {
+        remove_sockets(&self.paths);
+    }
+}

src/context.rs

@@ -17,7 +17,7 @@ use std::{
 };
 
 use crate::timesync::timesync_listener;
-use crate::virtio::{VsockAction, VsockConfig};
+use crate::virtio::{VirtioDeviceConfig, VsockAction, VsockConfig};
 use anyhow::{anyhow, Context};
 use env_logger::{Builder, Env, Target};
 
@@ -202,10 +202,7 @@ impl TryFrom<Args> for KrunContext {
         if let Some(timesync_port) = args.timesync {
             let vsock_config = VsockConfig {
                 port: timesync_port,
-                socket_url: PathBuf::from(format!(
-                    "/tmp/krunkit_timesync_{}.sock",
-                    std::process::id()
-                )),
+                socket_url: Self::timesync_socket_path(),
                 action: VsockAction::Connect,
             };
             unsafe { vsock_config.krun_ctx_set(id)? }
@@ -217,6 +214,30 @@ impl TryFrom<Args> for KrunContext {
 }
 
 impl KrunContext {
+    /// Generate the timesync socket path for this process.
+    fn timesync_socket_path() -> PathBuf {
+        PathBuf::from(format!(
+            "/tmp/krunkit_timesync_{}.sock",
+            std::process::id()
+        ))
+    }
+
+    /// Collect all vsock socket paths that should be cleaned up on exit.
+    pub fn vsock_socket_paths(&self) -> Vec<PathBuf> {
+        let mut paths = Vec::new();
+        for device in &self.args.devices {
+            if let VirtioDeviceConfig::Vsock(vsock) = device {
+                if vsock.action == VsockAction::Connect {
+                    paths.push(vsock.socket_url.clone());
+                }
+            }
+        }
+        if self.args.timesync.is_some() {
+            paths.push(Self::timesync_socket_path());
+        }
+        paths
+    }
+
     /// Spawn a thread to listen for shutdown requests and run the workload. If behaving properly,
     /// the main thread will never return from this function.
     pub fn run(&self) -> Result<(), anyhow::Error> {

src/main.rs

@@ -2,12 +2,14 @@
 
 #![allow(dead_code)]
 
+mod cleanup;
 mod cmdline;
 mod context;
 mod status;
 mod timesync;
 mod virtio;
 
+use cleanup::VsockCleanupGuard;
 use cmdline::Args;
 use context::KrunContext;
 
@@ -18,6 +20,9 @@ fn main() -> Result<(), anyhow::Error> {
     // accordingly.
     let ctx = KrunContext::try_from(Args::parse())?;
 
+    // The guard removes vsock socket files on drop (normal exit) and on Ctrl-C.
+    let _cleanup = VsockCleanupGuard::new(ctx.vsock_socket_paths())?;
+
     // Run the workload. If behaving properly, the main thread will not return from this
     // function.
     ctx.run()?;