1 /*
2 * Copyright 2016 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 * BandwidthControllerTest.cpp - unit tests for BandwidthController.cpp
17 */
18
19 #include <string>
20 #include <vector>
21
22 #include <inttypes.h>
23 #include <fcntl.h>
24 #include <unistd.h>
25 #include <sys/types.h>
26 #include <sys/socket.h>
27
28 #include <gtest/gtest.h>
29
30 #include <android-base/strings.h>
31 #include <android-base/stringprintf.h>
32
33 #include <netdutils/MockSyscalls.h>
34 #include "BandwidthController.h"
35 #include "Fwmark.h"
36 #include "IptablesBaseTest.h"
37 #include "mainline/XtBpfProgLocations.h"
38 #include "tun_interface.h"
39
40 using ::testing::_;
41 using ::testing::ByMove;
42 using ::testing::Invoke;
43 using ::testing::Return;
44 using ::testing::StrictMock;
45
46 using android::base::Join;
47 using android::base::StringPrintf;
48 using android::net::TunInterface;
49 using android::netdutils::UniqueFile;
50 using android::netdutils::status::ok;
51
52 class BandwidthControllerTest : public IptablesBaseTest {
53 protected:
BandwidthControllerTest()54 BandwidthControllerTest() {
55 BandwidthController::iptablesRestoreFunction = fakeExecIptablesRestoreWithOutput;
56 }
57 BandwidthController mBw;
58 TunInterface mTun;
59
SetUp()60 void SetUp() {
61 ASSERT_EQ(0, mTun.init());
62 }
63
TearDown()64 void TearDown() {
65 mTun.destroy();
66 }
67
expectSetupCommands(const std::string & expectedClean,const std::string & expectedAccounting)68 void expectSetupCommands(const std::string& expectedClean,
69 const std::string& expectedAccounting) {
70 std::string expectedList =
71 "*filter\n"
72 "-S\n"
73 "COMMIT\n";
74
75 std::string expectedFlush =
76 "*filter\n"
77 ":bw_INPUT -\n"
78 ":bw_OUTPUT -\n"
79 ":bw_FORWARD -\n"
80 ":bw_happy_box -\n"
81 ":bw_penalty_box -\n"
82 ":bw_data_saver -\n"
83 ":bw_costly_shared -\n"
84 ":bw_global_alert -\n"
85 "COMMIT\n"
86 "*raw\n"
87 ":bw_raw_PREROUTING -\n"
88 "COMMIT\n"
89 "*mangle\n"
90 ":bw_mangle_POSTROUTING -\n"
91 "COMMIT\n";
92
93 ExpectedIptablesCommands expected = {{ V4, expectedList }};
94 if (expectedClean.size()) {
95 expected.push_back({ V4V6, expectedClean });
96 }
97 expected.push_back({ V4V6, expectedFlush });
98 if (expectedAccounting.size()) {
99 expected.push_back({ V4V6, expectedAccounting });
100 }
101
102 expectIptablesRestoreCommands(expected);
103 }
104
105 using IptOp = BandwidthController::IptOp;
106
runIptablesAlertCmd(IptOp a,const char * b,int64_t c)107 int runIptablesAlertCmd(IptOp a, const char* b, int64_t c) {
108 return mBw.runIptablesAlertCmd(a, b, c);
109 }
110
setCostlyAlert(const std::string & a,int64_t b,int64_t * c)111 int setCostlyAlert(const std::string& a, int64_t b, int64_t* c) {
112 return mBw.setCostlyAlert(a, b, c);
113 }
114
removeCostlyAlert(const std::string & a,int64_t * b)115 int removeCostlyAlert(const std::string& a, int64_t* b) { return mBw.removeCostlyAlert(a, b); }
116
expectUpdateQuota(uint64_t quota)117 void expectUpdateQuota(uint64_t quota) {
118 uintptr_t dummy;
119 FILE* dummyFile = reinterpret_cast<FILE*>(&dummy);
120
121 EXPECT_CALL(mSyscalls, fopen(_, _)).WillOnce(Return(ByMove(UniqueFile(dummyFile))));
122 EXPECT_CALL(mSyscalls, vfprintf(dummyFile, _, _))
123 .WillOnce(Invoke([quota](FILE*, const std::string&, va_list ap) {
124 EXPECT_EQ(quota, va_arg(ap, uint64_t));
125 return 0;
126 }));
127 EXPECT_CALL(mSyscalls, fclose(dummyFile)).WillOnce(Return(ok));
128 }
129
130 StrictMock<android::netdutils::ScopedMockSyscalls> mSyscalls;
131 };
132
TEST_F(BandwidthControllerTest,TestSetupIptablesHooks)133 TEST_F(BandwidthControllerTest, TestSetupIptablesHooks) {
134 // Pretend some bw_costly_shared_<iface> rules already exist...
135 addIptablesRestoreOutput(
136 "-P OUTPUT ACCEPT\n"
137 "-N bw_costly_rmnet_data0\n"
138 "-N bw_costly_shared\n"
139 "-N unrelated\n"
140 "-N bw_costly_rmnet_data7\n");
141
142 // ... and expect that they be flushed and deleted.
143 std::string expectedCleanCmds =
144 "*filter\n"
145 ":bw_costly_rmnet_data0 -\n"
146 "-X bw_costly_rmnet_data0\n"
147 ":bw_costly_rmnet_data7 -\n"
148 "-X bw_costly_rmnet_data7\n"
149 "COMMIT\n";
150
151 mBw.setupIptablesHooks();
152 expectSetupCommands(expectedCleanCmds, "");
153 }
154
TEST_F(BandwidthControllerTest,TestCheckUidBillingMask)155 TEST_F(BandwidthControllerTest, TestCheckUidBillingMask) {
156 uint32_t uidBillingMask = Fwmark::getUidBillingMask();
157
158 // If mask is non-zero, and mask & mask-1 is equal to 0, then the mask is a power of two.
159 bool isPowerOfTwo = uidBillingMask && (uidBillingMask & (uidBillingMask - 1)) == 0;
160
161 // Must be exactly a power of two
162 EXPECT_TRUE(isPowerOfTwo);
163 }
164
TEST_F(BandwidthControllerTest,TestEnableBandwidthControl)165 TEST_F(BandwidthControllerTest, TestEnableBandwidthControl) {
166 // Pretend no bw_costly_shared_<iface> rules already exist...
167 addIptablesRestoreOutput(
168 "-P OUTPUT ACCEPT\n"
169 "-N bw_costly_shared\n"
170 "-N unrelated\n");
171
172 // ... so none are flushed or deleted.
173 // clang-format off
174 static const std::string expectedClean = "";
175 static const std::string expectedAccounting =
176 "*filter\n"
177 "-A bw_INPUT -j bw_global_alert\n"
178 "-A bw_INPUT -p esp -j RETURN\n"
179 "-A bw_INPUT -m mark --mark 0x100000/0x100000 -j RETURN\n"
180 "-A bw_INPUT -j MARK --or-mark 0x100000\n"
181 "-A bw_OUTPUT -j bw_global_alert\n"
182 "-A bw_costly_shared -j bw_penalty_box\n"
183 "-I bw_penalty_box -m bpf --object-pinned " XT_BPF_DENYLIST_PROG_PATH " -j REJECT\n"
184 "-A bw_penalty_box -j bw_happy_box\n"
185 "-A bw_happy_box -j bw_data_saver\n"
186 "-A bw_data_saver -j RETURN\n"
187 "-I bw_happy_box -m bpf --object-pinned " XT_BPF_ALLOWLIST_PROG_PATH " -j RETURN\n"
188 "COMMIT\n"
189 "*raw\n"
190 "-A bw_raw_PREROUTING -m mark --mark 0xdeadc1a7 -j DROP\n"
191 "-A bw_raw_PREROUTING -i ipsec+ -j RETURN\n"
192 "-A bw_raw_PREROUTING -m policy --pol ipsec --dir in -j RETURN\n"
193 "-A bw_raw_PREROUTING -m bpf --object-pinned " XT_BPF_INGRESS_PROG_PATH "\n"
194 "COMMIT\n"
195 "*mangle\n"
196 "-A bw_mangle_POSTROUTING -o ipsec+ -j RETURN\n"
197 "-A bw_mangle_POSTROUTING -m policy --pol ipsec --dir out -j RETURN\n"
198 "-A bw_mangle_POSTROUTING -j MARK --set-mark 0x0/0x100000\n"
199 "-A bw_mangle_POSTROUTING -m bpf --object-pinned " XT_BPF_EGRESS_PROG_PATH "\n"
200 "COMMIT\n";
201 // clang-format on
202
203 mBw.enableBandwidthControl();
204 expectSetupCommands(expectedClean, expectedAccounting);
205 }
206
TEST_F(BandwidthControllerTest,TestEnableDataSaver)207 TEST_F(BandwidthControllerTest, TestEnableDataSaver) {
208 mBw.enableDataSaver(true);
209 std::string expected4 =
210 "*filter\n"
211 ":bw_data_saver -\n"
212 "-A bw_data_saver -j REJECT\n"
213 "COMMIT\n";
214 std::string expected6 =
215 "*filter\n"
216 ":bw_data_saver -\n"
217 "-A bw_data_saver -p icmpv6 --icmpv6-type packet-too-big -j RETURN\n"
218 "-A bw_data_saver -p icmpv6 --icmpv6-type router-solicitation -j RETURN\n"
219 "-A bw_data_saver -p icmpv6 --icmpv6-type router-advertisement -j RETURN\n"
220 "-A bw_data_saver -p icmpv6 --icmpv6-type neighbour-solicitation -j RETURN\n"
221 "-A bw_data_saver -p icmpv6 --icmpv6-type neighbour-advertisement -j RETURN\n"
222 "-A bw_data_saver -p icmpv6 --icmpv6-type redirect -j RETURN\n"
223 "-A bw_data_saver -j REJECT\n"
224 "COMMIT\n";
225 expectIptablesRestoreCommands({
226 {V4, expected4},
227 {V6, expected6},
228 });
229
230 mBw.enableDataSaver(false);
231 std::string expected = {
232 "*filter\n"
233 ":bw_data_saver -\n"
234 "-A bw_data_saver -j RETURN\n"
235 "COMMIT\n"};
236 expectIptablesRestoreCommands({
237 {V4, expected},
238 {V6, expected},
239 });
240 }
241
makeInterfaceQuotaCommands(const std::string & iface,int ruleIndex,int64_t quota)242 const std::vector<std::string> makeInterfaceQuotaCommands(const std::string& iface, int ruleIndex,
243 int64_t quota) {
244 const std::string chain = "bw_costly_" + iface;
245 const char* c_chain = chain.c_str();
246 const char* c_iface = iface.c_str();
247 std::vector<std::string> cmds = {
248 "*filter",
249 StringPrintf(":%s -", c_chain),
250 StringPrintf("-A %s -j bw_penalty_box", c_chain),
251 StringPrintf("-I bw_INPUT %d -i %s -j %s", ruleIndex, c_iface, c_chain),
252 StringPrintf("-I bw_OUTPUT %d -o %s -j %s", ruleIndex, c_iface, c_chain),
253 StringPrintf("-A bw_FORWARD -i %s -j %s", c_iface, c_chain),
254 StringPrintf("-A bw_FORWARD -o %s -j %s", c_iface, c_chain),
255 StringPrintf("-A %s -m quota2 ! --quota %" PRIu64 " --name %s -j REJECT", c_chain,
256 quota, c_iface),
257 "COMMIT\n",
258 };
259 return {Join(cmds, "\n")};
260 }
261
removeInterfaceQuotaCommands(const std::string & iface)262 const std::vector<std::string> removeInterfaceQuotaCommands(const std::string& iface) {
263 const std::string chain = "bw_costly_" + iface;
264 const char* c_chain = chain.c_str();
265 const char* c_iface = iface.c_str();
266 std::vector<std::string> cmds = {
267 "*filter",
268 StringPrintf("-D bw_INPUT -i %s -j %s", c_iface, c_chain),
269 StringPrintf("-D bw_OUTPUT -o %s -j %s", c_iface, c_chain),
270 StringPrintf("-D bw_FORWARD -i %s -j %s", c_iface, c_chain),
271 StringPrintf("-D bw_FORWARD -o %s -j %s", c_iface, c_chain),
272 StringPrintf("-F %s", c_chain),
273 StringPrintf("-X %s", c_chain),
274 "COMMIT\n",
275 };
276 return {Join(cmds, "\n")};
277 }
278
TEST_F(BandwidthControllerTest,TestSetInterfaceQuota)279 TEST_F(BandwidthControllerTest, TestSetInterfaceQuota) {
280 constexpr uint64_t kOldQuota = 123456;
281 const std::string iface = mTun.name();
282 std::vector<std::string> expected = makeInterfaceQuotaCommands(iface, 1, kOldQuota);
283
284 EXPECT_EQ(0, mBw.setInterfaceQuota(iface, kOldQuota));
285 expectIptablesRestoreCommands(expected);
286
287 constexpr uint64_t kNewQuota = kOldQuota + 1;
288 expected = {};
289 expectUpdateQuota(kNewQuota);
290 EXPECT_EQ(0, mBw.setInterfaceQuota(iface, kNewQuota));
291 expectIptablesRestoreCommands(expected);
292
293 expected = removeInterfaceQuotaCommands(iface);
294 EXPECT_EQ(0, mBw.removeInterfaceQuota(iface));
295 expectIptablesRestoreCommands(expected);
296 }
297
makeInterfaceSharedQuotaCommands(const std::string & iface,int ruleIndex,int64_t quota,bool insertQuota)298 const std::vector<std::string> makeInterfaceSharedQuotaCommands(const std::string& iface,
299 int ruleIndex, int64_t quota,
300 bool insertQuota) {
301 const std::string chain = "bw_costly_shared";
302 const char* c_chain = chain.c_str();
303 const char* c_iface = iface.c_str();
304 std::vector<std::string> cmds = {
305 "*filter",
306 StringPrintf("-I bw_INPUT %d -i %s -j %s", ruleIndex, c_iface, c_chain),
307 StringPrintf("-I bw_OUTPUT %d -o %s -j %s", ruleIndex, c_iface, c_chain),
308 StringPrintf("-A bw_FORWARD -i %s -j %s", c_iface, c_chain),
309 StringPrintf("-A bw_FORWARD -o %s -j %s", c_iface, c_chain),
310 };
311 if (insertQuota) {
312 cmds.push_back(StringPrintf("-I %s -m quota2 ! --quota %" PRIu64 " --name shared -j REJECT",
313 c_chain, quota));
314 }
315 cmds.push_back("COMMIT\n");
316 return {Join(cmds, "\n")};
317 }
318
removeInterfaceSharedQuotaCommands(const std::string & iface,int64_t quota,bool deleteQuota)319 const std::vector<std::string> removeInterfaceSharedQuotaCommands(const std::string& iface,
320 int64_t quota, bool deleteQuota) {
321 const std::string chain = "bw_costly_shared";
322 const char* c_chain = chain.c_str();
323 const char* c_iface = iface.c_str();
324 std::vector<std::string> cmds = {
325 "*filter",
326 StringPrintf("-D bw_INPUT -i %s -j %s", c_iface, c_chain),
327 StringPrintf("-D bw_OUTPUT -o %s -j %s", c_iface, c_chain),
328 StringPrintf("-D bw_FORWARD -i %s -j %s", c_iface, c_chain),
329 StringPrintf("-D bw_FORWARD -o %s -j %s", c_iface, c_chain),
330 };
331 if (deleteQuota) {
332 cmds.push_back(StringPrintf("-D %s -m quota2 ! --quota %" PRIu64 " --name shared -j REJECT",
333 c_chain, quota));
334 }
335 cmds.push_back("COMMIT\n");
336 return {Join(cmds, "\n")};
337 }
338
TEST_F(BandwidthControllerTest,TestSetInterfaceSharedQuotaDuplicate)339 TEST_F(BandwidthControllerTest, TestSetInterfaceSharedQuotaDuplicate) {
340 constexpr uint64_t kQuota = 123456;
341 const std::string iface = mTun.name();
342 std::vector<std::string> expected = makeInterfaceSharedQuotaCommands(iface, 1, 123456, true);
343 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kQuota));
344 expectIptablesRestoreCommands(expected);
345
346 expected = {};
347 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kQuota));
348 expectIptablesRestoreCommands(expected);
349
350 expected = removeInterfaceSharedQuotaCommands(iface, kQuota, true);
351 EXPECT_EQ(0, mBw.removeInterfaceSharedQuota(iface));
352 expectIptablesRestoreCommands(expected);
353 }
354
TEST_F(BandwidthControllerTest,TestSetInterfaceSharedQuotaUpdate)355 TEST_F(BandwidthControllerTest, TestSetInterfaceSharedQuotaUpdate) {
356 constexpr uint64_t kOldQuota = 123456;
357 const std::string iface = mTun.name();
358 std::vector<std::string> expected = makeInterfaceSharedQuotaCommands(iface, 1, kOldQuota, true);
359 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kOldQuota));
360 expectIptablesRestoreCommands(expected);
361
362 constexpr uint64_t kNewQuota = kOldQuota + 1;
363 expected = {};
364 expectUpdateQuota(kNewQuota);
365 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kNewQuota));
366 expectIptablesRestoreCommands(expected);
367
368 expected = removeInterfaceSharedQuotaCommands(iface, kNewQuota, true);
369 EXPECT_EQ(0, mBw.removeInterfaceSharedQuota(iface));
370 expectIptablesRestoreCommands(expected);
371 }
372
TEST_F(BandwidthControllerTest,TestSetInterfaceSharedQuotaTwoInterfaces)373 TEST_F(BandwidthControllerTest, TestSetInterfaceSharedQuotaTwoInterfaces) {
374 constexpr uint64_t kQuota = 123456;
375 const std::vector<std::string> ifaces{
376 {"a" + mTun.name()},
377 {"b" + mTun.name()},
378 };
379
380 for (const auto& iface : ifaces) {
381 // Quota rule is only added when the total number of
382 // interfaces transitions from 0 -> 1.
383 bool first = (iface == ifaces[0]);
384 auto expected = makeInterfaceSharedQuotaCommands(iface, 1, kQuota, first);
385 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kQuota));
386 expectIptablesRestoreCommands(expected);
387 }
388
389 for (const auto& iface : ifaces) {
390 // Quota rule is only removed when the total number of
391 // interfaces transitions from 1 -> 0.
392 bool last = (iface == ifaces[1]);
393 auto expected = removeInterfaceSharedQuotaCommands(iface, kQuota, last);
394 EXPECT_EQ(0, mBw.removeInterfaceSharedQuota(iface));
395 expectIptablesRestoreCommands(expected);
396 }
397 }
398
TEST_F(BandwidthControllerTest,IptablesAlertCmd)399 TEST_F(BandwidthControllerTest, IptablesAlertCmd) {
400 std::vector<std::string> expected = {
401 "*filter\n"
402 "-I bw_global_alert -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
403 "COMMIT\n"};
404 EXPECT_EQ(0, runIptablesAlertCmd(IptOp::IptOpInsert, "MyWonderfulAlert", 123456));
405 expectIptablesRestoreCommands(expected);
406
407 expected = {
408 "*filter\n"
409 "-D bw_global_alert -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
410 "COMMIT\n"};
411 EXPECT_EQ(0, runIptablesAlertCmd(IptOp::IptOpDelete, "MyWonderfulAlert", 123456));
412 expectIptablesRestoreCommands(expected);
413 }
414
TEST_F(BandwidthControllerTest,CostlyAlert)415 TEST_F(BandwidthControllerTest, CostlyAlert) {
416 const int64_t kQuota = 123456;
417 int64_t alertBytes = 0;
418
419 std::vector<std::string> expected = {
420 "*filter\n"
421 "-A bw_costly_shared -m quota2 ! --quota 123456 --name sharedAlert\n"
422 "COMMIT\n"
423 };
424 EXPECT_EQ(0, setCostlyAlert("shared", kQuota, &alertBytes));
425 EXPECT_EQ(kQuota, alertBytes);
426 expectIptablesRestoreCommands(expected);
427
428 expected = {};
429 expectUpdateQuota(kQuota);
430 EXPECT_EQ(0, setCostlyAlert("shared", kQuota + 1, &alertBytes));
431 EXPECT_EQ(kQuota + 1, alertBytes);
432 expectIptablesRestoreCommands(expected);
433
434 expected = {
435 "*filter\n"
436 "-D bw_costly_shared -m quota2 ! --quota 123457 --name sharedAlert\n"
437 "COMMIT\n"
438 };
439 EXPECT_EQ(0, removeCostlyAlert("shared", &alertBytes));
440 EXPECT_EQ(0, alertBytes);
441 expectIptablesRestoreCommands(expected);
442 }
443
444