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