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   // Update the process environment with the one in the stored model.
  68   // Subsequent changes to the model do not alter the process environment. The
  69   // state of the process environment during execution of this method is not
  70   // defined. If the process environment is altered by another thread during the
  71   // execution of this method the results are not defined.
  72   void setAsCurrentEnvironment();
  73
  74   // Get a copy of the model environment in the form used by `folly::Subprocess`
  75   std::vector<std::string> toVector() const;
  76
  77   // Get a copy of the model environment in the form commonly used by C routines
  78   // such as execve, execle, etc. Example usage:
  79   //
  80   // EnvironmentState forChild{};
  81   // ... manipulate `forChild` as needed ...
  82   // execve("/bin/program",pArgs,forChild.toPointerArray().get());
  83   std::unique_ptr<char*, void (*)(char**)> toPointerArray() const;
  84
  85  private:
  86   EnvironmentState() {}
  87   EnvType env_;
  88 };
  89
  90 struct MalformedEnvironment : std::runtime_error {
  91   using std::runtime_error::runtime_error;
  92 };
  93 } // namespace experimental
  94
  95 namespace test {
  96 // RAII class allowing scoped changes to the process environment. The
  97 // environment state at the time of its construction is restored at the time
  98 // of its destruction.
  99 struct EnvVarSaver {
 100   EnvVarSaver()
 101       : state_(std::make_unique<experimental::EnvironmentState>(
 102             experimental::EnvironmentState::fromCurrentEnvironment())) {}
 103
 104   EnvVarSaver(EnvVarSaver&& other) noexcept : state_(std::move(other.state_)) {}
 105
 106   EnvVarSaver& operator=(EnvVarSaver&& other) noexcept {
 107     state_ = std::move(other.state_);
 108     return *this;
 109   }
 110
 111   ~EnvVarSaver() {
 112     if (state_) {
 113       state_->setAsCurrentEnvironment();
 114     }
 115   }
 116
 117  private:
 118   std::unique_ptr<experimental::EnvironmentState> state_;
 119 };
 120 } // namespace test
 121 } // namespace folly