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: 53 BandwidthControllerTest() { 54 BandwidthController::iptablesRestoreFunction = fakeExecIptablesRestoreWithOutput; 55 } 56 BandwidthController mBw; 57 TunInterface mTun; 58 59 void SetUp() { 60 ASSERT_EQ(0, mTun.init()); 61 } 62 63 void TearDown() { 64 mTun.destroy(); 65 } 66 67 void addIptablesRestoreOutput(std::string contents) { 68 sIptablesRestoreOutput.push_back(contents); 69 } 70 71 void addIptablesRestoreOutput(std::string contents1, std::string contents2) { 72 sIptablesRestoreOutput.push_back(contents1); 73 sIptablesRestoreOutput.push_back(contents2); 74 } 75 76 void clearIptablesRestoreOutput() { 77 sIptablesRestoreOutput.clear(); 78 } 79 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 117 int runIptablesAlertCmd(IptOp a, const char *b, int64_t c) { 118 return mBw.runIptablesAlertCmd(a, b, c); 119 } 120 121 int runIptablesAlertFwdCmd(IptOp a, const char *b, int64_t c) { 122 return mBw.runIptablesAlertFwdCmd(a, b, c); 123 } 124 125 int setCostlyAlert(const std::string a, int64_t b, int64_t *c) { 126 return mBw.setCostlyAlert(a, b, c); 127 } 128 129 int removeCostlyAlert(const std::string a, int64_t *b) { 130 return mBw.removeCostlyAlert(a, b); 131 } 132 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 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 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 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 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 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 286 void expectNoSocketClientResponse(int fd) { 287 char buf[64]; 288 EXPECT_EQ(-1, read(fd, buf, sizeof(buf))); 289 } 290 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 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 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 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 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 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 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 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 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 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 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 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 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
