1 /*
2  * Copyright 2021 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  * BpfHandlerTest.cpp - unit tests for BpfHandler.cpp
17  */
18 
19 #include <private/android_filesystem_config.h>
20 #include <sys/socket.h>
21 
22 #include <gtest/gtest.h>
23 
24 #define BPF_MAP_MAKE_VISIBLE_FOR_TESTING
25 #include "BpfHandler.h"
26 
27 using namespace android::bpf;  // NOLINT(google-build-using-namespace): exempted
28 
29 namespace android {
30 namespace net {
31 
32 using base::Result;
33 
34 constexpr int TEST_MAP_SIZE = 10;
35 constexpr int TEST_COOKIE = 1;
36 constexpr uid_t TEST_UID = 10086;
37 constexpr uid_t TEST_UID2 = 54321;
38 constexpr uint32_t TEST_TAG = 42;
39 constexpr uint32_t TEST_COUNTERSET = 1;
40 constexpr uint32_t TEST_PER_UID_STATS_ENTRIES_LIMIT = 3;
41 constexpr uint32_t TEST_TOTAL_UID_STATS_ENTRIES_LIMIT = 7;
42 
43 #define ASSERT_VALID(x) ASSERT_TRUE((x).isValid())
44 
45 class BpfHandlerTest : public ::testing::Test {
46   protected:
BpfHandlerTest()47     BpfHandlerTest()
48         : mBh(TEST_PER_UID_STATS_ENTRIES_LIMIT, TEST_TOTAL_UID_STATS_ENTRIES_LIMIT) {}
49     BpfHandler mBh;
50     BpfMap<uint64_t, UidTagValue> mFakeCookieTagMap;
51     BpfMap<StatsKey, StatsValue> mFakeStatsMapA;
52     BpfMap<uint32_t, uint32_t> mFakeConfigurationMap;
53     BpfMap<uint32_t, uint8_t> mFakeUidPermissionMap;
54 
SetUp()55     void SetUp() {
56         ASSERT_EQ(0, setrlimitForTest());
57 
58         mFakeCookieTagMap.resetMap(BPF_MAP_TYPE_HASH, TEST_MAP_SIZE);
59         ASSERT_VALID(mFakeCookieTagMap);
60 
61         mFakeStatsMapA.resetMap(BPF_MAP_TYPE_HASH, TEST_MAP_SIZE);
62         ASSERT_VALID(mFakeStatsMapA);
63 
64         mFakeConfigurationMap.resetMap(BPF_MAP_TYPE_ARRAY, CONFIGURATION_MAP_SIZE);
65         ASSERT_VALID(mFakeConfigurationMap);
66 
67         mFakeUidPermissionMap.resetMap(BPF_MAP_TYPE_HASH, TEST_MAP_SIZE);
68         ASSERT_VALID(mFakeUidPermissionMap);
69 
70         mBh.mCookieTagMap = mFakeCookieTagMap;
71         ASSERT_VALID(mBh.mCookieTagMap);
72         mBh.mStatsMapA = mFakeStatsMapA;
73         ASSERT_VALID(mBh.mStatsMapA);
74         mBh.mConfigurationMap = mFakeConfigurationMap;
75         ASSERT_VALID(mBh.mConfigurationMap);
76         // Always write to stats map A by default.
77         static_assert(SELECT_MAP_A == 0, "bpf map arrays are zero-initialized");
78 
79         mBh.mUidPermissionMap = mFakeUidPermissionMap;
80         ASSERT_VALID(mBh.mUidPermissionMap);
81     }
82 
setUpSocketAndTag(int protocol,uint64_t * cookie,uint32_t tag,uid_t uid,uid_t realUid)83     int setUpSocketAndTag(int protocol, uint64_t* cookie, uint32_t tag, uid_t uid,
84                           uid_t realUid) {
85         int sock = socket(protocol, SOCK_STREAM | SOCK_CLOEXEC, 0);
86         EXPECT_LE(0, sock);
87         *cookie = getSocketCookie(sock);
88         EXPECT_NE(NONEXISTENT_COOKIE, *cookie);
89         EXPECT_EQ(0, mBh.tagSocket(sock, tag, uid, realUid));
90         return sock;
91     }
92 
expectUidTag(uint64_t cookie,uid_t uid,uint32_t tag)93     void expectUidTag(uint64_t cookie, uid_t uid, uint32_t tag) {
94         Result<UidTagValue> tagResult = mFakeCookieTagMap.readValue(cookie);
95         ASSERT_RESULT_OK(tagResult);
96         EXPECT_EQ(uid, tagResult.value().uid);
97         EXPECT_EQ(tag, tagResult.value().tag);
98     }
99 
expectNoTag(uint64_t cookie)100     void expectNoTag(uint64_t cookie) { EXPECT_FALSE(mFakeCookieTagMap.readValue(cookie).ok()); }
101 
populateFakeStats(uint64_t cookie,uint32_t uid,uint32_t tag,StatsKey * key)102     void populateFakeStats(uint64_t cookie, uint32_t uid, uint32_t tag, StatsKey* key) {
103         UidTagValue cookieMapkey = {.uid = (uint32_t)uid, .tag = tag};
104         EXPECT_RESULT_OK(mFakeCookieTagMap.writeValue(cookie, cookieMapkey, BPF_ANY));
105         *key = {.uid = uid, .tag = tag, .counterSet = TEST_COUNTERSET, .ifaceIndex = 1};
106         StatsValue statsMapValue = {.rxPackets = 1, .rxBytes = 100};
107         EXPECT_RESULT_OK(mFakeStatsMapA.writeValue(*key, statsMapValue, BPF_ANY));
108         key->tag = 0;
109         EXPECT_RESULT_OK(mFakeStatsMapA.writeValue(*key, statsMapValue, BPF_ANY));
110         // put tag information back to statsKey
111         key->tag = tag;
112     }
113 
114     template <class Key, class Value>
expectMapEmpty(BpfMap<Key,Value> & map)115     void expectMapEmpty(BpfMap<Key, Value>& map) {
116         auto isEmpty = map.isEmpty();
117         EXPECT_RESULT_OK(isEmpty);
118         EXPECT_TRUE(isEmpty.value());
119     }
120 
expectTagSocketReachLimit(uint32_t tag,uint32_t uid)121     void expectTagSocketReachLimit(uint32_t tag, uint32_t uid) {
122         int sock = socket(AF_INET6, SOCK_STREAM | SOCK_CLOEXEC, 0);
123         EXPECT_LE(0, sock);
124         if (sock < 0) return;
125         uint64_t sockCookie = getSocketCookie(sock);
126         EXPECT_NE(NONEXISTENT_COOKIE, sockCookie);
127         EXPECT_EQ(-EMFILE, mBh.tagSocket(sock, tag, uid, uid));
128         expectNoTag(sockCookie);
129 
130         // Delete stats entries then tag socket success
131         StatsKey key = {.uid = uid, .tag = 0, .counterSet = TEST_COUNTERSET, .ifaceIndex = 1};
132         ASSERT_RESULT_OK(mFakeStatsMapA.deleteValue(key));
133         EXPECT_EQ(0, mBh.tagSocket(sock, tag, uid, uid));
134         expectUidTag(sockCookie, uid, tag);
135     }
136 };
137 
TEST_F(BpfHandlerTest,TestTagSocketV4)138 TEST_F(BpfHandlerTest, TestTagSocketV4) {
139     uint64_t sockCookie;
140     int v4socket = setUpSocketAndTag(AF_INET, &sockCookie, TEST_TAG, TEST_UID, TEST_UID);
141     expectUidTag(sockCookie, TEST_UID, TEST_TAG);
142     ASSERT_EQ(0, mBh.untagSocket(v4socket));
143     expectNoTag(sockCookie);
144     expectMapEmpty(mFakeCookieTagMap);
145 }
146 
TEST_F(BpfHandlerTest,TestReTagSocket)147 TEST_F(BpfHandlerTest, TestReTagSocket) {
148     uint64_t sockCookie;
149     int v4socket = setUpSocketAndTag(AF_INET, &sockCookie, TEST_TAG, TEST_UID, TEST_UID);
150     expectUidTag(sockCookie, TEST_UID, TEST_TAG);
151     ASSERT_EQ(0, mBh.tagSocket(v4socket, TEST_TAG + 1, TEST_UID + 1, TEST_UID + 1));
152     expectUidTag(sockCookie, TEST_UID + 1, TEST_TAG + 1);
153 }
154 
TEST_F(BpfHandlerTest,TestTagTwoSockets)155 TEST_F(BpfHandlerTest, TestTagTwoSockets) {
156     uint64_t sockCookie1;
157     uint64_t sockCookie2;
158     int v4socket1 = setUpSocketAndTag(AF_INET, &sockCookie1, TEST_TAG, TEST_UID, TEST_UID);
159     setUpSocketAndTag(AF_INET, &sockCookie2, TEST_TAG, TEST_UID, TEST_UID);
160     expectUidTag(sockCookie1, TEST_UID, TEST_TAG);
161     expectUidTag(sockCookie2, TEST_UID, TEST_TAG);
162     ASSERT_EQ(0, mBh.untagSocket(v4socket1));
163     expectNoTag(sockCookie1);
164     expectUidTag(sockCookie2, TEST_UID, TEST_TAG);
165     ASSERT_FALSE(mFakeCookieTagMap.getNextKey(sockCookie2).ok());
166 }
167 
TEST_F(BpfHandlerTest,TestTagSocketV6)168 TEST_F(BpfHandlerTest, TestTagSocketV6) {
169     uint64_t sockCookie;
170     int v6socket = setUpSocketAndTag(AF_INET6, &sockCookie, TEST_TAG, TEST_UID, TEST_UID);
171     expectUidTag(sockCookie, TEST_UID, TEST_TAG);
172     ASSERT_EQ(0, mBh.untagSocket(v6socket));
173     expectNoTag(sockCookie);
174     expectMapEmpty(mFakeCookieTagMap);
175 }
176 
TEST_F(BpfHandlerTest,TestTagInvalidSocket)177 TEST_F(BpfHandlerTest, TestTagInvalidSocket) {
178     int invalidSocket = -1;
179     ASSERT_GT(0, mBh.tagSocket(invalidSocket, TEST_TAG, TEST_UID, TEST_UID));
180     expectMapEmpty(mFakeCookieTagMap);
181 }
182 
TEST_F(BpfHandlerTest,TestTagSocketWithUnsupportedFamily)183 TEST_F(BpfHandlerTest, TestTagSocketWithUnsupportedFamily) {
184     int packetSocket = socket(AF_PACKET, SOCK_DGRAM | SOCK_CLOEXEC, 0);
185     EXPECT_LE(0, packetSocket);
186     EXPECT_NE(NONEXISTENT_COOKIE, getSocketCookie(packetSocket));
187     EXPECT_EQ(-EAFNOSUPPORT, mBh.tagSocket(packetSocket, TEST_TAG, TEST_UID, TEST_UID));
188 }
189 
TEST_F(BpfHandlerTest,TestTagSocketWithUnsupportedProtocol)190 TEST_F(BpfHandlerTest, TestTagSocketWithUnsupportedProtocol) {
191     int rawSocket = socket(AF_INET, SOCK_RAW | SOCK_CLOEXEC, IPPROTO_RAW);
192     EXPECT_LE(0, rawSocket);
193     EXPECT_NE(NONEXISTENT_COOKIE, getSocketCookie(rawSocket));
194     EXPECT_EQ(-EPROTONOSUPPORT, mBh.tagSocket(rawSocket, TEST_TAG, TEST_UID, TEST_UID));
195 }
196 
TEST_F(BpfHandlerTest,TestTagSocketWithoutPermission)197 TEST_F(BpfHandlerTest, TestTagSocketWithoutPermission) {
198     int sock = socket(AF_INET6, SOCK_STREAM | SOCK_CLOEXEC, 0);
199     ASSERT_NE(-1, sock);
200     ASSERT_EQ(-EPERM, mBh.tagSocket(sock, TEST_TAG, TEST_UID, TEST_UID2));
201     expectMapEmpty(mFakeCookieTagMap);
202 }
203 
TEST_F(BpfHandlerTest,TestTagSocketWithPermission)204 TEST_F(BpfHandlerTest, TestTagSocketWithPermission) {
205     // Grant permission to real uid. In practice, the uid permission map will be updated by
206     // TrafficController::setPermissionForUids().
207     uid_t realUid = TEST_UID2;
208     ASSERT_RESULT_OK(mFakeUidPermissionMap.writeValue(realUid,
209                      BPF_PERMISSION_UPDATE_DEVICE_STATS, BPF_ANY));
210 
211     // Tag a socket to a different uid other then realUid.
212     uint64_t sockCookie;
213     int v6socket = setUpSocketAndTag(AF_INET6, &sockCookie, TEST_TAG, TEST_UID, realUid);
214     expectUidTag(sockCookie, TEST_UID, TEST_TAG);
215     EXPECT_EQ(0, mBh.untagSocket(v6socket));
216     expectNoTag(sockCookie);
217     expectMapEmpty(mFakeCookieTagMap);
218 
219     // Tag a socket to AID_CLAT other then realUid.
220     int sock = socket(AF_INET6, SOCK_STREAM | SOCK_CLOEXEC, 0);
221     ASSERT_NE(-1, sock);
222     ASSERT_EQ(-EPERM, mBh.tagSocket(sock, TEST_TAG, AID_CLAT, realUid));
223     expectMapEmpty(mFakeCookieTagMap);
224 }
225 
TEST_F(BpfHandlerTest,TestUntagInvalidSocket)226 TEST_F(BpfHandlerTest, TestUntagInvalidSocket) {
227     int invalidSocket = -1;
228     ASSERT_GT(0, mBh.untagSocket(invalidSocket));
229     int v4socket = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
230     ASSERT_GT(0, mBh.untagSocket(v4socket));
231     expectMapEmpty(mFakeCookieTagMap);
232 }
233 
TEST_F(BpfHandlerTest,TestTagSocketReachLimitFail)234 TEST_F(BpfHandlerTest, TestTagSocketReachLimitFail) {
235     uid_t uid = TEST_UID;
236     StatsKey tagStatsMapKey[3];
237     for (int i = 0; i < 3; i++) {
238         uint64_t cookie = TEST_COOKIE + i;
239         uint32_t tag = TEST_TAG + i;
240         populateFakeStats(cookie, uid, tag, &tagStatsMapKey[i]);
241     }
242     expectTagSocketReachLimit(TEST_TAG, TEST_UID);
243 }
244 
TEST_F(BpfHandlerTest,TestTagSocketReachTotalLimitFail)245 TEST_F(BpfHandlerTest, TestTagSocketReachTotalLimitFail) {
246     StatsKey tagStatsMapKey[4];
247     for (int i = 0; i < 4; i++) {
248         uint64_t cookie = TEST_COOKIE + i;
249         uint32_t tag = TEST_TAG + i;
250         uid_t uid = TEST_UID + i;
251         populateFakeStats(cookie, uid, tag, &tagStatsMapKey[i]);
252     }
253     expectTagSocketReachLimit(TEST_TAG, TEST_UID);
254 }
255 
256 }  // namespace net
257 }  // namespace android
258