Expand environment-handling in folly

Summary: Extract the environment-as-STL-map logic out of `test::EnvVarSaver` and into a separate class `experimental::EnvironmentState` so that other code which needs to manipulate the environment can do so more easily. Add routines to set the process environment from the state of an `EnvironmentState`'s map, and to extract the environment in the forms required by `Subprocess` and the UNIX routines `execve` etc. Reviewed By: yfeldblum Differential Revision: D4713307 fbshipit-source-id: 6b1380dd29b9ba41c97b886814dd3eee91fc1c0f
mtabebe · Mar 24, 2017 · 55cd50f · 55cd50f
1 parent d91466d
commit 55cd50f
Show file tree

Hide file tree

Showing 4 changed files with 306 additions and 42 deletions.
diff --git a/folly/experimental/EnvUtil.cpp b/folly/experimental/EnvUtil.cpp
@@ -20,37 +20,72 @@
 #include <folly/portability/Stdlib.h>
 #include <folly/portability/Unistd.h>
 
-namespace folly {
-namespace test {
+using namespace folly;
+using namespace folly::experimental;
 
-static std::map<std::string, std::string> getEnvVarMap() {
-  std::map<std::string, std::string> data;
-  for (auto it = environ; *it != nullptr; ++it) {
+EnvironmentState EnvironmentState::fromCurrentEnvironment() {
+  std::unordered_map<std::string, std::string> data;
+  for (auto it = environ; it && *it; ++it) {
     std::string key, value;
-    split("=", *it, key, value);
-    if (key.empty()) {
-      continue;
+    folly::StringPiece entry(*it);
+    auto equalsPosition = entry.find('=');
+    if (equalsPosition == entry.npos) {
+      throw MalformedEnvironment{to<std::string>(
+          "Environment contains an non key-value-pair string \"", entry, "\"")};
     }
-    CHECK(!data.count(key)) << "already contains: " << key;
-    data.emplace(move(key), move(value));
+    key = entry.subpiece(0, equalsPosition).toString();
+    value = entry.subpiece(equalsPosition + 1).toString();
+    if (data.count(key)) {
+      throw MalformedEnvironment{to<std::string>(
+          "Environment contains duplicate value for \"", key, "\"")};
+    }
+    data.emplace(std::move(key), std::move(value));
   }
-  return data;
-}
-
-EnvVarSaver::EnvVarSaver() {
-  saved_ = getEnvVarMap();
+  return EnvironmentState{std::move(data)};
 }
 
-EnvVarSaver::~EnvVarSaver() {
-  for (const auto& kvp : getEnvVarMap()) {
-    if (saved_.count(kvp.first)) {
-      continue;
-    }
-    PCHECK(0 == unsetenv(kvp.first.c_str()));
-  }
-  for (const auto& kvp : saved_) {
+void EnvironmentState::setAsCurrentEnvironment() {
+  PCHECK(0 == clearenv());
+  for (const auto& kvp : env_) {
     PCHECK(0 == setenv(kvp.first.c_str(), kvp.second.c_str(), (int)true));
   }
 }
+
+std::vector<std::string> EnvironmentState::toVector() const {
+  std::vector<std::string> result;
+  for (auto const& pair : env_) {
+    result.emplace_back(to<std::string>(pair.first, "=", pair.second));
+  }
+  return result;
 }
+
+std::unique_ptr<char*, void (*)(char**)> EnvironmentState::toPointerArray()
+    const {
+  size_t totalStringLength{};
+  for (auto const& pair : env_) {
+    totalStringLength += pair.first.size() + pair.second.size() +
+        2 /* intermediate '=' and the terminating NUL */;
+  }
+  size_t allocationRequired =
+      (totalStringLength / sizeof(char*) + 1) + env_.size() + 1;
+  char** raw = new char*[allocationRequired];
+  char** ptrBase = raw;
+  char* stringBase = reinterpret_cast<char*>(&raw[env_.size() + 1]);
+  char* const stringEnd = reinterpret_cast<char*>(&raw[allocationRequired]);
+  for (auto const& pair : env_) {
+    std::string const& key = pair.first;
+    std::string const& value = pair.second;
+    *ptrBase = stringBase;
+    size_t lengthIncludingNullTerminator = key.size() + 1 + value.size() + 1;
+    CHECK_GT(stringEnd - lengthIncludingNullTerminator, stringBase);
+    memcpy(stringBase, key.c_str(), key.size());
+    stringBase += key.size();
+    *stringBase++ = '=';
+    memcpy(stringBase, value.c_str(), value.size() + 1);
+    stringBase += value.size() + 1;
+    ++ptrBase;
+  }
+  *ptrBase = nullptr;
+  CHECK_EQ(env_.size(), ptrBase - raw);
+  return {raw, [](char** ptr) { delete[] ptr; }};
 }
diff --git a/folly/experimental/EnvUtil.h b/folly/experimental/EnvUtil.h
@@ -16,19 +16,106 @@
 
 #pragma once
 
+#include <folly/Memory.h>
 #include <map>
 #include <string>
+#include <unordered_map>
+#include <vector>
 
 namespace folly {
+namespace experimental {
+
+// Class to model the process environment in idiomatic C++
+//
+// Changes to the modeled environment do not change the process environment
+// unless `setAsCurrentEnvironment()` is called.
+struct EnvironmentState {
+  using EnvType = std::unordered_map<std::string, std::string>;
+
+  // Returns an EnvironmentState containing a copy of the current process
+  // environment. Subsequent changes to the process environment do not
+  // alter the stored model. If the process environment is altered during the
+  // execution of this method the results are not defined.
+  //
+  // Throws MalformedEnvironment if the process environment cannot be modeled.
+  static EnvironmentState fromCurrentEnvironment();
+
+  // Returns an empty EnvironmentState
+  static EnvironmentState empty() {
+    return {};
+  }
+
+  explicit EnvironmentState(EnvType const& env) : env_(env) {}
+  explicit EnvironmentState(EnvType&& env) : env_(std::move(env)) {}
+
+  // Get the model environment for querying.
+  EnvType const& operator*() const {
+    return env_;
+  }
+  EnvType const* operator->() const {
+    return &env_;
+  }
+
+  // Get the model environment for mutation or querying.
+  EnvType& operator*() {
+    return env_;
+  }
+  EnvType* operator->() {
+    return &env_;
+  }
+
+  // Update the process environment with the one in the stored model.
+  // Subsequent changes to the model do not alter the process environment. The
+  // state of the process environment during execution of this method is not
+  // defined. If the process environment is altered by another thread during the
+  // execution of this method the results are not defined.
+  void setAsCurrentEnvironment();
+
+  // Get a copy of the model environment in the form used by `folly::Subprocess`
+  std::vector<std::string> toVector() const;
+
+  // Get a copy of the model environment in the form commonly used by C routines
+  // such as execve, execle, etc. Example usage:
+  //
+  // EnvironmentState forChild{};
+  // ... manipulate `forChild` as needed ...
+  // execve("/bin/program",pArgs,forChild.toPointerArray().get());
+  std::unique_ptr<char*, void (*)(char**)> toPointerArray() const;
+
+ private:
+  EnvironmentState() {}
+  EnvType env_;
+};
+
+struct MalformedEnvironment : std::runtime_error {
+  using std::runtime_error::runtime_error;
+};
+} // namespace experimental
+
 namespace test {
+// RAII class allowing scoped changes to the process environment. The
+// environment state at the time of its construction is restored at the time
+// of its destruction.
+struct EnvVarSaver {
+  EnvVarSaver()
+      : state_(make_unique<experimental::EnvironmentState>(
+            experimental::EnvironmentState::fromCurrentEnvironment())) {}
+
+  EnvVarSaver(EnvVarSaver&& other) noexcept : state_(std::move(other.state_)) {}
+
+  EnvVarSaver& operator=(EnvVarSaver&& other) noexcept {
+    state_ = std::move(other.state_);
+    return *this;
+  }
 
-class EnvVarSaver {
- public:
-  EnvVarSaver();
-  ~EnvVarSaver();
+  ~EnvVarSaver() {
+    if (state_) {
+      state_->setAsCurrentEnvironment();
+    }
+  }
 
  private:
-  std::map<std::string, std::string> saved_;
+  std::unique_ptr<experimental::EnvironmentState> state_;
 };
-}
-}
+} // namespace test
+} // namespace folly
diff --git a/folly/experimental/test/EnvUtilSubprocess.cpp b/folly/experimental/test/EnvUtilSubprocess.cpp
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2017 Facebook, Inc.
+ *
+ * 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 <stdlib.h>
+#include <string.h>
+
+int main() {
+  char* spork = getenv("spork");
+  if (!spork) {
+    return 1;
+  }
+  return strcmp("foon", spork);
+}