From 56d11ae733509d9be534591b4aa86f138d39bb45 Mon Sep 17 00:00:00 2001
From: Wiktor Garbacz <wiktorg@google.com>
Date: Fri, 18 Aug 2023 06:52:10 -0700
Subject: [PATCH] Client::PrepareEnvironment simplify by supporting just a
 single preserved fd

PiperOrigin-RevId: 558133382
Change-Id: I043985fcf331761b424ce720791711e5ea1f4fb9
---
 sandboxed_api/sandbox2/client.cc     | 12 +++++-------
 sandboxed_api/sandbox2/client.h      | 10 +++++-----
 sandboxed_api/sandbox2/forkserver.cc |  8 ++------
 3 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/sandboxed_api/sandbox2/client.cc b/sandboxed_api/sandbox2/client.cc
index 1137adb..4055994 100644
--- a/sandboxed_api/sandbox2/client.cc
+++ b/sandboxed_api/sandbox2/client.cc
@@ -148,8 +148,8 @@ std::string Client::GetFdMapEnvVar() const {
                       absl::StrJoin(fd_map_, ",", absl::PairFormatter(",")));
 }
 
-void Client::PrepareEnvironment(std::vector<int>* preserve_fds) {
-  SetUpIPC(preserve_fds);
+void Client::PrepareEnvironment(int* preserved_fd) {
+  SetUpIPC(preserved_fd);
   SetUpCwd();
 }
 
@@ -202,7 +202,7 @@ void Client::SetUpCwd() {
   }
 }
 
-void Client::SetUpIPC(std::vector<int>* preserve_fds) {
+void Client::SetUpIPC(int* preserved_fd) {
   uint32_t num_of_fd_pairs;
   SAPI_RAW_CHECK(comms_->RecvUint32(&num_of_fd_pairs),
                  "receiving number of fd pairs");
@@ -211,10 +211,8 @@ void Client::SetUpIPC(std::vector<int>* preserve_fds) {
   SAPI_RAW_VLOG(1, "Will receive %d file descriptor pairs", num_of_fd_pairs);
 
   absl::flat_hash_map<int, int*> preserve_fds_map;
-  if (preserve_fds) {
-    for (int& fd : *preserve_fds) {
-      preserve_fds_map.emplace(fd, &fd);
-    }
+  if (preserved_fd) {
+    preserve_fds_map.emplace(*preserved_fd, preserved_fd);
   }
 
   for (uint32_t i = 0; i < num_of_fd_pairs; ++i) {
diff --git a/sandboxed_api/sandbox2/client.h b/sandboxed_api/sandbox2/client.h
index ec25a28..02484ef 100644
--- a/sandboxed_api/sandbox2/client.h
+++ b/sandboxed_api/sandbox2/client.h
@@ -91,10 +91,10 @@ class Client {
   std::string GetFdMapEnvVar() const;
 
   // Sets up communication channels with the sandbox.
-  // preserve_fds contains file descriptors that should be kept open and alive.
-  // The FD numbers might be changed if needed and are updated in the vector.
-  // preserve_fds can be a nullptr, equivallent to an empty vector.
-  void SetUpIPC(std::vector<int>* preserve_fds);
+  // preserved_fd contains file descriptor that should be kept open and alive.
+  // The FD number might be changed if needed.
+  // preserved_fd can be a nullptr.
+  void SetUpIPC(int* preserved_fd);
 
   // Sets up the current working directory.
   void SetUpCwd();
@@ -105,7 +105,7 @@ class Client {
   // Applies sandbox-bpf policy, have limits applied on us, and become ptrace'd.
   void ApplyPolicyAndBecomeTracee();
 
-  void PrepareEnvironment(std::vector<int>* preserve_fds = nullptr);
+  void PrepareEnvironment(int* preserved_fd = nullptr);
   void EnableSandbox();
 };
 
diff --git a/sandboxed_api/sandbox2/forkserver.cc b/sandboxed_api/sandbox2/forkserver.cc
index 7505d38..fc96f95 100644
--- a/sandboxed_api/sandbox2/forkserver.cc
+++ b/sandboxed_api/sandbox2/forkserver.cc
@@ -287,8 +287,7 @@ void ForkServer::LaunchChild(const ForkRequest& request, int execve_fd,
   SAPI_RAW_CHECK(request.mode() != FORKSERVER_FORK_UNSPECIFIED,
                  "Forkserver mode is unspecified");
 
-  bool will_execve = (request.mode() == FORKSERVER_FORK_EXECVE ||
-                      request.mode() == FORKSERVER_FORK_EXECVE_SANDBOX);
+  const bool will_execve = execve_fd != -1;
 
   // Prepare the arguments before sandboxing (if needed), as doing it after
   // sandoxing can cause syscall violations (e.g. related to memory management).
@@ -355,10 +354,7 @@ void ForkServer::LaunchChild(const ForkRequest& request, int execve_fd,
     // The following client calls are basically SandboxMeHere. We split it so
     // that we can set up the envp after we received the file descriptors but
     // before we enable the syscall filter.
-    std::vector<int> preserved_fds;
-    preserved_fds.push_back(execve_fd);
-    c.PrepareEnvironment(&preserved_fds);
-    execve_fd = preserved_fds[0];
+    c.PrepareEnvironment(&execve_fd);
 
     if (client_comms.GetConnectionFD() != Comms::kSandbox2ClientCommsFD) {
       envs.push_back(absl::StrCat(Comms::kSandbox2CommsFDEnvVar, "=",