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