folly/experimental/EnvUtil.h

   1 /*
   2  * Copyright 2017 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
  17 #pragma once
  18
  19 #include <folly/Memory.h>
  20 #include <map>
  21 #include <string>
  22 #include <unordered_map>
  23 #include <vector>
  24
  25 namespace folly {
  26 namespace experimental {
  27
  28 // Class to model the process environment in idiomatic C++
  29 //
  30 // Changes to the modeled environment do not change the process environment
  31 // unless `setAsCurrentEnvironment()` is called.
  32 struct EnvironmentState {
  33   using EnvType = std::unordered_map<std::string, std::string>;
  34
  35   // Returns an EnvironmentState containing a copy of the current process
  36   // environment. Subsequent changes to the process environment do not
  37   // alter the stored model. If the process environment is altered during the
  38   // execution of this method the results are not defined.
  39   //
  40   // Throws MalformedEnvironment if the process environment cannot be modeled.
  41   static EnvironmentState fromCurrentEnvironment();
  42
  43   // Returns an empty EnvironmentState
  44   static EnvironmentState empty() {
  45     return {};
  46   }
  47
  48   explicit EnvironmentState(EnvType const& env) : env_(env) {}
  49   explicit EnvironmentState(EnvType&& env) : env_(std::move(env)) {}
  50
  51   // Get the model environment for querying.
  52   EnvType const& operator*() const {
  53     return env_;
  54   }
  55   EnvType const* operator->() const {
  56     return &env_;
  57   }
  58
  59   // Get the model environment for mutation or querying.
  60   EnvType& operator*() {
  61     return env_;
  62   }
  63   EnvType* operator->() {
  64     return &env_;
  65   }
  66
  67 #if __linux__ && !FOLLY_MOBILE
  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 #endif
  75
  76   // Get a copy of the model environment in the form used by `folly::Subprocess`
  77   std::vector<std::string> toVector() const;
  78
  79   // Get a copy of the model environment in the form commonly used by C routines
  80   // such as execve, execle, etc. Example usage:
  81   //
  82   // EnvironmentState forChild{};
  83   // ... manipulate `forChild` as needed ...
  84   // execve("/bin/program",pArgs,forChild.toPointerArray().get());
  85   std::unique_ptr<char*, void (*)(char**)> toPointerArray() const;
  86
  87  private:
  88   EnvironmentState() {}
  89   EnvType env_;
  90 };
  91
  92 struct MalformedEnvironment : std::runtime_error {
  93   using std::runtime_error::runtime_error;
  94 };
  95 } // namespace experimental
  96
  97 namespace test {
  98 // RAII class allowing scoped changes to the process environment. The
  99 // environment state at the time of its construction is restored at the time
 100 // of its destruction.
 101 struct EnvVarSaver {
 102   EnvVarSaver()
 103       : state_(make_unique<experimental::EnvironmentState>(
 104             experimental::EnvironmentState::fromCurrentEnvironment())) {}
 105
 106   EnvVarSaver(EnvVarSaver&& other) noexcept : state_(std::move(other.state_)) {}
 107
 108   EnvVarSaver& operator=(EnvVarSaver&& other) noexcept {
 109     state_ = std::move(other.state_);
 110     return *this;
 111   }
 112
 113   ~EnvVarSaver() {
 114     if (state_) {
 115       state_->setAsCurrentEnvironment();
 116     }
 117   }
 118
 119  private:
 120   std::unique_ptr<experimental::EnvironmentState> state_;
 121 };
 122 } // namespace test
 123 } // namespace folly