From 82316c3b40800ad8958bf64adcb928a4ee3bee5e Mon Sep 17 00:00:00 2001
From: Allen Webb <allenwebb@google.com>
Date: Fri, 1 Oct 2021 10:25:39 -0500
Subject: [PATCH] sys_util: Add wait_for_pid.

This adds a safe wrapper for libc::waitforpid that converts the *status
field to an enum so the various edgecases can be handled.

BUG=None
TEST=cargo build

Change-Id: Ic518e686b8ea60fa968f849f1cae571ebfe069e8
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/3199623
Reviewed-by: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Chirantan Ekbote <chirantan@chromium.org>
Auto-Submit: Allen Webb <allenwebb@google.com>
Tested-by: kokoro <noreply+kokoro@google.com>
Commit-Queue: Allen Webb <allenwebb@google.com>
---
 sys_util/src/lib.rs | 64 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/sys_util/src/lib.rs b/sys_util/src/lib.rs
index 822e8f143e..d8d293dfaf 100644
--- a/sys_util/src/lib.rs
+++ b/sys_util/src/lib.rs
@@ -89,6 +89,7 @@ pub use crate::signalfd::Error as SignalFdError;
 pub use crate::write_zeroes::{PunchHole, WriteZeroes, WriteZeroesAt};
 
 use std::cell::Cell;
+use std::convert::TryFrom;
 use std::ffi::CStr;
 use std::fs::{remove_file, File, OpenOptions};
 use std::mem;
@@ -274,6 +275,69 @@ pub fn fallocate(
     syscall!(unsafe { libc::fallocate64(file.as_raw_fd(), mode, offset, len) }).map(|_| ())
 }
 
+/// A trait used to abstract types that provide a process id that can be operated on.
+pub trait AsRawPid {
+    fn as_raw_pid(&self) -> Pid;
+}
+
+impl AsRawPid for Pid {
+    fn as_raw_pid(&self) -> Pid {
+        *self
+    }
+}
+
+impl AsRawPid for std::process::Child {
+    fn as_raw_pid(&self) -> Pid {
+        self.id() as Pid
+    }
+}
+
+/// A logical set of the values *status can take from libc::wait and libc::waitpid.
+pub enum WaitStatus {
+    Continued,
+    Exited(u8),
+    Running,
+    Signaled(Signal),
+    Stopped(Signal),
+}
+
+impl From<c_int> for WaitStatus {
+    fn from(status: c_int) -> WaitStatus {
+        use WaitStatus::*;
+        if libc::WIFEXITED(status) {
+            Exited(libc::WEXITSTATUS(status) as u8)
+        } else if libc::WIFSIGNALED(status) {
+            Signaled(Signal::try_from(libc::WTERMSIG(status)).unwrap())
+        } else if libc::WIFSTOPPED(status) {
+            Stopped(Signal::try_from(libc::WSTOPSIG(status)).unwrap())
+        } else if libc::WIFCONTINUED(status) {
+            Continued
+        } else {
+            Running
+        }
+    }
+}
+
+/// A safe wrapper around waitpid.
+///
+/// On success if a process was reaped, it will be returned as the first value.
+/// The second returned value is the WaitStatus from the libc::waitpid() call.
+///
+/// Note: this can block if libc::WNOHANG is not set and EINTR is not handled internally.
+pub fn wait_for_pid<A: AsRawPid>(pid: A, options: c_int) -> Result<(Option<Pid>, WaitStatus)> {
+    let pid = pid.as_raw_pid();
+    let mut status: c_int = 1;
+    // Safe because status is owned and the error is checked.
+    let ret = unsafe { libc::waitpid(pid, &mut status, options) };
+    if ret < 0 {
+        return errno_result();
+    }
+    Ok((
+        if ret == 0 { None } else { Some(ret) },
+        WaitStatus::from(status),
+    ))
+}
+
 /// Reaps a child process that has terminated.
 ///
 /// Returns `Ok(pid)` where `pid` is the process that was reaped or `Ok(0)` if none of the children