Utility to intercept fork() system calls

Summary: Creates a utility that will intercept any and all calls to the fork system method. This is a test only library.

Reviewed By: vitaut

Differential Revision: D79674149

fbshipit-source-id: a44da2794ed5e61d3b54848dc197d66c7a9510c6

Author: Charlie Cook

third-party/thrift/src/thrift/compiler/test/fork_intercept.cc

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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
+ *
+ *     http://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 <pthread.h>
+#include <folly/lang/Exception.h>
+#include <thrift/compiler/test/fork_intercept.h>
+
+bool fork_intercept::intercept = []() {
+  bool installed = false;
+#if FOLLY_HAVE_PTHREAD_ATFORK
+
+  pthread_atfork(
+      []() {
+        if (!fork_intercept::get_intercept()) {
+          // Call the real fork
+          return;
+        }
+        folly::terminate_with<std::runtime_error>(
+            "INTERCEPTED: Direct fork() override - intercepted fork() call\n");
+      },
+      nullptr,
+      nullptr);
+  installed = true;
+#endif
+  return installed;
+}();

third-party/thrift/src/thrift/compiler/test/fork_intercept.h

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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
+ *
+ *     http://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.
+ */
+
+#pragma once
+
+/**
+ * This class interfaces with the fork_intercept library to control if forking
+ * is allowed or not. By default it disallows forking and should remain so
+ * during normal operation. However, during testing, forking can be allowed for
+ * testing purposes.
+ */
+class fork_intercept {
+  static bool intercept;
+
+ public:
+  void static set_intercept(bool should_crash) { intercept = should_crash; }
+
+  bool static get_intercept() { return intercept; }
+};

third-party/thrift/src/thrift/compiler/test/fork_intercept_test.cc

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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
+ *
+ *     http://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 <gtest/gtest.h>
+
+#include <sys/types.h>
+#include <folly/Subprocess.h>
+#include <thrift/compiler/test/fork_intercept.h>
+
+namespace apache::thrift::compiler {
+namespace {
+
+TEST(ForkInterceptTest, DirectForkInterception) {
+  // The death test itself needs to fork, so we allow it here but the process
+  // executing the death test will reinstall the intercept
+  fork_intercept::set_intercept(false);
+  ASSERT_DEATH(
+      ([]() {
+        fork_intercept::set_intercept(true);
+        fork();
+      })(),
+      "INTERCEPTED: Direct fork\\(\\) override - intercepted fork\\(\\) call");
+}
+
+TEST(ForkInterceptTest, FollySubprocessInterception) {
+  // This should be allowed by the fork intercept override since
+  // folly::Subprocess does not use fork()
+  std::vector<std::string> args = {"/bin/echo", "Hello world!"};
+  auto proc = folly::Subprocess(args);
+  proc.wait();
+}
+
+} // namespace
+} // namespace apache::thrift::compiler