Switch BandwidthController per-UID commands to iptables-restore

(cherry picked from commit 6849001bc61f965896a2d0fa42b3915879ee0b5b)

Bug: 32073253
Test: bullhead builds and boots
Test: netd_{unit,integration}_test pass
Test: Turning datasaver on/off changes rules as expected
Test: Modifying datasaver whitelist changes rules as expected when datasaver is on
Test: Adding and removing cell data limits changes rules as expected
Test: No IptablesRestoreController methods in normal usage
Change-Id: I83723db6a539b641308ef0f74ac30b4db304295c
Merged-In: I83723db6a539b641308ef0f74ac30b4db304295c

2 files changed, 31 insertions, 75 deletions

diff --git a/server/BandwidthController.cpp b/server/BandwidthController.cpp
index 828ef3ad..52604872 100644
--- a/server/BandwidthController.cpp
+++ b/server/BandwidthController.cpp
@@ -148,7 +148,6 @@ const std::string GET_TETHER_STATS_COMMAND = StringPrintf(
  */
 
 const std::string COMMIT_AND_CLOSE = "COMMIT\n";
-const std::string DATA_SAVER_ENABLE_COMMAND = "-R bw_data_saver 1";
 const std::string HAPPY_BOX_WHITELIST_COMMAND = StringPrintf(
     "-I bw_happy_box -m owner --uid-owner %d-%d --jump RETURN", 0, MAX_SYSTEM_UID);
 
@@ -301,8 +300,11 @@ int BandwidthController::disableBandwidthControl(void) {
 }
 
 int BandwidthController::enableDataSaver(bool enable) {
-    return runIpxtablesCmd(DATA_SAVER_ENABLE_COMMAND.c_str(),
-                           enable ? IptJumpReject : IptJumpReturn, IptFailShow);
+    std::string cmd = StringPrintf(
+        "*filter\n"
+        "-R bw_data_saver 1%s\n"
+        "COMMIT\n", jumpToString(enable ? IptJumpReject : IptJumpReturn));
+    return iptablesRestoreFunction(V4V6, cmd, nullptr);
 }
 
 int BandwidthController::runCommands(int numCommands, const char *commands[],
@@ -321,25 +323,6 @@ int BandwidthController::runCommands(int numCommands, const char *commands[],
     return 0;
 }
 
-std::string BandwidthController::makeIptablesSpecialAppCmd(IptOp op, int uid, const char *chain) {
-    std::string res;
-    char *buff;
-    const char *opFlag;
-
-    switch (op) {
-    case IptOpInsert:
-        opFlag = "-I";
-        break;
-    case IptOpDelete:
-        opFlag = "-D";
-        break;
-    }
-    asprintf(&buff, "%s %s -m owner --uid-owner %d", opFlag, chain, uid);
-    res = buff;
-    free(buff);
-    return res;
-}
-
 int BandwidthController::addNaughtyApps(int numUids, char *appUids[]) {
     return manipulateNaughtyApps(numUids, appUids, IptOpInsert);
 }
@@ -364,51 +347,16 @@ int BandwidthController::manipulateNiceApps(int numUids, char *appStrUids[], Ipt
     return manipulateSpecialApps(numUids, appStrUids, "bw_happy_box", IptJumpReturn, op);
 }
 
-
 int BandwidthController::manipulateSpecialApps(int numUids, char *appStrUids[],
                                                const char *chain,
                                                IptJumpOp jumpHandling, IptOp op) {
-
-    int uidNum;
-    const char *failLogTemplate;
-    int appUids[numUids];
-    std::string iptCmd;
-
-    switch (op) {
-    case IptOpInsert:
-        failLogTemplate = "Failed to add app uid %s(%d) to %s.";
-        break;
-    case IptOpDelete:
-        failLogTemplate = "Failed to delete app uid %s(%d) from %s box.";
-        break;
+    std::string cmd = "*filter\n";
+    for (int uidNum = 0; uidNum < numUids; uidNum++) {
+        StringAppendF(&cmd, "%s %s -m owner --uid-owner %s%s\n", opToString(op), chain,
+                      appStrUids[uidNum], jumpToString(jumpHandling));
     }
-
-    for (uidNum = 0; uidNum < numUids; uidNum++) {
-        char *end;
-        appUids[uidNum] = strtoul(appStrUids[uidNum], &end, 0);
-        if (*end || !*appStrUids[uidNum]) {
-            ALOGE(failLogTemplate, appStrUids[uidNum], appUids[uidNum], chain);
-            goto fail_parse;
-        }
-    }
-
-    for (uidNum = 0; uidNum < numUids; uidNum++) {
-        int uid = appUids[uidNum];
-
-        iptCmd = makeIptablesSpecialAppCmd(op, uid, chain);
-        if (runIpxtablesCmd(iptCmd.c_str(), jumpHandling)) {
-            ALOGE(failLogTemplate, appStrUids[uidNum], uid, chain);
-            goto fail_with_uidNum;
-        }
-    }
-    return 0;
-
-fail_with_uidNum:
-    /* Try to remove the uid that failed in any case*/
-    iptCmd = makeIptablesSpecialAppCmd(IptOpDelete, appUids[uidNum], chain);
-    runIpxtablesCmd(iptCmd.c_str(), jumpHandling);
-fail_parse:
-    return -1;
+    StringAppendF(&cmd, "COMMIT\n");
+    return iptablesRestoreFunction(V4V6, cmd, nullptr);
 }
 
 std::string BandwidthController::makeIptablesQuotaCmd(IptFullOp op, const char *costName, int64_t quota) {
diff --git a/server/BandwidthControllerTest.cpp b/server/BandwidthControllerTest.cpp
index ad4745cc..954db575 100644
--- a/server/BandwidthControllerTest.cpp
+++ b/server/BandwidthControllerTest.cpp
@@ -190,15 +190,19 @@ TEST_F(BandwidthControllerTest, TestDisableBandwidthControl) {
 TEST_F(BandwidthControllerTest, TestEnableDataSaver) {
     mBw.enableDataSaver(true);
     std::vector<std::string> expected = {
-        "-R bw_data_saver 1 --jump REJECT",
+        "*filter\n"
+        "-R bw_data_saver 1 --jump REJECT\n"
+        "COMMIT\n"
     };
-    expectIptablesCommands(expected);
+    expectIptablesRestoreCommands(expected);
 
     mBw.enableDataSaver(false);
     expected = {
-        "-R bw_data_saver 1 --jump RETURN",
+        "*filter\n"
+        "-R bw_data_saver 1 --jump RETURN\n"
+        "COMMIT\n"
     };
-    expectIptablesCommands(expected);
+    expectIptablesRestoreCommands(expected);
 }
 
 std::string kIPv4TetherCounters = android::base::Join(std::vector<std::string> {
@@ -441,18 +445,22 @@ TEST_F(BandwidthControllerTest, ManipulateSpecialApps) {
     std::vector<const char *> appUids = { "1000", "1001", "10012" };
 
     std::vector<std::string> expected = {
-        "-I bw_happy_box -m owner --uid-owner 1000 --jump RETURN",
-        "-I bw_happy_box -m owner --uid-owner 1001 --jump RETURN",
-        "-I bw_happy_box -m owner --uid-owner 10012 --jump RETURN",
+        "*filter\n"
+        "-I bw_happy_box -m owner --uid-owner 1000 --jump RETURN\n"
+        "-I bw_happy_box -m owner --uid-owner 1001 --jump RETURN\n"
+        "-I bw_happy_box -m owner --uid-owner 10012 --jump RETURN\n"
+        "COMMIT\n"
     };
     EXPECT_EQ(0, mBw.addNiceApps(appUids.size(), const_cast<char**>(&appUids[0])));
-    expectIptablesCommands(expected);
+    expectIptablesRestoreCommands(expected);
 
     expected = {
-        "-D bw_penalty_box -m owner --uid-owner 1000 --jump REJECT",
-        "-D bw_penalty_box -m owner --uid-owner 1001 --jump REJECT",
-        "-D bw_penalty_box -m owner --uid-owner 10012 --jump REJECT",
+        "*filter\n"
+        "-D bw_penalty_box -m owner --uid-owner 1000 --jump REJECT\n"
+        "-D bw_penalty_box -m owner --uid-owner 1001 --jump REJECT\n"
+        "-D bw_penalty_box -m owner --uid-owner 10012 --jump REJECT\n"
+        "COMMIT\n"
     };
     EXPECT_EQ(0, mBw.removeNaughtyApps(appUids.size(), const_cast<char**>(&appUids[0])));
-    expectIptablesCommands(expected);
+    expectIptablesRestoreCommands(expected);
 }