From 5efae5cdf5bbfd6f1bf950424a07a66b11a7c7fe Mon Sep 17 00:00:00 2001
From: Wiktor Garbacz <wiktorg@google.com>
Date: Wed, 29 Mar 2023 02:21:31 -0700
Subject: [PATCH] Do not exit from within ForkServer to get more precise
 coverage data

PiperOrigin-RevId: 520273079
Change-Id: I3f37d9eacc2c284c45f37842e1e63364cf64faf2
---
 sandboxed_api/sandbox2/BUILD.bazel       |  1 +
 sandboxed_api/sandbox2/CMakeLists.txt    |  1 +
 sandboxed_api/sandbox2/forkingclient.cc  | 11 ++++++++++-
 sandboxed_api/sandbox2/forkserver.cc     |  5 +++--
 sandboxed_api/sandbox2/forkserver.h      |  3 +++
 sandboxed_api/sandbox2/forkserver_bin.cc |  5 +++--
 6 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/sandboxed_api/sandbox2/BUILD.bazel b/sandboxed_api/sandbox2/BUILD.bazel
index 0564e5d..39682ed 100644
--- a/sandboxed_api/sandbox2/BUILD.bazel
+++ b/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",
     ],
 )
diff --git a/sandboxed_api/sandbox2/CMakeLists.txt b/sandboxed_api/sandbox2/CMakeLists.txt
index e4f5584..de47282 100644
--- a/sandboxed_api/sandbox2/CMakeLists.txt
+++ b/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
diff --git a/sandboxed_api/sandbox2/forkingclient.cc b/sandboxed_api/sandbox2/forkingclient.cc
index 23c9b69..c7ce23e 100644
--- a/sandboxed_api/sandbox2/forkingclient.cc
+++ b/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
diff --git a/sandboxed_api/sandbox2/forkserver.cc b/sandboxed_api/sandbox2/forkserver.cc
index a11e593..78b94cf 100644
--- a/sandboxed_api/sandbox2/forkserver.cc
+++ b/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
diff --git a/sandboxed_api/sandbox2/forkserver.h b/sandboxed_api/sandbox2/forkserver.h
index 6603fe1..393ba9e 100644
--- a/sandboxed_api/sandbox2/forkserver.h
+++ b/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.
diff --git a/sandboxed_api/sandbox2/forkserver_bin.cc b/sandboxed_api/sandbox2/forkserver_bin.cc
index 7952aa5..f2208e4 100644
--- a/sandboxed_api/sandbox2/forkserver_bin.cc
+++ b/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");
 }