diff options
-rw-r--r-- | server/BandwidthController.cpp | 51 | ||||
-rw-r--r-- | server/BandwidthControllerTest.cpp | 12 |
2 files changed, 35 insertions, 28 deletions
diff --git a/server/BandwidthController.cpp b/server/BandwidthController.cpp index 3d1a56ca..50a87119 100644 --- a/server/BandwidthController.cpp +++ b/server/BandwidthController.cpp @@ -746,10 +746,7 @@ int BandwidthController::removeInterfaceAlert(const std::string& iface) { int BandwidthController::setCostlyAlert(const std::string& costName, int64_t bytes, int64_t* alertBytes) { - char *alertQuotaCmd; - char *chainName; int res = 0; - char *alertName; if (!isIfaceName(costName)) { ALOGE("setCostlyAlert: Invalid costName \"%s\"", costName.c_str()); @@ -760,27 +757,29 @@ int BandwidthController::setCostlyAlert(const std::string& costName, int64_t byt ALOGE("Invalid bytes value. 1..max_int64."); return -1; } - asprintf(&alertName, "%sAlert", costName.c_str()); + + std::string alertName = costName + "Alert"; + std::string chainName = "bw_costly_" + costName; if (*alertBytes) { res = updateQuota(alertName, *alertBytes); } else { - asprintf(&chainName, "bw_costly_%s", costName.c_str()); - asprintf(&alertQuotaCmd, ALERT_IPT_TEMPLATE, "-A", chainName, bytes, alertName); - res |= runIpxtablesCmd(alertQuotaCmd, IptJumpNoAdd); - free(alertQuotaCmd); - free(chainName); + std::vector<std::string> commands = { + "*filter\n", + StringPrintf(ALERT_IPT_TEMPLATE, "-A", chainName.c_str(), bytes, alertName.c_str()), + "COMMIT\n" + }; + res = iptablesRestoreFunction(V4V6, Join(commands, ""), nullptr); + if (res) { + ALOGE("Failed to set costly alert for %s", costName.c_str()); + } + } + if (res == 0) { + *alertBytes = bytes; } - *alertBytes = bytes; - free(alertName); return res; } int BandwidthController::removeCostlyAlert(const std::string& costName, int64_t* alertBytes) { - char *alertQuotaCmd; - char *chainName; - char *alertName; - int res = 0; - if (!isIfaceName(costName)) { ALOGE("removeCostlyAlert: Invalid costName \"%s\"", costName.c_str()); return -1; @@ -791,16 +790,20 @@ int BandwidthController::removeCostlyAlert(const std::string& costName, int64_t* return -1; } - asprintf(&alertName, "%sAlert", costName.c_str()); - asprintf(&chainName, "bw_costly_%s", costName.c_str()); - asprintf(&alertQuotaCmd, ALERT_IPT_TEMPLATE, "-D", chainName, *alertBytes, alertName); - res |= runIpxtablesCmd(alertQuotaCmd, IptJumpNoAdd); - free(alertQuotaCmd); - free(chainName); + std::string alertName = costName + "Alert"; + std::string chainName = "bw_costly_" + costName; + std::vector<std::string> commands = { + "*filter\n", + StringPrintf(ALERT_IPT_TEMPLATE, "-D", chainName.c_str(), *alertBytes, alertName.c_str()), + "COMMIT\n" + }; + if (iptablesRestoreFunction(V4V6, Join(commands, ""), nullptr) != 0) { + ALOGE("Failed to remove costly alert %s", costName.c_str()); + return -1; + } *alertBytes = 0; - free(alertName); - return res; + return 0; } void BandwidthController::addStats(TetherStatsList& statsList, const TetherStats& stats) { diff --git a/server/BandwidthControllerTest.cpp b/server/BandwidthControllerTest.cpp index becfe49b..066f9ebc 100644 --- a/server/BandwidthControllerTest.cpp +++ b/server/BandwidthControllerTest.cpp @@ -585,24 +585,28 @@ TEST_F(BandwidthControllerTest, CostlyAlert) { int64_t alertBytes = 0; std::vector<std::string> expected = { - "-A bw_costly_shared -m quota2 ! --quota 123456 --name sharedAlert\n", + "*filter\n" + "-A bw_costly_shared -m quota2 ! --quota 123456 --name sharedAlert\n" + "COMMIT\n" }; EXPECT_EQ(0, setCostlyAlert("shared", kQuota, &alertBytes)); EXPECT_EQ(kQuota, alertBytes); - expectIptablesCommands(expected); + expectIptablesRestoreCommands(expected); expected = {}; expectUpdateQuota(kQuota); EXPECT_EQ(0, setCostlyAlert("shared", kQuota + 1, &alertBytes)); EXPECT_EQ(kQuota + 1, alertBytes); - expectIptablesCommands(expected); + expectIptablesRestoreCommands(expected); expected = { + "*filter\n" "-D bw_costly_shared -m quota2 ! --quota 123457 --name sharedAlert\n" + "COMMIT\n" }; EXPECT_EQ(0, removeCostlyAlert("shared", &alertBytes)); EXPECT_EQ(0, alertBytes); - expectIptablesCommands(expected); + expectIptablesRestoreCommands(expected); } TEST_F(BandwidthControllerTest, ManipulateSpecialApps) { |