server/BandwidthControllerTest.cpp

   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 <gtest/gtest.h>
  23
  24 #include <android-base/strings.h>
  25
  26 #include "BandwidthController.h"
  27 #include "IptablesBaseTest.h"
  28
  29 FILE *fake_popen(const char *, const char *) {
  30     return NULL;
  31 };
  32
  33 class BandwidthControllerTest : public IptablesBaseTest {
  34 public:
  35     BandwidthControllerTest() {
  36         BandwidthController::execFunction = fake_android_fork_exec;
  37         BandwidthController::popenFunction = fake_popen;
  38         BandwidthController::iptablesRestoreFunction = fakeExecIptablesRestore;
  39     }
  40     BandwidthController mBw;
  41 };
  42
  43 TEST_F(BandwidthControllerTest, TestSetupIptablesHooks) {
  44     mBw.setupIptablesHooks();
  45     std::vector<std::string> expected = {
  46         "*filter\n"
  47         ":bw_INPUT -\n"
  48         ":bw_OUTPUT -\n"
  49         ":bw_FORWARD -\n"
  50         ":bw_happy_box -\n"
  51         ":bw_penalty_box -\n"
  52         ":bw_data_saver -\n"
  53         ":bw_costly_shared -\n"
  54         "COMMIT\n"
  55         "*raw\n"
  56         ":bw_raw_PREROUTING -\n"
  57         "COMMIT\n"
  58         "*mangle\n"
  59         ":bw_mangle_POSTROUTING -\n"
  60         "COMMIT\n\x04"
  61     };
  62     expectIptablesRestoreCommands(expected);
  63 }
  64
  65 TEST_F(BandwidthControllerTest, TestEnableBandwidthControl) {
  66     mBw.enableBandwidthControl(false);
  67     std::string expectedFlush =
  68         "*filter\n"
  69         ":bw_INPUT -\n"
  70         ":bw_OUTPUT -\n"
  71         ":bw_FORWARD -\n"
  72         ":bw_happy_box -\n"
  73         ":bw_penalty_box -\n"
  74         ":bw_data_saver -\n"
  75         ":bw_costly_shared -\n"
  76         "COMMIT\n"
  77         "*raw\n"
  78         ":bw_raw_PREROUTING -\n"
  79         "COMMIT\n"
  80         "*mangle\n"
  81         ":bw_mangle_POSTROUTING -\n"
  82         "COMMIT\n\x04";
  83      std::string expectedAccounting =
  84         "*filter\n"
  85         "-A bw_INPUT -m owner --socket-exists\n"
  86         "-A bw_OUTPUT -m owner --socket-exists\n"
  87         "-A bw_costly_shared --jump bw_penalty_box\n"
  88         "-A bw_penalty_box --jump bw_happy_box\n"
  89         "-A bw_happy_box --jump bw_data_saver\n"
  90         "-A bw_data_saver -j RETURN\n"
  91         "-I bw_happy_box -m owner --uid-owner 0-9999 --jump RETURN\n"
  92         "COMMIT\n"
  93         "*raw\n"
  94         "-A bw_raw_PREROUTING -m owner --socket-exists\n"
  95         "COMMIT\n"
  96         "*mangle\n"
  97         "-A bw_mangle_POSTROUTING -m owner --socket-exists\n"
  98         "COMMIT\n\x04";
  99
 100     expectIptablesRestoreCommands({ expectedFlush, expectedAccounting });
 101 }
 102
 103 TEST_F(BandwidthControllerTest, TestDisableBandwidthControl) {
 104     mBw.disableBandwidthControl();
 105     const std::string expected =
 106         "*filter\n"
 107         ":bw_INPUT -\n"
 108         ":bw_OUTPUT -\n"
 109         ":bw_FORWARD -\n"
 110         ":bw_happy_box -\n"
 111         ":bw_penalty_box -\n"
 112         ":bw_data_saver -\n"
 113         ":bw_costly_shared -\n"
 114         "COMMIT\n"
 115         "*raw\n"
 116         ":bw_raw_PREROUTING -\n"
 117         "COMMIT\n"
 118         "*mangle\n"
 119         ":bw_mangle_POSTROUTING -\n"
 120         "COMMIT\n\x04";
 121     expectIptablesRestoreCommands({ expected });
 122 }
 123
 124 TEST_F(BandwidthControllerTest, TestEnableDataSaver) {
 125     mBw.enableDataSaver(true);
 126     std::vector<std::string> expected = {
 127         "-R bw_data_saver 1 --jump REJECT",
 128     };
 129     expectIptablesCommands(expected);
 130
 131     mBw.enableDataSaver(false);
 132     expected = {
 133         "-R bw_data_saver 1 --jump RETURN",
 134     };
 135     expectIptablesCommands(expected);
 136 }