From 55cd50f37f6d8b137093fc976b780465442174fa Mon Sep 17 00:00:00 2001 From: Phil Willoughby Date: Fri, 24 Mar 2017 02:05:42 -0700 Subject: [PATCH] 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 --- folly/experimental/EnvUtil.cpp | 81 +++++++--- folly/experimental/EnvUtil.h | 101 ++++++++++++- folly/experimental/test/EnvUtilSubprocess.cpp | 26 ++++ folly/experimental/test/EnvUtilTest.cpp | 140 ++++++++++++++++-- 4 files changed, 306 insertions(+), 42 deletions(-) create mode 100644 folly/experimental/test/EnvUtilSubprocess.cpp diff --git a/folly/experimental/EnvUtil.cpp b/folly/experimental/EnvUtil.cpp index f31bd37d..5c16ee74 100644 --- a/folly/experimental/EnvUtil.cpp +++ b/folly/experimental/EnvUtil.cpp @@ -20,37 +20,72 @@ #include #include -namespace folly { -namespace test { +using namespace folly; +using namespace folly::experimental; -static std::map getEnvVarMap() { - std::map data; - for (auto it = environ; *it != nullptr; ++it) { +EnvironmentState EnvironmentState::fromCurrentEnvironment() { + std::unordered_map 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( + "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( + "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 EnvironmentState::toVector() const { + std::vector result; + for (auto const& pair : env_) { + result.emplace_back(to(pair.first, "=", pair.second)); + } + return result; } + +std::unique_ptr 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(&raw[env_.size() + 1]); + char* const stringEnd = reinterpret_cast(&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 index a44e07bc..3ab3b16a 100644 --- a/folly/experimental/EnvUtil.h +++ b/folly/experimental/EnvUtil.h @@ -16,19 +16,106 @@ #pragma once +#include #include #include +#include +#include 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; + + // 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 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 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::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 saved_; + std::unique_ptr state_; }; -} -} +} // namespace test +} // namespace folly diff --git a/folly/experimental/test/EnvUtilSubprocess.cpp b/folly/experimental/test/EnvUtilSubprocess.cpp new file mode 100644 index 00000000..6bde6635 --- /dev/null +++ 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 +#include + +int main() { + char* spork = getenv("spork"); + if (!spork) { + return 1; + } + return strcmp("foon", spork); +} diff --git a/folly/experimental/test/EnvUtilTest.cpp b/folly/experimental/test/EnvUtilTest.cpp index 2265e514..bcac455d 100644 --- a/folly/experimental/test/EnvUtilTest.cpp +++ b/folly/experimental/test/EnvUtilTest.cpp @@ -16,22 +16,27 @@ #include -#include - #include -#include - #include +#include #include #include #include +#include +#include +#include using namespace folly; -using namespace folly::test; +using folly::test::EnvVarSaver; +using folly::experimental::EnvironmentState; +using folly::experimental::MalformedEnvironment; -class EnvVarSaverTest : public testing::Test {}; +DEFINE_string( + env_util_subprocess_binary, + "./env_util_subprocess", + "Location of the `env_util_subprocess` test helper program"); -TEST_F(EnvVarSaverTest, ExampleNew) { +TEST(EnvVarSaverTest, ExampleNew) { auto key = "hahahahaha"; EXPECT_EQ(nullptr, getenv(key)); @@ -42,24 +47,135 @@ TEST_F(EnvVarSaverTest, ExampleNew) { auto saver = make_unique(); PCHECK(0 == setenv(key, "blah", true)); - EXPECT_EQ("blah", std::string{getenv(key)}); + EXPECT_STREQ("blah", getenv(key)); saver = nullptr; EXPECT_EQ(nullptr, getenv(key)); } -TEST_F(EnvVarSaverTest, ExampleExisting) { +TEST(EnvVarSaverTest, ExampleExisting) { auto key = "PATH"; EXPECT_NE(nullptr, getenv(key)); auto value = std::string{getenv(key)}; auto saver = make_unique(); PCHECK(0 == setenv(key, "blah", true)); - EXPECT_EQ("blah", std::string{getenv(key)}); + EXPECT_STREQ("blah", getenv(key)); saver = nullptr; - EXPECT_TRUE(value == getenv(key)); + EXPECT_EQ(value, getenv(key)); +} + +TEST(EnvVarSaverTest, Movable) { + Optional pSaver1; + pSaver1.emplace(); + auto key = "PATH"; + EXPECT_NE(nullptr, getenv(key)); + auto value = std::string{getenv(key)}; + Optional pSaver2; + pSaver2.emplace(std::move(*pSaver1)); + pSaver1.clear(); + PCHECK(0 == setenv(key, "blah", true)); + EXPECT_STREQ("blah", getenv(key)); + pSaver2.clear(); + EXPECT_EQ(value, getenv(key)); +} + +TEST(EnvironmentStateTest, FailOnEmptyString) { + EnvVarSaver saver{}; + char test[4] = "A=B"; + PCHECK(0 == putenv(test)); + auto okState = EnvironmentState::fromCurrentEnvironment(); + test[0] = 0; + EXPECT_THROW( + EnvironmentState::fromCurrentEnvironment(), MalformedEnvironment); +} + +TEST(EnvironmentStateTest, MovableAndCopyable) { + auto initialState = EnvironmentState::fromCurrentEnvironment(); + auto copiedState1 = EnvironmentState::empty(); + copiedState1.operator=(initialState); + EnvironmentState copiedState2{initialState}; + EXPECT_EQ(*initialState, *copiedState1); + EXPECT_EQ(*initialState, *copiedState2); + (*initialState)["foo"] = "bar"; + EXPECT_EQ(0, copiedState1->count("foo")); + EXPECT_EQ(0, copiedState2->count("foo")); + auto movedState1 = EnvironmentState::empty(); + movedState1.operator=(std::move(copiedState1)); + EnvironmentState movedState2{std::move(copiedState2)}; + EXPECT_EQ(0, movedState1->count("foo")); + EXPECT_EQ(0, movedState2->count("foo")); + initialState->erase("foo"); + EXPECT_EQ(*initialState, *movedState1); + EXPECT_EQ(*initialState, *movedState2); +} + +TEST(EnvironmentStateTest, FailOnDuplicate) { + EnvVarSaver saver{}; + char test[7] = "PATG=B"; + PCHECK(0 == putenv(test)); + auto okState = EnvironmentState::fromCurrentEnvironment(); + test[3] = 'H'; + EXPECT_THROW( + EnvironmentState::fromCurrentEnvironment(), MalformedEnvironment); +} + +TEST(EnvironmentStateTest, Separation) { + EnvVarSaver saver{}; + auto initialState = EnvironmentState::fromCurrentEnvironment(); + PCHECK(0 == setenv("spork", "foon", true)); + auto updatedState = EnvironmentState::fromCurrentEnvironment(); + EXPECT_EQ(0, initialState->count("spork")); + EXPECT_EQ(1, updatedState->count("spork")); + EXPECT_EQ("foon", (*updatedState)["spork"]); + updatedState->erase("spork"); + EXPECT_EQ(0, updatedState->count("spork")); + EXPECT_STREQ("foon", getenv("spork")); +} + +TEST(EnvironmentStateTest, Update) { + EnvVarSaver saver{}; + auto env = EnvironmentState::fromCurrentEnvironment(); + EXPECT_EQ(nullptr, getenv("spork")); + (*env)["spork"] = "foon"; + EXPECT_EQ(nullptr, getenv("spork")); + env.setAsCurrentEnvironment(); + EXPECT_STREQ("foon", getenv("spork")); +} + +TEST(EnvironmentStateTest, forSubprocess) { + auto env = EnvironmentState::empty(); + (*env)["spork"] = "foon"; + std::vector expected = {"spork=foon"}; + auto vec = env.toVector(); + EXPECT_EQ(expected, vec); + Subprocess subProcess{{fLS::FLAGS_env_util_subprocess_binary}, + {}, + fLS::FLAGS_env_util_subprocess_binary.c_str(), + &vec}; + EXPECT_EQ(0, subProcess.wait().exitStatus()); +} + +TEST(EnvironmentStateTest, forC) { + auto env = EnvironmentState::empty(); + (*env)["spork"] = "foon"; + EXPECT_STREQ("spork=foon", env.toPointerArray().get()[0]); + EXPECT_EQ(nullptr, env.toPointerArray().get()[1]); + char const* program = fLS::FLAGS_env_util_subprocess_binary.c_str(); + pid_t pid; + PCHECK( + 0 == posix_spawn( + &pid, + program, + nullptr, + nullptr, + nullptr, + env.toPointerArray().get())); + int result; + PCHECK(pid == waitpid(pid, &result, 0)); + EXPECT_EQ(0, result); } -TEST_F(EnvVarSaverTest, ExampleDeleting) { +TEST(EnvVarSaverTest, ExampleDeleting) { auto key = "PATH"; EXPECT_NE(nullptr, getenv(key)); auto value = std::string{getenv(key)}; -- 2.34.1