Do not exit from within ForkServer to get more precise coverage data

PiperOrigin-RevId: 520273079
Change-Id: I3f37d9eacc2c284c45f37842e1e63364cf64faf2

sandboxed_api/sandbox2/BUILD.bazel

@@ -722,6 +722,7 @@ cc_library(
         ":comms",
         ":forkserver",
         ":sanitizer",
+        "//sandboxed_api/util:raw_logging",
         "@com_google_absl//absl/log:check",
     ],
 )

sandboxed_api/sandbox2/CMakeLists.txt

@@ -638,6 +638,7 @@ target_link_libraries(sandbox2_forkingclient
           absl::log
           sandbox2::sanitizer
           sapi::base
+          sapi::raw_logging
   PUBLIC sandbox2::client
          sandbox2::comms
          sandbox2::forkserver

sandboxed_api/sandbox2/forkingclient.cc

@@ -14,12 +14,16 @@
 
 #include "sandboxed_api/sandbox2/forkingclient.h"
 
+#include <sys/types.h>
 #include <unistd.h>
 
+#include <cstdlib>
 #include <memory>
 
 #include "absl/log/check.h"
+#include "sandboxed_api/sandbox2/forkserver.h"
 #include "sandboxed_api/sandbox2/sanitizer.h"
+#include "sandboxed_api/util/raw_logging.h"
 
 namespace sandbox2 {
 
@@ -36,7 +40,12 @@ pid_t ForkingClient::WaitAndFork() {
                    << ") during sandbox2::Client::WaitAndFork()";
     fork_server_worker_ = std::make_unique<ForkServer>(comms_);
   }
-  return fork_server_worker_->ServeRequest();
+  pid_t pid = fork_server_worker_->ServeRequest();
+  if (pid == -1 && fork_server_worker_->IsTerminated()) {
+    SAPI_RAW_VLOG(1, "ForkServer Comms closed. Exiting");
+    exit(0);
+  }
+  return pid;
 }
 
 }  // namespace sandbox2

sandboxed_api/sandbox2/forkserver.cc

@@ -389,8 +389,7 @@ pid_t ForkServer::ServeRequest() {
   ForkRequest fork_request;
   if (!comms_->RecvProtoBuf(&fork_request)) {
     if (comms_->IsTerminated()) {
-      SAPI_RAW_VLOG(1, "ForkServer Comms closed. Exiting");
-      exit(0);
+      return -1;
     }
     SAPI_RAW_LOG(FATAL, "Failed to receive ForkServer request");
   }
@@ -533,6 +532,8 @@ pid_t ForkServer::ServeRequest() {
   return sandboxee_pid;
 }
 
+bool ForkServer::IsTerminated() const { return comms_->IsTerminated(); }
+
 bool ForkServer::Initialize() {
   // All processes spawned by the fork'd/execute'd process will see this process
   // as /sbin/init. Therefore it will receive (and ignore) their final status

sandboxed_api/sandbox2/forkserver.h

@@ -42,6 +42,9 @@ class ForkServer {
     }
   }
 
+  // Returns whether the connection with the forkserver was terminated.
+  bool IsTerminated() const;
+
   // Receives a fork request from the master process. The started process does
   // not need to be waited for (with waitid/waitpid/wait3/wait4) as the current
   // process will have the SIGCHLD set to sa_flags=SA_NOCLDWAIT.

sandboxed_api/sandbox2/forkserver_bin.cc

@@ -60,12 +60,13 @@ int main() {
   sandbox2::Comms comms(sandbox2::Comms::kDefaultConnection);
   sandbox2::ForkServer fork_server(&comms);
 
-  while (true) {
+  while (!fork_server.IsTerminated()) {
     pid_t child_pid = fork_server.ServeRequest();
-    if (!child_pid) {
+    if (child_pid == 0) {
       // FORKSERVER_FORK sent to the global forkserver. This case does not make
       // sense, we thus kill the process here.
       _Exit(0);
     }
   }
+  SAPI_RAW_VLOG(1, "ForkServer Comms closed. Exiting");
 }