folly/experimental/EnvUtil.h

   1 /*
   2  * Copyright 2017-present Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 #pragma once
  17
  18 #include <map>
  19 #include <string>
  20 #include <unordered_map>
  21 #include <vector>
  22
  23 #include <folly/CPortability.h>
  24 #include <folly/Memory.h>
  25
  26 namespace folly {
  27 namespace experimental {
  28
  29 // Class to model the process environment in idiomatic C++
  30 //
  31 // Changes to the modeled environment do not change the process environment
  32 // unless `setAsCurrentEnvironment()` is called.
  33 struct EnvironmentState {
  34   using EnvType = std::unordered_map<std::string, std::string>;
  35
  36   // Returns an EnvironmentState containing a copy of the current process
  37   // environment. Subsequent changes to the process environment do not
  38   // alter the stored model. If the process environment is altered during the
  39   // execution of this method the results are not defined.
  40   //
  41   // Throws MalformedEnvironment if the process environment cannot be modeled.
  42   static EnvironmentState fromCurrentEnvironment();
  43
  44   // Returns an empty EnvironmentState
  45   static EnvironmentState empty() {
  46     return {};
  47   }
  48
  49   explicit EnvironmentState(EnvType const& env) : env_(env) {}
  50   explicit EnvironmentState(EnvType&& env) : env_(std::move(env)) {}
  51
  52   // Get the model environment for querying.
  53   EnvType const& operator*() const {
  54     return env_;
  55   }
  56   EnvType const* operator->() const {
  57     return &env_;
  58   }
  59
  60   // Get the model environment for mutation or querying.
  61   EnvType& operator*() {
  62     return env_;
  63   }
  64   EnvType* operator->() {
  65     return &env_;
  66   }
  67
  68   // Update the process environment with the one in the stored model.
  69   // Subsequent changes to the model do not alter the process environment. The
  70   // state of the process environment during execution of this method is not
  71   // defined. If the process environment is altered by another thread during the
  72   // execution of this method the results are not defined.
  73   void setAsCurrentEnvironment();
  74
  75   // Get a copy of the model environment in the form used by `folly::Subprocess`
  76   std::vector<std::string> toVector() const;
  77
  78   // Get a copy of the model environment in the form commonly used by C routines
  79   // such as execve, execle, etc. Example usage:
  80   //
  81   // EnvironmentState forChild{};
  82   // ... manipulate `forChild` as needed ...
  83   // execve("/bin/program",pArgs,forChild.toPointerArray().get());
  84   std::unique_ptr<char*, void (*)(char**)> toPointerArray() const;
  85
  86  private:
  87   EnvironmentState() {}
  88   EnvType env_;
  89 };
  90
  91 struct FOLLY_EXPORT MalformedEnvironment : std::runtime_error {
  92   using std::runtime_error::runtime_error;
  93 };
  94 } // namespace experimental
  95
  96 namespace test {
  97 // RAII class allowing scoped changes to the process environment. The
  98 // environment state at the time of its construction is restored at the time
  99 // of its destruction.
 100 struct EnvVarSaver {
 101   EnvVarSaver()
 102       : state_(std::make_unique<experimental::EnvironmentState>(
 103             experimental::EnvironmentState::fromCurrentEnvironment())) {}
 104
 105   EnvVarSaver(EnvVarSaver&& other) noexcept : state_(std::move(other.state_)) {}
 106
 107   EnvVarSaver& operator=(EnvVarSaver&& other) noexcept {
 108     state_ = std::move(other.state_);
 109     return *this;
 110   }
 111
 112   ~EnvVarSaver() {
 113     if (state_) {
 114       state_->setAsCurrentEnvironment();
 115     }
 116   }
 117
 118  private:
 119   std::unique_ptr<experimental::EnvironmentState> state_;
 120 };
 121 } // namespace test
 122 } // namespace folly