Use absl::flat_hash_set + Status in favor of std::set in the sanitizer API

PiperOrigin-RevId: 395061068
Change-Id: I31548eb6fc9f27f55acf25bd6d3d0b941a529e63

sandboxed_api/sandbox2/BUILD.bazel

@@ -388,6 +388,7 @@ cc_library(
         "//sandboxed_api/util:strerror",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_glog//:glog",
@@ -774,6 +775,7 @@ cc_test(
         "//sandboxed_api:testing",
         "//sandboxed_api/sandbox2/util:bpf_helper",
         "//sandboxed_api/util:status_matchers",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",

sandboxed_api/sandbox2/forkserver.cc

@@ -556,11 +556,13 @@ void ForkServer::CreateInitialNamespaces() {
 void ForkServer::SanitizeEnvironment() {
   // Mark all file descriptors, except the standard ones (needed
   // for proper sandboxed process operations), as close-on-exec.
-  SAPI_RAW_CHECK(sanitizer::SanitizeCurrentProcess(
+  absl::Status status = sanitizer::SanitizeCurrentProcess(
-                     {STDIN_FILENO, STDOUT_FILENO, STDERR_FILENO,
+      {STDIN_FILENO, STDOUT_FILENO, STDERR_FILENO,
-                      Comms::kSandbox2ClientCommsFD},
+       Comms::kSandbox2ClientCommsFD},
-                     /* close_fds = */ false),
+      /* close_fds = */ false);
-                 "while sanitizing process");
+  SAPI_RAW_CHECK(
+      status.ok(),
+      absl::StrCat("while sanitizing process: ", status.message()).c_str());
 }
 
 void ForkServer::ExecuteProcess(int execve_fd, const char** argv,

sandboxed_api/sandbox2/forkserver_bin.cc

@@ -32,9 +32,13 @@ int main() {
   google::LogToStderr();
 
   // Close all non-essential FDs to keep newly opened FD numbers consistent.
-  sandbox2::sanitizer::CloseAllFDsExcept(
+  absl::Status status = sandbox2::sanitizer::CloseAllFDsExcept(
       {0, 1, 2, sandbox2::Comms::kSandbox2ClientCommsFD});
 
+  if (!status.ok()) {
+    SAPI_RAW_LOG(WARNING, "Closing non-essential FDs failed");
+  }
+
   // Make the process' name easily recognizable with ps/pstree.
   if (prctl(PR_SET_NAME, "S2-FORK-SERV", 0, 0, 0) != 0) {
     SAPI_RAW_PLOG(WARNING, "prctl(PR_SET_NAME, 'S2-FORK-SERV')");

sandboxed_api/sandbox2/monitor.cc

@@ -44,6 +44,7 @@
 
 #include <glog/logging.h>
 #include "absl/cleanup/cleanup.h"
+#include "absl/container/flat_hash_set.h"
 #include "sandboxed_api/util/flag.h"
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
@@ -641,9 +642,12 @@ bool Monitor::InitPtraceAttach() {
   sanitizer::WaitForTsan();
 
   // Get a list of tasks.
-  std::set<int> tasks;
+  absl::flat_hash_set<int> tasks;
-  if (!sanitizer::GetListOfTasks(pid_, &tasks)) {
+  if (auto task_list = sanitizer::GetListOfTasks(pid_); task_list.ok()) {
-    LOG(ERROR) << "Could not get list of tasks";
+    tasks = *std::move(task_list);
+  } else {
+    LOG(ERROR) << "Could not get list of tasks: "
+               << task_list.status().message();
     return false;
   }
 
@@ -662,13 +666,13 @@ bool Monitor::InitPtraceAttach() {
                  << ".";
   }
 
-  std::set<int> tasks_attached;
+  absl::flat_hash_set<int> tasks_attached;
   int retries = 0;
   absl::Time deadline = absl::Now() + absl::Seconds(2);
 
   // In some situations we allow ptrace to try again when it fails.
   while (!tasks.empty()) {
-    std::set<int> tasks_left;
+    absl::flat_hash_set<int> tasks_left;
     for (int task : tasks) {
       constexpr intptr_t options =
           PTRACE_O_TRACESYSGOOD | PTRACE_O_TRACEFORK | PTRACE_O_TRACEVFORK |
@@ -719,8 +723,11 @@ bool Monitor::InitPtraceAttach() {
   }
 
   // Get a list of tasks after attaching.
-  if (!sanitizer::GetListOfTasks(pid_, &tasks)) {
+  if (auto tasks_list = sanitizer::GetListOfTasks(pid_); tasks_list.ok()) {
-    LOG(ERROR) << "Could not get list of tasks";
+    tasks = *std::move(tasks_list);
+  } else {
+    LOG(ERROR) << "Could not get list of tasks: "
+               << tasks_list.status().message();
     return false;
   }
 

sandboxed_api/sandbox2/sanitizer.cc

@@ -42,6 +42,7 @@
 #include "sandboxed_api/sandbox2/util.h"
 #include "sandboxed_api/util/file_helpers.h"
 #include "sandboxed_api/util/fileops.h"
+#include "sandboxed_api/util/os_error.h"
 #include "sandboxed_api/util/raw_logging.h"
 #include "sandboxed_api/util/status_macros.h"
 #include "sandboxed_api/util/strerror.h"
@@ -91,41 +92,29 @@ absl::StatusOr<absl::flat_hash_set<int>> GetListOfFDs() {
   return fds;
 }
 
-bool GetListOfTasks(int pid, std::set<int>* tasks) {
+absl::StatusOr<absl::flat_hash_set<int>> GetListOfTasks(int pid) {
   const std::string task_dir = absl::StrCat("/proc/", pid, "/task");
-  auto task_entries = ListNumericalDirectoryEntries(task_dir);
+  return ListNumericalDirectoryEntries(task_dir);
-  if (!task_entries.ok()) {
-    return false;
-  }
-
-  tasks->clear();
-  tasks->insert(task_entries->begin(), task_entries->end());
-  return true;
 }
 
-bool CloseAllFDsExcept(const std::set<int>& fd_exceptions) {
+absl::Status CloseAllFDsExcept(const absl::flat_hash_set<int>& fd_exceptions) {
-  absl::StatusOr<absl::flat_hash_set<int>> fds = GetListOfFDs();
+  SAPI_ASSIGN_OR_RETURN(absl::flat_hash_set<int> fds, GetListOfFDs());
-  if (!fds.ok()) {
-    return false;
-  }
 
-  for (auto fd : *fds) {
+  for (const auto& fd : fds) {
     if (fd_exceptions.find(fd) != fd_exceptions.end()) {
       continue;
     }
     SAPI_RAW_VLOG(2, "Closing FD:%d", fd);
     close(fd);
   }
-  return true;
+  return absl::OkStatus();
 }
 
-bool MarkAllFDsAsCOEExcept(const std::set<int>& fd_exceptions) {
+absl::Status MarkAllFDsAsCOEExcept(
-  auto fds = GetListOfFDs();
+    const absl::flat_hash_set<int>& fd_exceptions) {
-  if (!fds.ok()) {
+  SAPI_ASSIGN_OR_RETURN(absl::flat_hash_set<int> fds, GetListOfFDs());
-    return false;
-  }
 
-  for (auto fd : *fds) {
+  for (const auto& fd : fds) {
     if (fd_exceptions.find(fd) != fd_exceptions.end()) {
       continue;
     }
@@ -134,17 +123,16 @@ bool MarkAllFDsAsCOEExcept(const std::set<int>& fd_exceptions) {
 
     int flags = fcntl(fd, F_GETFD);
     if (flags == -1) {
-      SAPI_RAW_PLOG(ERROR, "fcntl(%d, F_GETFD) failed", fd);
+      return absl::InternalError(
-      return false;
+          sapi::OsErrorMessage(errno, "fcntl(", fd, ", F_GETFD) failed"));
     }
     if (fcntl(fd, F_SETFD, flags | FD_CLOEXEC) == -1) {
-      SAPI_RAW_PLOG(ERROR, "fcntl(%d, F_SETFD, %x | FD_CLOEXEC) failed", fd,
+      return absl::InternalError(sapi::OsErrorMessage(
-                    flags);
+          errno, "fcntl(", fd, ", F_SETFD, ", flags, " | FD_CLOEXEC) failed"));
-      return false;
     }
   }
 
-  return true;
+  return absl::OkStatus();
 }
 
 int GetNumberOfThreads(int pid) {
@@ -180,8 +168,8 @@ void WaitForTsan() {
 #endif
 }
 
-bool SanitizeCurrentProcess(const std::set<int>& fd_exceptions,
+absl::Status SanitizeCurrentProcess(
-                            bool close_fds) {
+    const absl::flat_hash_set<int>& fd_exceptions, bool close_fds) {
   SAPI_RAW_VLOG(1, "Sanitizing PID: %zu, close_fds: %d", syscall(__NR_getpid),
                 close_fds);
 
@@ -190,23 +178,13 @@ bool SanitizeCurrentProcess(const std::set<int>& fd_exceptions,
 
   // If the parent goes down, so should we.
   if (prctl(PR_SET_PDEATHSIG, SIGKILL, 0, 0, 0) != 0) {
-    SAPI_RAW_PLOG(ERROR, "prctl(PR_SET_PDEATHSIG, SIGKILL) failed");
+    return absl::InternalError(
-    return false;
+        sapi::OsErrorMessage(errno, "prctl(PR_SET_PDEATHSIG, SIGKILL) failed"));
   }
 
   // Close or mark as close-on-exec open file descriptors.
-  if (close_fds) {
+  return close_fds ? CloseAllFDsExcept(fd_exceptions)
-    if (!CloseAllFDsExcept(fd_exceptions)) {
+                   : MarkAllFDsAsCOEExcept(fd_exceptions);
-      SAPI_RAW_LOG(ERROR, "Failed to close all fds");
-      return false;
-    }
-  } else {
-    if (!MarkAllFDsAsCOEExcept(fd_exceptions)) {
-      SAPI_RAW_LOG(ERROR, "Failed to mark all fds as closed");
-      return false;
-    }
-  }
-  return true;
 }
 
 }  // namespace sandbox2::sanitizer

sandboxed_api/sandbox2/sanitizer.h

@@ -18,10 +18,9 @@
 #ifndef SANDBOXED_API_SANDBOX2_SANITIZER_H_
 #define SANDBOXED_API_SANDBOX2_SANITIZER_H_
 
-#include <set>
-
 #include "absl/base/macros.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 
 namespace sandbox2 {
@@ -32,11 +31,12 @@ absl::StatusOr<absl::flat_hash_set<int>> GetListOfFDs();
 
 // Closes all file descriptors in the current process except the ones in
 // fd_exceptions.
-bool CloseAllFDsExcept(const std::set<int>& fd_exceptions);
+absl::Status CloseAllFDsExcept(const absl::flat_hash_set<int>& fd_exceptions);
 
 // Marks all file descriptors as close-on-exec, except the ones in
 // fd_exceptions.
-bool MarkAllFDsAsCOEExcept(const std::set<int>& fd_exceptions);
+absl::Status MarkAllFDsAsCOEExcept(
+    const absl::flat_hash_set<int>& fd_exceptions);
 
 // Returns the number of threads in the process 'pid'. Returns -1 in case of
 // errors.
@@ -53,10 +53,11 @@ void WaitForTsan();
 // Sanitizes current process (which will not execve a sandboxed binary).
 // File-descriptors in fd_exceptions will be either closed
 // (close_fds == true), or marked as close-on-exec (close_fds == false).
-bool SanitizeCurrentProcess(const std::set<int>& fd_exceptions, bool close_fds);
+absl::Status SanitizeCurrentProcess(
+    const absl::flat_hash_set<int>& fd_exceptions, bool close_fds);
 
 // Returns a list of tasks for a pid.
-bool GetListOfTasks(int pid, std::set<int>* tasks);
+absl::StatusOr<absl::flat_hash_set<int>> GetListOfTasks(int pid);
 
 }  // namespace sanitizer
 }  // namespace sandbox2

sandboxed_api/sandbox2/sanitizer_test.cc

@@ -27,6 +27,7 @@
 #include <glog/logging.h>
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/memory/memory.h"
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
@@ -89,13 +90,13 @@ TEST(SanitizerTest, TestMarkFDsAsCOE) {
   int null_fd = open("/dev/null", O_RDWR);
   ASSERT_THAT(null_fd, Ne(-1));
 
-  const std::set<int> exceptions = {STDIN_FILENO, STDOUT_FILENO, STDERR_FILENO,
+  const absl::flat_hash_set<int> keep = {STDIN_FILENO, STDOUT_FILENO,
-                                    null_fd};
+                                         STDERR_FILENO, null_fd};
-  ASSERT_THAT(sanitizer::MarkAllFDsAsCOEExcept(exceptions), IsTrue());
+  ASSERT_THAT(sanitizer::MarkAllFDsAsCOEExcept(keep), sapi::IsOk());
 
   const std::string path = GetTestSourcePath("sandbox2/testcases/sanitizer");
   std::vector<std::string> args;
-  for (auto fd : exceptions) {
+  for (auto fd : keep) {
     args.push_back(absl::StrCat(fd));
   }
   ASSERT_THAT(RunTestcase(path, args), Eq(0));