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 "IptablesBaseTest.h"
36 #include "tun_interface.h"
37
38 using ::testing::ByMove;
39 using ::testing::Invoke;
40 using ::testing::Return;
41 using ::testing::StrictMock;
42 using ::testing::Test;
43 using ::testing::_;
44
45 using android::base::Join;
46 using android::base::StringPrintf;
47 using android::net::TunInterface;
48 using android::netdutils::status::ok;
49 using android::netdutils::UniqueFile;
50
51 class BandwidthControllerTest : public IptablesBaseTest {
52 protected:
BandwidthControllerTest()53 BandwidthControllerTest() {
54 BandwidthController::iptablesRestoreFunction = fakeExecIptablesRestoreWithOutput;
55 }
56 BandwidthController mBw;
57 TunInterface mTun;
58
SetUp()59 void SetUp() {
60 ASSERT_EQ(0, mTun.init());
61 }
62
TearDown()63 void TearDown() {
64 mTun.destroy();
65 }
66
addIptablesRestoreOutput(std::string contents)67 void addIptablesRestoreOutput(std::string contents) {
68 sIptablesRestoreOutput.push_back(contents);
69 }
70
addIptablesRestoreOutput(std::string contents1,std::string contents2)71 void addIptablesRestoreOutput(std::string contents1, std::string contents2) {
72 sIptablesRestoreOutput.push_back(contents1);
73 sIptablesRestoreOutput.push_back(contents2);
74 }
75
clearIptablesRestoreOutput()76 void clearIptablesRestoreOutput() {
77 sIptablesRestoreOutput.clear();
78 }
79
expectSetupCommands(const std::string & expectedClean,std::string expectedAccounting)80 void expectSetupCommands(const std::string& expectedClean, std::string expectedAccounting) {
81 std::string expectedList =
82 "*filter\n"
83 "-S\n"
84 "COMMIT\n";
85
86 std::string expectedFlush =
87 "*filter\n"
88 ":bw_INPUT -\n"
89 ":bw_OUTPUT -\n"
90 ":bw_FORWARD -\n"
91 ":bw_happy_box -\n"
92 ":bw_penalty_box -\n"
93 ":bw_data_saver -\n"
94 ":bw_costly_shared -\n"
95 "COMMIT\n"
96 "*raw\n"
97 ":bw_raw_PREROUTING -\n"
98 "COMMIT\n"
99 "*mangle\n"
100 ":bw_mangle_POSTROUTING -\n"
101 "COMMIT\n";
102
103 ExpectedIptablesCommands expected = {{ V4, expectedList }};
104 if (expectedClean.size()) {
105 expected.push_back({ V4V6, expectedClean });
106 }
107 expected.push_back({ V4V6, expectedFlush });
108 if (expectedAccounting.size()) {
109 expected.push_back({ V4V6, expectedAccounting });
110 }
111
112 expectIptablesRestoreCommands(expected);
113 }
114
115 using IptOp = BandwidthController::IptOp;
116
runIptablesAlertCmd(IptOp a,const char * b,int64_t c)117 int runIptablesAlertCmd(IptOp a, const char *b, int64_t c) {
118 return mBw.runIptablesAlertCmd(a, b, c);
119 }
120
runIptablesAlertFwdCmd(IptOp a,const char * b,int64_t c)121 int runIptablesAlertFwdCmd(IptOp a, const char *b, int64_t c) {
122 return mBw.runIptablesAlertFwdCmd(a, b, c);
123 }
124
setCostlyAlert(const std::string a,int64_t b,int64_t * c)125 int setCostlyAlert(const std::string a, int64_t b, int64_t *c) {
126 return mBw.setCostlyAlert(a, b, c);
127 }
128
removeCostlyAlert(const std::string a,int64_t * b)129 int removeCostlyAlert(const std::string a, int64_t *b) {
130 return mBw.removeCostlyAlert(a, b);
131 }
132
expectUpdateQuota(uint64_t quota)133 void expectUpdateQuota(uint64_t quota) {
134 uintptr_t dummy;
135 FILE* dummyFile = reinterpret_cast<FILE*>(&dummy);
136
137 EXPECT_CALL(mSyscalls, fopen(_, _)).WillOnce(Return(ByMove(UniqueFile(dummyFile))));
138 EXPECT_CALL(mSyscalls, vfprintf(dummyFile, _, _))
139 .WillOnce(Invoke([quota](FILE*, const std::string&, va_list ap) {
140 EXPECT_EQ(quota, va_arg(ap, uint64_t));
141 return 0;
142 }));
143 EXPECT_CALL(mSyscalls, fclose(dummyFile)).WillOnce(Return(ok));
144 }
145
146 StrictMock<android::netdutils::ScopedMockSyscalls> mSyscalls;
147 };
148
TEST_F(BandwidthControllerTest,TestSetupIptablesHooks)149 TEST_F(BandwidthControllerTest, TestSetupIptablesHooks) {
150 // Pretend some bw_costly_shared_<iface> rules already exist...
151 addIptablesRestoreOutput(
152 "-P OUTPUT ACCEPT\n"
153 "-N bw_costly_rmnet_data0\n"
154 "-N bw_costly_shared\n"
155 "-N unrelated\n"
156 "-N bw_costly_rmnet_data7\n");
157
158 // ... and expect that they be flushed and deleted.
159 std::string expectedCleanCmds =
160 "*filter\n"
161 ":bw_costly_rmnet_data0 -\n"
162 "-X bw_costly_rmnet_data0\n"
163 ":bw_costly_rmnet_data7 -\n"
164 "-X bw_costly_rmnet_data7\n"
165 "COMMIT\n";
166
167 mBw.setupIptablesHooks();
168 expectSetupCommands(expectedCleanCmds, "");
169 }
170
TEST_F(BandwidthControllerTest,TestEnableBandwidthControl)171 TEST_F(BandwidthControllerTest, TestEnableBandwidthControl) {
172 // Pretend no bw_costly_shared_<iface> rules already exist...
173 addIptablesRestoreOutput(
174 "-P OUTPUT ACCEPT\n"
175 "-N bw_costly_shared\n"
176 "-N unrelated\n");
177
178 // ... so none are flushed or deleted.
179 std::string expectedClean = "";
180
181 std::string expectedAccounting =
182 "*filter\n"
183 "-A bw_INPUT -m owner --socket-exists\n"
184 "-A bw_OUTPUT -m owner --socket-exists\n"
185 "-A bw_costly_shared --jump bw_penalty_box\n"
186 "-A bw_penalty_box --jump bw_happy_box\n"
187 "-A bw_happy_box --jump bw_data_saver\n"
188 "-A bw_data_saver -j RETURN\n"
189 "-I bw_happy_box -m owner --uid-owner 0-9999 --jump RETURN\n"
190 "COMMIT\n"
191 "*raw\n"
192 "-A bw_raw_PREROUTING -m owner --socket-exists\n"
193 "COMMIT\n"
194 "*mangle\n"
195 "-A bw_mangle_POSTROUTING -m owner --socket-exists\n"
196 "COMMIT\n";
197
198 mBw.enableBandwidthControl(false);
199 expectSetupCommands(expectedClean, expectedAccounting);
200 }
201
TEST_F(BandwidthControllerTest,TestDisableBandwidthControl)202 TEST_F(BandwidthControllerTest, TestDisableBandwidthControl) {
203 // Pretend some bw_costly_shared_<iface> rules already exist...
204 addIptablesRestoreOutput(
205 "-P OUTPUT ACCEPT\n"
206 "-N bw_costly_rmnet_data0\n"
207 "-N bw_costly_shared\n"
208 "-N unrelated\n"
209 "-N bw_costly_rmnet_data7\n");
210
211 // ... and expect that they be flushed.
212 std::string expectedCleanCmds =
213 "*filter\n"
214 ":bw_costly_rmnet_data0 -\n"
215 ":bw_costly_rmnet_data7 -\n"
216 "COMMIT\n";
217
218 mBw.disableBandwidthControl();
219 expectSetupCommands(expectedCleanCmds, "");
220 }
221
TEST_F(BandwidthControllerTest,TestEnableDataSaver)222 TEST_F(BandwidthControllerTest, TestEnableDataSaver) {
223 mBw.enableDataSaver(true);
224 std::string expected4 =
225 "*filter\n"
226 ":bw_data_saver -\n"
227 "-A bw_data_saver --jump REJECT\n"
228 "COMMIT\n";
229 std::string expected6 =
230 "*filter\n"
231 ":bw_data_saver -\n"
232 "-A bw_data_saver -p icmpv6 --icmpv6-type packet-too-big -j RETURN\n"
233 "-A bw_data_saver -p icmpv6 --icmpv6-type router-solicitation -j RETURN\n"
234 "-A bw_data_saver -p icmpv6 --icmpv6-type router-advertisement -j RETURN\n"
235 "-A bw_data_saver -p icmpv6 --icmpv6-type neighbour-solicitation -j RETURN\n"
236 "-A bw_data_saver -p icmpv6 --icmpv6-type neighbour-advertisement -j RETURN\n"
237 "-A bw_data_saver -p icmpv6 --icmpv6-type redirect -j RETURN\n"
238 "-A bw_data_saver --jump REJECT\n"
239 "COMMIT\n";
240 expectIptablesRestoreCommands({
241 {V4, expected4},
242 {V6, expected6},
243 });
244
245 mBw.enableDataSaver(false);
246 std::string expected = {
247 "*filter\n"
248 ":bw_data_saver -\n"
249 "-A bw_data_saver --jump RETURN\n"
250 "COMMIT\n"
251 };
252 expectIptablesRestoreCommands({
253 {V4, expected},
254 {V6, expected},
255 });
256 }
257
258 std::string kIPv4TetherCounters = Join(std::vector<std::string> {
259 "Chain natctrl_tether_counters (4 references)",
260 " pkts bytes target prot opt in out source destination",
261 " 26 2373 RETURN all -- wlan0 rmnet0 0.0.0.0/0 0.0.0.0/0",
262 " 27 2002 RETURN all -- rmnet0 wlan0 0.0.0.0/0 0.0.0.0/0",
263 " 1040 107471 RETURN all -- bt-pan rmnet0 0.0.0.0/0 0.0.0.0/0",
264 " 1450 1708806 RETURN all -- rmnet0 bt-pan 0.0.0.0/0 0.0.0.0/0",
265 }, '\n');
266
267 std::string kIPv6TetherCounters = Join(std::vector<std::string> {
268 "Chain natctrl_tether_counters (2 references)",
269 " pkts bytes target prot opt in out source destination",
270 " 10000 10000000 RETURN all wlan0 rmnet0 ::/0 ::/0",
271 " 20000 20000000 RETURN all rmnet0 wlan0 ::/0 ::/0",
272 }, '\n');
273
readSocketClientResponse(int fd)274 std::string readSocketClientResponse(int fd) {
275 char buf[32768];
276 ssize_t bytesRead = read(fd, buf, sizeof(buf));
277 if (bytesRead < 0) {
278 return "";
279 }
280 for (int i = 0; i < bytesRead; i++) {
281 if (buf[i] == '\0') buf[i] = '\n';
282 }
283 return std::string(buf, bytesRead);
284 }
285
expectNoSocketClientResponse(int fd)286 void expectNoSocketClientResponse(int fd) {
287 char buf[64];
288 EXPECT_EQ(-1, read(fd, buf, sizeof(buf)));
289 }
290
TEST_F(BandwidthControllerTest,TestGetTetherStats)291 TEST_F(BandwidthControllerTest, TestGetTetherStats) {
292 int socketPair[2];
293 ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, socketPair));
294 ASSERT_EQ(0, fcntl(socketPair[0], F_SETFL, O_NONBLOCK | fcntl(socketPair[0], F_GETFL)));
295 ASSERT_EQ(0, fcntl(socketPair[1], F_SETFL, O_NONBLOCK | fcntl(socketPair[1], F_GETFL)));
296 SocketClient cli(socketPair[0], false);
297
298 std::string err;
299 BandwidthController::TetherStats filter;
300
301 // If no filter is specified, both IPv4 and IPv6 counters must have at least one interface pair.
302 addIptablesRestoreOutput(kIPv4TetherCounters);
303 ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
304 expectNoSocketClientResponse(socketPair[1]);
305 clearIptablesRestoreOutput();
306
307 addIptablesRestoreOutput(kIPv6TetherCounters);
308 ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
309 clearIptablesRestoreOutput();
310
311 // IPv4 and IPv6 counters are properly added together.
312 addIptablesRestoreOutput(kIPv4TetherCounters, kIPv6TetherCounters);
313 filter = BandwidthController::TetherStats();
314 std::string expected =
315 "114 wlan0 rmnet0 10002373 10026 20002002 20027\n"
316 "114 bt-pan rmnet0 107471 1040 1708806 1450\n"
317 "200 Tethering stats list completed\n";
318 ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
319 ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
320 expectNoSocketClientResponse(socketPair[1]);
321 clearIptablesRestoreOutput();
322
323 // Test filtering.
324 addIptablesRestoreOutput(kIPv4TetherCounters, kIPv6TetherCounters);
325 filter = BandwidthController::TetherStats("bt-pan", "rmnet0", -1, -1, -1, -1);
326 expected = "221 bt-pan rmnet0 107471 1040 1708806 1450\n";
327 ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
328 ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
329 expectNoSocketClientResponse(socketPair[1]);
330 clearIptablesRestoreOutput();
331
332 addIptablesRestoreOutput(kIPv4TetherCounters, kIPv6TetherCounters);
333 filter = BandwidthController::TetherStats("wlan0", "rmnet0", -1, -1, -1, -1);
334 expected = "221 wlan0 rmnet0 10002373 10026 20002002 20027\n";
335 ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
336 ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
337 clearIptablesRestoreOutput();
338
339 // Select nonexistent interfaces.
340 addIptablesRestoreOutput(kIPv4TetherCounters, kIPv6TetherCounters);
341 filter = BandwidthController::TetherStats("rmnet0", "foo0", -1, -1, -1, -1);
342 expected = "200 Tethering stats list completed\n";
343 ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
344 ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
345 clearIptablesRestoreOutput();
346
347 // No stats with a filter: no error.
348 addIptablesRestoreOutput("", "");
349 ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
350 ASSERT_EQ("200 Tethering stats list completed\n", readSocketClientResponse(socketPair[1]));
351 clearIptablesRestoreOutput();
352
353 addIptablesRestoreOutput("foo", "foo");
354 ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
355 ASSERT_EQ("200 Tethering stats list completed\n", readSocketClientResponse(socketPair[1]));
356 clearIptablesRestoreOutput();
357
358 // No stats and empty filter: error.
359 filter = BandwidthController::TetherStats();
360 addIptablesRestoreOutput("", kIPv6TetherCounters);
361 ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
362 expectNoSocketClientResponse(socketPair[1]);
363 clearIptablesRestoreOutput();
364
365 addIptablesRestoreOutput(kIPv4TetherCounters, "");
366 ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
367 expectNoSocketClientResponse(socketPair[1]);
368 clearIptablesRestoreOutput();
369
370 // Include only one pair of interfaces and things are fine.
371 std::vector<std::string> counterLines = android::base::Split(kIPv4TetherCounters, "\n");
372 std::vector<std::string> brokenCounterLines = counterLines;
373 counterLines.resize(4);
374 std::string counters = Join(counterLines, "\n") + "\n";
375 addIptablesRestoreOutput(counters, counters);
376 expected =
377 "114 wlan0 rmnet0 4746 52 4004 54\n"
378 "200 Tethering stats list completed\n";
379 ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
380 ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
381 clearIptablesRestoreOutput();
382
383 // But if interfaces aren't paired, it's always an error.
384 err = "";
385 counterLines.resize(3);
386 counters = Join(counterLines, "\n") + "\n";
387 addIptablesRestoreOutput(counters, counters);
388 ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
389 expectNoSocketClientResponse(socketPair[1]);
390 clearIptablesRestoreOutput();
391
392 // Token unit test of the fact that we return the stats in the error message which the caller
393 // ignores.
394 std::string expectedError = counters;
395 EXPECT_EQ(expectedError, err);
396
397 addIptablesRestoreOutput(kIPv4TetherCounters);
398 ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
399 expectNoSocketClientResponse(socketPair[1]);
400 clearIptablesRestoreOutput();
401 addIptablesRestoreOutput(kIPv6TetherCounters);
402 ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
403 expectNoSocketClientResponse(socketPair[1]);
404 clearIptablesRestoreOutput();
405 }
406
makeInterfaceQuotaCommands(const std::string & iface,int ruleIndex,int64_t quota)407 const std::vector<std::string> makeInterfaceQuotaCommands(const std::string& iface, int ruleIndex,
408 int64_t quota) {
409 const std::string chain = "bw_costly_" + iface;
410 const char* c_chain = chain.c_str();
411 const char* c_iface = iface.c_str();
412 std::vector<std::string> cmds = {
413 "*filter",
414 StringPrintf(":%s -", c_chain),
415 StringPrintf("-A %s -j bw_penalty_box", c_chain),
416 StringPrintf("-I bw_INPUT %d -i %s --jump %s", ruleIndex, c_iface, c_chain),
417 StringPrintf("-I bw_OUTPUT %d -o %s --jump %s", ruleIndex, c_iface, c_chain),
418 StringPrintf("-A bw_FORWARD -o %s --jump %s", c_iface, c_chain),
419 StringPrintf("-A %s -m quota2 ! --quota %" PRIu64 " --name %s --jump REJECT", c_chain,
420 quota, c_iface),
421 "COMMIT\n",
422 };
423 return {Join(cmds, "\n")};
424 }
425
removeInterfaceQuotaCommands(const std::string & iface)426 const std::vector<std::string> removeInterfaceQuotaCommands(const std::string& iface) {
427 const std::string chain = "bw_costly_" + iface;
428 const char* c_chain = chain.c_str();
429 const char* c_iface = iface.c_str();
430 std::vector<std::string> cmds = {
431 "*filter",
432 StringPrintf("-D bw_INPUT -i %s --jump %s", c_iface, c_chain),
433 StringPrintf("-D bw_OUTPUT -o %s --jump %s", c_iface, c_chain),
434 StringPrintf("-D bw_FORWARD -o %s --jump %s", c_iface, c_chain),
435 StringPrintf("-F %s", c_chain),
436 StringPrintf("-X %s", c_chain),
437 "COMMIT\n",
438 };
439 return {Join(cmds, "\n")};
440 }
441
TEST_F(BandwidthControllerTest,TestSetInterfaceQuota)442 TEST_F(BandwidthControllerTest, TestSetInterfaceQuota) {
443 constexpr uint64_t kOldQuota = 123456;
444 const std::string iface = mTun.name();
445 std::vector<std::string> expected = makeInterfaceQuotaCommands(iface, 1, kOldQuota);
446
447 EXPECT_EQ(0, mBw.setInterfaceQuota(iface, kOldQuota));
448 expectIptablesRestoreCommands(expected);
449
450 constexpr uint64_t kNewQuota = kOldQuota + 1;
451 expected = {};
452 expectUpdateQuota(kNewQuota);
453 EXPECT_EQ(0, mBw.setInterfaceQuota(iface, kNewQuota));
454 expectIptablesRestoreCommands(expected);
455
456 expected = removeInterfaceQuotaCommands(iface);
457 EXPECT_EQ(0, mBw.removeInterfaceQuota(iface));
458 expectIptablesRestoreCommands(expected);
459 }
460
makeInterfaceSharedQuotaCommands(const std::string & iface,int ruleIndex,int64_t quota,bool insertQuota)461 const std::vector<std::string> makeInterfaceSharedQuotaCommands(const std::string& iface,
462 int ruleIndex, int64_t quota,
463 bool insertQuota) {
464 const std::string chain = "bw_costly_shared";
465 const char* c_chain = chain.c_str();
466 const char* c_iface = iface.c_str();
467 std::vector<std::string> cmds = {
468 "*filter",
469 StringPrintf("-I bw_INPUT %d -i %s --jump %s", ruleIndex, c_iface, c_chain),
470 StringPrintf("-I bw_OUTPUT %d -o %s --jump %s", ruleIndex, c_iface, c_chain),
471 StringPrintf("-A bw_FORWARD -o %s --jump %s", c_iface, c_chain),
472 };
473 if (insertQuota) {
474 cmds.push_back(StringPrintf(
475 "-I %s -m quota2 ! --quota %" PRIu64 " --name shared --jump REJECT", c_chain, quota));
476 }
477 cmds.push_back("COMMIT\n");
478 return {Join(cmds, "\n")};
479 }
480
removeInterfaceSharedQuotaCommands(const std::string & iface,int64_t quota,bool deleteQuota)481 const std::vector<std::string> removeInterfaceSharedQuotaCommands(const std::string& iface,
482 int64_t quota, bool deleteQuota) {
483 const std::string chain = "bw_costly_shared";
484 const char* c_chain = chain.c_str();
485 const char* c_iface = iface.c_str();
486 std::vector<std::string> cmds = {
487 "*filter",
488 StringPrintf("-D bw_INPUT -i %s --jump %s", c_iface, c_chain),
489 StringPrintf("-D bw_OUTPUT -o %s --jump %s", c_iface, c_chain),
490 StringPrintf("-D bw_FORWARD -o %s --jump %s", c_iface, c_chain),
491 };
492 if (deleteQuota) {
493 cmds.push_back(StringPrintf(
494 "-D %s -m quota2 ! --quota %" PRIu64 " --name shared --jump REJECT", c_chain, quota));
495 }
496 cmds.push_back("COMMIT\n");
497 return {Join(cmds, "\n")};
498 }
499
TEST_F(BandwidthControllerTest,TestSetInterfaceSharedQuotaDuplicate)500 TEST_F(BandwidthControllerTest, TestSetInterfaceSharedQuotaDuplicate) {
501 constexpr uint64_t kQuota = 123456;
502 const std::string iface = mTun.name();
503 std::vector<std::string> expected = makeInterfaceSharedQuotaCommands(iface, 1, 123456, true);
504 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kQuota));
505 expectIptablesRestoreCommands(expected);
506
507 expected = {};
508 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kQuota));
509 expectIptablesRestoreCommands(expected);
510
511 expected = removeInterfaceSharedQuotaCommands(iface, kQuota, true);
512 EXPECT_EQ(0, mBw.removeInterfaceSharedQuota(iface));
513 expectIptablesRestoreCommands(expected);
514 }
515
TEST_F(BandwidthControllerTest,TestSetInterfaceSharedQuotaUpdate)516 TEST_F(BandwidthControllerTest, TestSetInterfaceSharedQuotaUpdate) {
517 constexpr uint64_t kOldQuota = 123456;
518 const std::string iface = mTun.name();
519 std::vector<std::string> expected = makeInterfaceSharedQuotaCommands(iface, 1, kOldQuota, true);
520 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kOldQuota));
521 expectIptablesRestoreCommands(expected);
522
523 constexpr uint64_t kNewQuota = kOldQuota + 1;
524 expected = {};
525 expectUpdateQuota(kNewQuota);
526 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kNewQuota));
527 expectIptablesRestoreCommands(expected);
528
529 expected = removeInterfaceSharedQuotaCommands(iface, kNewQuota, true);
530 EXPECT_EQ(0, mBw.removeInterfaceSharedQuota(iface));
531 expectIptablesRestoreCommands(expected);
532 }
533
TEST_F(BandwidthControllerTest,TestSetInterfaceSharedQuotaTwoInterfaces)534 TEST_F(BandwidthControllerTest, TestSetInterfaceSharedQuotaTwoInterfaces) {
535 constexpr uint64_t kQuota = 123456;
536 const std::vector<std::string> ifaces{
537 {"a" + mTun.name()},
538 {"b" + mTun.name()},
539 };
540
541 for (const auto& iface : ifaces) {
542 // Quota rule is only added when the total number of
543 // interfaces transitions from 0 -> 1.
544 bool first = (iface == ifaces[0]);
545 auto expected = makeInterfaceSharedQuotaCommands(iface, 1, kQuota, first);
546 EXPECT_EQ(0, mBw.setInterfaceSharedQuota(iface, kQuota));
547 expectIptablesRestoreCommands(expected);
548 }
549
550 for (const auto& iface : ifaces) {
551 // Quota rule is only removed when the total number of
552 // interfaces transitions from 1 -> 0.
553 bool last = (iface == ifaces[1]);
554 auto expected = removeInterfaceSharedQuotaCommands(iface, kQuota, last);
555 EXPECT_EQ(0, mBw.removeInterfaceSharedQuota(iface));
556 expectIptablesRestoreCommands(expected);
557 }
558 }
559
TEST_F(BandwidthControllerTest,IptablesAlertCmd)560 TEST_F(BandwidthControllerTest, IptablesAlertCmd) {
561 std::vector<std::string> expected = {
562 "*filter\n"
563 "-I bw_INPUT -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
564 "-I bw_OUTPUT -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
565 "COMMIT\n"
566 };
567 EXPECT_EQ(0, runIptablesAlertCmd(IptOp::IptOpInsert, "MyWonderfulAlert", 123456));
568 expectIptablesRestoreCommands(expected);
569
570 expected = {
571 "*filter\n"
572 "-D bw_INPUT -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
573 "-D bw_OUTPUT -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
574 "COMMIT\n"
575 };
576 EXPECT_EQ(0, runIptablesAlertCmd(IptOp::IptOpDelete, "MyWonderfulAlert", 123456));
577 expectIptablesRestoreCommands(expected);
578 }
579
TEST_F(BandwidthControllerTest,IptablesAlertFwdCmd)580 TEST_F(BandwidthControllerTest, IptablesAlertFwdCmd) {
581 std::vector<std::string> expected = {
582 "*filter\n"
583 "-I bw_FORWARD -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
584 "COMMIT\n"
585 };
586 EXPECT_EQ(0, runIptablesAlertFwdCmd(IptOp::IptOpInsert, "MyWonderfulAlert", 123456));
587 expectIptablesRestoreCommands(expected);
588
589 expected = {
590 "*filter\n"
591 "-D bw_FORWARD -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
592 "COMMIT\n"
593 };
594 EXPECT_EQ(0, runIptablesAlertFwdCmd(IptOp::IptOpDelete, "MyWonderfulAlert", 123456));
595 expectIptablesRestoreCommands(expected);
596 }
597
TEST_F(BandwidthControllerTest,CostlyAlert)598 TEST_F(BandwidthControllerTest, CostlyAlert) {
599 const int64_t kQuota = 123456;
600 int64_t alertBytes = 0;
601
602 std::vector<std::string> expected = {
603 "*filter\n"
604 "-A bw_costly_shared -m quota2 ! --quota 123456 --name sharedAlert\n"
605 "COMMIT\n"
606 };
607 EXPECT_EQ(0, setCostlyAlert("shared", kQuota, &alertBytes));
608 EXPECT_EQ(kQuota, alertBytes);
609 expectIptablesRestoreCommands(expected);
610
611 expected = {};
612 expectUpdateQuota(kQuota);
613 EXPECT_EQ(0, setCostlyAlert("shared", kQuota + 1, &alertBytes));
614 EXPECT_EQ(kQuota + 1, alertBytes);
615 expectIptablesRestoreCommands(expected);
616
617 expected = {
618 "*filter\n"
619 "-D bw_costly_shared -m quota2 ! --quota 123457 --name sharedAlert\n"
620 "COMMIT\n"
621 };
622 EXPECT_EQ(0, removeCostlyAlert("shared", &alertBytes));
623 EXPECT_EQ(0, alertBytes);
624 expectIptablesRestoreCommands(expected);
625 }
626
TEST_F(BandwidthControllerTest,ManipulateSpecialApps)627 TEST_F(BandwidthControllerTest, ManipulateSpecialApps) {
628 std::vector<const char *> appUids = { "1000", "1001", "10012" };
629
630 std::vector<std::string> expected = {
631 "*filter\n"
632 "-I bw_happy_box -m owner --uid-owner 1000 --jump RETURN\n"
633 "-I bw_happy_box -m owner --uid-owner 1001 --jump RETURN\n"
634 "-I bw_happy_box -m owner --uid-owner 10012 --jump RETURN\n"
635 "COMMIT\n"
636 };
637 EXPECT_EQ(0, mBw.addNiceApps(appUids.size(), const_cast<char**>(&appUids[0])));
638 expectIptablesRestoreCommands(expected);
639
640 expected = {
641 "*filter\n"
642 "-D bw_penalty_box -m owner --uid-owner 1000 --jump REJECT\n"
643 "-D bw_penalty_box -m owner --uid-owner 1001 --jump REJECT\n"
644 "-D bw_penalty_box -m owner --uid-owner 10012 --jump REJECT\n"
645 "COMMIT\n"
646 };
647 EXPECT_EQ(0, mBw.removeNaughtyApps(appUids.size(), const_cast<char**>(&appUids[0])));
648 expectIptablesRestoreCommands(expected);
649 }
650