an option for Subprocess child to print pid to a buffer

Summary:
Let the child process print its pid into a provided buffer.

The use-case is to that the child process can print its pid into zero'd storage in an environment variable, before the call to `execve`. Then the child process after `execve` will be able to see its own pid in that environment variable.

The caller is responsible for preparing the environment-variable vector with the buffer.

This is somewhat dangerous since the caller can provide any buffer, but it is simpler to implement the capability this way.

This follows a feature of SystemD, which spawns child processes with their pid's in the environment variable `SYSTEMD_EXEC_PID`.

Reviewed By: Gownta

Differential Revision: D70577643

fbshipit-source-id: 7d600b74ae239e74c3c57e10510d232818b43744

Author: yfeldblum

folly/Subprocess.cpp

@@ -136,6 +136,7 @@ static Ret subprocess_libc_load(
   X(fcntl, fcntl)                              \
   X(pthread_sigmask, pthread_sigmask)          \
   X(signal, signal)                            \
+  X(sprintf, sprintf)                          \
   X(strtol, strtol)                            \
   X(vfork, vfork)
 
@@ -190,9 +191,13 @@ __attribute__((constructor(101))) static void subprocess_libc_init() {
 struct Subprocess::SpawnRawArgs {
   struct Scratch {
     std::vector<std::pair<int, int>> fdActions;
+    std::vector<char*> setPrintPidToBuffer;
 
     explicit Scratch(Options const& options)
-        : fdActions{options.fdActions_.begin(), options.fdActions_.end()} {
+        : fdActions{options.fdActions_.begin(), options.fdActions_.end()},
+          setPrintPidToBuffer{
+              options.setPrintPidToBuffer_.begin(),
+              options.setPrintPidToBuffer_.end()} {
       std::sort(fdActions.begin(), fdActions.end());
     }
   };
@@ -209,6 +214,8 @@ struct Subprocess::SpawnRawArgs {
   int parentDeathSignal{};
   bool processGroupLeader{};
   bool usePath{};
+  char* const* setPrintPidToBufferData{};
+  size_t setPrintPidToBufferSize{};
 
   // assigned explicitly
   char const* const* argv{};
@@ -231,7 +238,9 @@ struct Subprocess::SpawnRawArgs {
         parentDeathSignal{options.parentDeathSignal_},
 #endif
         processGroupLeader{options.processGroupLeader_},
-        usePath{options.usePath_} {
+        usePath{options.usePath_},
+        setPrintPidToBufferData{scratch.setPrintPidToBuffer.data()},
+        setPrintPidToBufferSize{scratch.setPrintPidToBuffer.size()} {
     static_assert(std::is_standard_layout_v<Subprocess::SpawnRawArgs>);
     static_assert(std::is_trivially_destructible_v<Subprocess::SpawnRawArgs>);
   }
@@ -374,6 +383,14 @@ Subprocess::Options& Subprocess::Options::fd(int fd, int action) {
   return *this;
 }
 
+Subprocess::Options& Subprocess::Options::addPrintPidToBuffer(span<char> buf) {
+  if (buf.size() < kPidBufferMinSize) {
+    throw std::invalid_argument("buf size too small");
+  }
+  setPrintPidToBuffer_.insert(buf.data());
+  return *this;
+}
+
 Subprocess::Subprocess() = default;
 
 Subprocess::Subprocess(
@@ -833,6 +850,11 @@ int Subprocess::prepareChild(SpawnRawArgs const& args) {
     }
   }
 
+  for (size_t i = 0; i < args.setPrintPidToBufferSize; ++i) {
+    auto buf = args.setPrintPidToBufferData[i];
+    detail::subprocess_libc::sprintf(buf, "%d", getpid());
+  }
+
   // The user callback comes last, so that the child is otherwise all set up.
   if (args.dangerousPostForkPreExecCallback) {
     if (int error = (*args.dangerousPostForkPreExecCallback)()) {

folly/Subprocess.h

@@ -330,6 +330,15 @@ class Subprocess {
     friend class Subprocess;
 
    public:
+    // digits10 is the maximum number of decimal digits such that any number
+    // up to this many decimal digits can always be represented in the given
+    // integer type
+    // but we need to have storage for the decimal representation of any
+    // integer, so +1, and we need to have storage for the terminal null, so
+    // again +1.
+    static inline constexpr size_t kPidBufferMinSize =
+        std::numeric_limits<pid_t>::digits10 + 2;
+
     Options() {} // E.g. https://gcc.gnu.org/bugzilla/show_bug.cgi?id=58328
 
     /**
@@ -538,6 +547,8 @@ class Subprocess {
     }
 #endif
 
+    Options& addPrintPidToBuffer(span<char> buf);
+
    private:
     typedef boost::container::flat_map<int, int> FdMap;
     FdMap fdActions_;
@@ -559,6 +570,7 @@ class Subprocess {
 #if defined(__linux__)
     Optional<cpu_set_t> cpuSet_;
 #endif
+    std::unordered_set<char*> setPrintPidToBuffer_;
   };
 
   // Non-copyable, but movable

folly/test/BUCK

@@ -1625,10 +1625,12 @@ cpp_unittest(
         "//folly:format",
         "//folly:string",
         "//folly:subprocess",
+        "//folly/container:span",
         "//folly/experimental/io:fs_util",
         "//folly/gen:base",
         "//folly/gen:file",
         "//folly/gen:string",
+        "//folly/portability:gmock",
         "//folly/portability:gtest",
         "//folly/portability:unistd",
         "//folly/testing:test_util",

folly/test/SubprocessTest.cpp

@@ -19,6 +19,9 @@
 #include <sys/types.h>
 
 #include <chrono>
+#include <string>
+#include <string_view>
+#include <vector>
 
 #include <boost/container/flat_set.hpp>
 #include <glog/logging.h>
@@ -27,10 +30,12 @@
 #include <folly/FileUtil.h>
 #include <folly/Format.h>
 #include <folly/String.h>
+#include <folly/container/span.h>
 #include <folly/experimental/io/FsUtil.h>
 #include <folly/gen/Base.h>
 #include <folly/gen/File.h>
 #include <folly/gen/String.h>
+#include <folly/portability/GMock.h>
 #include <folly/portability/GTest.h>
 #include <folly/portability/Unistd.h>
 #include <folly/testing/TestUtil.h>
@@ -39,6 +44,7 @@ FOLLY_GNU_DISABLE_WARNING("-Wdeprecated-declarations")
 
 using namespace folly;
 using namespace std::chrono_literals;
+using namespace std::string_view_literals;
 
 namespace std::chrono {
 template <typename Rep, typename Period>
@@ -920,3 +926,47 @@ TEST(KeepFileOpenSubprocessTest, KeepsFileOpen) {
   std::sort(std::begin(fds), std::end(fds));
   EXPECT_EQ(fmt::format("{}\n", fmt::join(fds, "\n")), p.first);
 }
+
+static_assert(
+    Subprocess::Options::kPidBufferMinSize ==
+    std::numeric_limits<pid_t>::digits10 + 2);
+
+TEST(WritePidIntoBufTest, WritesPidIntoBufTooSmall) {
+  constexpr size_t size = Subprocess::Options::kPidBufferMinSize;
+  char buf[size - 1] = {};
+  auto options = Subprocess::Options();
+  EXPECT_THROW(options.addPrintPidToBuffer(buf), std::invalid_argument);
+}
+
+TEST(WritePidIntoBufTest, WritesPidIntoBuf) {
+  constexpr size_t size = Subprocess::Options::kPidBufferMinSize;
+  char buf[size] = {};
+  std::memset(buf, 0xA5, size);
+  auto options = Subprocess::Options().addPrintPidToBuffer(buf);
+  Subprocess proc(std::vector<std::string>{"/bin/true"}, options);
+  EXPECT_EQ(fmt::format("{}", proc.pid()), buf);
+  proc.wait();
+}
+
+TEST(WritePidIntoBufTest, WritesPidIntoBufExampleEnvVar) {
+  // this test effectively duplicates WritesPidIntoBuf but may serve as a
+  // reference for how to use this feature with environment-variable storage
+  //
+  // systemd does something like this:
+  // https://www.freedesktop.org/software/systemd/man/latest/systemd.exec.html#%24SYSTEMD_EXEC_PID
+  constexpr size_t size = Subprocess::Options::kPidBufferMinSize;
+  constexpr auto prefix = "FOLLY_TEST_SUBPROCESS_PID="sv;
+  std::vector<std::string> env;
+  env.emplace_back("FOLLY_TEST_GREETING=hello world");
+  auto& var = env.emplace_back(prefix);
+  var.resize(prefix.size() + size); // must be stable! no more changes to env!
+  auto buf = folly::span{var}.subspan(prefix.size());
+  auto options = Subprocess::Options().pipeStdout().addPrintPidToBuffer(buf);
+  Subprocess proc(std::vector<std::string>{"/bin/env"}, options, nullptr, &env);
+  auto pid = proc.pid();
+  auto p = proc.communicate();
+  proc.wait();
+  std::vector<std::string_view> lines;
+  folly::split('\n', p.first, lines);
+  EXPECT_THAT(lines, testing::Contains(fmt::format("{}{}", prefix, pid)));
+}