forked from microsoft/SysmonCommon
-
Notifications
You must be signed in to change notification settings - Fork 0
/
environmentvar.cpp
190 lines (157 loc) · 8.58 KB
/
environmentvar.cpp
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
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
// UnitTests for environment variable caching.
#include "test.h"
#if defined _WIN64 || defined _WIN32
using ::testing::_;
using ::testing::Expectation;
using ::testing::Invoke;
using ::testing::AnyNumber;
#define AppDataPath L"%appdata%\\file.txt"
// InvalidArguments check for failure on invalid input.
TEST(EnvCache, InvalidArguments) {
WCHAR ruleData[] = L"input buffer";
WCHAR expandVariable[MAX_PATH];
LPWSTR nullRuleData = nullptr;
LPWSTR realRuleData = ruleData;
EXPECT_FALSE(ExpandEnvironmentVariable(1, nullptr, expandVariable, _countof(expandVariable)));
EXPECT_FALSE(ExpandEnvironmentVariable(1, &nullRuleData, expandVariable, _countof(expandVariable)));
EXPECT_FALSE(ExpandEnvironmentVariable(1, &realRuleData, nullptr, _countof(expandVariable)));
EXPECT_FALSE(ExpandEnvironmentVariable(1, &realRuleData, expandVariable, 0));
}
// NoVariable checks that a string without environment variables would work even if that case should not happen.
TEST(EnvCache, NoVariable)
{
WCHAR ruleData[] = L"a simple string";
LPWSTR inputRuleData = ruleData;
WCHAR expandVariable[MAX_PATH];
MockEnvironmentVariableCache cache;
// Expect only a call to GetUserToken with the mockSessionId.
EXPECT_CALL(cache, GetUserToken(mockSessionId, _))
.Times(AnyNumber())
.WillRepeatedly(Invoke(MockEnvironmentVariableCache::MockGetUserToken));
ASSERT_TRUE(cache.Expand(mockSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
}
// SessionZero tests that a session without token will successfully default to system environment variables.
TEST(EnvCache, SessionZero)
{
WCHAR ruleData[] = L"%systemroot%";
LPWSTR inputRuleData = ruleData;
WCHAR expandVariable[MAX_PATH];
MockEnvironmentVariableCache cache;
// Simulate GetUserToken error on session zero.
EXPECT_CALL(cache, GetUserToken(0, _))
.Times(AnyNumber())
.WillRepeatedly(Invoke([](ULONG sessionId, PHANDLE token){
SetLastError(ERROR_NO_TOKEN);
return FALSE;
}));
EXPECT_CALL(cache, AddValueToCache(0, _, _)).Times(1);
// Session zero call is successful.
ASSERT_TRUE(cache.Expand(0, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
}
// Storage verifies that the cache stores the entry only when previously unknown.
TEST(EnvCache, Storage)
{
WCHAR ruleData[] = AppDataPath;
LPWSTR inputRuleData = ruleData;
WCHAR expandVariable[MAX_PATH];
MockEnvironmentVariableCache cache;
// Expect only a call to GetUserToken with the mockSessionId.
EXPECT_CALL(cache, GetUserToken(mockSessionId, _))
.Times(AnyNumber())
.WillRepeatedly(Invoke(MockEnvironmentVariableCache::MockGetUserToken));
EXPECT_CALL(cache, AddValueToCache(mockSessionId, _, _)).Times(1);
// First call fills the cache.
ASSERT_TRUE(cache.Expand(mockSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
std::wstring expandVariableLower = expandVariable;
std::transform(expandVariableLower.begin(), expandVariableLower.end(), expandVariableLower.begin(), ::towlower);
EXPECT_STREQ(expandVariableLower.c_str(), expandVariable) << "The expanded variable should be lower case";
// Second call will use the cache, no call to AddValueToCache.
ZeroMemory(expandVariable, sizeof(expandVariable));
inputRuleData = ruleData;
ASSERT_TRUE(cache.Expand(mockSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
}
// Logoff checks that a session is correctly removed for logoff.
TEST(EnvCache, Logoff)
{
WCHAR ruleData[] = AppDataPath;
LPWSTR inputRuleData = ruleData;
WCHAR expandVariable[MAX_PATH];
MockEnvironmentVariableCache cache;
// Support two mock session ids.
EXPECT_CALL(cache, GetUserToken(mockSessionId, _))
.Times(AnyNumber())
.WillRepeatedly(Invoke(MockEnvironmentVariableCache::MockGetUserToken));
EXPECT_CALL(cache, GetUserToken(mockSecondSessionId, _))
.Times(AnyNumber())
.WillRepeatedly(Invoke(MockEnvironmentVariableCache::MockGetUserToken));
// Values are added to the cache twice for the first mock, once only for the second as it is not cleared.
EXPECT_CALL(cache, AddValueToCache(mockSessionId, _, _)).Times(2);
EXPECT_CALL(cache, AddValueToCache(mockSecondSessionId, _, _)).Times(1);
// First calls fill the cache for both sessions.
ASSERT_TRUE(cache.Expand(mockSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
ZeroMemory(expandVariable, sizeof(expandVariable));
inputRuleData = ruleData;
ASSERT_TRUE(cache.Expand(mockSecondSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
cache.CleanSessionCache(mockSessionId);
// Second call fills the cache again but only for main session.
ZeroMemory(expandVariable, sizeof(expandVariable));
inputRuleData = ruleData;
ASSERT_TRUE(cache.Expand(mockSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
ZeroMemory(expandVariable, sizeof(expandVariable));
inputRuleData = ruleData;
ASSERT_TRUE(cache.Expand(mockSecondSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
}
// ConfigChange checks that sessions are correctly cleaned on config update.
TEST(EnvCache, ConfigChange)
{
WCHAR ruleData[] = AppDataPath;
LPWSTR inputRuleData = ruleData;
WCHAR expandVariable[MAX_PATH];
MockEnvironmentVariableCache cache;
// Support two mock session ids.
EXPECT_CALL(cache, GetUserToken(mockSessionId, _))
.Times(AnyNumber())
.WillRepeatedly(Invoke(MockEnvironmentVariableCache::MockGetUserToken));
EXPECT_CALL(cache, GetUserToken(mockSecondSessionId, _))
.Times(AnyNumber())
.WillRepeatedly(Invoke(MockEnvironmentVariableCache::MockGetUserToken));
// Values are added to the cache twice after the clean.
EXPECT_CALL(cache, AddValueToCache(mockSessionId, _, _)).Times(2);
EXPECT_CALL(cache, AddValueToCache(mockSecondSessionId, _, _)).Times(2);
// First calls fill the cache for both sessions.
ASSERT_TRUE(cache.Expand(mockSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
ZeroMemory(expandVariable, sizeof(expandVariable));
inputRuleData = ruleData;
ASSERT_TRUE(cache.Expand(mockSecondSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
cache.CleanCache();
// Second call fills the cache again but only for main session.
ZeroMemory(expandVariable, sizeof(expandVariable));
inputRuleData = ruleData;
ASSERT_TRUE(cache.Expand(mockSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
ZeroMemory(expandVariable, sizeof(expandVariable));
inputRuleData = ruleData;
ASSERT_TRUE(cache.Expand(mockSecondSessionId, &inputRuleData, expandVariable, _countof(expandVariable)));
EXPECT_EQ(inputRuleData, expandVariable) << "On expansion, inputRuleData should be changed to point to expandVariable";
EXPECT_STRNE(ruleData, expandVariable);
}
#endif