1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
|
#include "env.h"
#include <util/generic/maybe.h>
#include <util/generic/string.h>
#include <util/generic/yexception.h>
#ifdef _win_
#include <util/generic/scope.h>
#include <util/generic/vector.h>
#include <util/system/yassert.h>
#include "winint.h"
#else
#ifndef _linux_
#include <util/generic/vector.h> // ClearEnv impl
#endif
#include <cerrno>
#include <cstdlib>
extern char** environ;
#endif
/**
* On Windows there may be many copies of environment variables, there at least two known, one is
* manipulated by Win32 API, another by C runtime, so we must be consistent in the choice of
* functions used to manipulate them.
*
* Relevant links:
* - http://bugs.python.org/issue16633
* - https://a.yandex-team.ru/review/108892/details
*/
TMaybe<TString> TryGetEnv(const TString& key) {
#ifdef _win_
size_t len = GetEnvironmentVariableA(key.data(), nullptr, 0);
if (len == 0) {
if (GetLastError() == ERROR_ENVVAR_NOT_FOUND) {
return Nothing();
}
return TString{};
}
TVector<char> buffer(len);
size_t bufferSize;
do {
bufferSize = buffer.size();
len = GetEnvironmentVariableA(key.data(), buffer.data(), static_cast<DWORD>(bufferSize));
if (len > bufferSize) {
buffer.resize(len);
}
} while (len > bufferSize);
return TString(buffer.data(), len);
#else
const char* env = getenv(key.data());
if (!env) {
return Nothing();
}
return TString(env);
#endif
}
TString GetEnv(const TString& key, const TString& def) {
TMaybe<TString> value = TryGetEnv(key);
if (value.Defined()) {
return *std::move(value);
}
return def;
}
void SetEnv(const TString& key, const TString& value) {
bool isOk = false;
#ifdef _win_
isOk = SetEnvironmentVariable(key.data(), value.data());
#else
isOk = (0 == setenv(key.data(), value.data(), true /*replace*/));
#endif
Y_ENSURE_EX(isOk, TSystemError() << "failed to SetEnv");
}
void UnsetEnv(const TString& key) {
bool notFound = false;
#ifdef _win_
bool ok = SetEnvironmentVariable(key.c_str(), NULL);
notFound = !ok && (GetLastError() == ERROR_ENVVAR_NOT_FOUND);
#else
bool ok = (0 == unsetenv(key.c_str()));
#if defined(_darwin_)
notFound = !ok && (errno == EINVAL);
#endif
#endif
Y_ENSURE_EX(ok || notFound, TSystemError() << "failed to unset environment variable " << key.Quote());
}
void IterateEnv(const std::function<void(TStringBuf, TStringBuf)>& f, bool ignoreMalformedStrings) {
#ifdef _win_
// Env block format:
// Var1=Value1\0
// Var2=Value2\0
// ...
// VarN=ValueN\0\0
// Or "\0" if empty
auto envBlock = GetEnvironmentStringsA();
if (!envBlock) {
ythrow TSystemError() << "failed to get environment variables";
}
Y_DEFER {
bool ok = FreeEnvironmentStringsA(envBlock);
Y_ABORT_UNLESS(ok, "failed to free env block"); // ¯\_(ツ)_/¯
};
const char* charEnv = envBlock;
while (*charEnv) {
TStringBuf varStr(charEnv);
TStringBuf name, value;
if (varStr.TrySplit('=', name, value)) { // Not optimal (we could scan for null byte and '=' in one pass), but less bugprone
f(name, value);
} else {
Y_ENSURE(ignoreMalformedStrings, "Environment contains a string without '=': \"" << varStr << '"');
}
charEnv = varStr.data() + varStr.size() + 1;
}
#else
if (!environ) { // may be null after clearenv
return;
}
for (char** var = environ; *var; ++var) {
TStringBuf varStr(*var);
TStringBuf name, value;
if (varStr.TrySplit('=', name, value)) {
f(name, value);
} else {
Y_ENSURE(ignoreMalformedStrings, "Environment contains a string without '=': \"" << varStr << '"');
}
}
#endif
}
void ClearEnv() {
#if defined(_win_)
wchar_t emptyEnv[] = {L'\0'};
// At least on Wine, SetEnvironmentStringsA expects a multi-byte string that is internally converted to UTF-16.
// So it's easier to use SetEnvironmentStringsW directly.
bool ok = SetEnvironmentStringsW(emptyEnv);
Y_ENSURE_EX(ok, TSystemError() << "failed to clear environment");
#elif defined(_linux_)
bool ok = !clearenv();
Y_ENSURE_EX(ok, TSystemError() << "failed to clear environment");
#else
// Darwin or other unix platform that may not have clearenv
TVector<TString> keys;
IterateEnv(
[&keys](TStringBuf name, TStringBuf) {
keys.emplace_back(name);
},
true // Ignore malformed strings - they don't seem to be accessible anyway.
// IterateEnv won't expose these strings, and GetEnv/SetEnv/UnsetEnv ignore them (at least in glibc).
);
for (const auto& key : keys) {
UnsetEnv(key); // duplicate UnsetEnv is ok.
}
#endif
}
|
