1 /*
2  * Copyright 2017 The Android Open Source Project
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  * TrafficControllerTest.cpp - unit tests for TrafficController.cpp
17  */
18 
19 #include <cstdint>
20 #include <string>
21 #include <vector>
22 
23 #include <fcntl.h>
24 #include <inttypes.h>
25 #include <linux/inet_diag.h>
26 #include <linux/sock_diag.h>
27 #include <sys/socket.h>
28 #include <sys/types.h>
29 #include <unistd.h>
30 
31 #include <gtest/gtest.h>
32 
33 #include <android-base/stringprintf.h>
34 #include <android-base/strings.h>
35 
36 #include <netdutils/MockSyscalls.h>
37 
38 #include "FirewallController.h"
39 #include "TrafficController.h"
40 #include "bpf/BpfUtils.h"
41 
42 using namespace android::bpf;  // NOLINT(google-build-using-namespace): grandfathered
43 
44 namespace android {
45 namespace net {
46 
47 using base::Result;
48 using netdutils::isOk;
49 
50 constexpr int TEST_MAP_SIZE = 10;
51 constexpr int TEST_COOKIE = 1;
52 constexpr uid_t TEST_UID = 10086;
53 constexpr uid_t TEST_UID2 = 54321;
54 constexpr uid_t TEST_UID3 = 98765;
55 constexpr uint32_t TEST_TAG = 42;
56 constexpr uint32_t TEST_COUNTERSET = 1;
57 constexpr uint32_t DEFAULT_COUNTERSET = 0;
58 constexpr uint32_t TEST_PER_UID_STATS_ENTRIES_LIMIT = 3;
59 constexpr uint32_t TEST_TOTAL_UID_STATS_ENTRIES_LIMIT = 7;
60 
61 #define ASSERT_VALID(x) ASSERT_TRUE((x).isValid())
62 
63 class TrafficControllerTest : public ::testing::Test {
64   protected:
TrafficControllerTest()65     TrafficControllerTest()
66         : mTc(TEST_PER_UID_STATS_ENTRIES_LIMIT, TEST_TOTAL_UID_STATS_ENTRIES_LIMIT) {}
67     TrafficController mTc;
68     BpfMap<uint64_t, UidTagValue> mFakeCookieTagMap;
69     BpfMap<uint32_t, uint8_t> mFakeUidCounterSetMap;
70     BpfMap<uint32_t, StatsValue> mFakeAppUidStatsMap;
71     BpfMap<StatsKey, StatsValue> mFakeStatsMapA;
72     BpfMap<uint32_t, uint8_t> mFakeConfigurationMap;
73     BpfMap<uint32_t, UidOwnerValue> mFakeUidOwnerMap;
74     BpfMap<uint32_t, uint8_t> mFakeUidPermissionMap;
75 
SetUp()76     void SetUp() {
77         std::lock_guard guard(mTc.mMutex);
78         ASSERT_EQ(0, setrlimitForTest());
79 
80         mFakeCookieTagMap.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(uint64_t), sizeof(UidTagValue),
81                                           TEST_MAP_SIZE, 0));
82         ASSERT_VALID(mFakeCookieTagMap);
83 
84         mFakeUidCounterSetMap.reset(
85             createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), TEST_MAP_SIZE, 0));
86         ASSERT_VALID(mFakeUidCounterSetMap);
87 
88         mFakeAppUidStatsMap.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(StatsValue),
89                                             TEST_MAP_SIZE, 0));
90         ASSERT_VALID(mFakeAppUidStatsMap);
91 
92         mFakeStatsMapA.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(StatsKey), sizeof(StatsValue),
93                                        TEST_MAP_SIZE, 0));
94         ASSERT_VALID(mFakeStatsMapA);
95 
96         mFakeConfigurationMap.reset(
97                 createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), 1, 0));
98         ASSERT_VALID(mFakeConfigurationMap);
99 
100         mFakeUidOwnerMap.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(UidOwnerValue),
101                                          TEST_MAP_SIZE, 0));
102         ASSERT_VALID(mFakeUidOwnerMap);
103         mFakeUidPermissionMap.reset(
104                 createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), TEST_MAP_SIZE, 0));
105         ASSERT_VALID(mFakeUidPermissionMap);
106 
107         mTc.mCookieTagMap.reset(dupFd(mFakeCookieTagMap.getMap()));
108         ASSERT_VALID(mTc.mCookieTagMap);
109         mTc.mUidCounterSetMap.reset(dupFd(mFakeUidCounterSetMap.getMap()));
110         ASSERT_VALID(mTc.mUidCounterSetMap);
111         mTc.mAppUidStatsMap.reset(dupFd(mFakeAppUidStatsMap.getMap()));
112         ASSERT_VALID(mTc.mAppUidStatsMap);
113         mTc.mStatsMapA.reset(dupFd(mFakeStatsMapA.getMap()));
114         ASSERT_VALID(mTc.mStatsMapA);
115         mTc.mConfigurationMap.reset(dupFd(mFakeConfigurationMap.getMap()));
116         ASSERT_VALID(mTc.mConfigurationMap);
117 
118         // Always write to stats map A by default.
119         ASSERT_RESULT_OK(mTc.mConfigurationMap.writeValue(CURRENT_STATS_MAP_CONFIGURATION_KEY,
120                                                           SELECT_MAP_A, BPF_ANY));
121         mTc.mUidOwnerMap.reset(dupFd(mFakeUidOwnerMap.getMap()));
122         ASSERT_VALID(mTc.mUidOwnerMap);
123         mTc.mUidPermissionMap.reset(dupFd(mFakeUidPermissionMap.getMap()));
124         ASSERT_VALID(mTc.mUidPermissionMap);
125         mTc.mPrivilegedUser.clear();
126     }
127 
dupFd(const android::base::unique_fd & mapFd)128     int dupFd(const android::base::unique_fd& mapFd) {
129         return fcntl(mapFd.get(), F_DUPFD_CLOEXEC, 0);
130     }
131 
setUpSocketAndTag(int protocol,uint64_t * cookie,uint32_t tag,uid_t uid,uid_t callingUid)132     int setUpSocketAndTag(int protocol, uint64_t* cookie, uint32_t tag, uid_t uid,
133                           uid_t callingUid) {
134         int sock = socket(protocol, SOCK_STREAM | SOCK_CLOEXEC, 0);
135         EXPECT_LE(0, sock);
136         *cookie = getSocketCookie(sock);
137         EXPECT_NE(NONEXISTENT_COOKIE, *cookie);
138         EXPECT_EQ(0, mTc.tagSocket(sock, tag, uid, callingUid));
139         return sock;
140     }
141 
expectUidTag(uint64_t cookie,uid_t uid,uint32_t tag)142     void expectUidTag(uint64_t cookie, uid_t uid, uint32_t tag) {
143         Result<UidTagValue> tagResult = mFakeCookieTagMap.readValue(cookie);
144         ASSERT_RESULT_OK(tagResult);
145         EXPECT_EQ(uid, tagResult.value().uid);
146         EXPECT_EQ(tag, tagResult.value().tag);
147     }
148 
expectNoTag(uint64_t cookie)149     void expectNoTag(uint64_t cookie) { EXPECT_FALSE(mFakeCookieTagMap.readValue(cookie).ok()); }
150 
populateFakeStats(uint64_t cookie,uint32_t uid,uint32_t tag,StatsKey * key)151     void populateFakeStats(uint64_t cookie, uint32_t uid, uint32_t tag, StatsKey* key) {
152         UidTagValue cookieMapkey = {.uid = (uint32_t)uid, .tag = tag};
153         EXPECT_RESULT_OK(mFakeCookieTagMap.writeValue(cookie, cookieMapkey, BPF_ANY));
154         *key = {.uid = uid, .tag = tag, .counterSet = TEST_COUNTERSET, .ifaceIndex = 1};
155         StatsValue statsMapValue = {.rxPackets = 1, .rxBytes = 100};
156         uint8_t counterSet = TEST_COUNTERSET;
157         EXPECT_RESULT_OK(mFakeUidCounterSetMap.writeValue(uid, counterSet, BPF_ANY));
158         EXPECT_RESULT_OK(mFakeStatsMapA.writeValue(*key, statsMapValue, BPF_ANY));
159         key->tag = 0;
160         EXPECT_RESULT_OK(mFakeStatsMapA.writeValue(*key, statsMapValue, BPF_ANY));
161         EXPECT_RESULT_OK(mFakeAppUidStatsMap.writeValue(uid, statsMapValue, BPF_ANY));
162         // put tag information back to statsKey
163         key->tag = tag;
164     }
165 
checkUidOwnerRuleForChain(ChildChain chain,UidOwnerMatchType match)166     void checkUidOwnerRuleForChain(ChildChain chain, UidOwnerMatchType match) {
167         uint32_t uid = TEST_UID;
168         EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, DENY, DENYLIST));
169         Result<UidOwnerValue> value = mFakeUidOwnerMap.readValue(uid);
170         EXPECT_RESULT_OK(value);
171         EXPECT_TRUE(value.value().rule & match);
172 
173         uid = TEST_UID2;
174         EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, ALLOW, ALLOWLIST));
175         value = mFakeUidOwnerMap.readValue(uid);
176         EXPECT_RESULT_OK(value);
177         EXPECT_TRUE(value.value().rule & match);
178 
179         EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, DENY, ALLOWLIST));
180         value = mFakeUidOwnerMap.readValue(uid);
181         EXPECT_FALSE(value.ok());
182         EXPECT_EQ(ENOENT, value.error().code());
183 
184         uid = TEST_UID;
185         EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, ALLOW, DENYLIST));
186         value = mFakeUidOwnerMap.readValue(uid);
187         EXPECT_FALSE(value.ok());
188         EXPECT_EQ(ENOENT, value.error().code());
189 
190         uid = TEST_UID3;
191         EXPECT_EQ(-ENOENT, mTc.changeUidOwnerRule(chain, uid, ALLOW, DENYLIST));
192         value = mFakeUidOwnerMap.readValue(uid);
193         EXPECT_FALSE(value.ok());
194         EXPECT_EQ(ENOENT, value.error().code());
195     }
196 
checkEachUidValue(const std::vector<int32_t> & uids,UidOwnerMatchType match)197     void checkEachUidValue(const std::vector<int32_t>& uids, UidOwnerMatchType match) {
198         for (uint32_t uid : uids) {
199             Result<UidOwnerValue> value = mFakeUidOwnerMap.readValue(uid);
200             EXPECT_RESULT_OK(value);
201             EXPECT_TRUE(value.value().rule & match);
202         }
203         std::set<uint32_t> uidSet(uids.begin(), uids.end());
204         const auto checkNoOtherUid = [&uidSet](const int32_t& key,
205                                                const BpfMap<uint32_t, UidOwnerValue>&) {
206             EXPECT_NE(uidSet.end(), uidSet.find(key));
207             return Result<void>();
208         };
209         EXPECT_RESULT_OK(mFakeUidOwnerMap.iterate(checkNoOtherUid));
210     }
211 
checkUidMapReplace(const std::string & name,const std::vector<int32_t> & uids,UidOwnerMatchType match)212     void checkUidMapReplace(const std::string& name, const std::vector<int32_t>& uids,
213                             UidOwnerMatchType match) {
214         bool isAllowlist = true;
215         EXPECT_EQ(0, mTc.replaceUidOwnerMap(name, isAllowlist, uids));
216         checkEachUidValue(uids, match);
217 
218         isAllowlist = false;
219         EXPECT_EQ(0, mTc.replaceUidOwnerMap(name, isAllowlist, uids));
220         checkEachUidValue(uids, match);
221     }
222 
expectUidOwnerMapValues(const std::vector<uint32_t> & appUids,uint8_t expectedRule,uint32_t expectedIif)223     void expectUidOwnerMapValues(const std::vector<uint32_t>& appUids, uint8_t expectedRule,
224                                  uint32_t expectedIif) {
225         for (uint32_t uid : appUids) {
226             Result<UidOwnerValue> value = mFakeUidOwnerMap.readValue(uid);
227             EXPECT_RESULT_OK(value);
228             EXPECT_EQ(expectedRule, value.value().rule)
229                     << "Expected rule for UID " << uid << " to be " << expectedRule << ", but was "
230                     << value.value().rule;
231             EXPECT_EQ(expectedIif, value.value().iif)
232                     << "Expected iif for UID " << uid << " to be " << expectedIif << ", but was "
233                     << value.value().iif;
234         }
235     }
236 
237     template <class Key, class Value>
expectMapEmpty(BpfMap<Key,Value> & map)238     void expectMapEmpty(BpfMap<Key, Value>& map) {
239         auto isEmpty = map.isEmpty();
240         EXPECT_RESULT_OK(isEmpty);
241         EXPECT_TRUE(isEmpty.value());
242     }
243 
expectUidPermissionMapValues(const std::vector<uid_t> & appUids,uint8_t expectedValue)244     void expectUidPermissionMapValues(const std::vector<uid_t>& appUids, uint8_t expectedValue) {
245         for (uid_t uid : appUids) {
246             Result<uint8_t> value = mFakeUidPermissionMap.readValue(uid);
247             EXPECT_RESULT_OK(value);
248             EXPECT_EQ(expectedValue, value.value())
249                     << "Expected value for UID " << uid << " to be " << expectedValue
250                     << ", but was " << value.value();
251         }
252     }
253 
expectPrivilegedUserSet(const std::vector<uid_t> & appUids)254     void expectPrivilegedUserSet(const std::vector<uid_t>& appUids) {
255         std::lock_guard guard(mTc.mMutex);
256         EXPECT_EQ(appUids.size(), mTc.mPrivilegedUser.size());
257         for (uid_t uid : appUids) {
258             EXPECT_NE(mTc.mPrivilegedUser.end(), mTc.mPrivilegedUser.find(uid));
259         }
260     }
261 
expectPrivilegedUserSetEmpty()262     void expectPrivilegedUserSetEmpty() {
263         std::lock_guard guard(mTc.mMutex);
264         EXPECT_TRUE(mTc.mPrivilegedUser.empty());
265     }
266 
addPrivilegedUid(uid_t uid)267     void addPrivilegedUid(uid_t uid) {
268         std::vector privilegedUid = {uid};
269         mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, privilegedUid);
270     }
271 
removePrivilegedUid(uid_t uid)272     void removePrivilegedUid(uid_t uid) {
273         std::vector privilegedUid = {uid};
274         mTc.setPermissionForUids(INetd::PERMISSION_NONE, privilegedUid);
275     }
276 
expectFakeStatsUnchanged(uint64_t cookie,uint32_t tag,uint32_t uid,StatsKey tagStatsMapKey)277     void expectFakeStatsUnchanged(uint64_t cookie, uint32_t tag, uint32_t uid,
278                                   StatsKey tagStatsMapKey) {
279         Result<UidTagValue> cookieMapResult = mFakeCookieTagMap.readValue(cookie);
280         EXPECT_RESULT_OK(cookieMapResult);
281         EXPECT_EQ(uid, cookieMapResult.value().uid);
282         EXPECT_EQ(tag, cookieMapResult.value().tag);
283         Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid);
284         EXPECT_RESULT_OK(counterSetResult);
285         EXPECT_EQ(TEST_COUNTERSET, counterSetResult.value());
286         Result<StatsValue> statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey);
287         EXPECT_RESULT_OK(statsMapResult);
288         EXPECT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
289         EXPECT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
290         tagStatsMapKey.tag = 0;
291         statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey);
292         EXPECT_RESULT_OK(statsMapResult);
293         EXPECT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
294         EXPECT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
295         auto appStatsResult = mFakeAppUidStatsMap.readValue(uid);
296         EXPECT_RESULT_OK(appStatsResult);
297         EXPECT_EQ((uint64_t)1, appStatsResult.value().rxPackets);
298         EXPECT_EQ((uint64_t)100, appStatsResult.value().rxBytes);
299     }
300 
expectTagSocketReachLimit(uint32_t tag,uint32_t uid)301     void expectTagSocketReachLimit(uint32_t tag, uint32_t uid) {
302         int sock = socket(AF_INET6, SOCK_STREAM | SOCK_CLOEXEC, 0);
303         EXPECT_LE(0, sock);
304         if (sock < 0) return;
305         uint64_t sockCookie = getSocketCookie(sock);
306         EXPECT_NE(NONEXISTENT_COOKIE, sockCookie);
307         EXPECT_EQ(-EMFILE, mTc.tagSocket(sock, tag, uid, uid));
308         expectNoTag(sockCookie);
309 
310         // Delete stats entries then tag socket success
311         EXPECT_EQ(0, mTc.deleteTagData(0, uid, 0));
312         EXPECT_EQ(0, mTc.tagSocket(sock, tag, uid, uid));
313         expectUidTag(sockCookie, uid, tag);
314     }
315 };
316 
TEST_F(TrafficControllerTest,TestTagSocketV4)317 TEST_F(TrafficControllerTest, TestTagSocketV4) {
318     uint64_t sockCookie;
319     int v4socket = setUpSocketAndTag(AF_INET, &sockCookie, TEST_TAG, TEST_UID, TEST_UID);
320     expectUidTag(sockCookie, TEST_UID, TEST_TAG);
321     ASSERT_EQ(0, mTc.untagSocket(v4socket));
322     expectNoTag(sockCookie);
323     expectMapEmpty(mFakeCookieTagMap);
324 }
325 
TEST_F(TrafficControllerTest,TestReTagSocket)326 TEST_F(TrafficControllerTest, TestReTagSocket) {
327     uint64_t sockCookie;
328     int v4socket = setUpSocketAndTag(AF_INET, &sockCookie, TEST_TAG, TEST_UID, TEST_UID);
329     expectUidTag(sockCookie, TEST_UID, TEST_TAG);
330     ASSERT_EQ(0, mTc.tagSocket(v4socket, TEST_TAG + 1, TEST_UID + 1, TEST_UID + 1));
331     expectUidTag(sockCookie, TEST_UID + 1, TEST_TAG + 1);
332 }
333 
TEST_F(TrafficControllerTest,TestTagTwoSockets)334 TEST_F(TrafficControllerTest, TestTagTwoSockets) {
335     uint64_t sockCookie1;
336     uint64_t sockCookie2;
337     int v4socket1 = setUpSocketAndTag(AF_INET, &sockCookie1, TEST_TAG, TEST_UID, TEST_UID);
338     setUpSocketAndTag(AF_INET, &sockCookie2, TEST_TAG, TEST_UID, TEST_UID);
339     expectUidTag(sockCookie1, TEST_UID, TEST_TAG);
340     expectUidTag(sockCookie2, TEST_UID, TEST_TAG);
341     ASSERT_EQ(0, mTc.untagSocket(v4socket1));
342     expectNoTag(sockCookie1);
343     expectUidTag(sockCookie2, TEST_UID, TEST_TAG);
344     ASSERT_FALSE(mFakeCookieTagMap.getNextKey(sockCookie2).ok());
345 }
346 
TEST_F(TrafficControllerTest,TestTagSocketV6)347 TEST_F(TrafficControllerTest, TestTagSocketV6) {
348     uint64_t sockCookie;
349     int v6socket = setUpSocketAndTag(AF_INET6, &sockCookie, TEST_TAG, TEST_UID, TEST_UID);
350     expectUidTag(sockCookie, TEST_UID, TEST_TAG);
351     ASSERT_EQ(0, mTc.untagSocket(v6socket));
352     expectNoTag(sockCookie);
353     expectMapEmpty(mFakeCookieTagMap);
354 }
355 
TEST_F(TrafficControllerTest,TestTagInvalidSocket)356 TEST_F(TrafficControllerTest, TestTagInvalidSocket) {
357     int invalidSocket = -1;
358     ASSERT_GT(0, mTc.tagSocket(invalidSocket, TEST_TAG, TEST_UID, TEST_UID));
359     expectMapEmpty(mFakeCookieTagMap);
360 }
361 
TEST_F(TrafficControllerTest,TestTagSocketWithoutPermission)362 TEST_F(TrafficControllerTest, TestTagSocketWithoutPermission) {
363     int sock = socket(AF_INET6, SOCK_STREAM | SOCK_CLOEXEC, 0);
364     ASSERT_NE(-1, sock);
365     ASSERT_EQ(-EPERM, mTc.tagSocket(sock, TEST_TAG, TEST_UID, TEST_UID2));
366     expectMapEmpty(mFakeCookieTagMap);
367 }
368 
TEST_F(TrafficControllerTest,TestTagSocketWithPermission)369 TEST_F(TrafficControllerTest, TestTagSocketWithPermission) {
370     // Grant permission to calling uid.
371     std::vector<uid_t> callingUid = {TEST_UID2};
372     mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, callingUid);
373 
374     // Tag a socket to a different uid other then callingUid.
375     uint64_t sockCookie;
376     int v6socket = setUpSocketAndTag(AF_INET6, &sockCookie, TEST_TAG, TEST_UID, TEST_UID2);
377     expectUidTag(sockCookie, TEST_UID, TEST_TAG);
378     EXPECT_EQ(0, mTc.untagSocket(v6socket));
379     expectNoTag(sockCookie);
380     expectMapEmpty(mFakeCookieTagMap);
381 
382     // Clean up the permission
383     mTc.setPermissionForUids(INetd::PERMISSION_NONE, callingUid);
384     expectPrivilegedUserSetEmpty();
385 }
386 
TEST_F(TrafficControllerTest,TestUntagInvalidSocket)387 TEST_F(TrafficControllerTest, TestUntagInvalidSocket) {
388     int invalidSocket = -1;
389     ASSERT_GT(0, mTc.untagSocket(invalidSocket));
390     int v4socket = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
391     ASSERT_GT(0, mTc.untagSocket(v4socket));
392     expectMapEmpty(mFakeCookieTagMap);
393 }
394 
TEST_F(TrafficControllerTest,TestTagSocketReachLimitFail)395 TEST_F(TrafficControllerTest, TestTagSocketReachLimitFail) {
396     uid_t uid = TEST_UID;
397     StatsKey tagStatsMapKey[4];
398     for (int i = 0; i < 3; i++) {
399         uint64_t cookie = TEST_COOKIE + i;
400         uint32_t tag = TEST_TAG + i;
401         populateFakeStats(cookie, uid, tag, &tagStatsMapKey[i]);
402     }
403     expectTagSocketReachLimit(TEST_TAG, TEST_UID);
404 }
405 
TEST_F(TrafficControllerTest,TestTagSocketReachTotalLimitFail)406 TEST_F(TrafficControllerTest, TestTagSocketReachTotalLimitFail) {
407     StatsKey tagStatsMapKey[4];
408     for (int i = 0; i < 4; i++) {
409         uint64_t cookie = TEST_COOKIE + i;
410         uint32_t tag = TEST_TAG + i;
411         uid_t uid = TEST_UID + i;
412         populateFakeStats(cookie, uid, tag, &tagStatsMapKey[i]);
413     }
414     expectTagSocketReachLimit(TEST_TAG, TEST_UID);
415 }
416 
TEST_F(TrafficControllerTest,TestSetCounterSet)417 TEST_F(TrafficControllerTest, TestSetCounterSet) {
418     uid_t callingUid = TEST_UID2;
419     addPrivilegedUid(callingUid);
420     ASSERT_EQ(0, mTc.setCounterSet(TEST_COUNTERSET, TEST_UID, callingUid));
421     uid_t uid = TEST_UID;
422     Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid);
423     ASSERT_RESULT_OK(counterSetResult);
424     ASSERT_EQ(TEST_COUNTERSET, counterSetResult.value());
425     ASSERT_EQ(0, mTc.setCounterSet(DEFAULT_COUNTERSET, TEST_UID, callingUid));
426     ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid).ok());
427     expectMapEmpty(mFakeUidCounterSetMap);
428 }
429 
TEST_F(TrafficControllerTest,TestSetCounterSetWithoutPermission)430 TEST_F(TrafficControllerTest, TestSetCounterSetWithoutPermission) {
431     ASSERT_EQ(-EPERM, mTc.setCounterSet(TEST_COUNTERSET, TEST_UID, TEST_UID2));
432     uid_t uid = TEST_UID;
433     ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid).ok());
434     expectMapEmpty(mFakeUidCounterSetMap);
435 }
436 
TEST_F(TrafficControllerTest,TestSetInvalidCounterSet)437 TEST_F(TrafficControllerTest, TestSetInvalidCounterSet) {
438     uid_t callingUid = TEST_UID2;
439     addPrivilegedUid(callingUid);
440     ASSERT_GT(0, mTc.setCounterSet(OVERFLOW_COUNTERSET, TEST_UID, callingUid));
441     uid_t uid = TEST_UID;
442     ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid).ok());
443     expectMapEmpty(mFakeUidCounterSetMap);
444 }
445 
TEST_F(TrafficControllerTest,TestDeleteTagDataWithoutPermission)446 TEST_F(TrafficControllerTest, TestDeleteTagDataWithoutPermission) {
447     uint64_t cookie = 1;
448     uid_t uid = TEST_UID;
449     uint32_t tag = TEST_TAG;
450     StatsKey tagStatsMapKey;
451     populateFakeStats(cookie, uid, tag, &tagStatsMapKey);
452     ASSERT_EQ(-EPERM, mTc.deleteTagData(0, TEST_UID, TEST_UID2));
453 
454     expectFakeStatsUnchanged(cookie, tag, uid, tagStatsMapKey);
455 }
456 
TEST_F(TrafficControllerTest,TestDeleteTagData)457 TEST_F(TrafficControllerTest, TestDeleteTagData) {
458     uid_t callingUid = TEST_UID2;
459     addPrivilegedUid(callingUid);
460     uint64_t cookie = 1;
461     uid_t uid = TEST_UID;
462     uint32_t tag = TEST_TAG;
463     StatsKey tagStatsMapKey;
464     populateFakeStats(cookie, uid, tag, &tagStatsMapKey);
465     ASSERT_EQ(0, mTc.deleteTagData(TEST_TAG, TEST_UID, callingUid));
466     ASSERT_FALSE(mFakeCookieTagMap.readValue(cookie).ok());
467     Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid);
468     ASSERT_RESULT_OK(counterSetResult);
469     ASSERT_EQ(TEST_COUNTERSET, counterSetResult.value());
470     ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey).ok());
471     tagStatsMapKey.tag = 0;
472     Result<StatsValue> statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey);
473     ASSERT_RESULT_OK(statsMapResult);
474     ASSERT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
475     ASSERT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
476     auto appStatsResult = mFakeAppUidStatsMap.readValue(TEST_UID);
477     ASSERT_RESULT_OK(appStatsResult);
478     ASSERT_EQ((uint64_t)1, appStatsResult.value().rxPackets);
479     ASSERT_EQ((uint64_t)100, appStatsResult.value().rxBytes);
480 }
481 
TEST_F(TrafficControllerTest,TestDeleteAllUidData)482 TEST_F(TrafficControllerTest, TestDeleteAllUidData) {
483     uid_t callingUid = TEST_UID2;
484     addPrivilegedUid(callingUid);
485     uint64_t cookie = 1;
486     uid_t uid = TEST_UID;
487     uint32_t tag = TEST_TAG;
488     StatsKey tagStatsMapKey;
489     populateFakeStats(cookie, uid, tag, &tagStatsMapKey);
490     ASSERT_EQ(0, mTc.deleteTagData(0, TEST_UID, callingUid));
491     ASSERT_FALSE(mFakeCookieTagMap.readValue(cookie).ok());
492     ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid).ok());
493     ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey).ok());
494     tagStatsMapKey.tag = 0;
495     ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey).ok());
496     ASSERT_FALSE(mFakeAppUidStatsMap.readValue(TEST_UID).ok());
497 }
498 
TEST_F(TrafficControllerTest,TestDeleteDataWithTwoTags)499 TEST_F(TrafficControllerTest, TestDeleteDataWithTwoTags) {
500     uid_t callingUid = TEST_UID2;
501     addPrivilegedUid(callingUid);
502     uint64_t cookie1 = 1;
503     uint64_t cookie2 = 2;
504     uid_t uid = TEST_UID;
505     uint32_t tag1 = TEST_TAG;
506     uint32_t tag2 = TEST_TAG + 1;
507     StatsKey tagStatsMapKey1;
508     StatsKey tagStatsMapKey2;
509     populateFakeStats(cookie1, uid, tag1, &tagStatsMapKey1);
510     populateFakeStats(cookie2, uid, tag2, &tagStatsMapKey2);
511     ASSERT_EQ(0, mTc.deleteTagData(TEST_TAG, TEST_UID, callingUid));
512     ASSERT_FALSE(mFakeCookieTagMap.readValue(cookie1).ok());
513     Result<UidTagValue> cookieMapResult = mFakeCookieTagMap.readValue(cookie2);
514     ASSERT_RESULT_OK(cookieMapResult);
515     ASSERT_EQ(TEST_UID, cookieMapResult.value().uid);
516     ASSERT_EQ(TEST_TAG + 1, cookieMapResult.value().tag);
517     Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid);
518     ASSERT_RESULT_OK(counterSetResult);
519     ASSERT_EQ(TEST_COUNTERSET, counterSetResult.value());
520     ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey1).ok());
521     Result<StatsValue> statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey2);
522     ASSERT_RESULT_OK(statsMapResult);
523     ASSERT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
524     ASSERT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
525 }
526 
TEST_F(TrafficControllerTest,TestDeleteDataWithTwoUids)527 TEST_F(TrafficControllerTest, TestDeleteDataWithTwoUids) {
528     uid_t callingUid = TEST_UID2;
529     addPrivilegedUid(callingUid);
530     uint64_t cookie1 = 1;
531     uint64_t cookie2 = 2;
532     uid_t uid1 = TEST_UID;
533     uid_t uid2 = TEST_UID + 1;
534     uint32_t tag = TEST_TAG;
535     StatsKey tagStatsMapKey1;
536     StatsKey tagStatsMapKey2;
537     populateFakeStats(cookie1, uid1, tag, &tagStatsMapKey1);
538     populateFakeStats(cookie2, uid2, tag, &tagStatsMapKey2);
539 
540     // Delete the stats of one of the uid. Check if it is properly collected by
541     // removedStats.
542     ASSERT_EQ(0, mTc.deleteTagData(0, uid2, callingUid));
543     ASSERT_FALSE(mFakeCookieTagMap.readValue(cookie2).ok());
544     Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid1);
545     ASSERT_RESULT_OK(counterSetResult);
546     ASSERT_EQ(TEST_COUNTERSET, counterSetResult.value());
547     ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid2).ok());
548     ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey2).ok());
549     tagStatsMapKey2.tag = 0;
550     ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey2).ok());
551     ASSERT_FALSE(mFakeAppUidStatsMap.readValue(uid2).ok());
552     tagStatsMapKey1.tag = 0;
553     Result<StatsValue> statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey1);
554     ASSERT_RESULT_OK(statsMapResult);
555     ASSERT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
556     ASSERT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
557     auto appStatsResult = mFakeAppUidStatsMap.readValue(uid1);
558     ASSERT_RESULT_OK(appStatsResult);
559     ASSERT_EQ((uint64_t)1, appStatsResult.value().rxPackets);
560     ASSERT_EQ((uint64_t)100, appStatsResult.value().rxBytes);
561 
562     // Delete the stats of the other uid.
563     ASSERT_EQ(0, mTc.deleteTagData(0, uid1, callingUid));
564     ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey1).ok());
565     ASSERT_FALSE(mFakeAppUidStatsMap.readValue(uid1).ok());
566 }
567 
TEST_F(TrafficControllerTest,TestUpdateOwnerMapEntry)568 TEST_F(TrafficControllerTest, TestUpdateOwnerMapEntry) {
569     uint32_t uid = TEST_UID;
570     ASSERT_TRUE(isOk(mTc.updateOwnerMapEntry(STANDBY_MATCH, uid, DENY, DENYLIST)));
571     Result<UidOwnerValue> value = mFakeUidOwnerMap.readValue(uid);
572     ASSERT_RESULT_OK(value);
573     ASSERT_TRUE(value.value().rule & STANDBY_MATCH);
574 
575     ASSERT_TRUE(isOk(mTc.updateOwnerMapEntry(DOZABLE_MATCH, uid, ALLOW, ALLOWLIST)));
576     value = mFakeUidOwnerMap.readValue(uid);
577     ASSERT_RESULT_OK(value);
578     ASSERT_TRUE(value.value().rule & DOZABLE_MATCH);
579 
580     ASSERT_TRUE(isOk(mTc.updateOwnerMapEntry(DOZABLE_MATCH, uid, DENY, ALLOWLIST)));
581     value = mFakeUidOwnerMap.readValue(uid);
582     ASSERT_RESULT_OK(value);
583     ASSERT_FALSE(value.value().rule & DOZABLE_MATCH);
584 
585     ASSERT_TRUE(isOk(mTc.updateOwnerMapEntry(STANDBY_MATCH, uid, ALLOW, DENYLIST)));
586     ASSERT_FALSE(mFakeUidOwnerMap.readValue(uid).ok());
587 
588     uid = TEST_UID2;
589     ASSERT_FALSE(isOk(mTc.updateOwnerMapEntry(STANDBY_MATCH, uid, ALLOW, DENYLIST)));
590     ASSERT_FALSE(mFakeUidOwnerMap.readValue(uid).ok());
591 }
592 
TEST_F(TrafficControllerTest,TestChangeUidOwnerRule)593 TEST_F(TrafficControllerTest, TestChangeUidOwnerRule) {
594     checkUidOwnerRuleForChain(DOZABLE, DOZABLE_MATCH);
595     checkUidOwnerRuleForChain(STANDBY, STANDBY_MATCH);
596     checkUidOwnerRuleForChain(POWERSAVE, POWERSAVE_MATCH);
597     checkUidOwnerRuleForChain(RESTRICTED, RESTRICTED_MATCH);
598     ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(NONE, TEST_UID, ALLOW, ALLOWLIST));
599     ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(INVALID_CHAIN, TEST_UID, ALLOW, ALLOWLIST));
600 }
601 
TEST_F(TrafficControllerTest,TestReplaceUidOwnerMap)602 TEST_F(TrafficControllerTest, TestReplaceUidOwnerMap) {
603     std::vector<int32_t> uids = {TEST_UID, TEST_UID2, TEST_UID3};
604     checkUidMapReplace("fw_dozable", uids, DOZABLE_MATCH);
605     checkUidMapReplace("fw_standby", uids, STANDBY_MATCH);
606     checkUidMapReplace("fw_powersave", uids, POWERSAVE_MATCH);
607     checkUidMapReplace("fw_restricted", uids, RESTRICTED_MATCH);
608     ASSERT_EQ(-EINVAL, mTc.replaceUidOwnerMap("unknow", true, uids));
609 }
610 
TEST_F(TrafficControllerTest,TestReplaceSameChain)611 TEST_F(TrafficControllerTest, TestReplaceSameChain) {
612     std::vector<int32_t> uids = {TEST_UID, TEST_UID2, TEST_UID3};
613     checkUidMapReplace("fw_dozable", uids, DOZABLE_MATCH);
614     std::vector<int32_t> newUids = {TEST_UID2, TEST_UID3};
615     checkUidMapReplace("fw_dozable", newUids, DOZABLE_MATCH);
616 }
617 
TEST_F(TrafficControllerTest,TestDenylistUidMatch)618 TEST_F(TrafficControllerTest, TestDenylistUidMatch) {
619     std::vector<uint32_t> appUids = {1000, 1001, 10012};
620     ASSERT_TRUE(isOk(
621             mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, BandwidthController::IptOpInsert)));
622     expectUidOwnerMapValues(appUids, PENALTY_BOX_MATCH, 0);
623     ASSERT_TRUE(isOk(
624             mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, BandwidthController::IptOpDelete)));
625     expectMapEmpty(mFakeUidOwnerMap);
626 }
627 
TEST_F(TrafficControllerTest,TestAllowlistUidMatch)628 TEST_F(TrafficControllerTest, TestAllowlistUidMatch) {
629     std::vector<uint32_t> appUids = {1000, 1001, 10012};
630     ASSERT_TRUE(isOk(
631             mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, BandwidthController::IptOpInsert)));
632     expectUidOwnerMapValues(appUids, HAPPY_BOX_MATCH, 0);
633     ASSERT_TRUE(isOk(
634             mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, BandwidthController::IptOpDelete)));
635     expectMapEmpty(mFakeUidOwnerMap);
636 }
637 
TEST_F(TrafficControllerTest,TestReplaceMatchUid)638 TEST_F(TrafficControllerTest, TestReplaceMatchUid) {
639     std::vector<uint32_t> appUids = {1000, 1001, 10012};
640     // Add appUids to the denylist and expect that their values are all PENALTY_BOX_MATCH.
641     ASSERT_TRUE(isOk(
642             mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, BandwidthController::IptOpInsert)));
643     expectUidOwnerMapValues(appUids, PENALTY_BOX_MATCH, 0);
644 
645     // Add the same UIDs to the allowlist and expect that we get PENALTY_BOX_MATCH |
646     // HAPPY_BOX_MATCH.
647     ASSERT_TRUE(isOk(
648             mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, BandwidthController::IptOpInsert)));
649     expectUidOwnerMapValues(appUids, HAPPY_BOX_MATCH | PENALTY_BOX_MATCH, 0);
650 
651     // Remove the same UIDs from the allowlist and check the PENALTY_BOX_MATCH is still there.
652     ASSERT_TRUE(isOk(
653             mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, BandwidthController::IptOpDelete)));
654     expectUidOwnerMapValues(appUids, PENALTY_BOX_MATCH, 0);
655 
656     // Remove the same UIDs from the denylist and check the map is empty.
657     ASSERT_TRUE(isOk(
658             mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, BandwidthController::IptOpDelete)));
659     ASSERT_FALSE(mFakeUidOwnerMap.getFirstKey().ok());
660 }
661 
TEST_F(TrafficControllerTest,TestDeleteWrongMatchSilentlyFails)662 TEST_F(TrafficControllerTest, TestDeleteWrongMatchSilentlyFails) {
663     std::vector<uint32_t> appUids = {1000, 1001, 10012};
664     // If the uid does not exist in the map, trying to delete a rule about it will fail.
665     ASSERT_FALSE(isOk(
666             mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, BandwidthController::IptOpDelete)));
667     expectMapEmpty(mFakeUidOwnerMap);
668 
669     // Add denylist rules for appUids.
670     ASSERT_TRUE(isOk(
671             mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, BandwidthController::IptOpInsert)));
672     expectUidOwnerMapValues(appUids, HAPPY_BOX_MATCH, 0);
673 
674     // Delete (non-existent) denylist rules for appUids, and check that this silently does
675     // nothing if the uid is in the map but does not have denylist match. This is required because
676     // NetworkManagementService will try to remove a uid from denylist after adding it to the
677     // allowlist and if the remove fails it will not update the uid status.
678     ASSERT_TRUE(isOk(
679             mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, BandwidthController::IptOpDelete)));
680     expectUidOwnerMapValues(appUids, HAPPY_BOX_MATCH, 0);
681 }
682 
TEST_F(TrafficControllerTest,TestAddUidInterfaceFilteringRules)683 TEST_F(TrafficControllerTest, TestAddUidInterfaceFilteringRules) {
684     int iif0 = 15;
685     ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif0, {1000, 1001})));
686     expectUidOwnerMapValues({1000, 1001}, IIF_MATCH, iif0);
687 
688     // Add some non-overlapping new uids. They should coexist with existing rules
689     int iif1 = 16;
690     ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif1, {2000, 2001})));
691     expectUidOwnerMapValues({1000, 1001}, IIF_MATCH, iif0);
692     expectUidOwnerMapValues({2000, 2001}, IIF_MATCH, iif1);
693 
694     // Overwrite some existing uids
695     int iif2 = 17;
696     ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif2, {1000, 2000})));
697     expectUidOwnerMapValues({1001}, IIF_MATCH, iif0);
698     expectUidOwnerMapValues({2001}, IIF_MATCH, iif1);
699     expectUidOwnerMapValues({1000, 2000}, IIF_MATCH, iif2);
700 }
701 
TEST_F(TrafficControllerTest,TestRemoveUidInterfaceFilteringRules)702 TEST_F(TrafficControllerTest, TestRemoveUidInterfaceFilteringRules) {
703     int iif0 = 15;
704     int iif1 = 16;
705     ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif0, {1000, 1001})));
706     ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif1, {2000, 2001})));
707     expectUidOwnerMapValues({1000, 1001}, IIF_MATCH, iif0);
708     expectUidOwnerMapValues({2000, 2001}, IIF_MATCH, iif1);
709 
710     // Rmove some uids
711     ASSERT_TRUE(isOk(mTc.removeUidInterfaceRules({1001, 2001})));
712     expectUidOwnerMapValues({1000}, IIF_MATCH, iif0);
713     expectUidOwnerMapValues({2000}, IIF_MATCH, iif1);
714     checkEachUidValue({1000, 2000}, IIF_MATCH);  // Make sure there are only two uids remaining
715 
716     // Remove non-existent uids shouldn't fail
717     ASSERT_TRUE(isOk(mTc.removeUidInterfaceRules({2000, 3000})));
718     expectUidOwnerMapValues({1000}, IIF_MATCH, iif0);
719     checkEachUidValue({1000}, IIF_MATCH);  // Make sure there are only one uid remaining
720 
721     // Remove everything
722     ASSERT_TRUE(isOk(mTc.removeUidInterfaceRules({1000})));
723     expectMapEmpty(mFakeUidOwnerMap);
724 }
725 
TEST_F(TrafficControllerTest,TestUidInterfaceFilteringRulesCoexistWithExistingMatches)726 TEST_F(TrafficControllerTest, TestUidInterfaceFilteringRulesCoexistWithExistingMatches) {
727     // Set up existing PENALTY_BOX_MATCH rules
728     ASSERT_TRUE(isOk(mTc.updateUidOwnerMap({1000, 1001, 10012}, PENALTY_BOX_MATCH,
729                                            BandwidthController::IptOpInsert)));
730     expectUidOwnerMapValues({1000, 1001, 10012}, PENALTY_BOX_MATCH, 0);
731 
732     // Add some partially-overlapping uid owner rules and check result
733     int iif1 = 32;
734     ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif1, {10012, 10013, 10014})));
735     expectUidOwnerMapValues({1000, 1001}, PENALTY_BOX_MATCH, 0);
736     expectUidOwnerMapValues({10012}, PENALTY_BOX_MATCH | IIF_MATCH, iif1);
737     expectUidOwnerMapValues({10013, 10014}, IIF_MATCH, iif1);
738 
739     // Removing some PENALTY_BOX_MATCH rules should not change uid interface rule
740     ASSERT_TRUE(isOk(mTc.updateUidOwnerMap({1001, 10012}, PENALTY_BOX_MATCH,
741                                            BandwidthController::IptOpDelete)));
742     expectUidOwnerMapValues({1000}, PENALTY_BOX_MATCH, 0);
743     expectUidOwnerMapValues({10012, 10013, 10014}, IIF_MATCH, iif1);
744 
745     // Remove all uid interface rules
746     ASSERT_TRUE(isOk(mTc.removeUidInterfaceRules({10012, 10013, 10014})));
747     expectUidOwnerMapValues({1000}, PENALTY_BOX_MATCH, 0);
748     // Make sure these are the only uids left
749     checkEachUidValue({1000}, PENALTY_BOX_MATCH);
750 }
751 
TEST_F(TrafficControllerTest,TestUidInterfaceFilteringRulesCoexistWithNewMatches)752 TEST_F(TrafficControllerTest, TestUidInterfaceFilteringRulesCoexistWithNewMatches) {
753     int iif1 = 56;
754     // Set up existing uid interface rules
755     ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif1, {10001, 10002})));
756     expectUidOwnerMapValues({10001, 10002}, IIF_MATCH, iif1);
757 
758     // Add some partially-overlapping doze rules
759     EXPECT_EQ(0, mTc.replaceUidOwnerMap("fw_dozable", true, {10002, 10003}));
760     expectUidOwnerMapValues({10001}, IIF_MATCH, iif1);
761     expectUidOwnerMapValues({10002}, DOZABLE_MATCH | IIF_MATCH, iif1);
762     expectUidOwnerMapValues({10003}, DOZABLE_MATCH, 0);
763 
764     // Introduce a third rule type (powersave) on various existing UIDs
765     EXPECT_EQ(0, mTc.replaceUidOwnerMap("fw_powersave", true, {10000, 10001, 10002, 10003}));
766     expectUidOwnerMapValues({10000}, POWERSAVE_MATCH, 0);
767     expectUidOwnerMapValues({10001}, POWERSAVE_MATCH | IIF_MATCH, iif1);
768     expectUidOwnerMapValues({10002}, POWERSAVE_MATCH | DOZABLE_MATCH | IIF_MATCH, iif1);
769     expectUidOwnerMapValues({10003}, POWERSAVE_MATCH | DOZABLE_MATCH, 0);
770 
771     // Remove all doze rules
772     EXPECT_EQ(0, mTc.replaceUidOwnerMap("fw_dozable", true, {}));
773     expectUidOwnerMapValues({10000}, POWERSAVE_MATCH, 0);
774     expectUidOwnerMapValues({10001}, POWERSAVE_MATCH | IIF_MATCH, iif1);
775     expectUidOwnerMapValues({10002}, POWERSAVE_MATCH | IIF_MATCH, iif1);
776     expectUidOwnerMapValues({10003}, POWERSAVE_MATCH, 0);
777 
778     // Remove all powersave rules, expect ownerMap to only have uid interface rules left
779     EXPECT_EQ(0, mTc.replaceUidOwnerMap("fw_powersave", true, {}));
780     expectUidOwnerMapValues({10001, 10002}, IIF_MATCH, iif1);
781     // Make sure these are the only uids left
782     checkEachUidValue({10001, 10002}, IIF_MATCH);
783 }
784 
TEST_F(TrafficControllerTest,TestGrantInternetPermission)785 TEST_F(TrafficControllerTest, TestGrantInternetPermission) {
786     std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
787 
788     mTc.setPermissionForUids(INetd::PERMISSION_INTERNET, appUids);
789     expectMapEmpty(mFakeUidPermissionMap);
790     expectPrivilegedUserSetEmpty();
791 }
792 
TEST_F(TrafficControllerTest,TestRevokeInternetPermission)793 TEST_F(TrafficControllerTest, TestRevokeInternetPermission) {
794     std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
795 
796     mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
797     expectUidPermissionMapValues(appUids, INetd::PERMISSION_NONE);
798 }
799 
TEST_F(TrafficControllerTest,TestPermissionUninstalled)800 TEST_F(TrafficControllerTest, TestPermissionUninstalled) {
801     std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
802 
803     mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, appUids);
804     expectUidPermissionMapValues(appUids, INetd::PERMISSION_UPDATE_DEVICE_STATS);
805     expectPrivilegedUserSet(appUids);
806 
807     std::vector<uid_t> uidToRemove = {TEST_UID};
808     mTc.setPermissionForUids(INetd::PERMISSION_UNINSTALLED, uidToRemove);
809 
810     std::vector<uid_t> uidRemain = {TEST_UID3, TEST_UID2};
811     expectUidPermissionMapValues(uidRemain, INetd::PERMISSION_UPDATE_DEVICE_STATS);
812     expectPrivilegedUserSet(uidRemain);
813 
814     mTc.setPermissionForUids(INetd::PERMISSION_UNINSTALLED, uidRemain);
815     expectMapEmpty(mFakeUidPermissionMap);
816     expectPrivilegedUserSetEmpty();
817 }
818 
TEST_F(TrafficControllerTest,TestGrantUpdateStatsPermission)819 TEST_F(TrafficControllerTest, TestGrantUpdateStatsPermission) {
820     std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
821 
822     mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, appUids);
823     expectUidPermissionMapValues(appUids, INetd::PERMISSION_UPDATE_DEVICE_STATS);
824     expectPrivilegedUserSet(appUids);
825 
826     mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
827     expectPrivilegedUserSetEmpty();
828     expectUidPermissionMapValues(appUids, INetd::PERMISSION_NONE);
829 }
830 
TEST_F(TrafficControllerTest,TestRevokeUpdateStatsPermission)831 TEST_F(TrafficControllerTest, TestRevokeUpdateStatsPermission) {
832     std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
833 
834     mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, appUids);
835     expectPrivilegedUserSet(appUids);
836 
837     std::vector<uid_t> uidToRemove = {TEST_UID};
838     mTc.setPermissionForUids(INetd::PERMISSION_NONE, uidToRemove);
839 
840     std::vector<uid_t> uidRemain = {TEST_UID3, TEST_UID2};
841     expectPrivilegedUserSet(uidRemain);
842 
843     mTc.setPermissionForUids(INetd::PERMISSION_NONE, uidRemain);
844     expectPrivilegedUserSetEmpty();
845 }
846 
TEST_F(TrafficControllerTest,TestGrantWrongPermission)847 TEST_F(TrafficControllerTest, TestGrantWrongPermission) {
848     std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
849 
850     mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
851     expectPrivilegedUserSetEmpty();
852     expectUidPermissionMapValues(appUids, INetd::PERMISSION_NONE);
853 }
854 
TEST_F(TrafficControllerTest,TestGrantDuplicatePermissionSlientlyFail)855 TEST_F(TrafficControllerTest, TestGrantDuplicatePermissionSlientlyFail) {
856     std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
857 
858     mTc.setPermissionForUids(INetd::PERMISSION_INTERNET, appUids);
859     expectMapEmpty(mFakeUidPermissionMap);
860 
861     std::vector<uid_t> uidToAdd = {TEST_UID};
862     mTc.setPermissionForUids(INetd::PERMISSION_INTERNET, uidToAdd);
863 
864     expectPrivilegedUserSetEmpty();
865 
866     mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
867     expectUidPermissionMapValues(appUids, INetd::PERMISSION_NONE);
868 
869     mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, appUids);
870     expectPrivilegedUserSet(appUids);
871 
872     mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, uidToAdd);
873     expectPrivilegedUserSet(appUids);
874 
875     mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
876     expectPrivilegedUserSetEmpty();
877 }
878 
879 }  // namespace net
880 }  // namespace android
881