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