Use xt_bpf programs to do bandwidth control

To completely move away from xt_qtaguid module, the bandwidth controller
should not using it for uid owner match any more. Instead, it can use a
eBPF map to store the uid need to be matched and use two eBPF program
running on the xt_bpf hooks to filter out the packet.

Bug: 80649292
Test: ./netd_unit_test
Change-Id: I8e9c7cb3371aae0c24ccc6f64e05e6cbd4f78aae
diff --git a/server/BandwidthController.cpp b/server/BandwidthController.cpp
index 7cd1598..c85d102 100644
--- a/server/BandwidthController.cpp
+++ b/server/BandwidthController.cpp
@@ -56,6 +56,7 @@
 #include "FirewallController.h" /* For makeCriticalCommands */
 #include "Fwmark.h"
 #include "NetdConstants.h"
+#include "TrafficController.h"
 #include "bpf/BpfUtils.h"
 
 /* Alphabetical */
@@ -71,9 +72,13 @@
 using android::base::Join;
 using android::base::StringAppendF;
 using android::base::StringPrintf;
+using android::bpf::XT_BPF_BLACKLIST_PROG_PATH;
 using android::bpf::XT_BPF_EGRESS_PROG_PATH;
 using android::bpf::XT_BPF_INGRESS_PROG_PATH;
+using android::bpf::XT_BPF_WHITELIST_PROG_PATH;
+using android::net::gCtls;
 using android::netdutils::StatusOr;
+using android::netdutils::Status;
 using android::netdutils::UniqueFile;
 
 namespace {
@@ -149,6 +154,11 @@
 const std::string COMMIT_AND_CLOSE = "COMMIT\n";
 const std::string HAPPY_BOX_WHITELIST_COMMAND = StringPrintf(
     "-I bw_happy_box -m owner --uid-owner %d-%d --jump RETURN", 0, MAX_SYSTEM_UID);
+const std::string BPF_HAPPY_BOX_WHITELIST_COMMAND =
+    StringPrintf("-I bw_happy_box -m bpf --object-pinned %s -j RETURN", XT_BPF_WHITELIST_PROG_PATH);
+const std::string BPF_PENALTY_BOX_BLACKLIST_COMMAND =
+    StringPrintf("-I bw_penalty_box -m bpf --object-pinned %s -j REJECT",
+                 XT_BPF_BLACKLIST_PROG_PATH);
 
 static const std::vector<std::string> IPT_FLUSH_COMMANDS = {
     /*
@@ -206,8 +216,7 @@
  * See go/ipsec-data-accounting for more information.
  */
 
-const std::vector<std::string> getBasicAccountingCommands() {
-    bool useBpf = BandwidthController::getBpfStatsStatus();
+const std::vector<std::string> getBasicAccountingCommands(const bool useBpf) {
     const std::vector<std::string> ipt_basic_accounting_commands = {
         "*filter",
         // Prevents IPSec double counting (ESP and UDP-encap-ESP respectively)
@@ -224,10 +233,11 @@
         "-A bw_OUTPUT -m owner --socket-exists", /* This is a tracking rule. */
 
         "-A bw_costly_shared --jump bw_penalty_box",
+        useBpf ? BPF_PENALTY_BOX_BLACKLIST_COMMAND : "",
         "-A bw_penalty_box --jump bw_happy_box",
         "-A bw_happy_box --jump bw_data_saver",
         "-A bw_data_saver -j RETURN",
-        HAPPY_BOX_WHITELIST_COMMAND,
+        useBpf ? BPF_HAPPY_BOX_WHITELIST_COMMAND : HAPPY_BOX_WHITELIST_COMMAND,
         "COMMIT",
 
         "*raw",
@@ -237,7 +247,7 @@
         "-A bw_raw_PREROUTING -m policy --pol ipsec --dir in -j RETURN",
         "-A bw_raw_PREROUTING -m owner --socket-exists", /* This is a tracking rule. */
         useBpf ? StringPrintf("-A bw_raw_PREROUTING -m bpf --object-pinned %s",
-                              XT_BPF_INGRESS_PROG_PATH):"",
+                              XT_BPF_INGRESS_PROG_PATH) : "",
         "COMMIT",
 
         "*mangle",
@@ -249,7 +259,7 @@
         StringPrintf("-A bw_mangle_POSTROUTING -j MARK --set-mark 0x0/0x%x",
                      uidBillingMask), // Clear the mark before sending this packet
         useBpf ? StringPrintf("-A bw_mangle_POSTROUTING -m bpf --object-pinned %s",
-                              XT_BPF_EGRESS_PROG_PATH):"",
+                              XT_BPF_EGRESS_PROG_PATH) : "",
         COMMIT_AND_CLOSE
     };
     return ipt_basic_accounting_commands;
@@ -266,9 +276,11 @@
 
 }  // namespace
 
-bool BandwidthController::getBpfStatsStatus() {
+bool BandwidthController::getBpfStatus() {
     return (access(XT_BPF_INGRESS_PROG_PATH, F_OK) != -1) &&
-           (access(XT_BPF_EGRESS_PROG_PATH, F_OK) != -1);
+           (access(XT_BPF_EGRESS_PROG_PATH, F_OK) != -1) &&
+           (access(XT_BPF_WHITELIST_PROG_PATH, F_OK) != -1) &&
+           (access(XT_BPF_BLACKLIST_PROG_PATH, F_OK) != -1);
 }
 
 BandwidthController::BandwidthController() {
@@ -306,7 +318,8 @@
 
     flushCleanTables(false);
 
-    std::string commands = Join(getBasicAccountingCommands(), '\n');
+    mBpfSupported = getBpfStatus();
+    std::string commands = Join(getBasicAccountingCommands(mBpfSupported), '\n');
     return iptablesRestoreFunction(V4V6, commands, nullptr);
 }
 
@@ -360,6 +373,13 @@
 int BandwidthController::manipulateSpecialApps(const std::vector<std::string>& appStrUids,
                                                const std::string& chain, IptJumpOp jumpHandling,
                                                IptOp op) {
+    if (mBpfSupported) {
+      Status status = gCtls->trafficCtrl.updateBandwidthUidMap(appStrUids, jumpHandling, op);
+      if (!isOk(status)) {
+          ALOGE("unable to update the Bandwidth Uid Map: %s", toString(status).c_str());
+      }
+      return status.code();
+    }
     std::string cmd = "*filter\n";
     for (const auto& appStrUid : appStrUids) {
         StringAppendF(&cmd, "%s %s -m owner --uid-owner %s%s\n", opToString(op), chain.c_str(),
diff --git a/server/BandwidthController.h b/server/BandwidthController.h
index efacdce..3f6c0ad 100644
--- a/server/BandwidthController.h
+++ b/server/BandwidthController.h
@@ -33,7 +33,7 @@
     BandwidthController();
 
     int setupIptablesHooks();
-    static bool getBpfStatsStatus();
+    static bool getBpfStatus();
 
     int enableBandwidthControl(bool force);
     int disableBandwidthControl();
@@ -69,6 +69,9 @@
     static const char LOCAL_RAW_PREROUTING[];
     static const char LOCAL_MANGLE_POSTROUTING[];
 
+    enum IptJumpOp { IptJumpReject, IptJumpReturn, IptJumpNoAdd };
+    enum IptOp { IptOpInsert, IptOpDelete };
+
   private:
     struct QuotaInfo {
         int64_t quota;
@@ -77,8 +80,6 @@
 
     enum IptIpVer { IptIpV4, IptIpV6 };
     enum IptFullOp { IptFullOpInsert, IptFullOpDelete, IptFullOpAppend };
-    enum IptJumpOp { IptJumpReject, IptJumpReturn, IptJumpNoAdd };
-    enum IptOp { IptOpInsert, IptOpDelete };
     enum QuotaType { QuotaUnique, QuotaShared };
     enum RunCmdErrHandling { RunCmdFailureBad, RunCmdFailureOk };
 #if LOG_NDEBUG
@@ -125,6 +126,8 @@
     static const char *opToString(IptOp op);
     static const char *jumpToString(IptJumpOp jumpHandling);
 
+    bool mBpfSupported;
+
     int64_t mSharedQuotaBytes = 0;
     int64_t mSharedAlertBytes = 0;
     int64_t mGlobalAlertBytes = 0;
diff --git a/server/BandwidthControllerTest.cpp b/server/BandwidthControllerTest.cpp
index 938cea0..95af148 100644
--- a/server/BandwidthControllerTest.cpp
+++ b/server/BandwidthControllerTest.cpp
@@ -46,8 +46,10 @@
 
 using android::base::Join;
 using android::base::StringPrintf;
+using android::bpf::XT_BPF_BLACKLIST_PROG_PATH;
 using android::bpf::XT_BPF_EGRESS_PROG_PATH;
 using android::bpf::XT_BPF_INGRESS_PROG_PATH;
+using android::bpf::XT_BPF_WHITELIST_PROG_PATH;
 using android::net::TunInterface;
 using android::netdutils::status::ok;
 using android::netdutils::UniqueFile;
@@ -180,7 +182,7 @@
     std::string expectedClean = "";
 
     uint32_t uidBillingMask = Fwmark::getUidBillingMask();
-    bool useBpf = BandwidthController::getBpfStatsStatus();
+    bool useBpf = BandwidthController::getBpfStatus();
     std::string expectedAccounting =
         "*filter\n"
         "-A bw_INPUT -p esp -j RETURN\n" +
@@ -191,11 +193,24 @@
         "-A bw_OUTPUT -o " IPSEC_IFACE_PREFIX "+ -j RETURN\n"
         "-A bw_OUTPUT -m policy --pol ipsec --dir out -j RETURN\n"
         "-A bw_OUTPUT -m owner --socket-exists\n"
-        "-A bw_costly_shared --jump bw_penalty_box\n"
-        "-A bw_penalty_box --jump bw_happy_box\n"
-        "-A bw_happy_box --jump bw_data_saver\n"
-        "-A bw_data_saver -j RETURN\n"
-        "-I bw_happy_box -m owner --uid-owner 0-9999 --jump RETURN\n"
+        "-A bw_costly_shared --jump bw_penalty_box\n";
+    if (useBpf) {
+        expectedAccounting +=
+            StringPrintf("-I bw_penalty_box -m bpf --object-pinned %s -j REJECT\n",
+                         XT_BPF_BLACKLIST_PROG_PATH) +
+            "-A bw_penalty_box --jump bw_happy_box\n" +
+            "-A bw_happy_box --jump bw_data_saver\n"
+            "-A bw_data_saver -j RETURN\n" +
+            StringPrintf("-I bw_happy_box -m bpf --object-pinned %s -j RETURN\n",
+                         XT_BPF_WHITELIST_PROG_PATH);
+    } else {
+        expectedAccounting +=
+            "\n-A bw_penalty_box --jump bw_happy_box\n"
+            "-A bw_happy_box --jump bw_data_saver\n"
+            "-A bw_data_saver -j RETURN\n"
+            "-I bw_happy_box -m owner --uid-owner 0-9999 --jump RETURN\n";
+    }
+    expectedAccounting +=
         "COMMIT\n"
         "*raw\n"
         "-A bw_raw_PREROUTING -i " IPSEC_IFACE_PREFIX "+ -j RETURN\n"
@@ -506,6 +521,7 @@
 }
 
 TEST_F(BandwidthControllerTest, ManipulateSpecialApps) {
+    if (BandwidthController::getBpfStatus()) return;
     std::vector<const char *> appUids = { "1000", "1001", "10012" };
 
     std::vector<std::string> expected = {
diff --git a/server/TrafficController.cpp b/server/TrafficController.cpp
index d6a6480..18c3689 100644
--- a/server/TrafficController.cpp
+++ b/server/TrafficController.cpp
@@ -45,7 +45,6 @@
 #include <netdutils/Syscalls.h>
 #include "TrafficController.h"
 #include "bpf/BpfMap.h"
-#include "bpf/bpf_shared.h"
 
 #include "DumpWriter.h"
 #include "FirewallController.h"
@@ -117,19 +116,7 @@
 }
 
 Status initialOwnerMap(BpfMap<uint32_t, uint8_t>& map) {
-    // Check and delete all the entries from the map in case it is a runtime
-    // restart
-    const auto deleteAllEntries = [](const uint32_t& key, BpfMap<uint32_t, uint8_t>& map) {
-        Status res = map.deleteValue(key);
-        if (!isOk(res) && (res.code() == ENOENT)) {
-            ALOGE("Failed to delete data(uid=%u): %s\n", key, strerror(res.code()));
-        }
-        return netdutils::status::ok;
-    };
-    // It is safe to delete from this map because nothing will concurrently iterate over it:
-    // - Nothing in netd will iterate over it because we're holding mOwnerMatchMutex.
-    // - Nothing outside netd iterates over it.
-    map.iterate(deleteAllEntries);
+    map.clear();
     uint32_t mapSettingKey = UID_MAP_ENABLED;
     uint8_t defaultMapState = 0;
     return map.writeValue(mapSettingKey, defaultMapState, BPF_NOEXIST);
@@ -187,6 +174,11 @@
         mIfaceStatsMap.getOrCreate(IFACE_STATS_MAP_SIZE, IFACE_STATS_MAP_PATH, BPF_MAP_TYPE_HASH));
     RETURN_IF_NOT_OK(changeOwnerAndMode(IFACE_STATS_MAP_PATH, AID_NET_BW_STATS, "IfaceStatsMap",
                                         false));
+
+    RETURN_IF_NOT_OK(mBandwidthUidMap.getOrCreate(UID_OWNER_MAP_SIZE, BANDWIDTH_UID_MAP_PATH,
+                                                  BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(changeOwnerAndMode(BANDWIDTH_UID_MAP_PATH, AID_ROOT, "BandwidthUidMap", true));
+    mBandwidthUidMap.clear();
     return netdutils::status::ok;
 }
 
@@ -265,6 +257,14 @@
     if (ret != 0 && errno == ENOENT) {
         prog_args.push_back((char*)"-m");
     }
+    ret = access(XT_BPF_WHITELIST_PROG_PATH, R_OK);
+    if (ret != 0 && errno == ENOENT) {
+        prog_args.push_back((char*)"-w");
+    }
+    ret = access(XT_BPF_BLACKLIST_PROG_PATH, R_OK);
+    if (ret != 0 && errno == ENOENT) {
+        prog_args.push_back((char*)"-b");
+    }
 
     if (prog_args.size() == 1) {
         // all program are loaded already.
@@ -447,6 +447,59 @@
     return netdutils::status::ok;
 }
 
+BandwithMatchType TrafficController::jumpOpToMatch(BandwidthController::IptJumpOp jumpHandling) {
+    switch (jumpHandling) {
+        case BandwidthController::IptJumpReject:
+            return BLACKLISTMATCH;
+        case BandwidthController::IptJumpReturn:
+            return WHITELISTMATCH;
+        case BandwidthController::IptJumpNoAdd:
+            return NO_MATCH;
+    }
+}
+
+Status TrafficController::updateBandwidthUidMap(const std::vector<std::string>& appStrUids,
+                                                BandwidthController::IptJumpOp jumpHandling,
+                                                BandwidthController::IptOp op) {
+    uint8_t command = jumpOpToMatch(jumpHandling);
+    if (command == NO_MATCH) {
+        return statusFromErrno(EINVAL, StringPrintf("invalid IptJumpOp: %d, command: %d",
+                                                    jumpHandling, command));
+    }
+    for (const auto& appStrUid : appStrUids) {
+        char* endPtr;
+        long uid = strtol(appStrUid.c_str(), &endPtr, 10);
+        if ((errno == ERANGE && (uid == LONG_MAX || uid == LONG_MIN)) ||
+            (endPtr == appStrUid.c_str()) || (*endPtr != '\0')) {
+               return statusFromErrno(errno, "invalid uid string:" + appStrUid);
+        }
+
+        auto match = mBandwidthUidMap.readValue(uid);
+        if (op == BandwidthController::IptOpDelete) {
+            if(!isOk(match)) {
+                return statusFromErrno(EINVAL, StringPrintf("uid(%ld): %s does not exist in map",
+                                                            uid, appStrUid.c_str()));
+            }
+            uint8_t newValue = match.value() & ~command;
+            if (newValue == 0) {
+                RETURN_IF_NOT_OK(mBandwidthUidMap.deleteValue((uint32_t)uid));
+            } else {
+                RETURN_IF_NOT_OK(mBandwidthUidMap.writeValue((uint32_t)uid, newValue, BPF_ANY));
+            }
+        } else if (op == BandwidthController::IptOpInsert) {
+            uint8_t newValue = command;
+            if (isOk(match)) {
+                newValue |= match.value();
+            }
+            RETURN_IF_NOT_OK(mBandwidthUidMap.writeValue((uint32_t)uid, newValue, BPF_ANY));
+        } else {
+            // Cannot happen.
+            return statusFromErrno(EINVAL, StringPrintf("invalid IptOp: %d, %d", op, command));
+        }
+    }
+    return netdutils::status::ok;
+}
+
 int TrafficController::changeUidOwnerRule(ChildChain chain, uid_t uid, FirewallRule rule,
                                           FirewallType type) {
     std::lock_guard<std::mutex> guard(mOwnerMatchMutex);
@@ -622,6 +675,8 @@
                getMapStatus(mStandbyUidMap.getMap(), STANDBY_UID_MAP_PATH).c_str());
     dw.println("mPowerSaveUidMap status: %s",
                getMapStatus(mPowerSaveUidMap.getMap(), POWERSAVE_UID_MAP_PATH).c_str());
+    dw.println("mBandwidthUidMap status: %s",
+               getMapStatus(mBandwidthUidMap.getMap(), BANDWIDTH_UID_MAP_PATH).c_str());
 
     dw.blankline();
     dw.println("Cgroup ingress program status: %s",
@@ -631,6 +686,10 @@
                getProgramStatus(XT_BPF_INGRESS_PROG_PATH).c_str());
     dw.println("xt_bpf egress program status: %s",
                getProgramStatus(XT_BPF_EGRESS_PROG_PATH).c_str());
+    dw.println("xt_bpf bandwidth whitelist program status: %s",
+               getProgramStatus(XT_BPF_WHITELIST_PROG_PATH).c_str());
+    dw.println("xt_bpf bandwidth blacklist program status: %s",
+               getProgramStatus(XT_BPF_BLACKLIST_PROG_PATH).c_str());
 
     if(!verbose) return;
 
@@ -756,6 +815,11 @@
         dw.println("mDozableUidMap print end with error: %s", res.msg().c_str());
     }
 
+    dumpBpfMap("mBandwidthUidMap", dw, "");
+    res = mBandwidthUidMap.iterateWithValue(printUidInfo);
+    if (!isOk(res)) {
+        dw.println("mBandwidthUidMap print end with error: %s", res.msg().c_str());
+    }
     dw.decIndent();
 
     dw.decIndent();
diff --git a/server/TrafficController.h b/server/TrafficController.h
index 79f7d14..e4e32f4 100644
--- a/server/TrafficController.h
+++ b/server/TrafficController.h
@@ -20,12 +20,14 @@
 #include <linux/bpf.h>
 
 #include <netdutils/StatusOr.h>
+#include "BandwidthController.h"
 #include "FirewallController.h"
 #include "NetlinkListener.h"
 #include "Network.h"
 #include "android-base/thread_annotations.h"
 #include "android-base/unique_fd.h"
 #include "bpf/BpfMap.h"
+#include "bpf/bpf_shared.h"
 
 using android::bpf::BpfMap;
 using android::bpf::IfaceValue;
@@ -106,6 +108,9 @@
                                        const std::vector<int32_t>& uids, FirewallRule rule,
                                        FirewallType type);
 
+    netdutils::Status updateBandwidthUidMap(const std::vector<std::string>& appStrUids,
+                                            BandwidthController::IptJumpOp jumpHandling,
+                                            BandwidthController::IptOp op);
     static const String16 DUMP_KEYWORD;
 
     int toggleUidOwnerMap(ChildChain chain, bool enable);
@@ -192,6 +197,11 @@
      */
     BpfMap<uint32_t, uint8_t> mPowerSaveUidMap GUARDED_BY(mOwnerMatchMutex);
 
+    /*
+     * mBandwidthUidMap: Store uids that are used for bandwidth control uid match.
+     */
+    BpfMap<uint32_t, uint8_t> mBandwidthUidMap;
+
     std::unique_ptr<NetlinkListenerInterface> mSkDestroyListener;
 
     bool ebpfSupported;
@@ -202,6 +212,8 @@
                                            base::unique_fd& cg_fd);
 
     netdutils::Status initMaps();
+
+    BandwithMatchType jumpOpToMatch(BandwidthController::IptJumpOp jumpHandling);
     // For testing
     friend class TrafficControllerTest;
 };
diff --git a/server/TrafficControllerTest.cpp b/server/TrafficControllerTest.cpp
index a354f83..fee5ee4 100644
--- a/server/TrafficControllerTest.cpp
+++ b/server/TrafficControllerTest.cpp
@@ -78,6 +78,7 @@
     BpfMap<uint32_t, uint8_t> mFakeDozableUidMap;
     BpfMap<uint32_t, uint8_t> mFakeStandbyUidMap;
     BpfMap<uint32_t, uint8_t> mFakePowerSaveUidMap;
+    BpfMap<uint32_t, uint8_t> mFakeBandwidthUidMap;
 
     void SetUp() {
         std::lock_guard<std::mutex> ownerGuard(mTc.mOwnerMatchMutex);
@@ -114,6 +115,10 @@
         mFakePowerSaveUidMap.reset(
             createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), TEST_MAP_SIZE, 0));
         ASSERT_LE(0, mFakePowerSaveUidMap.getMap());
+
+        mFakeBandwidthUidMap.reset(
+            createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakeBandwidthUidMap.getMap());
         // Make sure trafficController use the eBPF code path.
         mTc.ebpfSupported = true;
 
@@ -125,6 +130,7 @@
         mTc.mDozableUidMap.reset(mFakeDozableUidMap.getMap());
         mTc.mStandbyUidMap.reset(mFakeStandbyUidMap.getMap());
         mTc.mPowerSaveUidMap.reset(mFakePowerSaveUidMap.getMap());
+        mTc.mBandwidthUidMap.reset(mFakeBandwidthUidMap.getMap());
     }
 
     int setUpSocketAndTag(int protocol, uint64_t* cookie, uint32_t tag, uid_t uid) {
@@ -220,6 +226,18 @@
         checkEachUidValue(uids, BPF_DROP, targetMap);
     }
 
+    void expectBandwidthMapValues(const std::vector<std::string>& appStrUids,
+                                  uint8_t expectedValue) {
+        for (std::string strUid : appStrUids) {
+            uint32_t uid = stoi(strUid);
+            StatusOr<uint8_t> value = mFakeBandwidthUidMap.readValue(uid);
+            EXPECT_TRUE(isOk(value));
+            EXPECT_EQ(expectedValue, value.value()) <<
+                "Expected value for UID " << uid << " to be " << expectedValue <<
+                ", but was " << value.value();
+        }
+    }
+
     void TearDown() {
         std::lock_guard<std::mutex> ownerGuard(mTc.mOwnerMatchMutex);
         mFakeCookieTagMap.reset();
@@ -478,5 +496,72 @@
     ASSERT_EQ(-EINVAL, mTc.replaceUidOwnerMap("unknow", true, uids));
 }
 
+TEST_F(TrafficControllerTest, TestBlacklistUidMatch) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    std::vector<std::string> appStrUids = {"1000", "1001", "10012"};
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReject,
+                                               BandwidthController::IptOpInsert)));
+    expectBandwidthMapValues(appStrUids, BLACKLISTMATCH);
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReject,
+                                               BandwidthController::IptOpDelete)));
+    ASSERT_FALSE(isOk(mFakeBandwidthUidMap.getFirstKey()));
+}
+
+TEST_F(TrafficControllerTest, TestWhitelistUidMatch) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    std::vector<std::string> appStrUids = {"1000", "1001", "10012"};
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReturn,
+                                               BandwidthController::IptOpInsert)));
+    expectBandwidthMapValues(appStrUids, WHITELISTMATCH);
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReturn,
+                                               BandwidthController::IptOpDelete)));
+    ASSERT_FALSE(isOk(mFakeBandwidthUidMap.getFirstKey()));
+}
+
+TEST_F(TrafficControllerTest, TestReplaceMatchUid) {
+    std::vector<std::string> appStrUids = {"1000", "1001", "10012"};
+    // Add appStrUids to the blacklist and expect that their values are all BLACKLISTMATCH.
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReject,
+                                               BandwidthController::IptOpInsert)));
+    expectBandwidthMapValues(appStrUids, BLACKLISTMATCH);
+
+    // Add the same UIDs to the whitelist and expect that we get BLACKLISTMATCH | WHITELISTMATCH.
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReturn,
+                                               BandwidthController::IptOpInsert)));
+    expectBandwidthMapValues(appStrUids, WHITELISTMATCH | BLACKLISTMATCH);
+
+    // Remove the same UIDs from the whitelist and check the BLACKLISTMATCH is still there.
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReturn,
+                                               BandwidthController::IptOpDelete)));
+    expectBandwidthMapValues(appStrUids, BLACKLISTMATCH);
+
+    // Remove the same UIDs from the blacklist and check the map is empty.
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReject,
+                                               BandwidthController::IptOpDelete)));
+    ASSERT_FALSE(isOk(mFakeBandwidthUidMap.getFirstKey()));
+}
+
+TEST_F(TrafficControllerTest, TestDeleteWrongMatchSilentlyFails) {
+    std::vector<std::string> appStrUids = {"1000", "1001", "10012"};
+    // If the uid does not exist in the map, trying to delete a rule about it will fail.
+    ASSERT_FALSE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReject,
+                                                BandwidthController::IptOpDelete)));
+    ASSERT_FALSE(isOk(mFakeBandwidthUidMap.getFirstKey()));
+
+    // Add blacklist rules for appStrUids.
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReturn,
+                                               BandwidthController::IptOpInsert)));
+    expectBandwidthMapValues(appStrUids, WHITELISTMATCH);
+
+    // Delete (non-existent) blacklist rules for appStrUids, and check that this silently does
+    // nothing if the uid is in the map but does not have blacklist match. This is required because
+    // NetworkManagementService will try to remove a uid from blacklist after adding it to the
+    // whitelist and if the remove fails it will not update the uid status.
+    ASSERT_TRUE(isOk(mTc.updateBandwidthUidMap(appStrUids, BandwidthController::IptJumpReject,
+                                                BandwidthController::IptOpDelete)));
+    expectBandwidthMapValues(appStrUids, WHITELISTMATCH);
+}
 }  // namespace net
 }  // namespace android