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