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/bpfloader/BpfLoader.cpp b/bpfloader/BpfLoader.cpp
index 7cb4616..9d75077 100644
--- a/bpfloader/BpfLoader.cpp
+++ b/bpfloader/BpfLoader.cpp
@@ -44,8 +44,7 @@
 #include <netdutils/Misc.h>
 #include <netdutils/Slice.h>
 #include "bpf/BpfUtils.h"
-
-#include "bpf_shared.h"
+#include "bpf/bpf_shared.h"
 
 using android::base::unique_fd;
 using android::netdutils::Slice;
@@ -84,29 +83,36 @@
         (x)[4], (x)[5], (x)[6], (x)[7]    \
     }
 
+#define DECLARE_MAP(_mapFd, _mapPath)                             \
+    unique_fd _mapFd(android::bpf::mapRetrieve((_mapPath), 0));   \
+    if (_mapFd < 0) {                                             \
+        FAIL("Failed to get map from %s", (_mapPath));            \
+    }
+
 #define MAP_CMD_SIZE 16
 #define LOG_BUF_SIZE 65536
 
 namespace android {
 namespace bpf {
 
-void makeFdReplacePattern(uint64_t code, uint64_t mapFd, char* pattern, char* cmd) {
-    char mapCode[sizeof(uint64_t)];
-    char mapCmd[sizeof(uint64_t)];
-    // The byte order is little endian for arm devices.
-    for (uint32_t i = 0; i < sizeof(uint64_t); i++) {
-        mapCode[i] = (code >> (i * 8)) & 0xFF;
-        mapCmd[i] = (mapFd >> (i * 8)) & 0xFF;
+struct ReplacePattern {
+    std::array<uint8_t, MAP_CMD_SIZE> search;
+    std::array<uint8_t, MAP_CMD_SIZE> replace;
+
+    ReplacePattern(uint64_t dummyFd, int realFd) {
+        // Ensure that the fd numbers are big-endian.
+        uint8_t beDummyFd[sizeof(uint64_t)];
+        uint8_t beRealFd[sizeof(uint64_t)];
+        for (size_t i = 0; i < sizeof(uint64_t); i++) {
+            beDummyFd[i] = (dummyFd >> (i * 8)) & 0xFF;
+            beRealFd[i] = (realFd >> (i * 8)) & 0xFF;
+        }
+        search = MAP_SEARCH_PATTERN(beDummyFd);
+        replace = MAP_REPLACE_PATTERN(beRealFd);
     }
+};
 
-    char tmpPattern[] = MAP_SEARCH_PATTERN(mapCode);
-    memcpy(pattern, tmpPattern, MAP_CMD_SIZE);
-    char tmpCmd[] = MAP_REPLACE_PATTERN(mapCmd);
-    memcpy(cmd, tmpCmd, MAP_CMD_SIZE);
-}
-
-int loadProg(const char* path, int cookieTagMap, int uidStatsMap, int tagStatsMap,
-             int uidCounterSetMap, int ifaceStatsMap) {
+int loadProg(const char* path, const std::vector<ReplacePattern> &mapPatterns) {
     int fd = open(path, O_RDONLY);
     if (fd == -1) {
         FAIL("Failed to open %s program: %s", path, strerror(errno));
@@ -159,28 +165,6 @@
     char* prog = new char[progSize]();
     memcpy(prog, progSection, progSize);
 
-    char cookieTagMapFdPattern[MAP_CMD_SIZE];
-    char cookieTagMapFdLoadByte[MAP_CMD_SIZE];
-    makeFdReplacePattern(COOKIE_TAG_MAP, cookieTagMap, cookieTagMapFdPattern,
-                         cookieTagMapFdLoadByte);
-
-    char uidCounterSetMapFdPattern[MAP_CMD_SIZE];
-    char uidCounterSetMapFdLoadByte[MAP_CMD_SIZE];
-    makeFdReplacePattern(UID_COUNTERSET_MAP, uidCounterSetMap, uidCounterSetMapFdPattern,
-                         uidCounterSetMapFdLoadByte);
-
-    char tagStatsMapFdPattern[MAP_CMD_SIZE];
-    char tagStatsMapFdLoadByte[MAP_CMD_SIZE];
-    makeFdReplacePattern(TAG_STATS_MAP, tagStatsMap, tagStatsMapFdPattern, tagStatsMapFdLoadByte);
-
-    char uidStatsMapFdPattern[MAP_CMD_SIZE];
-    char uidStatsMapFdLoadByte[MAP_CMD_SIZE];
-    makeFdReplacePattern(UID_STATS_MAP, uidStatsMap, uidStatsMapFdPattern, uidStatsMapFdLoadByte);
-
-    char ifaceStatsMapFdPattern[MAP_CMD_SIZE];
-    char ifaceStatsMapFdLoadByte[MAP_CMD_SIZE];
-    makeFdReplacePattern(IFACE_STATS_MAP, ifaceStatsMap, ifaceStatsMapFdPattern,
-                         ifaceStatsMapFdLoadByte);
 
     char* mapHead = prog;
     while ((uint64_t)(mapHead - prog + MAP_CMD_SIZE) < progSize) {
@@ -191,22 +175,9 @@
         // replace, and if so, replace it.
         mapHead = (char*)memchr(mapHead, MAP_LD_CMD_HEAD, progSize);
         if (!mapHead) break;
-        if ((uint64_t)(mapHead - prog + MAP_CMD_SIZE) < progSize) {
-            if (!memcmp(mapHead, cookieTagMapFdPattern, MAP_CMD_SIZE)) {
-                memcpy(mapHead, cookieTagMapFdLoadByte, MAP_CMD_SIZE);
-                mapHead += MAP_CMD_SIZE;
-            } else if (!memcmp(mapHead, uidCounterSetMapFdPattern, MAP_CMD_SIZE)) {
-                memcpy(mapHead, uidCounterSetMapFdLoadByte, MAP_CMD_SIZE);
-                mapHead += MAP_CMD_SIZE;
-            } else if (!memcmp(mapHead, tagStatsMapFdPattern, MAP_CMD_SIZE)) {
-                memcpy(mapHead, tagStatsMapFdLoadByte, MAP_CMD_SIZE);
-                mapHead += MAP_CMD_SIZE;
-            } else if (!memcmp(mapHead, uidStatsMapFdPattern, MAP_CMD_SIZE)) {
-                memcpy(mapHead, uidStatsMapFdLoadByte, MAP_CMD_SIZE);
-                mapHead += MAP_CMD_SIZE;
-            } else if (!memcmp(mapHead, ifaceStatsMapFdPattern, MAP_CMD_SIZE)) {
-                memcpy(mapHead, ifaceStatsMapFdLoadByte, MAP_CMD_SIZE);
-                mapHead += MAP_CMD_SIZE;
+        for (const auto& pattern : mapPatterns) {
+            if (!memcmp(mapHead, pattern.search.data(), MAP_CMD_SIZE)) {
+                memcpy(mapHead, pattern.replace.data(), MAP_CMD_SIZE);
             }
         }
         mapHead++;
@@ -221,9 +192,7 @@
 }
 
 int loadAndAttachProgram(bpf_attach_type type, const char* path, const char* name,
-                         const unique_fd& cookieTagMap, const unique_fd& uidCounterSetMap,
-                         const unique_fd& uidStatsMap, const unique_fd& tagStatsMap,
-                         const unique_fd& ifaceStatsMap) {
+                         std::vector<ReplacePattern> mapPatterns) {
     unique_fd cg_fd(open(CGROUP_ROOT_PATH, O_DIRECTORY | O_RDONLY | O_CLOEXEC));
     if (cg_fd < 0) {
         FAIL("Failed to open the cgroup directory");
@@ -231,17 +200,13 @@
 
     unique_fd fd;
     if (type == BPF_CGROUP_INET_EGRESS) {
-        fd.reset(loadProg(INGRESS_PROG, cookieTagMap.get(), uidStatsMap.get(), tagStatsMap.get(),
-                          uidCounterSetMap.get(), ifaceStatsMap.get()));
+        fd.reset(loadProg(INGRESS_PROG, mapPatterns));
     } else if (type == BPF_CGROUP_INET_INGRESS) {
-        fd.reset(loadProg(EGRESS_PROG, cookieTagMap.get(), uidStatsMap.get(), tagStatsMap.get(),
-                          uidCounterSetMap.get(), ifaceStatsMap.get()));
+        fd.reset(loadProg(EGRESS_PROG, mapPatterns));
     } else if (!strcmp(name, "xt_bpf_ingress_prog")) {
-        fd.reset(loadProg(XT_BPF_INGRESS_PROG, cookieTagMap.get(), uidStatsMap.get(),
-                          tagStatsMap.get(), uidCounterSetMap.get(), ifaceStatsMap.get()));
+        fd.reset(loadProg(XT_BPF_INGRESS_PROG, mapPatterns));
     } else if (!strcmp(name, "xt_bpf_egress_prog")) {
-        fd.reset(loadProg(XT_BPF_EGRESS_PROG, cookieTagMap.get(), uidStatsMap.get(),
-                          tagStatsMap.get(), uidCounterSetMap.get(), ifaceStatsMap.get()));
+        fd.reset(loadProg(XT_BPF_EGRESS_PROG, mapPatterns));
     } else {
         FAIL("Unrecognized program type: %s", name);
     }
@@ -269,13 +234,17 @@
 
 using android::bpf::BPF_EGRESS_PROG_PATH;
 using android::bpf::BPF_INGRESS_PROG_PATH;
-using android::bpf::COOKIE_UID_MAP_PATH;
+using android::bpf::COOKIE_TAG_MAP_PATH;
+using android::bpf::DOZABLE_UID_MAP_PATH;
 using android::bpf::IFACE_STATS_MAP_PATH;
+using android::bpf::POWERSAVE_UID_MAP_PATH;
+using android::bpf::STANDBY_UID_MAP_PATH;
 using android::bpf::TAG_STATS_MAP_PATH;
 using android::bpf::UID_COUNTERSET_MAP_PATH;
 using android::bpf::UID_STATS_MAP_PATH;
 using android::bpf::XT_BPF_EGRESS_PROG_PATH;
 using android::bpf::XT_BPF_INGRESS_PROG_PATH;
+using android::bpf::ReplacePattern;
 
 static void usage(void) {
     ALOGE( "Usage: ./bpfloader [-i] [-e]\n"
@@ -287,30 +256,25 @@
 
 int main(int argc, char** argv) {
     int ret = 0;
-    unique_fd cookieTagMap(android::bpf::mapRetrieve(COOKIE_UID_MAP_PATH, 0));
-    if (cookieTagMap < 0) {
-        FAIL("Failed to get cookieTagMap");
-    }
+    DECLARE_MAP(cookieTagMap, COOKIE_TAG_MAP_PATH);
+    DECLARE_MAP(uidCounterSetMap, UID_COUNTERSET_MAP_PATH);
+    DECLARE_MAP(uidStatsMap, UID_STATS_MAP_PATH);
+    DECLARE_MAP(tagStatsMap, TAG_STATS_MAP_PATH);
+    DECLARE_MAP(ifaceStatsMap, IFACE_STATS_MAP_PATH);
+    DECLARE_MAP(dozableUidMap, DOZABLE_UID_MAP_PATH);
+    DECLARE_MAP(standbyUidMap, STANDBY_UID_MAP_PATH);
+    DECLARE_MAP(powerSaveUidMap, POWERSAVE_UID_MAP_PATH);
 
-    unique_fd uidCounterSetMap(android::bpf::mapRetrieve(UID_COUNTERSET_MAP_PATH, 0));
-    if (uidCounterSetMap < 0) {
-        FAIL("Failed to get uidCounterSetMap");
-    }
-
-    unique_fd uidStatsMap(android::bpf::mapRetrieve(UID_STATS_MAP_PATH, 0));
-    if (uidStatsMap < 0) {
-        FAIL("Failed to get uidStatsMap");
-    }
-
-    unique_fd tagStatsMap(android::bpf::mapRetrieve(TAG_STATS_MAP_PATH, 0));
-    if (tagStatsMap < 0) {
-        FAIL("Failed to get tagStatsMap");
-    }
-
-    unique_fd ifaceStatsMap(android::bpf::mapRetrieve(IFACE_STATS_MAP_PATH, 0));
-    if (ifaceStatsMap < 0) {
-        FAIL("Failed to get ifaceStatsMap");
-    }
+    const std::vector<ReplacePattern> mapPatterns = {
+        ReplacePattern(COOKIE_TAG_MAP, cookieTagMap.get()),
+        ReplacePattern(UID_COUNTERSET_MAP, uidCounterSetMap.get()),
+        ReplacePattern(UID_STATS_MAP, uidStatsMap.get()),
+        ReplacePattern(TAG_STATS_MAP, tagStatsMap.get()),
+        ReplacePattern(IFACE_STATS_MAP, ifaceStatsMap.get()),
+        ReplacePattern(DOZABLE_UID_MAP, dozableUidMap.get()),
+        ReplacePattern(STANDBY_UID_MAP, standbyUidMap.get()),
+        ReplacePattern(POWERSAVE_UID_MAP, powerSaveUidMap.get()),
+    };
 
     int opt;
     bool doIngress = false, doEgress = false, doPrerouting = false, doMangle = false;
@@ -335,32 +299,28 @@
     }
     if (doIngress) {
         ret = android::bpf::loadAndAttachProgram(BPF_CGROUP_INET_INGRESS, BPF_INGRESS_PROG_PATH,
-                                                 "ingress_prog", cookieTagMap, uidCounterSetMap,
-                                                 uidStatsMap, tagStatsMap, ifaceStatsMap);
+                                                 "ingress_prog", mapPatterns);
         if (ret) {
             FAIL("Failed to set up ingress program");
         }
     }
     if (doEgress) {
         ret = android::bpf::loadAndAttachProgram(BPF_CGROUP_INET_EGRESS, BPF_EGRESS_PROG_PATH,
-                                                 "egress_prog", cookieTagMap, uidCounterSetMap,
-                                                 uidStatsMap, tagStatsMap, ifaceStatsMap);
+                                                 "egress_prog", mapPatterns);
         if (ret) {
             FAIL("Failed to set up ingress program");
         }
     }
     if (doPrerouting) {
         ret = android::bpf::loadAndAttachProgram(
-            MAX_BPF_ATTACH_TYPE, XT_BPF_INGRESS_PROG_PATH, "xt_bpf_ingress_prog", cookieTagMap,
-            uidCounterSetMap, uidStatsMap, tagStatsMap, ifaceStatsMap);
+            MAX_BPF_ATTACH_TYPE, XT_BPF_INGRESS_PROG_PATH, "xt_bpf_ingress_prog", mapPatterns);
         if (ret) {
             FAIL("Failed to set up xt_bpf program");
         }
     }
     if (doMangle) {
         ret = android::bpf::loadAndAttachProgram(
-            MAX_BPF_ATTACH_TYPE, XT_BPF_EGRESS_PROG_PATH, "xt_bpf_egress_prog", cookieTagMap,
-            uidCounterSetMap, uidStatsMap, tagStatsMap, ifaceStatsMap);
+            MAX_BPF_ATTACH_TYPE, XT_BPF_EGRESS_PROG_PATH, "xt_bpf_egress_prog", mapPatterns);
         if (ret) {
             FAIL("Failed to set up xt_bpf program");
         }
diff --git a/bpfloader/bpf_egress.c b/bpfloader/bpf_egress.c
index 9c2cc5f..5cf648c 100644
--- a/bpfloader/bpf_egress.c
+++ b/bpfloader/bpf_egress.c
@@ -15,11 +15,7 @@
  */
 
 #include <linux/bpf.h>
-#include <linux/if_ether.h>
-#include <linux/if_packet.h>
-#include <linux/ip.h>
 #include "bpf_kern.h"
-#include "bpf_shared.h"
 
 ELF_SEC(BPF_PROG_SEC_NAME)
 int bpf_cgroup_egress(struct __sk_buff* skb) {
diff --git a/bpfloader/bpf_ingress.c b/bpfloader/bpf_ingress.c
index d25877e..4bd5170 100644
--- a/bpfloader/bpf_ingress.c
+++ b/bpfloader/bpf_ingress.c
@@ -15,11 +15,7 @@
  */
 
 #include <linux/bpf.h>
-#include <linux/if_ether.h>
-#include <linux/if_packet.h>
-#include <linux/ip.h>
 #include "bpf_kern.h"
-#include "bpf_shared.h"
 
 ELF_SEC(BPF_PROG_SEC_NAME)
 int bpf_cgroup_ingress(struct __sk_buff* skb) {
diff --git a/bpfloader/bpf_kern.h b/bpfloader/bpf_kern.h
index 806b45a..83ad5b1 100644
--- a/bpfloader/bpf_kern.h
+++ b/bpfloader/bpf_kern.h
@@ -15,8 +15,13 @@
  */
 
 #include <linux/bpf.h>
+#include <linux/if_ether.h>
+#include <linux/in.h>
+#include <linux/in6.h>
+#include <linux/ip.h>
+#include <linux/ipv6.h>
 #include <stdint.h>
-#include "bpf_shared.h"
+#include "bpf/bpf_shared.h"
 
 #define ELF_SEC(NAME) __attribute__((section(NAME), used))
 
@@ -48,12 +53,17 @@
 static uint32_t (*get_socket_uid)(struct __sk_buff* skb) = (void*)BPF_FUNC_get_socket_uid;
 static int (*bpf_skb_load_bytes)(struct __sk_buff* skb, int off, void* to,
                                  int len) = (void*)BPF_FUNC_skb_load_bytes;
-
 #define BPF_PASS 1
 #define BPF_DROP 0
 #define BPF_EGRESS 0
 #define BPF_INGRESS 1
 
+#define IP_PROTO_OFF offsetof(struct iphdr, protocol)
+#define IPV6_PROTO_OFF offsetof(struct ipv6hdr, nexthdr)
+#define IPPROTO_IHL_OFF 0
+#define TCP_FLAG_OFF 13
+#define RST_OFFSET 2
+
 static __always_inline int xt_bpf_count(struct __sk_buff* skb, int type) {
     uint32_t key = skb->ifindex;
     struct stats_value* value;
@@ -96,6 +106,71 @@
     }
 }
 
+static inline int bpf_owner_match(struct __sk_buff* skb, uint32_t uid) {
+    int offset = -1;
+    int ret = 0;
+    if (skb->protocol == ETH_P_IP) {
+        offset = IP_PROTO_OFF;
+        uint8_t proto, ihl;
+        uint16_t flag;
+        ret = bpf_skb_load_bytes(skb, offset, &proto, 1);
+        if (!ret) {
+            if (proto == IPPROTO_ESP) {
+                return 1;
+            } else if (proto == IPPROTO_TCP) {
+                ret = bpf_skb_load_bytes(skb, IPPROTO_IHL_OFF, &ihl, 1);
+                ihl = ihl & 0x0F;
+                ret = bpf_skb_load_bytes(skb, ihl * 4 + TCP_FLAG_OFF, &flag, 1);
+                if (ret == 0 && (flag >> RST_OFFSET & 1)) {
+                    return BPF_PASS;
+                }
+            }
+        }
+    } else if (skb->protocol == ETH_P_IPV6) {
+        offset = IPV6_PROTO_OFF;
+        uint8_t proto;
+        ret = bpf_skb_load_bytes(skb, offset, &proto, 1);
+        if (!ret) {
+            if (proto == IPPROTO_ESP) {
+                return BPF_PASS;
+            } else if (proto == IPPROTO_TCP) {
+                uint16_t flag;
+                ret = bpf_skb_load_bytes(skb, sizeof(struct ipv6hdr) + TCP_FLAG_OFF, &flag, 1);
+                if (ret == 0 && (flag >> RST_OFFSET & 1)) {
+                    return BPF_PASS;
+                }
+            }
+        }
+    }
+
+    if ((uid <= MAX_SYSTEM_UID) && (uid >= MIN_SYSTEM_UID)) return BPF_PASS;
+
+    // In each of these maps, the entry with key UID_MAP_ENABLED tells us whether that
+    // map is enabled or not.
+    // TODO: replace this with a map of size one that contains a config structure defined in
+    // bpf_shared.h that can be written by userspace and read here.
+    uint32_t mapSettingKey = UID_MAP_ENABLED;
+    uint8_t* ownerMatch;
+    uint8_t* mapEnabled = find_map_entry(DOZABLE_UID_MAP, &mapSettingKey);
+    if (mapEnabled && *mapEnabled) {
+        ownerMatch = find_map_entry(DOZABLE_UID_MAP, &uid);
+        if (ownerMatch) return *ownerMatch;
+        return BPF_DROP;
+    }
+    mapEnabled = find_map_entry(STANDBY_UID_MAP, &mapSettingKey);
+    if (mapEnabled && *mapEnabled) {
+        ownerMatch = find_map_entry(STANDBY_UID_MAP, &uid);
+        if (ownerMatch) return *ownerMatch;
+    }
+    mapEnabled = find_map_entry(POWERSAVE_UID_MAP, &mapSettingKey);
+    if (mapEnabled && *mapEnabled) {
+        ownerMatch = find_map_entry(POWERSAVE_UID_MAP, &uid);
+        if (ownerMatch) return *ownerMatch;
+        return BPF_DROP;
+    }
+    return BPF_PASS;
+}
+
 static __always_inline inline int bpf_traffic_account(struct __sk_buff* skb, int direction) {
     uint64_t cookie = get_socket_cookie(skb);
     struct uid_tag* utag = find_map_entry(COOKIE_TAG_MAP, &cookie);
@@ -121,5 +196,5 @@
 
     key.tag = 0;
     bpf_update_stats(skb, UID_STATS_MAP, direction, key);
-    return BPF_PASS;
+    return bpf_owner_match(skb, uid);
 }
diff --git a/bpfloader/cgroup_bpf_egress_prog.o b/bpfloader/cgroup_bpf_egress_prog.o
index ab29366..4575693 100644
--- a/bpfloader/cgroup_bpf_egress_prog.o
+++ b/bpfloader/cgroup_bpf_egress_prog.o
Binary files differ
diff --git a/bpfloader/cgroup_bpf_ingress_prog.o b/bpfloader/cgroup_bpf_ingress_prog.o
index da7ccca..b989bd3 100644
--- a/bpfloader/cgroup_bpf_ingress_prog.o
+++ b/bpfloader/cgroup_bpf_ingress_prog.o
Binary files differ
diff --git a/bpfloader/xt_bpf_egress_prog.c b/bpfloader/xt_bpf_egress_prog.c
index 8109030..ae4dc31 100644
--- a/bpfloader/xt_bpf_egress_prog.c
+++ b/bpfloader/xt_bpf_egress_prog.c
@@ -15,11 +15,7 @@
  */
 
 #include <linux/bpf.h>
-#include <linux/if_ether.h>
-#include <linux/if_packet.h>
-#include <linux/ip.h>
 #include "bpf_kern.h"
-#include "bpf_shared.h"
 
 ELF_SEC(BPF_PROG_SEC_NAME)
 int xt_bpf_egress_prog(struct __sk_buff* skb) {
diff --git a/bpfloader/xt_bpf_ingress_prog.c b/bpfloader/xt_bpf_ingress_prog.c
index f9852b2..dc823c3 100644
--- a/bpfloader/xt_bpf_ingress_prog.c
+++ b/bpfloader/xt_bpf_ingress_prog.c
@@ -15,11 +15,7 @@
  */
 
 #include <linux/bpf.h>
-#include <linux/if_ether.h>
-#include <linux/if_packet.h>
-#include <linux/ip.h>
 #include "bpf_kern.h"
-#include "bpf_shared.h"
 
 ELF_SEC(BPF_PROG_SEC_NAME)
 int xt_bpf_ingress_prog(struct __sk_buff* skb) {
diff --git a/libbpf/BpfNetworkStats.cpp b/libbpf/BpfNetworkStats.cpp
index e9ec0c9..22a1e79 100644
--- a/libbpf/BpfNetworkStats.cpp
+++ b/libbpf/BpfNetworkStats.cpp
@@ -297,7 +297,7 @@
 }
 
 int cleanStatsMap() {
-    base::unique_fd cookieTagMap(bpf::mapRetrieve(COOKIE_UID_MAP_PATH, BPF_OPEN_FLAGS));
+    base::unique_fd cookieTagMap(bpf::mapRetrieve(COOKIE_TAG_MAP_PATH, BPF_OPEN_FLAGS));
     int ret = 0;
     if (cookieTagMap < 0) {
         ret = -errno;
diff --git a/libbpf/include/bpf/BpfUtils.h b/libbpf/include/bpf/BpfUtils.h
index 344727c..47f3ad4 100644
--- a/libbpf/include/bpf/BpfUtils.h
+++ b/libbpf/include/bpf/BpfUtils.h
@@ -29,6 +29,9 @@
 #define ptr_to_u64(x) ((uint64_t)(uintptr_t)x)
 #define DEFAULT_LOG_LEVEL 1
 
+#define BPF_PASS 1
+#define BPF_DROP 0
+
 /* instruction set for bpf program */
 
 #define MEM_LD(SIZE) (BPF_LDX | BPF_SIZE(SIZE) | BPF_MEM)
@@ -106,17 +109,21 @@
 
 constexpr const char* CGROUP_ROOT_PATH = "/dev/cg2_bpf";
 
-constexpr const char* COOKIE_UID_MAP_PATH = BPF_PATH "/traffic_cookie_uid_map";
+constexpr const char* COOKIE_TAG_MAP_PATH = BPF_PATH "/traffic_cookie_tag_map";
 constexpr const char* UID_COUNTERSET_MAP_PATH = BPF_PATH "/traffic_uid_counterSet_map";
 constexpr const char* UID_STATS_MAP_PATH = BPF_PATH "/traffic_uid_stats_map";
 constexpr const char* TAG_STATS_MAP_PATH = BPF_PATH "/traffic_tag_stats_map";
 constexpr const char* IFACE_INDEX_NAME_MAP_PATH = BPF_PATH "/traffic_iface_index_name_map";
 constexpr const char* IFACE_STATS_MAP_PATH = BPF_PATH "/traffic_iface_stats_map";
+constexpr const char* DOZABLE_UID_MAP_PATH = BPF_PATH "/traffic_dozable_uid_map";
+constexpr const char* STANDBY_UID_MAP_PATH = BPF_PATH "/traffic_standby_uid_map";
+constexpr const char* POWERSAVE_UID_MAP_PATH = BPF_PATH "/traffic_powersave_uid_map";
 
 const StatsKey NONEXISTENT_STATSKEY = {
     .uid = DEFAULT_OVERFLOWUID,
 };
 
+const uint32_t NONEXISTENT_UID = DEFAULT_OVERFLOWUID;
 const uint32_t NONEXISTENT_IFACE_STATS_KEY = 0;
 
 int createMap(bpf_map_type map_type, uint32_t key_size, uint32_t value_size,
diff --git a/bpfloader/bpf_shared.h b/libbpf/include/bpf/bpf_shared.h
similarity index 66%
rename from bpfloader/bpf_shared.h
rename to libbpf/include/bpf/bpf_shared.h
index 448867a..2c9d329 100644
--- a/bpfloader/bpf_shared.h
+++ b/libbpf/include/bpf/bpf_shared.h
@@ -23,3 +23,13 @@
 #define UID_STATS_MAP 0xbfdaafffffffffff
 #define TAG_STATS_MAP 0xbfaaafffffffffff
 #define IFACE_STATS_MAP 0xbf1faceaafffffff
+#define DOZABLE_UID_MAP 0Xbfd0ab1e1dafffff
+#define STANDBY_UID_MAP 0Xbfadb1daffffffff
+#define POWERSAVE_UID_MAP 0Xbf0eae1dafffffff
+// These are also defined in NetdConstants.h, but we want to minimize the number of headers
+// included by the BPF kernel program.
+// TODO: refactor the the following constant into a seperate file so
+// NetdConstants.h can also include it from there.
+#define MIN_SYSTEM_UID 0
+#define MAX_SYSTEM_UID 9999
+#define UID_MAP_ENABLED UINT32_MAX
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