diff --git a/include/gz/utils/Environment.hh b/include/gz/utils/Environment.hh index 9a7a062..3302f88 100644 --- a/include/gz/utils/Environment.hh +++ b/include/gz/utils/Environment.hh @@ -23,6 +23,7 @@ #include #include +#include namespace gz { @@ -78,6 +79,24 @@ bool GZ_UTILS_VISIBLE clearenv(); /// \brief Type alias for a collection of environment variables using EnvironmentMap = std::unordered_map; +/// \brief Type alias for a collection of environment variables +/// Each entry is of the form KEY=VAL +using EnvironmentStrings = std::vector; + +/// \brief Convert a vector of environment variables to a map +/// +/// \param[in] _envStrings Vector collection of environment variables +/// \return Mapped collection of environment variables. +EnvironmentMap GZ_UTILS_VISIBLE envStringsToMap( + const EnvironmentStrings &_envStrings); + +/// \brief Convert a map of environment variables to a vector +/// +/// \param[in] _envMap Collection of mapped environment variables. +/// \return Vector collection of environment variables. +EnvironmentStrings GZ_UTILS_VISIBLE envMapToStrings( + const EnvironmentMap &_envMap); + /// \brief Retrieve all current environment variables /// /// Note: This function is not thread-safe and should not be called diff --git a/include/gz/utils/Subprocess.hh b/include/gz/utils/Subprocess.hh index 787daf5..f7f9fab 100644 --- a/include/gz/utils/Subprocess.hh +++ b/include/gz/utils/Subprocess.hh @@ -19,8 +19,12 @@ #define GZ_UTILS__SUBPROCESS_HH_ #include "detail/subprocess.h" +#include "gz/utils/Environment.hh" + +#include #include #include +#include #include @@ -32,46 +36,95 @@ namespace gz namespace utils { +/// \brief Create a RAII-type object that encapsulates a subprocess. class Subprocess { + /// \brief Constructor + /// + /// This variant will spawn a subprocess that inherits the environment + /// from the calling process. + /// + /// \param[in] _commandLine set of arguments starting with an executable + /// used to spawn the subprocess + public: explicit Subprocess(const std::vector &_commandLine): + commandLine(_commandLine), + environment({}), + inheritEnvironment(true) + { + this->Create(); + } + + /// \brief Constructor + /// + /// This variant will spawn a subprocess that uses the user-specified + /// environment + /// + /// \param[in] _commandLine set of arguments starting with an executable + /// used to spawn the subprocess + /// \param[in] _environment environment variables to set in the spawned + /// subprocess public: Subprocess(const std::vector &_commandLine, - const std::vector &_environment = {}): + gz::utils::EnvironmentMap _environment): commandLine(_commandLine), - environment(_environment) + environment(std::move(_environment)), + inheritEnvironment(false) { this->Create(); } + /// \brief Constructor + /// + /// This variant will spawn a subprocess that uses the user-specified + /// environment + /// + /// \param[in] _commandLine set of arguments starting with an executable + /// used to spawn the subprocess + /// \param[in] _environment environment variables to set in the spawned + /// subprocess + public: Subprocess(const std::vector &_commandLine, + const gz::utils::EnvironmentStrings &_environment): + Subprocess(_commandLine, gz::utils::envStringsToMap(_environment)) + { + } + + private: void Create() { - if (this->process) + if (this->process != nullptr) return; this->process = new subprocess_s; + auto environmentStr = gz::utils::envMapToStrings(this->environment); + std::vector environmentCstr; std::vector commandLineCstr; + for (const auto &val : this->commandLine) { commandLineCstr.push_back(val.c_str()); } commandLineCstr.push_back(nullptr); - std::vector environmentCstr; - for (const auto &val : this->environment) + if (!this->inheritEnvironment) { - environmentCstr.push_back(val.c_str()); + for (const auto &val : environmentStr) + { + environmentCstr.push_back(val.c_str()); + } + environmentCstr.push_back(nullptr); } - environmentCstr.push_back(nullptr); int ret = -1; - if (this->environment.size()) + if (!this->inheritEnvironment) { ret = subprocess_create_ex(commandLineCstr.data(), 0, environmentCstr.data(), this->process); } else { - ret = subprocess_create(commandLineCstr.data(), 0, this->process); + ret = subprocess_create(commandLineCstr.data(), + subprocess_option_inherit_environment, + this->process); } if (0 != ret) @@ -84,21 +137,21 @@ class Subprocess public: ~Subprocess() { - if (this->process) + if (this->process != nullptr) subprocess_destroy(this->process); delete this->process; } public: std::string Stdout() { - std::string result{""}; - if (this->process) + std::string result; + if (this->process != nullptr) { - auto p_stdout = subprocess_stdout(this->process); + auto *p_stdout = subprocess_stdout(this->process); char buffer[128]; while (!feof(p_stdout)) { - if (fgets(buffer, 128, p_stdout) != NULL) + if (fgets(buffer, 128, p_stdout) != nullptr) result += buffer; } } @@ -107,14 +160,14 @@ class Subprocess public: std::string Stderr() { - std::string result{""}; - if (this->process) + std::string result; + if (this->process != nullptr) { - auto p_stdout = subprocess_stderr(this->process); + auto *p_stdout = subprocess_stderr(this->process); char buffer[128]; while (!feof(p_stdout)) { - if (fgets(buffer, 128, p_stdout) != NULL) + if (fgets(buffer, 128, p_stdout) != nullptr) result += buffer; } } @@ -124,7 +177,7 @@ class Subprocess public: bool Alive() { - if (this->process) + if (this->process != nullptr) return subprocess_alive(this->process); else return false; @@ -132,7 +185,7 @@ class Subprocess public: bool Terminate() { - if (this->process) + if (this->process != nullptr) return subprocess_terminate(this->process) != 0; else return false; @@ -141,7 +194,7 @@ class Subprocess public: int Join() { int return_code = -1; - if (this->process) + if (this->process != nullptr) { auto ret = subprocess_join(this->process, &return_code); if (ret != 0) @@ -155,7 +208,9 @@ class Subprocess protected: std::vector commandLine; - protected: std::vector environment; + protected: EnvironmentMap environment; + + protected: bool inheritEnvironment; protected: subprocess_s * process {nullptr}; }; diff --git a/src/Environment.cc b/src/Environment.cc index 4d86e0b..d477414 100644 --- a/src/Environment.cc +++ b/src/Environment.cc @@ -128,7 +128,6 @@ bool clearenv() } #endif return success; - } ///////////////////////////////////////////////// @@ -136,14 +135,6 @@ EnvironmentMap env() { EnvironmentMap ret; - // Helper function to split KEY=VAL - auto split = [](const std::string &_inp) - { - return std::make_pair( - _inp.substr(0, _inp.find('=')), - _inp.substr(_inp.find('=') + 1)); - }; - char **currentEnv = nullptr; #ifdef _WIN32 currentEnv = *__p__environ(); @@ -155,11 +146,12 @@ EnvironmentMap env() if (currentEnv == nullptr) return {}; + std::vector envStrings; for (; *currentEnv; ++currentEnv) { - ret.emplace(split(*currentEnv)); + envStrings.emplace_back(*currentEnv); } - return ret; + return envStringsToMap(envStrings); } ///////////////////////////////////////////////// @@ -174,20 +166,41 @@ bool setenv(const EnvironmentMap &_vars) } ///////////////////////////////////////////////// -std::string printenv() +EnvironmentMap envStringsToMap(const EnvironmentStrings &_envStrings) { - std::string ret; - // Variables are in an unordered_map as we generally don't - // care, but for printing sort for consistent display - auto currentEnv = env(); + EnvironmentMap ret; + for (const auto &pair : _envStrings) + { + auto eqPos = pair.find('='); + if (eqPos != std::string::npos) + { + ret.emplace(pair.substr(0, eqPos), pair.substr(eqPos + 1)); + } + } + return ret; +} + +///////////////////////////////////////////////// +EnvironmentStrings envMapToStrings(const EnvironmentMap &_envMap) +{ + EnvironmentStrings ret; auto sorted = std::vector>( - currentEnv.begin(), currentEnv.end()); + _envMap.begin(), _envMap.end()); std::sort(sorted.begin(), sorted.end()); - for (const auto &[key, value] : sorted) + for (auto [key, value] : sorted) + { + ret.push_back(key + "=" + value); + } + return ret; +} + +///////////////////////////////////////////////// +std::string printenv() +{ + std::string ret; + for (const auto &entry : envMapToStrings(env())) { - ret.append(key); - ret.append("="); - ret.append(value); + ret.append(entry); ret.append("\n"); } return ret; diff --git a/src/Environment_TEST.cc b/src/Environment_TEST.cc index 6d09923..496dddb 100644 --- a/src/Environment_TEST.cc +++ b/src/Environment_TEST.cc @@ -175,3 +175,39 @@ TEST(Environment, printenv) EXPECT_EQ(gz::utils::printenv(), "GZ_BAR_KEY=GZ_BAR_VAL\nGZ_BAZ_KEY=GZ_BAZ_VAL\nGZ_FOO_KEY=GZ_FOO_VAL\n"); } + +///////////////////////////////////////////////// +TEST(Environment, envStringsToMap) +{ + gz::utils::EnvironmentStrings strings; + strings.emplace_back("GZ_FOO_KEY=GZ_FOO_VAL"); + strings.emplace_back("GZ_BAR_KEY=GZ_BAR_VAL"); + strings.emplace_back("GZ_BAZ_KEY=GZ_BAZ_VAL"); + strings.emplace_back("BAD_KEY"); + + { + auto envMap = gz::utils::envStringsToMap(strings); + EXPECT_EQ(3u, envMap.size()); + EXPECT_EQ("GZ_FOO_VAL", envMap["GZ_FOO_KEY"]); + EXPECT_EQ("GZ_BAR_VAL", envMap["GZ_BAR_KEY"]); + EXPECT_EQ("GZ_BAZ_VAL", envMap["GZ_BAZ_KEY"]); + } +} + +///////////////////////////////////////////////// +TEST(Environment, envMapToStrings) +{ + gz::utils::EnvironmentMap envMap; + envMap.insert({{"GZ_FOO_KEY"}, {"GZ_FOO_VAL"}}); + envMap.insert({{"GZ_BAR_KEY"}, {"GZ_BAR_VAL"}}); + envMap.insert({{"GZ_BAZ_KEY"}, {"GZ_BAZ_VAL"}}); + + { + auto envStrings = gz::utils::envMapToStrings(envMap); + + EXPECT_EQ(3u, envStrings.size()); + EXPECT_EQ("GZ_BAR_KEY=GZ_BAR_VAL", envStrings[0]); + EXPECT_EQ("GZ_BAZ_KEY=GZ_BAZ_VAL", envStrings[1]); + EXPECT_EQ("GZ_FOO_KEY=GZ_FOO_VAL", envStrings[2]); + } +} diff --git a/test/integration/subprocess/subprocess_main.cc b/test/integration/subprocess/subprocess_main.cc index 436172e..d3afa05 100644 --- a/test/integration/subprocess/subprocess_main.cc +++ b/test/integration/subprocess/subprocess_main.cc @@ -24,8 +24,8 @@ class OutputSink { - public: OutputSink(const std::string &_dest): - dest(_dest) + public: explicit OutputSink(std::string _dest): + dest(std::move(_dest)) { } @@ -34,13 +34,11 @@ class OutputSink if (dest == "cout" || dest == "both") { std::cout << val << std::endl; - } else if (dest == "cerr" || dest == "both") { std::cerr << val << std::endl; } - return; } private: std::string dest; @@ -60,6 +58,10 @@ int main(int argc, char **argv) int iter_ms = 0; app.add_option("--iteration-ms", iter_ms, "length of one iteration"); + bool environment = false; + app.add_flag("--environment", environment, + "print the environment variables"); + CLI11_PARSE(app, argc, argv); auto sink = OutputSink(output); @@ -69,10 +71,9 @@ int main(int argc, char **argv) std::this_thread::sleep_for(std::chrono::milliseconds(iter_ms)); } - std::string env_var; - if(gz::utils::env("ENV_VAR", env_var)) + if (environment) { - sink.Write("ENV_VAR=" + env_var); + sink.Write(gz::utils::printenv()); } return 0; } diff --git a/test/integration/subprocess_TEST.cc b/test/integration/subprocess_TEST.cc index 32c6ed4..47c8976 100644 --- a/test/integration/subprocess_TEST.cc +++ b/test/integration/subprocess_TEST.cc @@ -21,6 +21,7 @@ #include #include +#include "gz/utils/Environment.hh" using Subprocess = gz::utils::Subprocess; @@ -33,6 +34,28 @@ TEST(Subprocess, CreateInvalid) EXPECT_FALSE(proc.Alive()); } +///////////////////////////////////////////////// +TEST(Subprocess, CreateInvalidSpaces) +{ + // Test if a user passes a string with spaces, rather than + // a vector of strings + std::string command(kExecutablePath); + command.append(" --help"); + auto proc = Subprocess({command}); + + // Block until the executable is done + auto ret = proc.Join(); + EXPECT_EQ(-1, ret); + + EXPECT_FALSE(proc.Alive()); + + auto cout = proc.Stdout(); + auto cerr = proc.Stdout(); + + EXPECT_TRUE(cout.empty()); + EXPECT_TRUE(cerr.empty()); +} + ///////////////////////////////////////////////// TEST(Subprocess, CreateValid) { @@ -94,38 +117,88 @@ TEST(Subprocess, Cerr) ///////////////////////////////////////////////// TEST(Subprocess, Environment) { + ASSERT_TRUE(gz::utils::clearenv()); + ASSERT_TRUE(gz::utils::setenv("GZ_FOO_KEY", "GZ_FOO_VAL")); + ASSERT_TRUE(gz::utils::setenv("GZ_BAR_KEY", "GZ_BAR_VAL")); + ASSERT_TRUE(gz::utils::setenv("GZ_BAZ_KEY", "GZ_BAZ_VAL")); + + { + // Default behavior is to inherit the environment + auto proc = Subprocess({kExecutablePath, "--output=cout", "--environment"}); + // Block until the executable is done + auto ret = proc.Join(); + EXPECT_EQ(0u, ret); + + auto cout = proc.Stdout(); + EXPECT_NE(std::string::npos, cout.find("GZ_FOO_KEY=GZ_FOO_VAL")); + EXPECT_NE(std::string::npos, cout.find("GZ_BAR_KEY=GZ_BAR_VAL")); + EXPECT_NE(std::string::npos, cout.find("GZ_BAZ_KEY=GZ_BAZ_VAL")); + } + { - auto proc = Subprocess({kExecutablePath, "--output=cout"}, - {"ENV_VAR=foobar"}); + // Passing an empty map as the second arg clears the environment + auto proc = Subprocess( + {kExecutablePath, "--output=cout", "--environment"}, + gz::utils::EnvironmentMap()); // Block until the executable is done auto ret = proc.Join(); EXPECT_EQ(0u, ret); auto cout = proc.Stdout(); - EXPECT_NE(std::string::npos, cout.find("ENV_VAR=foobar")); + EXPECT_EQ(std::string::npos, cout.find("GZ_FOO_KEY=GZ_FOO_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAR_KEY=GZ_BAR_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAZ_KEY=GZ_BAZ_VAL")); } { - auto proc = Subprocess({kExecutablePath, "--output=cerr"}, - {"ENV_VAR=foobar2"}); + // Passing an empty vector as the second arg clears the environment + auto proc = Subprocess( + {kExecutablePath, "--output=cout", "--environment"}, + gz::utils::EnvironmentStrings()); // Block until the executable is done auto ret = proc.Join(); EXPECT_EQ(0u, ret); - auto cerr = proc.Stderr(); - EXPECT_NE(std::string::npos, cerr.find("ENV_VAR=foobar2")); + auto cout = proc.Stdout(); + EXPECT_EQ(std::string::npos, cout.find("GZ_FOO_KEY=GZ_FOO_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAR_KEY=GZ_BAR_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAZ_KEY=GZ_BAZ_VAL")); } { - auto proc = Subprocess({kExecutablePath}, - {"ENV_VAR=foobar3"}); + // Passing a map sets those variables, clearing the rest + auto proc = Subprocess( + {kExecutablePath, "--output=cout", "--environment"}, + gz::utils::EnvironmentMap{ + {"QUX_KEY", "QUX_VAL"} + }); // Block until the executable is done auto ret = proc.Join(); EXPECT_EQ(0u, ret); auto cout = proc.Stdout(); - auto cerr = proc.Stderr(); - EXPECT_EQ(std::string::npos, cerr.find("ENV_VAR=foobar3")); - EXPECT_EQ(std::string::npos, cout.find("ENV_VAR=foobar3")); + EXPECT_EQ(std::string::npos, cout.find("GZ_FOO_KEY=GZ_FOO_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAR_KEY=GZ_BAR_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAZ_KEY=GZ_BAZ_VAL")); + EXPECT_NE(std::string::npos, cout.find("QUX_KEY=QUX_VAL")); } + + { + // Passing a map sets those variables, clearing the rest + auto proc = Subprocess( + {kExecutablePath, "--output=cout", "--environment"}, + gz::utils::EnvironmentStrings{ + {"QUX_KEY=QUX_VAL"} + }); + // Block until the executable is done + auto ret = proc.Join(); + EXPECT_EQ(0u, ret); + + auto cout = proc.Stdout(); + EXPECT_EQ(std::string::npos, cout.find("GZ_FOO_KEY=GZ_FOO_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAR_KEY=GZ_BAR_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAZ_KEY=GZ_BAZ_VAL")); + EXPECT_NE(std::string::npos, cout.find("QUX_KEY=QUX_VAL")); + } + }