Add xt_owner module support in trafficController

Add bpf maps for recording rules about socket owner uid filtering.
Modified the bpf program so that packets with uid listed in the
the uidOwnerMap will get handled according to userspace settings

Test: bpf program can be loaded and attached when boot
Bug: 72381727 30950746
Change-Id: I39497334fcb5e200dbf07a0046b85c227d59e2d7
diff --git a/server/FirewallController.cpp b/server/FirewallController.cpp
index 26c2126..2c59817 100644
--- a/server/FirewallController.cpp
+++ b/server/FirewallController.cpp
@@ -29,12 +29,18 @@
 #include <android-base/stringprintf.h>
 #include <cutils/log.h>
 
-#include "NetdConstants.h"
+#include "Controllers.h"
 #include "FirewallController.h"
+#include "NetdConstants.h"
+#include "bpf/BpfUtils.h"
 
 using android::base::Join;
 using android::base::StringAppendF;
 using android::base::StringPrintf;
+using android::bpf::DOZABLE_UID_MAP_PATH;
+using android::bpf::POWERSAVE_UID_MAP_PATH;
+using android::bpf::STANDBY_UID_MAP_PATH;
+using android::net::gCtls;
 
 auto FirewallController::execIptablesRestore = ::execIptablesRestore;
 
@@ -60,6 +66,10 @@
     "redirect",
 };
 
+bool getBpfOwnerStatus() {
+    return gCtls->trafficCtrl.checkBpfStatsEnable();
+}
+
 FirewallController::FirewallController(void) {
     // If no rules are set, it's in BLACKLIST mode
     mFirewallType = BLACKLIST;
@@ -68,6 +78,10 @@
 
 int FirewallController::setupIptablesHooks(void) {
     int res = 0;
+    mUseBpfOwnerMatch = getBpfOwnerStatus();
+    if (mUseBpfOwnerMatch) {
+        return res;
+    }
     res |= createChain(LOCAL_DOZABLE, getFirewallType(DOZABLE));
     res |= createChain(LOCAL_STANDBY, getFirewallType(STANDBY));
     res |= createChain(LOCAL_POWERSAVE, getFirewallType(POWERSAVE));
@@ -129,6 +143,10 @@
             return res;
     }
 
+    if (mUseBpfOwnerMatch) {
+        return gCtls->trafficCtrl.toggleUidOwnerMap(chain, enable);
+    }
+
     std::string command = "*filter\n";
     for (const char *parent : { LOCAL_INPUT, LOCAL_OUTPUT }) {
         StringAppendF(&command, "%s %s -j %s\n", (enable ? "-A" : "-D"), parent, name);
@@ -224,6 +242,9 @@
             ALOGW("Unknown child chain: %d", chain);
             return -1;
     }
+    if (mUseBpfOwnerMatch) {
+        return gCtls->trafficCtrl.changeUidOwnerRule(chain, uid, rule, firewallType);
+    }
 
     std::string command = "*filter\n";
     for (std::string chainName : chainNames) {
@@ -309,6 +330,9 @@
 
 int FirewallController::replaceUidChain(
         const std::string &name, bool isWhitelist, const std::vector<int32_t>& uids) {
+   if (mUseBpfOwnerMatch) {
+       return gCtls->trafficCtrl.replaceUidOwnerMap(name, isWhitelist, uids);
+   }
    std::string commands4 = makeUidRules(V4, name.c_str(), isWhitelist, uids);
    std::string commands6 = makeUidRules(V6, name.c_str(), isWhitelist, uids);
    return execIptablesRestore(V4, commands4.c_str()) | execIptablesRestore(V6, commands6.c_str());
diff --git a/server/FirewallController.h b/server/FirewallController.h
index c7928b9..8a222d3 100644
--- a/server/FirewallController.h
+++ b/server/FirewallController.h
@@ -85,6 +85,7 @@
 
 private:
     FirewallType mFirewallType;
+    bool mUseBpfOwnerMatch;
     std::set<std::string> mIfaceRules;
     int attachChain(const char*, const char*);
     int detachChain(const char*, const char*);
diff --git a/server/FirewallControllerTest.cpp b/server/FirewallControllerTest.cpp
index c1f43eb..952b21e 100644
--- a/server/FirewallControllerTest.cpp
+++ b/server/FirewallControllerTest.cpp
@@ -35,6 +35,7 @@
 protected:
     FirewallControllerTest() {
         FirewallController::execIptablesRestore = fakeExecIptablesRestore;
+        mFw.mUseBpfOwnerMatch = false;
     }
     FirewallController mFw;
 
diff --git a/server/IptablesRestoreControllerTest.cpp b/server/IptablesRestoreControllerTest.cpp
index 017870f..8194f58 100644
--- a/server/IptablesRestoreControllerTest.cpp
+++ b/server/IptablesRestoreControllerTest.cpp
@@ -32,6 +32,7 @@
 #include "IptablesRestoreController.h"
 #include "NetdConstants.h"
 #include "Stopwatch.h"
+#include "bpf/BpfUtils.h"
 
 #define XT_LOCK_NAME "/system/etc/xtables.lock"
 #define XT_LOCK_ATTEMPTS 10
@@ -39,6 +40,9 @@
 
 using android::base::Join;
 using android::base::StringPrintf;
+using android::bpf::DOZABLE_UID_MAP_PATH;
+using android::bpf::STANDBY_UID_MAP_PATH;
+using android::bpf::POWERSAVE_UID_MAP_PATH;
 using android::netdutils::ScopedMockSyscalls;
 using testing::Return;
 using testing::StrictMock;
@@ -250,13 +254,16 @@
   EXPECT_EQ(expected, output);
 }
 
+
 TEST_F(IptablesRestoreControllerTest, TestUidRuleBenchmark) {
     const std::vector<int> ITERATIONS = { 1, 5, 10 };
 
     const std::string IPTABLES_RESTORE_ADD =
-            "*filter\n-I fw_powersave -m owner --uid-owner 2000000000 -j RETURN\nCOMMIT\n";
+            StringPrintf("*filter\n-I %s -m owner --uid-owner 2000000000 -j RETURN\nCOMMIT\n",
+                         mChainName.c_str());
     const std::string IPTABLES_RESTORE_DEL =
-            "*filter\n-D fw_powersave -m owner --uid-owner 2000000000 -j RETURN\nCOMMIT\n";
+            StringPrintf("*filter\n-D %s -m owner --uid-owner 2000000000 -j RETURN\nCOMMIT\n",
+                         mChainName.c_str());
 
     for (const int iterations : ITERATIONS) {
         Stopwatch s;
diff --git a/server/TrafficController.cpp b/server/TrafficController.cpp
index 25142ee..443cefa 100644
--- a/server/TrafficController.cpp
+++ b/server/TrafficController.cpp
@@ -31,6 +31,7 @@
 #include <sys/types.h>
 #include <sys/utsname.h>
 #include <sys/wait.h>
+#include <mutex>
 #include <unordered_set>
 #include <vector>
 
@@ -43,7 +44,9 @@
 #include <netdutils/Syscalls.h>
 #include "TrafficController.h"
 #include "bpf/BpfUtils.h"
+#include "bpf/bpf_shared.h"
 
+#include "FirewallController.h"
 #include "InterfaceController.h"
 #include "NetlinkListener.h"
 #include "qtaguid/qtaguid.h"
@@ -89,8 +92,100 @@
     return listener;
 }
 
-Status TrafficController::start() {
+Status changeOwnerAndMode(const char* path, gid_t group, const char* debugName, bool netdOnly) {
+    int ret = chown(path, AID_ROOT, group);
+    if (ret != 0) return statusFromErrno(errno, StringPrintf("change %s group failed", debugName));
+
+    if (netdOnly) {
+        ret = chmod(path, S_IRWXU);
+    } else {
+        // Allow both netd and system server to obtain map fd from the path.
+        // chmod doesn't grant permission to all processes in that group to
+        // read/write the bpf map. They still need correct sepolicy to
+        // read/write the map.
+        ret = chmod(path, S_IRWXU | S_IRGRP);
+    }
+    if (ret != 0) return statusFromErrno(errno, StringPrintf("change %s mode failed", debugName));
+    return netdutils::status::ok;
+}
+
+TrafficController::TrafficController() {
     ebpfSupported = hasBpfSupport();
+}
+
+Status initialOwnerMap(const unique_fd& map_fd, const std::string mapName) {
+    uint32_t mapSettingKey = UID_MAP_ENABLED;
+    uint32_t defaultMapState = 0;
+    int ret = writeToMapEntry(map_fd, &mapSettingKey, &defaultMapState, BPF_NOEXIST);
+    // If it is already exist, it might be a runtime restart and we just keep
+    // the old state.
+    if (ret && errno != EEXIST) {
+        return statusFromErrno(errno, "Fail to set the initial state of " + mapName);
+    }
+    return netdutils::status::ok;
+}
+
+Status TrafficController::initMaps() {
+    std::lock_guard<std::mutex> guard(mOwnerMatchMutex);
+    ASSIGN_OR_RETURN(mCookieTagMap,
+                     setUpBPFMap(sizeof(uint64_t), sizeof(struct UidTag), COOKIE_UID_MAP_SIZE,
+                                 COOKIE_TAG_MAP_PATH, BPF_MAP_TYPE_HASH));
+
+    RETURN_IF_NOT_OK(changeOwnerAndMode(COOKIE_TAG_MAP_PATH, AID_NET_BW_ACCT, "CookieTagMap",
+                                        false));
+
+    ASSIGN_OR_RETURN(mUidCounterSetMap,
+                     setUpBPFMap(sizeof(uint32_t), sizeof(uint32_t), UID_COUNTERSET_MAP_SIZE,
+                                 UID_COUNTERSET_MAP_PATH, BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(changeOwnerAndMode(UID_COUNTERSET_MAP_PATH, AID_NET_BW_ACCT,
+                                        "UidCounterSetMap", false));
+
+    ASSIGN_OR_RETURN(mUidStatsMap,
+                     setUpBPFMap(sizeof(struct StatsKey), sizeof(struct StatsValue),
+                                 UID_STATS_MAP_SIZE, UID_STATS_MAP_PATH, BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(changeOwnerAndMode(UID_STATS_MAP_PATH, AID_NET_BW_STATS, "UidStatsMap",
+                                        false));
+
+    ASSIGN_OR_RETURN(mTagStatsMap,
+                     setUpBPFMap(sizeof(struct StatsKey), sizeof(struct StatsValue),
+                                 TAG_STATS_MAP_SIZE, TAG_STATS_MAP_PATH, BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(changeOwnerAndMode(TAG_STATS_MAP_PATH, AID_NET_BW_STATS, "TagStatsMap",
+                                        false));
+
+    ASSIGN_OR_RETURN(mIfaceIndexNameMap,
+                     setUpBPFMap(sizeof(uint32_t), IFNAMSIZ, IFACE_INDEX_NAME_MAP_SIZE,
+                                 IFACE_INDEX_NAME_MAP_PATH, BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(changeOwnerAndMode(IFACE_INDEX_NAME_MAP_PATH, AID_NET_BW_STATS,
+                                        "IfaceIndexNameMap", false));
+
+    ASSIGN_OR_RETURN(mDozableUidMap,
+                     setUpBPFMap(sizeof(uint32_t), sizeof(uint8_t), UID_OWNER_MAP_SIZE,
+                                 DOZABLE_UID_MAP_PATH, BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(changeOwnerAndMode(DOZABLE_UID_MAP_PATH, AID_ROOT, "DozableUidMap", true));
+    RETURN_IF_NOT_OK(initialOwnerMap(mDozableUidMap, "DozableUidMap"));
+
+    ASSIGN_OR_RETURN(mStandbyUidMap,
+                     setUpBPFMap(sizeof(uint32_t), sizeof(uint8_t), UID_OWNER_MAP_SIZE,
+                                 STANDBY_UID_MAP_PATH, BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(changeOwnerAndMode(STANDBY_UID_MAP_PATH, AID_ROOT, "StandbyUidMap", true));
+    RETURN_IF_NOT_OK(initialOwnerMap(mStandbyUidMap, "StandbyUidMap"));
+
+    ASSIGN_OR_RETURN(mPowerSaveUidMap,
+                     setUpBPFMap(sizeof(uint32_t), sizeof(uint8_t), UID_OWNER_MAP_SIZE,
+                                 POWERSAVE_UID_MAP_PATH, BPF_MAP_TYPE_HASH));
+    RETURN_IF_NOT_OK(changeOwnerAndMode(POWERSAVE_UID_MAP_PATH, AID_ROOT, "PowerSaveUidMap", true));
+    RETURN_IF_NOT_OK(initialOwnerMap(mPowerSaveUidMap, "PowerSaveUidMap"));
+
+    ASSIGN_OR_RETURN(mIfaceStatsMap,
+                     setUpBPFMap(sizeof(uint32_t), sizeof(struct StatsValue), 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 netdutils::status::ok;
+}
+
+Status TrafficController::start() {
+
     if (!ebpfSupported) {
         return netdutils::status::ok;
     }
@@ -103,75 +198,7 @@
      * if the socket no longer exist
      */
 
-    ALOGI("START to load TrafficController");
-
-    ASSIGN_OR_RETURN(mCookieTagMap,
-                     setUpBPFMap(sizeof(uint64_t), sizeof(struct UidTag), COOKIE_UID_MAP_SIZE,
-                                 COOKIE_UID_MAP_PATH, BPF_MAP_TYPE_HASH));
-
-    // Allow both netd and system server to obtain map fd from the path. Chown the group to
-    // net_bw_acct does not grant all process in that group the permission to access bpf maps. They
-    // still need correct sepolicy to read/write the map. And only system_server and netd have that
-    // permission for now.
-    int ret = chown(COOKIE_UID_MAP_PATH, AID_ROOT, AID_NET_BW_ACCT);
-    if (ret) {
-        return statusFromErrno(errno, "change cookieTagMap group failed.");
-    }
-    ret = chmod(COOKIE_UID_MAP_PATH, S_IRWXU | S_IRGRP | S_IWGRP);
-    if (ret) {
-        return statusFromErrno(errno, "change cookieTagMap mode failed.");
-    }
-
-    ASSIGN_OR_RETURN(mUidCounterSetMap,
-                     setUpBPFMap(sizeof(uint32_t), sizeof(uint32_t), UID_COUNTERSET_MAP_SIZE,
-                                 UID_COUNTERSET_MAP_PATH, BPF_MAP_TYPE_HASH));
-    // Only netd can access the file.
-    ret = chmod(UID_COUNTERSET_MAP_PATH, S_IRWXU);
-    if (ret) {
-        return statusFromErrno(errno, "change uidCounterSetMap mode failed.");
-    }
-
-    ASSIGN_OR_RETURN(mUidStatsMap,
-                     setUpBPFMap(sizeof(struct StatsKey), sizeof(struct StatsValue),
-                                 UID_STATS_MAP_SIZE, UID_STATS_MAP_PATH, BPF_MAP_TYPE_HASH));
-    // Change the file mode of pinned map so both netd and system server can get the map fd
-    // from it.
-    ret = chown(UID_STATS_MAP_PATH, AID_ROOT, AID_NET_BW_ACCT);
-    if (ret) {
-        return statusFromErrno(errno, "change uidStatsMap group failed.");
-    }
-    ret = chmod(UID_STATS_MAP_PATH, S_IRWXU | S_IRGRP | S_IWGRP);
-    if (ret) {
-        return statusFromErrno(errno, "change uidStatsMap mode failed.");
-    }
-
-    ASSIGN_OR_RETURN(mTagStatsMap,
-                     setUpBPFMap(sizeof(struct StatsKey), sizeof(struct StatsValue),
-                                 TAG_STATS_MAP_SIZE, TAG_STATS_MAP_PATH, BPF_MAP_TYPE_HASH));
-    // Change the file mode of pinned map so both netd and system server can get the map fd
-    // from the path.
-    ret = chown(TAG_STATS_MAP_PATH, AID_ROOT, AID_NET_BW_STATS);
-    if (ret) {
-        return statusFromErrno(errno, "change tagStatsMap group failed.");
-    }
-    ret = chmod(TAG_STATS_MAP_PATH, S_IRWXU | S_IRGRP | S_IWGRP);
-    if (ret) {
-        return statusFromErrno(errno, "change tagStatsMap mode failed.");
-    }
-
-    ASSIGN_OR_RETURN(mIfaceIndexNameMap,
-                     setUpBPFMap(sizeof(uint32_t), IFNAMSIZ, IFACE_INDEX_NAME_MAP_SIZE,
-                                 IFACE_INDEX_NAME_MAP_PATH, BPF_MAP_TYPE_HASH));
-    // Change the file mode of pinned map so both netd and system server can get the map fd
-    // from the path.
-    ret = chown(IFACE_INDEX_NAME_MAP_PATH, AID_ROOT, AID_NET_BW_STATS);
-    if (ret) {
-        return statusFromErrno(errno, "change ifaceIndexNameMap group failed.");
-    }
-    ret = chmod(IFACE_INDEX_NAME_MAP_PATH, S_IRWXU | S_IRGRP | S_IWGRP);
-    if (ret) {
-        return statusFromErrno(errno, "change ifaceIndexNameMap mode failed.");
-    }
+    RETURN_IF_NOT_OK(initMaps());
 
     // Fetch the list of currently-existing interfaces. At this point NetlinkHandler is
     // already running, so it will call addInterface() when any new interface appears.
@@ -181,19 +208,6 @@
         addInterface(ifacePair.first.c_str(), ifacePair.second);
     }
 
-    ASSIGN_OR_RETURN(mIfaceStatsMap,
-                     setUpBPFMap(sizeof(uint32_t), sizeof(struct StatsValue), IFACE_STATS_MAP_SIZE,
-                                 IFACE_STATS_MAP_PATH, BPF_MAP_TYPE_HASH));
-    // Change the file mode of pinned map so both netd and system server can get the map fd
-    // from the path.
-    ret = chown(IFACE_STATS_MAP_PATH, AID_ROOT, AID_NET_BW_STATS);
-    if (ret) {
-        return statusFromErrno(errno, "change ifaceStatsMap group failed.");
-    }
-    ret = chmod(IFACE_STATS_MAP_PATH, S_IRWXU | S_IRGRP);
-    if (ret) {
-        return statusFromErrno(errno, "change ifaceStatsMap mode failed.");
-    }
 
     auto result = makeSkDestroyListener();
     if (!isOk(result)) {
@@ -229,7 +243,7 @@
     std::vector<const char*> prog_args{
         "/system/bin/bpfloader",
     };
-    ret = access(BPF_INGRESS_PROG_PATH, R_OK);
+    int ret = access(BPF_INGRESS_PROG_PATH, R_OK);
     if (ret != 0 && errno == ENOENT) {
         prog_args.push_back((char*)"-i");
     }
@@ -456,6 +470,151 @@
     return res;
 }
 
+int TrafficController::updateOwnerMapEntry(const base::unique_fd& map_fd, uid_t uid,
+                                           FirewallRule rule, FirewallType type) {
+    int res = 0;
+    if (uid == NONEXISTENT_UID) {
+        ALOGE("This uid should never exist in maps: %u", uid);
+        return -EINVAL;
+    }
+
+    if (uid == UID_MAP_ENABLED) {
+        ALOGE("This uid is reserved for map state");
+        return -EINVAL;
+    }
+
+    if ((rule == ALLOW && type == WHITELIST) || (rule == DENY && type == BLACKLIST)) {
+        uint8_t flag = (type == WHITELIST) ? BPF_PASS : BPF_DROP;
+        res = writeToMapEntry(map_fd, &uid, &flag, BPF_ANY);
+        if (res) {
+            res = -errno;
+            ALOGE("Failed to add owner rule(uid: %u): %s", uid, strerror(errno));
+        }
+    } else if ((rule == ALLOW && type == BLACKLIST) || (rule == DENY && type == WHITELIST)) {
+        res = deleteMapEntry(map_fd, &uid);
+        if (res) {
+            res = -errno;
+            ALOGE("Failed to delete owner rule(uid: %u): %s", uid, strerror(errno));
+        }
+    } else {
+        //Cannot happen.
+        return -EINVAL;
+    }
+    return res;
+}
+
+int TrafficController::changeUidOwnerRule(ChildChain chain, uid_t uid, FirewallRule rule,
+                                          FirewallType type) {
+    std::lock_guard<std::mutex> guard(mOwnerMatchMutex);
+    if (!ebpfSupported) {
+        ALOGE("bpf is not set up, should use iptables rule");
+        return -ENOSYS;
+    }
+    switch (chain) {
+        case DOZABLE:
+            return updateOwnerMapEntry(mDozableUidMap, uid, rule, type);
+        case STANDBY:
+            return updateOwnerMapEntry(mStandbyUidMap, uid, rule, type);
+        case POWERSAVE:
+            return updateOwnerMapEntry(mPowerSaveUidMap, uid, rule, type);
+        case NONE:
+        default:
+            return -EINVAL;
+    }
+}
+
+int TrafficController::replaceUidsInMap(const base::unique_fd& map_fd,
+                                        const std::vector<int32_t> &uids,
+                                        FirewallRule rule, FirewallType type) {
+    std::set<int32_t> uidSet(uids.begin(), uids.end());
+    std::vector<uint32_t> uidsToDelete;
+    auto getUidsToDelete = [&uidsToDelete, &uidSet](void *key, const base::unique_fd&) {
+        uint32_t uid = *(uint32_t *)key;
+        if (uid != UID_MAP_ENABLED && uidSet.find((int32_t)uid) == uidSet.end()) {
+            uidsToDelete.push_back(uid);
+        }
+        return 0;
+    };
+    uint32_t nonExistentKey = NONEXISTENT_UID;
+    uint8_t dummyValue;
+    int ret = bpfIterateMap(nonExistentKey, dummyValue, map_fd, getUidsToDelete);
+
+    if (ret)  return ret;
+
+    for(auto uid : uidsToDelete) {
+        if(deleteMapEntry(map_fd, &uid)) {
+            ret = -errno;
+            ALOGE("Delete uid(%u) from owner Map %d failed: %s", uid, map_fd.get(),
+                  strerror(errno));
+            return -errno;
+        }
+    }
+
+    for (auto uid : uids) {
+        ret = updateOwnerMapEntry(map_fd, uid, rule, type);
+        if (ret) {
+            ALOGE("Failed to add owner rule(uid: %u, map: %d)", uid, map_fd.get());
+            return ret;
+        }
+    }
+    return 0;
+}
+
+int TrafficController::replaceUidOwnerMap(const std::string& name, bool isWhitelist,
+                                          const std::vector<int32_t>& uids) {
+    std::lock_guard<std::mutex> guard(mOwnerMatchMutex);
+    FirewallRule rule;
+    FirewallType type;
+    if (isWhitelist) {
+        type = WHITELIST;
+        rule = ALLOW;
+    } else {
+        type = BLACKLIST;
+        rule = DENY;
+    }
+    int ret;
+    if (!name.compare(FirewallController::LOCAL_DOZABLE)) {
+        ret = replaceUidsInMap(mDozableUidMap, uids, rule, type);
+    } else if (!name.compare(FirewallController::LOCAL_STANDBY)) {
+        ret = replaceUidsInMap(mStandbyUidMap, uids, rule, type);
+    } else if (!name.compare(FirewallController::LOCAL_POWERSAVE)) {
+        ret = replaceUidsInMap(mPowerSaveUidMap, uids, rule, type);
+    } else {
+        ALOGE("unknown chain name: %s", name.c_str());
+        return -EINVAL;
+    }
+    if (ret) {
+        ALOGE("Failed to clean up chain: %s: %s", name.c_str(), strerror(ret));
+        return ret;
+    }
+    return 0;
+}
+
+int TrafficController::toggleUidOwnerMap(ChildChain chain, bool enable) {
+    std::lock_guard<std::mutex> guard(mOwnerMatchMutex);
+    uint32_t keyUid = UID_MAP_ENABLED;
+    uint8_t mapState = enable ? 1 : 0;
+    int ret;
+    switch (chain) {
+        case DOZABLE:
+            ret = writeToMapEntry(mDozableUidMap, &keyUid, &mapState, BPF_EXIST);
+            break;
+        case STANDBY:
+            ret = writeToMapEntry(mStandbyUidMap, &keyUid, &mapState, BPF_EXIST);
+            break;
+        case POWERSAVE:
+            ret = writeToMapEntry(mPowerSaveUidMap, &keyUid, &mapState, BPF_EXIST);
+            break;
+        default:
+            return -EINVAL;
+    }
+    if (ret) {
+        ret = -errno;
+        ALOGE("Failed to toggleUidOwnerMap(%d): %s", chain, strerror(errno));
+    }
+    return ret;
+}
+
 bool TrafficController::checkBpfStatsEnable() {
     return ebpfSupported;
 }
diff --git a/server/TrafficController.h b/server/TrafficController.h
index 0991288..bc16d39 100644
--- a/server/TrafficController.h
+++ b/server/TrafficController.h
@@ -20,8 +20,10 @@
 #include <linux/bpf.h>
 
 #include <netdutils/StatusOr.h>
+#include "FirewallController.h"
 #include "NetlinkListener.h"
 #include "Network.h"
+#include "android-base/thread_annotations.h"
 #include "android-base/unique_fd.h"
 
 // Since we cannot garbage collect the stats map since device boot, we need to make these maps as
@@ -51,6 +53,7 @@
 
 class TrafficController {
   public:
+    TrafficController();
     /*
      * Initialize the whole controller
      */
@@ -100,6 +103,21 @@
      */
     int addInterface(const char* name, uint32_t ifaceIndex);
 
+    int changeUidOwnerRule(ChildChain chain, const uid_t uid, FirewallRule rule, FirewallType type);
+
+    int removeUidOwnerRule(const uid_t uid);
+
+    int replaceUidOwnerMap(const std::string& name, bool isWhitelist,
+                           const std::vector<int32_t>& uids);
+
+    int updateOwnerMapEntry(const base::unique_fd& map_fd, uid_t uid, FirewallRule rule,
+                            FirewallType type);
+
+    int replaceUidsInMap(const base::unique_fd& map_fd, const std::vector<int32_t> &uids,
+                         FirewallRule rule, FirewallType type);
+
+    int toggleUidOwnerMap(ChildChain chain, bool enable);
+
   private:
     /*
      * mCookieTagMap: Store the corresponding tag and uid for a specific socket.
@@ -152,13 +170,34 @@
      */
     base::unique_fd mIfaceStatsMap;
 
+    /*
+     * mDozableUidMap: Store uids that have related rules in dozable mode owner match
+     * chain.
+     */
+    base::unique_fd mDozableUidMap GUARDED_BY(mOwnerMatchMutex);
+
+    /*
+     * mStandbyUidMap: Store uids that have related rules in standby mode owner match
+     * chain.
+     */
+    base::unique_fd mStandbyUidMap GUARDED_BY(mOwnerMatchMutex);
+
+    /*
+     * mPowerSaveUidMap: Store uids that have related rules in power save mode owner match
+     * chain.
+     */
+    base::unique_fd mPowerSaveUidMap GUARDED_BY(mOwnerMatchMutex);
+
     std::unique_ptr<NetlinkListenerInterface> mSkDestroyListener;
 
     bool ebpfSupported;
 
+    std::mutex mOwnerMatchMutex;
+
     netdutils::Status loadAndAttachProgram(bpf_attach_type type, const char* path, const char* name,
                                            base::unique_fd& cg_fd);
 
+    netdutils::Status initMaps();
     // For testing
     friend class TrafficControllerTest;
 };
diff --git a/server/TrafficControllerTest.cpp b/server/TrafficControllerTest.cpp
index 1ad3915..2d47985 100644
--- a/server/TrafficControllerTest.cpp
+++ b/server/TrafficControllerTest.cpp
@@ -33,6 +33,7 @@
 #include <android-base/strings.h>
 
 #include <netdutils/MockSyscalls.h>
+#include "FirewallController.h"
 #include "TrafficController.h"
 #include "bpf/BpfUtils.h"
 
@@ -53,6 +54,8 @@
 
 constexpr int TEST_MAP_SIZE = 10;
 constexpr uid_t TEST_UID = 10086;
+constexpr uid_t TEST_UID2 = 54321;
+constexpr uid_t TEST_UID3 = 98765;
 constexpr uint32_t TEST_TAG = 42;
 constexpr int TEST_COUNTERSET = 1;
 constexpr int DEFAULT_COUNTERSET = 0;
@@ -70,8 +73,12 @@
     unique_fd mFakeUidCounterSetMap;
     unique_fd mFakeUidStatsMap;
     unique_fd mFakeTagStatsMap;
+    unique_fd mFakeDozableUidMap;
+    unique_fd mFakeStandbyUidMap;
+    unique_fd mFakePowerSaveUidMap;
 
     void SetUp() {
+        std::lock_guard<std::mutex> guard(mTc.mOwnerMatchMutex);
         SKIP_IF_BPF_NOT_SUPPORTED;
 
         mFakeCookieTagMap = unique_fd(createMap(BPF_MAP_TYPE_HASH, sizeof(uint64_t),
@@ -90,6 +97,17 @@
                                                sizeof(struct StatsValue), TEST_MAP_SIZE, 0));
         ASSERT_LE(0, mFakeTagStatsMap);
 
+        mFakeDozableUidMap = unique_fd(
+            createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint32_t), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakeDozableUidMap);
+
+        mFakeStandbyUidMap = unique_fd(
+            createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint32_t), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakeStandbyUidMap);
+
+        mFakePowerSaveUidMap = unique_fd(
+            createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint32_t), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakePowerSaveUidMap);
         // Make sure trafficController use the eBPF code path.
         mTc.ebpfSupported = true;
 
@@ -97,6 +115,9 @@
         mTc.mUidCounterSetMap.reset(mFakeUidCounterSetMap);
         mTc.mUidStatsMap.reset(mFakeUidStatsMap);
         mTc.mTagStatsMap.reset(mFakeTagStatsMap);
+        mTc.mDozableUidMap.reset(mFakeDozableUidMap);
+        mTc.mStandbyUidMap.reset(mFakeStandbyUidMap);
+        mTc.mPowerSaveUidMap.reset(mFakePowerSaveUidMap);
     }
 
     int setUpSocketAndTag(int protocol, uint64_t* cookie, uint32_t tag, uid_t uid) {
@@ -140,6 +161,62 @@
         key->tag = tag;
     }
 
+    void checkUidOwnerRuleForChain(ChildChain chain, const unique_fd& targetMap) {
+        uint32_t uid = TEST_UID;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, DENY, BLACKLIST));
+        uint8_t value;
+        EXPECT_EQ(0, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ((uint8_t)BPF_DROP, value);
+
+        uid = TEST_UID2;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, ALLOW, WHITELIST));
+        EXPECT_EQ(0, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ((uint8_t)BPF_PASS, value);
+
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, DENY, WHITELIST));
+        EXPECT_EQ(-1, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ(ENOENT, errno);
+
+        uid = TEST_UID;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, ALLOW, BLACKLIST));
+        EXPECT_EQ(-1, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ(ENOENT, errno);
+
+        uid = TEST_UID3;
+        EXPECT_EQ(-ENOENT, mTc.changeUidOwnerRule(chain, uid, ALLOW, BLACKLIST));
+        EXPECT_EQ(-1, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ(ENOENT, errno);
+    }
+
+    void checkEachUidValue(const std::vector<int32_t>& uids, const uint8_t expectValue,
+                           const unique_fd& targetMap) {
+        uint8_t value;
+        for (auto uid : uids) {
+            EXPECT_EQ(0, findMapEntry(targetMap, &uid, &value));
+            EXPECT_EQ((uint8_t)expectValue, value);
+        }
+        std::set<uint32_t> uidSet(uids.begin(), uids.end());
+        auto checkNoOtherUid = [&uidSet](void *key, const base::unique_fd&) {
+            int32_t uid = *(int32_t *)key;
+            EXPECT_NE(uidSet.end(), uidSet.find(uid));
+            return 0;
+        };
+        uint32_t nonExistentKey = NONEXISTENT_UID;
+        uint8_t dummyValue;
+        EXPECT_EQ(0, bpfIterateMap(nonExistentKey, dummyValue, targetMap, checkNoOtherUid));
+    }
+
+    void checkUidMapReplace(const std::string& name, const std::vector<int32_t>& uids,
+                            const unique_fd& targetMap) {
+        bool isWhitelist = true;
+        EXPECT_EQ(0, mTc.replaceUidOwnerMap(name, isWhitelist, uids));
+        checkEachUidValue(uids, BPF_PASS, targetMap);
+
+        isWhitelist = false;
+        EXPECT_EQ(0, mTc.replaceUidOwnerMap(name, isWhitelist, uids));
+        checkEachUidValue(uids, BPF_DROP, targetMap);
+    }
+
     void TearDown() {
         mFakeCookieTagMap.reset();
         mFakeUidCounterSetMap.reset();
@@ -357,5 +434,54 @@
     ASSERT_EQ((uint64_t)200, statsMapResult.rxBytes);
 }
 
+TEST_F(TrafficControllerTest, TestUpdateOwnerMapEntry) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    uint32_t uid = TEST_UID;
+    ASSERT_EQ(0, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, DENY, BLACKLIST));
+    uint8_t value;
+    ASSERT_EQ(0, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ((uint8_t)BPF_DROP, value);
+
+    uid = TEST_UID2;
+    ASSERT_EQ(0, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, ALLOW, WHITELIST));
+    ASSERT_EQ(0, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ((uint8_t)BPF_PASS, value);
+
+    ASSERT_EQ(0, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, DENY, WHITELIST));
+    ASSERT_EQ(-1, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ(ENOENT, errno);
+
+    uid = TEST_UID;
+    ASSERT_EQ(0, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, ALLOW, BLACKLIST));
+    ASSERT_EQ(-1, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ(ENOENT, errno);
+
+    uid = TEST_UID3;
+    ASSERT_EQ(-ENOENT, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, ALLOW, BLACKLIST));
+    ASSERT_EQ(-1, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ(ENOENT, errno);
+}
+
+TEST_F(TrafficControllerTest, TestChangeUidOwnerRule) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    checkUidOwnerRuleForChain(DOZABLE, mFakeDozableUidMap);
+    checkUidOwnerRuleForChain(STANDBY, mFakeStandbyUidMap);
+    checkUidOwnerRuleForChain(POWERSAVE, mFakePowerSaveUidMap);
+    ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(NONE, TEST_UID, ALLOW, WHITELIST));
+    ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(INVALID_CHAIN, TEST_UID, ALLOW, WHITELIST));
+}
+
+TEST_F(TrafficControllerTest, TestReplaceUidOwnerMap) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    std::vector<int32_t> uids = {TEST_UID, TEST_UID2, TEST_UID3};
+    checkUidMapReplace("fw_dozable", uids, mFakeDozableUidMap);
+    checkUidMapReplace("fw_standby", uids, mFakeStandbyUidMap);
+    checkUidMapReplace("fw_powersave", uids, mFakePowerSaveUidMap);
+    ASSERT_EQ(-EINVAL, mTc.replaceUidOwnerMap("unknow", true, uids));
+}
+
 }  // namespace net
 }  // namespace android