daemon: gracefully shutdown if proxy wrapped command exits

src/daemon.rs

@@ -1,4 +1,5 @@
 use anyhow::{Context, Result};
+use futures::future::OptionFuture;
 use futures::TryStreamExt;
 use std::collections::BTreeSet;
 use std::fs;
@@ -118,12 +119,13 @@ pub(crate) async fn run(state: State, on_ready: Option<Command>) -> Result<()> {
     let mut admin_handle = tokio::spawn(admin_server.run(grpc, adhoc, cancel.clone()));
     let mut proxy_handle = tokio::spawn(proxy.run(cancel.clone()));
 
-    if let Some(mut cmd) = on_ready {
-        tokio::spawn(async move {
-            let status = cmd.status().await;
-            process::exit(status.map_or(1, |s| s.code().unwrap_or(1)));
-        });
-    }
+    let mut task_status: Option<i32> = None;
+    // If there is no wrapped command, this variable will be None.
+    // The corresponding arm in tokio::select! will only be polled
+    // once, then never again since the pattern will not match.
+    let command_handle = on_ready.map(|mut cmd| tokio::spawn(async move { cmd.status().await }));
+    // Convert Option to OptionFuture to fit in tokio::select! arm
+    let mut command_handle = OptionFuture::from(command_handle);
 
     // Wait for shutdown signal or unexpected task exit. Signal results are
     // ignored - recv() only returns None if "already received" which can't
@@ -138,6 +140,22 @@ pub(crate) async fn run(state: State, on_ready: Option<Command>) -> Result<()> {
         r = &mut gossip_handle => error!(?r, "gossip loop exited"),
         r = &mut admin_handle => error!(?r, "admin server exited"),
         r = &mut proxy_handle => error!(?r, "proxy exited"),
+        Some(r) = &mut command_handle => {
+            match r {
+                Err(e) => {
+                    task_status = Some(1);
+                    error!(?e, "command failed");
+                },
+                Ok(Err(e)) => {
+                    task_status = Some(1);
+                    error!(?e, "command exited");
+                },
+                Ok(Ok(s)) => {
+                    task_status = Some(s.code().unwrap_or(1));
+                    info!("command exited with code {}", task_status.assert());
+                },
+            }
+        },
     }
 
     info!("shutting down");
@@ -166,5 +184,11 @@ pub(crate) async fn run(state: State, on_ready: Option<Command>) -> Result<()> {
     }
 
     info!("shutdown complete");
+
+    // if proxy command exited cleanly, exit with its status code.
+    if let Some(status) = task_status {
+        process::exit(status)
+    }
+
     Ok(())
 }

tests/e2e.rs

@@ -422,6 +422,97 @@ fn test_graceful_shutdown() {
     );
 }
 
+#[test]
+#[ignore = "e2e test requiring docker"]
+fn test_proxy_command_exit() {
+    let harness = Harness::new();
+
+    let proxy_cmd = r"intermesh proxy -- sleep 1";
+    let node = harness.launch("root", proxy_cmd).assert();
+
+    // Verify daemon's exit code
+    assert_eq!(node.exit_code().assert(), 0);
+
+    // Start a new container
+    let cmd = r"tail -f /dev/null";
+    let node = harness.launch("root", cmd).assert();
+    let _ = node.exec(&["intermesh", "proxy", "--", "sleep", "1"]);
+
+    // Wait for admin socket to disappear.
+    wait(|| node.exec(&["test", "!", "-S", "/var/run/intermesh/admin.sock"])).assert();
+
+    // Verify nftables rules were cleaned up
+    let result = node.exec(&["nft", "list", "table", "ip", "intermesh"]);
+    assert!(
+        result.is_err(),
+        "nftables table should be removed on shutdown"
+    );
+}
+
+#[test]
+#[ignore = "e2e test requiring docker"]
+fn test_proxy_command_error() {
+    let harness = Harness::new();
+
+    let proxy_cmd = r#"intermesh proxy -- sh -c "exit 123""#;
+    let node = harness.launch("root", proxy_cmd).assert();
+
+    // Verify daemon's exit code
+    assert_eq!(node.exit_code().assert(), 123);
+
+    // Start a new container
+    let cmd = r"tail -f /dev/null";
+    let node = harness.launch("root", cmd).assert();
+    let _ = node.exec(&["intermesh", "proxy", "--", "exit", "123"]);
+
+    // Wait for admin socket to disappear.
+    wait(|| node.exec(&["test", "!", "-S", "/var/run/intermesh/admin.sock"])).assert();
+
+    // Verify nftables rules were cleaned up
+    let result = node.exec(&["nft", "list", "table", "ip", "intermesh"]);
+    assert!(
+        result.is_err(),
+        "nftables table should be removed on shutdown"
+    );
+}
+
+#[test]
+#[ignore = "e2e test requiring docker"]
+fn test_proxy_command_killed() {
+    let harness = Harness::new();
+
+    let proxy_cmd = r"intermesh proxy -- sleep inf";
+    let node = harness.launch("root", proxy_cmd).assert();
+
+    // Verify nftables rules were injected
+    wait(|| node.exec(&["nft", "list", "table", "ip", "intermesh"])).assert();
+
+    // Kill wrapped command via SIGTERM
+    node.exec(&["pkill", "-TERM", "sleep"]).assert();
+
+    // Verify daemon's exit code
+    assert_eq!(node.exit_code().assert(), 1);
+
+    // Start a new container
+    let cmd = r"tail -f /dev/null";
+    let node = harness.launch("root", cmd).assert();
+    node.start_daemon(&["--intercept", "--", "sleep", "inf"])
+        .assert();
+
+    // Kill wrapped command via SIGTERM
+    node.exec(&["pkill", "-TERM", "sleep"]).assert();
+
+    // Wait for admin socket to disappear.
+    wait(|| node.exec(&["test", "!", "-S", "/var/run/intermesh/admin.sock"])).assert();
+
+    // Verify nftables rules were cleaned up
+    let result = node.exec(&["nft", "list", "table", "ip", "intermesh"]);
+    assert!(
+        result.is_err(),
+        "nftables table should be removed on shutdown"
+    );
+}
+
 #[test]
 #[ignore = "e2e test requiring docker"]
 fn test_runtime_endorsements_are_saved() {

tests/helpers/docker.rs

@@ -196,6 +196,18 @@ impl Node<'_> {
         Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
     }
 
+    pub(crate) fn exit_code(&self) -> Result<u32> {
+        let output = Command::new("docker")
+            .args(["wait", &self.container])
+            .output()?;
+
+        if !output.status.success() {
+            bail!("{}", command_error(&output, &self.container))
+        }
+
+        Ok(String::from_utf8(output.stdout)?.trim().parse::<u32>()?)
+    }
+
     pub(crate) fn get_dump(&self) -> Result<StateDump> {
         let dump_toml = self.exec(&["intermesh", "debug", "dump", "--toml"])?;
         toml::from_str(&dump_toml).map_err(Into::into)