Allow forkserver to use waitpid as alternative to sa_nochldwait

PiperOrigin-RevId: 529074278
Change-Id: If63015586673610e111ee589995e5264523be7a7

sandboxed_api/sandbox2/BUILD.bazel

@@ -247,6 +247,7 @@ cc_binary(
         ":util",
         "//sandboxed_api/util:raw_logging",
         "//sandboxed_api/util:strerror",
+        "@com_google_absl//absl/flags:parse",
         "@com_google_absl//absl/log:globals",
     ],
 )
@@ -606,6 +607,7 @@ cc_library(
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
@@ -848,12 +850,12 @@ cc_test(
     data = ["//sandboxed_api/sandbox2/testcases:minimal"],
     tags = ["no_qemu_user_mode"],
     deps = [
-        ":forkserver",
         ":forkserver_cc_proto",
         ":global_forkserver",
         ":sandbox2",
         "//sandboxed_api:testing",
         "//sandboxed_api/util:raw_logging",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",

sandboxed_api/sandbox2/CMakeLists.txt

@@ -206,6 +206,8 @@ set_target_properties(sandbox2_forkserver_bin PROPERTIES
 add_executable(sandbox2::forkserver_bin ALIAS sandbox2_forkserver_bin)
 target_link_libraries(sandbox2_forkserver_bin PRIVATE
   absl::log_globals
+  absl::flags
+  absl::flags_parse
   sandbox2::client
   sandbox2::comms
   sandbox2::forkserver
@@ -914,6 +916,7 @@ if(BUILD_TESTING AND SAPI_BUILD_TESTING)
   )
   target_link_libraries(sandbox2_forkserver_test PRIVATE
     absl::check
+    absl::log
     absl::strings
     sandbox2::forkserver
     sandbox2::forkserver_proto

sandboxed_api/sandbox2/forkserver.cc

@@ -39,6 +39,7 @@
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/flags/flag.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/match.h"
@@ -61,6 +62,13 @@
 #include "sandboxed_api/util/raw_logging.h"
 #include "sandboxed_api/util/strerror.h"
 
+// Make the forkserver use a signal handler for SIGCHLD instead of setting
+// `SA_NOCLDWAIT`. This is needed in certain cases where the process need
+// to be explicitly reaped by waitpid (see getrusage).
+ABSL_FLAG(
+    bool, sandbox2_forkserver_use_waitpid, false,
+    "Use waitpid to reap child processes instead of relying on SA_NOCLDWAIT");
+
 namespace sandbox2 {
 
 namespace file_util = ::sapi::file_util;
@@ -69,6 +77,18 @@ namespace {
 
 using ::sapi::StrError;
 
+void SigChldHandler(int signal) {
+  if (signal != SIGCHLD) {
+    return;
+  }
+
+  pid_t pid;
+  int wstatus;
+  do {
+    pid = waitpid(-1, &wstatus, WNOHANG);
+  } while (pid > 0);
+}
+
 // "Moves" FDs in move_fds from current to target FD number while keeping FDs
 // in keep_fds open - potentially moving them to another FD number as well in
 // case of colisions.
@@ -547,11 +567,20 @@ bool ForkServer::Initialize() {
   // Don't convert terminated child processes into zombies. It's up to the
   // sandbox (Monitor) to track them and receive/report their final status.
   struct sigaction sa;
+  if (absl::GetFlag(FLAGS_sandbox2_forkserver_use_waitpid)) {
+    sa.sa_handler = SigChldHandler;
+    sa.sa_flags = 0;
+  } else {
     sa.sa_handler = SIG_DFL;
     sa.sa_flags = SA_NOCLDWAIT;
+  }
   sigemptyset(&sa.sa_mask);
   if (sigaction(SIGCHLD, &sa, nullptr) == -1) {
+    if (absl::GetFlag(FLAGS_sandbox2_forkserver_use_waitpid)) {
+      SAPI_RAW_PLOG(ERROR, "sigaction(SIGCHLD, sa_handler=SigChldHandler)");
+    } else {
       SAPI_RAW_PLOG(ERROR, "sigaction(SIGCHLD, flags=SA_NOCLDWAIT)");
+    }
     return false;
   }
   return true;

sandboxed_api/sandbox2/forkserver_bin.cc

@@ -19,16 +19,20 @@
 #include <csignal>
 #include <cstdlib>
 
+#include "absl/flags/parse.h"
 #include "absl/log/globals.h"
 #include "sandboxed_api/sandbox2/comms.h"
 #include "sandboxed_api/sandbox2/forkserver.h"
 #include "sandboxed_api/sandbox2/sanitizer.h"
 #include "sandboxed_api/util/raw_logging.h"
 
-int main() {
+int main(int argc, char* argv[]) {
   // Make sure the logs go stderr.
   absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo);
 
+  // Parse command line arguments.
+  absl::ParseCommandLine(argc, argv);
+
   // Close all non-essential FDs to keep newly opened FD numbers consistent.
   absl::Status status = sandbox2::sanitizer::CloseAllFDsExcept(
       {0, 1, 2, sandbox2::Comms::kSandbox2ClientCommsFD});

sandboxed_api/sandbox2/forkserver_test.cc

@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "sandboxed_api/sandbox2/forkserver.h"
-
 #include <fcntl.h>
 #include <sys/socket.h>
 #include <syscall.h>
@@ -23,6 +21,7 @@
 
 #include "gtest/gtest.h"
 #include "absl/log/check.h"
+#include "absl/log/log.h"
 #include "absl/strings/str_cat.h"
 #include "sandboxed_api/sandbox2/forkserver.pb.h"
 #include "sandboxed_api/sandbox2/global_forkclient.h"

sandboxed_api/sandbox2/global_forkclient.cc

@@ -53,6 +53,10 @@
 #include "sandboxed_api/util/fileops.h"
 #include "sandboxed_api/util/raw_logging.h"
 
+ABSL_FLAG(
+    bool, sandbox2_forkserver_use_waitpid, false,
+    "Use waitpid to reap child processes instead of relying on SA_NOCLDWAIT");
+
 namespace sandbox2 {
 
 namespace file_util = ::sapi::file_util;
@@ -124,6 +128,7 @@ GlobalForkserverStartModeSet GetForkserverStartMode() {
 struct ForkserverArgs {
   int exec_fd;
   int comms_fd;
+  bool use_waitpid;
 };
 
 int LaunchForkserver(void* vargs) {
@@ -142,7 +147,11 @@ int LaunchForkserver(void* vargs) {
                   "duping comms fd failed");
 
   char proc_name[] = "S2-FORK-SERV";
-  char* const argv[] = {proc_name, nullptr};
+  char use_waitpid[] = "--sandbox2_forkserver_use_waitpid";
+  char* argv[] = {proc_name, nullptr, nullptr};
+  if (args->use_waitpid) {
+    argv[1] = use_waitpid;
+  }
   util::Execveat(args->exec_fd, "", argv, environ, AT_EMPTY_PATH);
   SAPI_RAW_PLOG(FATAL, "Could not launch forkserver binary");
 }
@@ -196,6 +205,7 @@ absl::StatusOr<std::unique_ptr<GlobalForkClient>> StartGlobalForkServer() {
   ForkserverArgs args = {
       .exec_fd = exec_fd,
       .comms_fd = sv[0],
+      .use_waitpid = absl::GetFlag(FLAGS_sandbox2_forkserver_use_waitpid),
   };
   pid_t pid = clone(LaunchForkserver, &stack[stack_size], clone_flags, &args,
                     nullptr, nullptr, nullptr);