Home | History | Annotate | Download | only in server 
      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 "BandwidthController.h"
     34 #include "IptablesBaseTest.h"
     35 #include "tun_interface.h"
     36 
     37 using android::base::StringPrintf;
     38 using android::net::TunInterface;
     39 
     40 class BandwidthControllerTest : public IptablesBaseTest {
     41 public:
     42     BandwidthControllerTest() {
     43         BandwidthController::execFunction = fake_android_fork_exec;
     44         BandwidthController::popenFunction = fake_popen;
     45         BandwidthController::iptablesRestoreFunction = fakeExecIptablesRestoreWithOutput;
     46     }
     47     BandwidthController mBw;
     48     TunInterface mTun;
     49 
     50     void SetUp() {
     51         ASSERT_EQ(0, mTun.init());
     52     }
     53 
     54     void TearDown() {
     55         mTun.destroy();
     56     }
     57 
     58     void addIptablesRestoreOutput(std::string contents) {
     59         sIptablesRestoreOutput.push_back(contents);
     60     }
     61 
     62     void addIptablesRestoreOutput(std::string contents1, std::string contents2) {
     63         sIptablesRestoreOutput.push_back(contents1);
     64         sIptablesRestoreOutput.push_back(contents2);
     65     }
     66 
     67     void clearIptablesRestoreOutput() {
     68         sIptablesRestoreOutput.clear();
     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 
    117 TEST_F(BandwidthControllerTest, TestSetupIptablesHooks) {
    118     // Pretend some bw_costly_shared_<iface> rules already exist...
    119     addIptablesRestoreOutput(
    120         "-P OUTPUT ACCEPT\n"
    121         "-N bw_costly_rmnet_data0\n"
    122         "-N bw_costly_shared\n"
    123         "-N unrelated\n"
    124         "-N bw_costly_rmnet_data7\n");
    125 
    126     // ... and expect that they be flushed and deleted.
    127     std::string expectedCleanCmds =
    128         "*filter\n"
    129         ":bw_costly_rmnet_data0 -\n"
    130         "-X bw_costly_rmnet_data0\n"
    131         ":bw_costly_rmnet_data7 -\n"
    132         "-X bw_costly_rmnet_data7\n"
    133         "COMMIT\n";
    134 
    135     mBw.setupIptablesHooks();
    136     expectSetupCommands(expectedCleanCmds, "");
    137 }
    138 
    139 TEST_F(BandwidthControllerTest, TestEnableBandwidthControl) {
    140     // Pretend no bw_costly_shared_<iface> rules already exist...
    141     addIptablesRestoreOutput(
    142         "-P OUTPUT ACCEPT\n"
    143         "-N bw_costly_shared\n"
    144         "-N unrelated\n");
    145 
    146     // ... so none are flushed or deleted.
    147     std::string expectedClean = "";
    148 
    149     std::string expectedAccounting =
    150         "*filter\n"
    151         "-A bw_INPUT -m owner --socket-exists\n"
    152         "-A bw_OUTPUT -m owner --socket-exists\n"
    153         "-A bw_costly_shared --jump bw_penalty_box\n"
    154         "-A bw_penalty_box --jump bw_happy_box\n"
    155         "-A bw_happy_box --jump bw_data_saver\n"
    156         "-A bw_data_saver -j RETURN\n"
    157         "-I bw_happy_box -m owner --uid-owner 0-9999 --jump RETURN\n"
    158         "COMMIT\n"
    159         "*raw\n"
    160         "-A bw_raw_PREROUTING -m owner --socket-exists\n"
    161         "COMMIT\n"
    162         "*mangle\n"
    163         "-A bw_mangle_POSTROUTING -m owner --socket-exists\n"
    164         "COMMIT\n";
    165 
    166     mBw.enableBandwidthControl(false);
    167     expectSetupCommands(expectedClean, expectedAccounting);
    168 }
    169 
    170 TEST_F(BandwidthControllerTest, TestDisableBandwidthControl) {
    171     // Pretend some bw_costly_shared_<iface> rules already exist...
    172     addIptablesRestoreOutput(
    173         "-P OUTPUT ACCEPT\n"
    174         "-N bw_costly_rmnet_data0\n"
    175         "-N bw_costly_shared\n"
    176         "-N unrelated\n"
    177         "-N bw_costly_rmnet_data7\n");
    178 
    179     // ... and expect that they be flushed.
    180     std::string expectedCleanCmds =
    181         "*filter\n"
    182         ":bw_costly_rmnet_data0 -\n"
    183         ":bw_costly_rmnet_data7 -\n"
    184         "COMMIT\n";
    185 
    186     mBw.disableBandwidthControl();
    187     expectSetupCommands(expectedCleanCmds, "");
    188 }
    189 
    190 TEST_F(BandwidthControllerTest, TestEnableDataSaver) {
    191     mBw.enableDataSaver(true);
    192     std::vector<std::string> expected = {
    193         "*filter\n"
    194         "-R bw_data_saver 1 --jump REJECT\n"
    195         "COMMIT\n"
    196     };
    197     expectIptablesRestoreCommands(expected);
    198 
    199     mBw.enableDataSaver(false);
    200     expected = {
    201         "*filter\n"
    202         "-R bw_data_saver 1 --jump RETURN\n"
    203         "COMMIT\n"
    204     };
    205     expectIptablesRestoreCommands(expected);
    206 }
    207 
    208 std::string kIPv4TetherCounters = android::base::Join(std::vector<std::string> {
    209     "Chain natctrl_tether_counters (4 references)",
    210     "    pkts      bytes target     prot opt in     out     source               destination",
    211     "      26     2373 RETURN     all  --  wlan0  rmnet0  0.0.0.0/0            0.0.0.0/0",
    212     "      27     2002 RETURN     all  --  rmnet0 wlan0   0.0.0.0/0            0.0.0.0/0",
    213     "    1040   107471 RETURN     all  --  bt-pan rmnet0  0.0.0.0/0            0.0.0.0/0",
    214     "    1450  1708806 RETURN     all  --  rmnet0 bt-pan  0.0.0.0/0            0.0.0.0/0",
    215 }, '\n');
    216 
    217 std::string kIPv6TetherCounters = android::base::Join(std::vector<std::string> {
    218     "Chain natctrl_tether_counters (2 references)",
    219     "    pkts      bytes target     prot opt in     out     source               destination",
    220     "   10000 10000000 RETURN     all      wlan0  rmnet0  ::/0                 ::/0",
    221     "   20000 20000000 RETURN     all      rmnet0 wlan0   ::/0                 ::/0",
    222 }, '\n');
    223 
    224 std::string readSocketClientResponse(int fd) {
    225     char buf[32768];
    226     ssize_t bytesRead = read(fd, buf, sizeof(buf));
    227     if (bytesRead < 0) {
    228         return "";
    229     }
    230     for (int i = 0; i < bytesRead; i++) {
    231         if (buf[i] == '\0') buf[i] = '\n';
    232     }
    233     return std::string(buf, bytesRead);
    234 }
    235 
    236 void expectNoSocketClientResponse(int fd) {
    237     char buf[64];
    238     EXPECT_EQ(-1, read(fd, buf, sizeof(buf)));
    239 }
    240 
    241 TEST_F(BandwidthControllerTest, TestGetTetherStats) {
    242     int socketPair[2];
    243     ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, socketPair));
    244     ASSERT_EQ(0, fcntl(socketPair[0], F_SETFL, O_NONBLOCK | fcntl(socketPair[0], F_GETFL)));
    245     ASSERT_EQ(0, fcntl(socketPair[1], F_SETFL, O_NONBLOCK | fcntl(socketPair[1], F_GETFL)));
    246     SocketClient cli(socketPair[0], false);
    247 
    248     std::string err;
    249     BandwidthController::TetherStats filter;
    250 
    251     // If no filter is specified, both IPv4 and IPv6 counters must have at least one interface pair.
    252     addIptablesRestoreOutput(kIPv4TetherCounters);
    253     ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
    254     expectNoSocketClientResponse(socketPair[1]);
    255     clearIptablesRestoreOutput();
    256 
    257     addIptablesRestoreOutput(kIPv6TetherCounters);
    258     ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
    259     clearIptablesRestoreOutput();
    260 
    261     // IPv4 and IPv6 counters are properly added together.
    262     addIptablesRestoreOutput(kIPv4TetherCounters, kIPv6TetherCounters);
    263     filter = BandwidthController::TetherStats();
    264     std::string expected =
    265             "114 wlan0 rmnet0 10002373 10026 20002002 20027\n"
    266             "114 bt-pan rmnet0 107471 1040 1708806 1450\n"
    267             "200 Tethering stats list completed\n";
    268     ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
    269     ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
    270     expectNoSocketClientResponse(socketPair[1]);
    271     clearIptablesRestoreOutput();
    272 
    273     // Test filtering.
    274     addIptablesRestoreOutput(kIPv4TetherCounters, kIPv6TetherCounters);
    275     filter = BandwidthController::TetherStats("bt-pan", "rmnet0", -1, -1, -1, -1);
    276     expected = "221 bt-pan rmnet0 107471 1040 1708806 1450\n";
    277     ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
    278     ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
    279     expectNoSocketClientResponse(socketPair[1]);
    280     clearIptablesRestoreOutput();
    281 
    282     addIptablesRestoreOutput(kIPv4TetherCounters, kIPv6TetherCounters);
    283     filter = BandwidthController::TetherStats("wlan0", "rmnet0", -1, -1, -1, -1);
    284     expected = "221 wlan0 rmnet0 10002373 10026 20002002 20027\n";
    285     ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
    286     ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
    287     clearIptablesRestoreOutput();
    288 
    289     // Select nonexistent interfaces.
    290     addIptablesRestoreOutput(kIPv4TetherCounters, kIPv6TetherCounters);
    291     filter = BandwidthController::TetherStats("rmnet0", "foo0", -1, -1, -1, -1);
    292     expected = "200 Tethering stats list completed\n";
    293     ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
    294     ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
    295     clearIptablesRestoreOutput();
    296 
    297     // No stats with a filter: no error.
    298     addIptablesRestoreOutput("", "");
    299     ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
    300     ASSERT_EQ("200 Tethering stats list completed\n", readSocketClientResponse(socketPair[1]));
    301     clearIptablesRestoreOutput();
    302 
    303     addIptablesRestoreOutput("foo", "foo");
    304     ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
    305     ASSERT_EQ("200 Tethering stats list completed\n", readSocketClientResponse(socketPair[1]));
    306     clearIptablesRestoreOutput();
    307 
    308     // No stats and empty filter: error.
    309     filter = BandwidthController::TetherStats();
    310     addIptablesRestoreOutput("", kIPv6TetherCounters);
    311     ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
    312     expectNoSocketClientResponse(socketPair[1]);
    313     clearIptablesRestoreOutput();
    314 
    315     addIptablesRestoreOutput(kIPv4TetherCounters, "");
    316     ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
    317     expectNoSocketClientResponse(socketPair[1]);
    318     clearIptablesRestoreOutput();
    319 
    320     // Include only one pair of interfaces and things are fine.
    321     std::vector<std::string> counterLines = android::base::Split(kIPv4TetherCounters, "\n");
    322     std::vector<std::string> brokenCounterLines = counterLines;
    323     counterLines.resize(4);
    324     std::string counters = android::base::Join(counterLines, "\n") + "\n";
    325     addIptablesRestoreOutput(counters, counters);
    326     expected =
    327             "114 wlan0 rmnet0 4746 52 4004 54\n"
    328             "200 Tethering stats list completed\n";
    329     ASSERT_EQ(0, mBw.getTetherStats(&cli, filter, err));
    330     ASSERT_EQ(expected, readSocketClientResponse(socketPair[1]));
    331     clearIptablesRestoreOutput();
    332 
    333     // But if interfaces aren't paired, it's always an error.
    334     err = "";
    335     counterLines.resize(3);
    336     counters = android::base::Join(counterLines, "\n") + "\n";
    337     addIptablesRestoreOutput(counters, counters);
    338     ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
    339     expectNoSocketClientResponse(socketPair[1]);
    340     clearIptablesRestoreOutput();
    341 
    342     // Token unit test of the fact that we return the stats in the error message which the caller
    343     // ignores.
    344     std::string expectedError = counters;
    345     EXPECT_EQ(expectedError, err);
    346 
    347     // popen() failing is always an error.
    348     addIptablesRestoreOutput(kIPv4TetherCounters);
    349     ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
    350     expectNoSocketClientResponse(socketPair[1]);
    351     clearIptablesRestoreOutput();
    352     addIptablesRestoreOutput(kIPv6TetherCounters);
    353     ASSERT_EQ(-1, mBw.getTetherStats(&cli, filter, err));
    354     expectNoSocketClientResponse(socketPair[1]);
    355     clearIptablesRestoreOutput();
    356 }
    357 
    358 const std::vector<std::string> makeInterfaceQuotaCommands(const char *iface, int ruleIndex,
    359                                                           int64_t quota) {
    360     std::vector<std::string> cmds = {
    361         StringPrintf("-F bw_costly_%s", iface),
    362         StringPrintf("-N bw_costly_%s", iface),
    363         StringPrintf("-A bw_costly_%s -j bw_penalty_box", iface),
    364         StringPrintf("-D bw_INPUT -i %s --jump bw_costly_%s", iface, iface),
    365         StringPrintf("-I bw_INPUT %d -i %s --jump bw_costly_%s", ruleIndex, iface, iface),
    366         StringPrintf("-D bw_OUTPUT -o %s --jump bw_costly_%s", iface, iface),
    367         StringPrintf("-I bw_OUTPUT %d -o %s --jump bw_costly_%s", ruleIndex, iface, iface),
    368         StringPrintf("-D bw_FORWARD -o %s --jump bw_costly_%s", iface, iface),
    369         StringPrintf("-A bw_FORWARD -o %s --jump bw_costly_%s", iface, iface),
    370         StringPrintf("-A bw_costly_%s -m quota2 ! --quota %" PRIu64 " --name %s --jump REJECT",
    371                      iface, quota, iface),
    372     };
    373     return cmds;
    374 }
    375 
    376 const std::vector<std::string> removeInterfaceQuotaCommands(const char *iface) {
    377     std::vector<std::string> cmds = {
    378         StringPrintf("-D bw_INPUT -i %s --jump bw_costly_%s", iface, iface),
    379         StringPrintf("-D bw_OUTPUT -o %s --jump bw_costly_%s", iface, iface),
    380         StringPrintf("-D bw_FORWARD -o %s --jump bw_costly_%s", iface, iface),
    381         StringPrintf("-F bw_costly_%s", iface),
    382         StringPrintf("-X bw_costly_%s", iface),
    383     };
    384     return cmds;
    385 }
    386 
    387 TEST_F(BandwidthControllerTest, TestSetInterfaceQuota) {
    388     const char *iface = mTun.name().c_str();
    389     std::vector<std::string> expected = makeInterfaceQuotaCommands(iface, 1, 123456);
    390 
    391     // prepCostlyInterface assumes that exactly one of the "-F chain" and "-N chain" commands fails.
    392     // So pretend that the first two commands (the IPv4 -F and the IPv6 -F) fail.
    393     std::deque<int> returnValues(expected.size() * 2, 0);
    394     returnValues[0] = 1;
    395     returnValues[1] = 1;
    396     setReturnValues(returnValues);
    397 
    398     EXPECT_EQ(0, mBw.setInterfaceQuota(iface, 123456));
    399     expectIptablesCommands(expected);
    400 
    401     expected = removeInterfaceQuotaCommands(iface);
    402     EXPECT_EQ(0, mBw.removeInterfaceQuota(iface));
    403     expectIptablesCommands(expected);
    404 }
    405 
    406 TEST_F(BandwidthControllerTest, IptablesAlertCmd) {
    407     std::vector<std::string> expected = {
    408         "*filter\n"
    409         "-I bw_INPUT -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
    410         "-I bw_OUTPUT -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
    411         "COMMIT\n"
    412     };
    413     EXPECT_EQ(0, runIptablesAlertCmd(IptOp::IptOpInsert, "MyWonderfulAlert", 123456));
    414     expectIptablesRestoreCommands(expected);
    415 
    416     expected = {
    417         "*filter\n"
    418         "-D bw_INPUT -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
    419         "-D bw_OUTPUT -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
    420         "COMMIT\n"
    421     };
    422     EXPECT_EQ(0, runIptablesAlertCmd(IptOp::IptOpDelete, "MyWonderfulAlert", 123456));
    423     expectIptablesRestoreCommands(expected);
    424 }
    425 
    426 TEST_F(BandwidthControllerTest, IptablesAlertFwdCmd) {
    427     std::vector<std::string> expected = {
    428         "*filter\n"
    429         "-I bw_FORWARD -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
    430         "COMMIT\n"
    431     };
    432     EXPECT_EQ(0, runIptablesAlertFwdCmd(IptOp::IptOpInsert, "MyWonderfulAlert", 123456));
    433     expectIptablesRestoreCommands(expected);
    434 
    435     expected = {
    436         "*filter\n"
    437         "-D bw_FORWARD -m quota2 ! --quota 123456 --name MyWonderfulAlert\n"
    438         "COMMIT\n"
    439     };
    440     EXPECT_EQ(0, runIptablesAlertFwdCmd(IptOp::IptOpDelete, "MyWonderfulAlert", 123456));
    441     expectIptablesRestoreCommands(expected);
    442 }
    443 
    444 TEST_F(BandwidthControllerTest, ManipulateSpecialApps) {
    445     std::vector<const char *> appUids = { "1000", "1001", "10012" };
    446 
    447     std::vector<std::string> expected = {
    448         "*filter\n"
    449         "-I bw_happy_box -m owner --uid-owner 1000 --jump RETURN\n"
    450         "-I bw_happy_box -m owner --uid-owner 1001 --jump RETURN\n"
    451         "-I bw_happy_box -m owner --uid-owner 10012 --jump RETURN\n"
    452         "COMMIT\n"
    453     };
    454     EXPECT_EQ(0, mBw.addNiceApps(appUids.size(), const_cast<char**>(&appUids[0])));
    455     expectIptablesRestoreCommands(expected);
    456 
    457     expected = {
    458         "*filter\n"
    459         "-D bw_penalty_box -m owner --uid-owner 1000 --jump REJECT\n"
    460         "-D bw_penalty_box -m owner --uid-owner 1001 --jump REJECT\n"
    461         "-D bw_penalty_box -m owner --uid-owner 10012 --jump REJECT\n"
    462         "COMMIT\n"
    463     };
    464     EXPECT_EQ(0, mBw.removeNaughtyApps(appUids.size(), const_cast<char**>(&appUids[0])));
    465     expectIptablesRestoreCommands(expected);
    466 }
    467