From ccdb1f3e7e8da9a8ecf4bf7b51c5cf632ec1a1a6 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jake@hillion.co.uk>
Date: Thu, 19 May 2022 16:42:08 +0100
Subject: [PATCH] wait for children and report return code

---
 src/lib.rs  | 42 +++++++++++++++++++++++++++++++++++++++---
 src/main.rs | 15 ++++++++++-----
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0733e5e..7561300 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -17,17 +17,19 @@ use std::path::Path;
 
 use nix::fcntl::OFlag;
 use nix::sys::socket;
-use nix::unistd;
+use nix::sys::wait::{waitpid, WaitPidFlag, WaitStatus};
+use nix::unistd::{self, Pid};
 
 pub struct RunArgs<'a> {
     pub spec: Option<&'a Path>,
     pub debug: bool,
+    pub daemon: bool,
 
     pub binary: &'a Path,
     pub binary_args: Vec<&'a str>,
 }
 
-pub fn run(args: &RunArgs) -> Result<()> {
+pub fn run(args: &RunArgs) -> Result<i32> {
     // parse the specification
     let spec: Specification = if let Some(m) = args.spec {
         if m.extension().map(|e| e == "json") == Some(true) {
@@ -62,7 +64,41 @@ pub fn run(args: &RunArgs) -> Result<()> {
     }
     .spawn()?;
 
-    Ok(())
+    if args.daemon {
+        return Ok(exitcode::OK);
+    }
+
+    info!("spawned successfully, awaiting children exiting...");
+    let mut exit_code = exitcode::OK;
+
+    loop {
+        let status = match waitpid(Some(Pid::from_raw(-1)), Some(WaitPidFlag::WEXITED)) {
+            Ok(v) => Ok(v),
+            Err(nix::Error::ECHILD) => {
+                info!("all child processes have exited, exiting...");
+                break;
+            }
+            Err(e) => Err(Error::Nix {
+                msg: "waitpid",
+                src: e,
+            }),
+        }?;
+
+        match status {
+            WaitStatus::Exited(pid, code) => {
+                if code != exitcode::OK {
+                    exit_code = code;
+                }
+                debug!("child {} exited with code {}", pid, code);
+            }
+            WaitStatus::Signaled(pid, sig, _coredump) => {
+                debug!("child {} was terminated with signal {}", pid, sig);
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    Ok(exit_code)
 }
 
 fn create_pipes(names: Vec<&str>) -> Result<HashMap<String, PipePair>> {
diff --git a/src/main.rs b/src/main.rs
index f6c0b92..07f7404 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,4 @@
-use log::{error, info};
+use log::error;
 
 use clone_shim::{run, RunArgs};
 
@@ -34,6 +34,13 @@ fn main() {
                 .help("Stop each spawned application process so that it can be attached to.")
                 .takes_value(false),
         )
+        .arg(
+            Arg::new("daemon")
+                .long("daemon")
+                .short('D')
+                .help("Detach the shim from all child processes and exit immediately.")
+                .takes_value(false),
+        )
         .arg(
             Arg::new("binary")
                 .index(1)
@@ -69,15 +76,13 @@ fn main() {
         let args = RunArgs {
             spec: matches.value_of("spec").map(Path::new),
             debug: matches.is_present("debug"),
+            daemon: matches.is_present("daemon"),
             binary,
             binary_args,
         };
 
         match run(&args) {
-            Ok(_) => {
-                info!("launched successfully");
-                exitcode::OK
-            }
+            Ok(_) => exitcode::OK,
             Err(e) => {
                 error!("error: {}", e);
                 -1