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 |
123 | |