Add a test for custom forkserver

PiperOrigin-RevId: 590187497
Change-Id: I9e2d4a2ed585a78bd3cb44b3f78d91afd527f6ab

sandboxed_api/sandbox2/BUILD.bazel

@@ -957,6 +957,7 @@ cc_test(
     copts = sapi_platform_copts(),
     data = [
         "//sandboxed_api/sandbox2/testcases:abort",
+        "//sandboxed_api/sandbox2/testcases:custom_fork",
         "//sandboxed_api/sandbox2/testcases:minimal",
         "//sandboxed_api/sandbox2/testcases:sleep",
         "//sandboxed_api/sandbox2/testcases:starve",
@@ -967,6 +968,7 @@ cc_test(
         "no_qemu_user_mode",
     ],
     deps = [
+        ":fork_client",
         ":sandbox2",
         "//sandboxed_api:config",
         "//sandboxed_api:testing",

sandboxed_api/sandbox2/CMakeLists.txt

@@ -1039,6 +1039,7 @@ if(BUILD_TESTING AND SAPI_BUILD_TESTING)
   )
   add_dependencies(sandbox2_sandbox2_test
     sandbox2::testcase_abort
+    sandbox2::testcase_custom_fork
     sandbox2::testcase_minimal
     sandbox2::testcase_sleep
     sandbox2::testcase_tsync
@@ -1050,6 +1051,7 @@ if(BUILD_TESTING AND SAPI_BUILD_TESTING)
     absl::synchronization
     absl::time
     sapi::config
+    sandbox2::fork_client
     sandbox2::sandbox2
     sapi::testing
     sapi::status_matchers

sandboxed_api/sandbox2/sandbox2_test.cc

@@ -35,6 +35,7 @@
 #include "absl/time/time.h"
 #include "sandboxed_api/config.h"
 #include "sandboxed_api/sandbox2/executor.h"
+#include "sandboxed_api/sandbox2/fork_client.h"
 #include "sandboxed_api/sandbox2/policy.h"
 #include "sandboxed_api/sandbox2/policybuilder.h"
 #include "sandboxed_api/sandbox2/result.h"
@@ -51,6 +52,7 @@ using ::testing::Eq;
 using ::testing::IsEmpty;
 using ::testing::IsTrue;
 using ::testing::Lt;
+using ::testing::Ne;
 
 class Sandbox2Test : public ::testing::TestWithParam<bool> {
  public:
@@ -195,6 +197,23 @@ TEST_P(Sandbox2Test, SandboxeeNotKilledWhenStartingThreadFinishes) {
   EXPECT_EQ(result.final_status(), Result::OK);
 }
 
+TEST_P(Sandbox2Test, CustomForkserverWorks) {
+  const std::string path = GetTestSourcePath("sandbox2/testcases/custom_fork");
+  std::vector<std::string> args = {path};
+  auto fork_executor = std::make_unique<Executor>(path, args);
+  std::unique_ptr<ForkClient> fork_client = fork_executor->StartForkServer();
+  ASSERT_THAT(fork_client.get(), Ne(nullptr));
+
+  SAPI_ASSERT_OK_AND_ASSIGN(auto policy,
+                            CreateDefaultTestPolicy(path).TryBuild());
+
+  Sandbox2 sandbox(std::make_unique<Executor>(fork_client.get()),
+                   std::move(policy));
+  ASSERT_THAT(SetUpSandbox(&sandbox), IsOk());
+  Result result = sandbox.Run();
+  EXPECT_EQ(result.final_status(), Result::OK);
+}
+
 TEST(StarvationTest, MonitorIsNotStarvedByTheSandboxee) {
   const std::string path = GetTestSourcePath("sandbox2/testcases/starve");
 

sandboxed_api/sandbox2/testcases/BUILD.bazel

@@ -261,3 +261,16 @@ cc_binary(
         "@com_google_absl//absl/strings:string_view",
     ],
 )
+
+cc_binary(
+    name = "custom_fork",
+    testonly = True,
+    srcs = ["custom_fork.cc"],
+    copts = sapi_platform_copts(),
+    features = ["fully_static_link"],
+    deps = [
+        "//sandboxed_api/sandbox2:comms",
+        "//sandboxed_api/sandbox2:forkingclient",
+        "//sandboxed_api/util:raw_logging",
+    ],
+)

sandboxed_api/sandbox2/testcases/CMakeLists.txt

@@ -322,3 +322,19 @@ target_link_libraries(sandbox2_testcase_network_proxy PRIVATE
   sapi::status
 )
 
+# sandboxed_api/sandbox2/testcases:custom_fork
+add_executable(sandbox2_testcase_custom_fork
+  custom_fork.cc
+)
+add_executable(sandbox2::testcase_custom_fork ALIAS sandbox2_testcase_custom_fork)
+set_target_properties(sandbox2_testcase_custom_fork PROPERTIES
+  OUTPUT_NAME custom_fork
+)
+target_link_libraries(sandbox2_testcase_custom_fork PRIVATE
+  sandbox2::comms
+  sandbox2::forkingclient
+  sapi::base
+  sapi::raw_logging
+)
+
+

sandboxed_api/sandbox2/testcases/custom_fork.cc

@@ -0,0 +1,34 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "sandboxed_api/sandbox2/comms.h"
+#include "sandboxed_api/sandbox2/forkingclient.h"
+#include "sandboxed_api/util/raw_logging.h"
+
+int main(int argc, char* argv[]) {
+  sandbox2::Comms comms(sandbox2::Comms::kDefaultConnection);
+  sandbox2::ForkingClient s2client(&comms);
+
+  for (;;) {
+    pid_t pid = s2client.WaitAndFork();
+    if (pid == -1) {
+      SAPI_RAW_LOG(FATAL, "Could not spawn a new sandboxee");
+    }
+    if (pid == 0) {
+      // Start sandboxing here
+      s2client.SandboxMeHere();
+      return 0;
+    }
+  }
+}