diff --git a/service/native/TrafficController.cpp b/service/native/TrafficController.cpp
new file mode 100644
index 0000000..c24a41b
--- /dev/null
+++ b/service/native/TrafficController.cpp
@@ -0,0 +1,910 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#define LOG_TAG "TrafficController"
+#include <inttypes.h>
+#include <linux/if_ether.h>
+#include <linux/in.h>
+#include <linux/inet_diag.h>
+#include <linux/netlink.h>
+#include <linux/sock_diag.h>
+#include <linux/unistd.h>
+#include <net/if.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/socket.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <sys/utsname.h>
+#include <sys/wait.h>
+#include <mutex>
+#include <unordered_set>
+#include <vector>
+
+#include <android-base/stringprintf.h>
+#include <android-base/strings.h>
+#include <android-base/unique_fd.h>
+#include <netdutils/StatusOr.h>
+
+#include <netdutils/Syscalls.h>
+#include <netdutils/Utils.h>
+#include "TrafficController.h"
+#include "bpf/BpfMap.h"
+
+#include "netdutils/DumpWriter.h"
+
+namespace android {
+namespace net {
+
+using base::StringPrintf;
+using base::unique_fd;
+using bpf::BpfMap;
+using bpf::OVERFLOW_COUNTERSET;
+using bpf::synchronizeKernelRCU;
+using netdutils::DumpWriter;
+using netdutils::getIfaceList;
+using netdutils::NetlinkListener;
+using netdutils::NetlinkListenerInterface;
+using netdutils::ScopedIndent;
+using netdutils::Slice;
+using netdutils::sSyscalls;
+using netdutils::Status;
+using netdutils::statusFromErrno;
+using netdutils::StatusOr;
+using netdutils::status::ok;
+
+constexpr int kSockDiagMsgType = SOCK_DIAG_BY_FAMILY;
+constexpr int kSockDiagDoneMsgType = NLMSG_DONE;
+
+const char* TrafficController::LOCAL_DOZABLE = "fw_dozable";
+const char* TrafficController::LOCAL_STANDBY = "fw_standby";
+const char* TrafficController::LOCAL_POWERSAVE = "fw_powersave";
+const char* TrafficController::LOCAL_RESTRICTED = "fw_restricted";
+
+static_assert(BPF_PERMISSION_INTERNET == INetd::PERMISSION_INTERNET,
+              "Mismatch between BPF and AIDL permissions: PERMISSION_INTERNET");
+static_assert(BPF_PERMISSION_UPDATE_DEVICE_STATS == INetd::PERMISSION_UPDATE_DEVICE_STATS,
+              "Mismatch between BPF and AIDL permissions: PERMISSION_UPDATE_DEVICE_STATS");
+
+#define FLAG_MSG_TRANS(result, flag, value) \
+    do {                                    \
+        if ((value) & (flag)) {             \
+            (result).append(" " #flag);     \
+            (value) &= ~(flag);             \
+        }                                   \
+    } while (0)
+
+const std::string uidMatchTypeToString(uint8_t match) {
+    std::string matchType;
+    FLAG_MSG_TRANS(matchType, HAPPY_BOX_MATCH, match);
+    FLAG_MSG_TRANS(matchType, PENALTY_BOX_MATCH, match);
+    FLAG_MSG_TRANS(matchType, DOZABLE_MATCH, match);
+    FLAG_MSG_TRANS(matchType, STANDBY_MATCH, match);
+    FLAG_MSG_TRANS(matchType, POWERSAVE_MATCH, match);
+    FLAG_MSG_TRANS(matchType, RESTRICTED_MATCH, match);
+    FLAG_MSG_TRANS(matchType, IIF_MATCH, match);
+    if (match) {
+        return StringPrintf("Unknown match: %u", match);
+    }
+    return matchType;
+}
+
+bool TrafficController::hasUpdateDeviceStatsPermission(uid_t uid) {
+    // This implementation is the same logic as method ActivityManager#checkComponentPermission.
+    // It implies that the calling uid can never be the same as PER_USER_RANGE.
+    uint32_t appId = uid % PER_USER_RANGE;
+    return ((appId == AID_ROOT) || (appId == AID_SYSTEM) ||
+            mPrivilegedUser.find(appId) != mPrivilegedUser.end());
+}
+
+const std::string UidPermissionTypeToString(int permission) {
+    if (permission == INetd::PERMISSION_NONE) {
+        return "PERMISSION_NONE";
+    }
+    if (permission == INetd::PERMISSION_UNINSTALLED) {
+        // This should never appear in the map, complain loudly if it does.
+        return "PERMISSION_UNINSTALLED error!";
+    }
+    std::string permissionType;
+    FLAG_MSG_TRANS(permissionType, BPF_PERMISSION_INTERNET, permission);
+    FLAG_MSG_TRANS(permissionType, BPF_PERMISSION_UPDATE_DEVICE_STATS, permission);
+    if (permission) {
+        return StringPrintf("Unknown permission: %u", permission);
+    }
+    return permissionType;
+}
+
+StatusOr<std::unique_ptr<NetlinkListenerInterface>> TrafficController::makeSkDestroyListener() {
+    const auto& sys = sSyscalls.get();
+    ASSIGN_OR_RETURN(auto event, sys.eventfd(0, EFD_CLOEXEC));
+    const int domain = AF_NETLINK;
+    const int type = SOCK_DGRAM | SOCK_CLOEXEC | SOCK_NONBLOCK;
+    const int protocol = NETLINK_INET_DIAG;
+    ASSIGN_OR_RETURN(auto sock, sys.socket(domain, type, protocol));
+
+    // TODO: if too many sockets are closed too quickly, we can overflow the socket buffer, and
+    // some entries in mCookieTagMap will not be freed. In order to fix this we would need to
+    // periodically dump all sockets and remove the tag entries for sockets that have been closed.
+    // For now, set a large-enough buffer that we can close hundreds of sockets without getting
+    // ENOBUFS and leaking mCookieTagMap entries.
+    int rcvbuf = 512 * 1024;
+    auto ret = sys.setsockopt(sock, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf));
+    if (!ret.ok()) {
+        ALOGW("Failed to set SkDestroyListener buffer size to %d: %s", rcvbuf, ret.msg().c_str());
+    }
+
+    sockaddr_nl addr = {
+        .nl_family = AF_NETLINK,
+        .nl_groups = 1 << (SKNLGRP_INET_TCP_DESTROY - 1) | 1 << (SKNLGRP_INET_UDP_DESTROY - 1) |
+                     1 << (SKNLGRP_INET6_TCP_DESTROY - 1) | 1 << (SKNLGRP_INET6_UDP_DESTROY - 1)};
+    RETURN_IF_NOT_OK(sys.bind(sock, addr));
+
+    const sockaddr_nl kernel = {.nl_family = AF_NETLINK};
+    RETURN_IF_NOT_OK(sys.connect(sock, kernel));
+
+    std::unique_ptr<NetlinkListenerInterface> listener =
+            std::make_unique<NetlinkListener>(std::move(event), std::move(sock), "SkDestroyListen");
+
+    return listener;
+}
+
+Status TrafficController::initMaps() {
+    std::lock_guard guard(mMutex);
+
+    RETURN_IF_NOT_OK(mCookieTagMap.init(COOKIE_TAG_MAP_PATH));
+    RETURN_IF_NOT_OK(mUidCounterSetMap.init(UID_COUNTERSET_MAP_PATH));
+    RETURN_IF_NOT_OK(mAppUidStatsMap.init(APP_UID_STATS_MAP_PATH));
+    RETURN_IF_NOT_OK(mStatsMapA.init(STATS_MAP_A_PATH));
+    RETURN_IF_NOT_OK(mStatsMapB.init(STATS_MAP_B_PATH));
+    RETURN_IF_NOT_OK(mIfaceIndexNameMap.init(IFACE_INDEX_NAME_MAP_PATH));
+    RETURN_IF_NOT_OK(mIfaceStatsMap.init(IFACE_STATS_MAP_PATH));
+
+    RETURN_IF_NOT_OK(mConfigurationMap.init(CONFIGURATION_MAP_PATH));
+    RETURN_IF_NOT_OK(
+            mConfigurationMap.writeValue(UID_RULES_CONFIGURATION_KEY, DEFAULT_CONFIG, BPF_ANY));
+    RETURN_IF_NOT_OK(mConfigurationMap.writeValue(CURRENT_STATS_MAP_CONFIGURATION_KEY, SELECT_MAP_A,
+                                                  BPF_ANY));
+
+    RETURN_IF_NOT_OK(mUidOwnerMap.init(UID_OWNER_MAP_PATH));
+    RETURN_IF_NOT_OK(mUidOwnerMap.clear());
+    RETURN_IF_NOT_OK(mUidPermissionMap.init(UID_PERMISSION_MAP_PATH));
+
+    return netdutils::status::ok;
+}
+
+Status TrafficController::start() {
+    /* When netd restarts from a crash without total system reboot, the program
+     * is still attached to the cgroup, detach it so the program can be freed
+     * and we can load and attach new program into the target cgroup.
+     *
+     * TODO: Scrape existing socket when run-time restart and clean up the map
+     * if the socket no longer exist
+     */
+
+    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.
+    std::map<std::string, uint32_t> ifacePairs;
+    ASSIGN_OR_RETURN(ifacePairs, getIfaceList());
+    for (const auto& ifacePair:ifacePairs) {
+        addInterface(ifacePair.first.c_str(), ifacePair.second);
+    }
+
+    auto result = makeSkDestroyListener();
+    if (!isOk(result)) {
+        ALOGE("Unable to create SkDestroyListener: %s", toString(result).c_str());
+    } else {
+        mSkDestroyListener = std::move(result.value());
+    }
+    // Rx handler extracts nfgenmsg looks up and invokes registered dispatch function.
+    const auto rxHandler = [this](const nlmsghdr&, const Slice msg) {
+        std::lock_guard guard(mMutex);
+        inet_diag_msg diagmsg = {};
+        if (extract(msg, diagmsg) < sizeof(inet_diag_msg)) {
+            ALOGE("Unrecognized netlink message: %s", toString(msg).c_str());
+            return;
+        }
+        uint64_t sock_cookie = static_cast<uint64_t>(diagmsg.id.idiag_cookie[0]) |
+                               (static_cast<uint64_t>(diagmsg.id.idiag_cookie[1]) << 32);
+
+        Status s = mCookieTagMap.deleteValue(sock_cookie);
+        if (!isOk(s) && s.code() != ENOENT) {
+            ALOGE("Failed to delete cookie %" PRIx64 ": %s", sock_cookie, toString(s).c_str());
+            return;
+        }
+    };
+    expectOk(mSkDestroyListener->subscribe(kSockDiagMsgType, rxHandler));
+
+    // In case multiple netlink message comes in as a stream, we need to handle the rxDone message
+    // properly.
+    const auto rxDoneHandler = [](const nlmsghdr&, const Slice msg) {
+        // Ignore NLMSG_DONE  messages
+        inet_diag_msg diagmsg = {};
+        extract(msg, diagmsg);
+    };
+    expectOk(mSkDestroyListener->subscribe(kSockDiagDoneMsgType, rxDoneHandler));
+
+    return netdutils::status::ok;
+}
+
+int TrafficController::setCounterSet(int counterSetNum, uid_t uid, uid_t callingUid) {
+    if (counterSetNum < 0 || counterSetNum >= OVERFLOW_COUNTERSET) return -EINVAL;
+
+    std::lock_guard guard(mMutex);
+    if (!hasUpdateDeviceStatsPermission(callingUid)) return -EPERM;
+
+    // The default counter set for all uid is 0, so deleting the current counterset for that uid
+    // will automatically set it to 0.
+    if (counterSetNum == 0) {
+        Status res = mUidCounterSetMap.deleteValue(uid);
+        if (isOk(res) || (!isOk(res) && res.code() == ENOENT)) {
+            return 0;
+        } else {
+            ALOGE("Failed to delete the counterSet: %s\n", strerror(res.code()));
+            return -res.code();
+        }
+    }
+    uint8_t tmpCounterSetNum = (uint8_t)counterSetNum;
+    Status res = mUidCounterSetMap.writeValue(uid, tmpCounterSetNum, BPF_ANY);
+    if (!isOk(res)) {
+        ALOGE("Failed to set the counterSet: %s, fd: %d", strerror(res.code()),
+              mUidCounterSetMap.getMap().get());
+        return -res.code();
+    }
+    return 0;
+}
+
+// This method only get called by system_server when an app get uinstalled, it
+// is called inside removeUidsLocked() while holding mStatsLock. So it is safe
+// to iterate and modify the stats maps.
+int TrafficController::deleteTagData(uint32_t tag, uid_t uid, uid_t callingUid) {
+    std::lock_guard guard(mMutex);
+    if (!hasUpdateDeviceStatsPermission(callingUid)) return -EPERM;
+
+    // First we go through the cookieTagMap to delete the target uid tag combination. Or delete all
+    // the tags related to the uid if the tag is 0.
+    const auto deleteMatchedCookieEntries = [uid, tag](const uint64_t& key,
+                                                       const UidTagValue& value,
+                                                       BpfMap<uint64_t, UidTagValue>& map) {
+        if (value.uid == uid && (value.tag == tag || tag == 0)) {
+            auto res = map.deleteValue(key);
+            if (res.ok() || (res.error().code() == ENOENT)) {
+                return base::Result<void>();
+            }
+            ALOGE("Failed to delete data(cookie = %" PRIu64 "): %s\n", key,
+                  strerror(res.error().code()));
+        }
+        // Move forward to next cookie in the map.
+        return base::Result<void>();
+    };
+    mCookieTagMap.iterateWithValue(deleteMatchedCookieEntries);
+    // Now we go through the Tag stats map and delete the data entry with correct uid and tag
+    // combination. Or all tag stats under that uid if the target tag is 0.
+    const auto deleteMatchedUidTagEntries = [uid, tag](const StatsKey& key,
+                                                       BpfMap<StatsKey, StatsValue>& map) {
+        if (key.uid == uid && (key.tag == tag || tag == 0)) {
+            auto res = map.deleteValue(key);
+            if (res.ok() || (res.error().code() == ENOENT)) {
+                //Entry is deleted, use the current key to get a new nextKey;
+                return base::Result<void>();
+            }
+            ALOGE("Failed to delete data(uid=%u, tag=%u): %s\n", key.uid, key.tag,
+                  strerror(res.error().code()));
+        }
+        return base::Result<void>();
+    };
+    mStatsMapB.iterate(deleteMatchedUidTagEntries);
+    mStatsMapA.iterate(deleteMatchedUidTagEntries);
+    // If the tag is not zero, we already deleted all the data entry required. If tag is 0, we also
+    // need to delete the stats stored in uidStatsMap and counterSet map.
+    if (tag != 0) return 0;
+
+    auto res = mUidCounterSetMap.deleteValue(uid);
+    if (!res.ok() && res.error().code() != ENOENT) {
+        ALOGE("Failed to delete counterSet data(uid=%u, tag=%u): %s\n", uid, tag,
+              strerror(res.error().code()));
+    }
+
+    auto deleteAppUidStatsEntry = [uid](const uint32_t& key,
+                                        BpfMap<uint32_t, StatsValue>& map) -> base::Result<void> {
+        if (key == uid) {
+            auto res = map.deleteValue(key);
+            if (res.ok() || (res.error().code() == ENOENT)) {
+                return {};
+            }
+            ALOGE("Failed to delete data(uid=%u): %s", key, strerror(res.error().code()));
+        }
+        return {};
+    };
+    mAppUidStatsMap.iterate(deleteAppUidStatsEntry);
+    return 0;
+}
+
+int TrafficController::addInterface(const char* name, uint32_t ifaceIndex) {
+    IfaceValue iface;
+    if (ifaceIndex == 0) {
+        ALOGE("Unknown interface %s(%d)", name, ifaceIndex);
+        return -1;
+    }
+
+    strlcpy(iface.name, name, sizeof(IfaceValue));
+    Status res = mIfaceIndexNameMap.writeValue(ifaceIndex, iface, BPF_ANY);
+    if (!isOk(res)) {
+        ALOGE("Failed to add iface %s(%d): %s", name, ifaceIndex, strerror(res.code()));
+        return -res.code();
+    }
+    return 0;
+}
+
+Status TrafficController::updateOwnerMapEntry(UidOwnerMatchType match, uid_t uid, FirewallRule rule,
+                                              FirewallType type) {
+    std::lock_guard guard(mMutex);
+    if ((rule == ALLOW && type == ALLOWLIST) || (rule == DENY && type == DENYLIST)) {
+        RETURN_IF_NOT_OK(addRule(uid, match));
+    } else if ((rule == ALLOW && type == DENYLIST) || (rule == DENY && type == ALLOWLIST)) {
+        RETURN_IF_NOT_OK(removeRule(uid, match));
+    } else {
+        //Cannot happen.
+        return statusFromErrno(EINVAL, "");
+    }
+    return netdutils::status::ok;
+}
+
+Status TrafficController::removeRule(uint32_t uid, UidOwnerMatchType match) {
+    auto oldMatch = mUidOwnerMap.readValue(uid);
+    if (oldMatch.ok()) {
+        UidOwnerValue newMatch = {
+                .iif = (match == IIF_MATCH) ? 0 : oldMatch.value().iif,
+                .rule = static_cast<uint8_t>(oldMatch.value().rule & ~match),
+        };
+        if (newMatch.rule == 0) {
+            RETURN_IF_NOT_OK(mUidOwnerMap.deleteValue(uid));
+        } else {
+            RETURN_IF_NOT_OK(mUidOwnerMap.writeValue(uid, newMatch, BPF_ANY));
+        }
+    } else {
+        return statusFromErrno(ENOENT, StringPrintf("uid: %u does not exist in map", uid));
+    }
+    return netdutils::status::ok;
+}
+
+Status TrafficController::addRule(uint32_t uid, UidOwnerMatchType match, uint32_t iif) {
+    // iif should be non-zero if and only if match == MATCH_IIF
+    if (match == IIF_MATCH && iif == 0) {
+        return statusFromErrno(EINVAL, "Interface match must have nonzero interface index");
+    } else if (match != IIF_MATCH && iif != 0) {
+        return statusFromErrno(EINVAL, "Non-interface match must have zero interface index");
+    }
+    auto oldMatch = mUidOwnerMap.readValue(uid);
+    if (oldMatch.ok()) {
+        UidOwnerValue newMatch = {
+                .iif = iif ? iif : oldMatch.value().iif,
+                .rule = static_cast<uint8_t>(oldMatch.value().rule | match),
+        };
+        RETURN_IF_NOT_OK(mUidOwnerMap.writeValue(uid, newMatch, BPF_ANY));
+    } else {
+        UidOwnerValue newMatch = {
+                .iif = iif,
+                .rule = static_cast<uint8_t>(match),
+        };
+        RETURN_IF_NOT_OK(mUidOwnerMap.writeValue(uid, newMatch, BPF_ANY));
+    }
+    return netdutils::status::ok;
+}
+
+Status TrafficController::updateUidOwnerMap(const std::vector<uint32_t>& appUids,
+                                            UidOwnerMatchType matchType, IptOp op) {
+    std::lock_guard guard(mMutex);
+    for (uint32_t uid : appUids) {
+        if (op == IptOpDelete) {
+            RETURN_IF_NOT_OK(removeRule(uid, matchType));
+        } else if (op == IptOpInsert) {
+            RETURN_IF_NOT_OK(addRule(uid, matchType));
+        } else {
+            // Cannot happen.
+            return statusFromErrno(EINVAL, StringPrintf("invalid IptOp: %d, %d", op, matchType));
+        }
+    }
+    return netdutils::status::ok;
+}
+
+FirewallType TrafficController::getFirewallType(ChildChain chain) {
+    switch (chain) {
+        case DOZABLE:
+            return ALLOWLIST;
+        case STANDBY:
+            return DENYLIST;
+        case POWERSAVE:
+            return ALLOWLIST;
+        case RESTRICTED:
+            return ALLOWLIST;
+        case NONE:
+        default:
+            return DENYLIST;
+    }
+}
+
+int TrafficController::changeUidOwnerRule(ChildChain chain, uid_t uid, FirewallRule rule,
+                                          FirewallType type) {
+    Status res;
+    switch (chain) {
+        case DOZABLE:
+            res = updateOwnerMapEntry(DOZABLE_MATCH, uid, rule, type);
+            break;
+        case STANDBY:
+            res = updateOwnerMapEntry(STANDBY_MATCH, uid, rule, type);
+            break;
+        case POWERSAVE:
+            res = updateOwnerMapEntry(POWERSAVE_MATCH, uid, rule, type);
+            break;
+        case RESTRICTED:
+            res = updateOwnerMapEntry(RESTRICTED_MATCH, uid, rule, type);
+            break;
+        case NONE:
+        default:
+            ALOGW("Unknown child chain: %d", chain);
+            return -EINVAL;
+    }
+    if (!isOk(res)) {
+        ALOGE("change uid(%u) rule of %d failed: %s, rule: %d, type: %d", uid, chain,
+              res.msg().c_str(), rule, type);
+        return -res.code();
+    }
+    return 0;
+}
+
+Status TrafficController::replaceRulesInMap(const UidOwnerMatchType match,
+                                            const std::vector<int32_t>& uids) {
+    std::lock_guard guard(mMutex);
+    std::set<int32_t> uidSet(uids.begin(), uids.end());
+    std::vector<uint32_t> uidsToDelete;
+    auto getUidsToDelete = [&uidsToDelete, &uidSet](const uint32_t& key,
+                                                    const BpfMap<uint32_t, UidOwnerValue>&) {
+        if (uidSet.find((int32_t) key) == uidSet.end()) {
+            uidsToDelete.push_back(key);
+        }
+        return base::Result<void>();
+    };
+    RETURN_IF_NOT_OK(mUidOwnerMap.iterate(getUidsToDelete));
+
+    for(auto uid : uidsToDelete) {
+        RETURN_IF_NOT_OK(removeRule(uid, match));
+    }
+
+    for (auto uid : uids) {
+        RETURN_IF_NOT_OK(addRule(uid, match));
+    }
+    return netdutils::status::ok;
+}
+
+Status TrafficController::addUidInterfaceRules(const int iif,
+                                               const std::vector<int32_t>& uidsToAdd) {
+    if (!iif) {
+        return statusFromErrno(EINVAL, "Interface rule must specify interface");
+    }
+    std::lock_guard guard(mMutex);
+
+    for (auto uid : uidsToAdd) {
+        netdutils::Status result = addRule(uid, IIF_MATCH, iif);
+        if (!isOk(result)) {
+            ALOGW("addRule failed(%d): uid=%d iif=%d", result.code(), uid, iif);
+        }
+    }
+    return netdutils::status::ok;
+}
+
+Status TrafficController::removeUidInterfaceRules(const std::vector<int32_t>& uidsToDelete) {
+    std::lock_guard guard(mMutex);
+
+    for (auto uid : uidsToDelete) {
+        netdutils::Status result = removeRule(uid, IIF_MATCH);
+        if (!isOk(result)) {
+            ALOGW("removeRule failed(%d): uid=%d", result.code(), uid);
+        }
+    }
+    return netdutils::status::ok;
+}
+
+int TrafficController::replaceUidOwnerMap(const std::string& name, bool isAllowlist __unused,
+                                          const std::vector<int32_t>& uids) {
+    // FirewallRule rule = isAllowlist ? ALLOW : DENY;
+    // FirewallType type = isAllowlist ? ALLOWLIST : DENYLIST;
+    Status res;
+    if (!name.compare(LOCAL_DOZABLE)) {
+        res = replaceRulesInMap(DOZABLE_MATCH, uids);
+    } else if (!name.compare(LOCAL_STANDBY)) {
+        res = replaceRulesInMap(STANDBY_MATCH, uids);
+    } else if (!name.compare(LOCAL_POWERSAVE)) {
+        res = replaceRulesInMap(POWERSAVE_MATCH, uids);
+    } else if (!name.compare(LOCAL_RESTRICTED)) {
+        res = replaceRulesInMap(RESTRICTED_MATCH, uids);
+    } else {
+        ALOGE("unknown chain name: %s", name.c_str());
+        return -EINVAL;
+    }
+    if (!isOk(res)) {
+        ALOGE("Failed to clean up chain: %s: %s", name.c_str(), res.msg().c_str());
+        return -res.code();
+    }
+    return 0;
+}
+
+int TrafficController::toggleUidOwnerMap(ChildChain chain, bool enable) {
+    std::lock_guard guard(mMutex);
+    uint32_t key = UID_RULES_CONFIGURATION_KEY;
+    auto oldConfiguration = mConfigurationMap.readValue(key);
+    if (!oldConfiguration.ok()) {
+        ALOGE("Cannot read the old configuration from map: %s",
+              oldConfiguration.error().message().c_str());
+        return -oldConfiguration.error().code();
+    }
+    Status res;
+    BpfConfig newConfiguration;
+    uint8_t match;
+    switch (chain) {
+        case DOZABLE:
+            match = DOZABLE_MATCH;
+            break;
+        case STANDBY:
+            match = STANDBY_MATCH;
+            break;
+        case POWERSAVE:
+            match = POWERSAVE_MATCH;
+            break;
+        case RESTRICTED:
+            match = RESTRICTED_MATCH;
+            break;
+        default:
+            return -EINVAL;
+    }
+    newConfiguration =
+            enable ? (oldConfiguration.value() | match) : (oldConfiguration.value() & (~match));
+    res = mConfigurationMap.writeValue(key, newConfiguration, BPF_EXIST);
+    if (!isOk(res)) {
+        ALOGE("Failed to toggleUidOwnerMap(%d): %s", chain, res.msg().c_str());
+    }
+    return -res.code();
+}
+
+Status TrafficController::swapActiveStatsMap() {
+    std::lock_guard guard(mMutex);
+
+    uint32_t key = CURRENT_STATS_MAP_CONFIGURATION_KEY;
+    auto oldConfiguration = mConfigurationMap.readValue(key);
+    if (!oldConfiguration.ok()) {
+        ALOGE("Cannot read the old configuration from map: %s",
+              oldConfiguration.error().message().c_str());
+        return Status(oldConfiguration.error().code(), oldConfiguration.error().message());
+    }
+
+    // Write to the configuration map to inform the kernel eBPF program to switch
+    // from using one map to the other. Use flag BPF_EXIST here since the map should
+    // be already populated in initMaps.
+    uint8_t newConfigure = (oldConfiguration.value() == SELECT_MAP_A) ? SELECT_MAP_B : SELECT_MAP_A;
+    auto res = mConfigurationMap.writeValue(CURRENT_STATS_MAP_CONFIGURATION_KEY, newConfigure,
+                                            BPF_EXIST);
+    if (!res.ok()) {
+        ALOGE("Failed to toggle the stats map: %s", strerror(res.error().code()));
+        return res;
+    }
+    // After changing the config, we need to make sure all the current running
+    // eBPF programs are finished and all the CPUs are aware of this config change
+    // before we modify the old map. So we do a special hack here to wait for
+    // the kernel to do a synchronize_rcu(). Once the kernel called
+    // synchronize_rcu(), the config we just updated will be available to all cores
+    // and the next eBPF programs triggered inside the kernel will use the new
+    // map configuration. So once this function returns we can safely modify the
+    // old stats map without concerning about race between the kernel and
+    // userspace.
+    int ret = synchronizeKernelRCU();
+    if (ret) {
+        ALOGE("map swap synchronize_rcu() ended with failure: %s", strerror(-ret));
+        return statusFromErrno(-ret, "map swap synchronize_rcu() failed");
+    }
+    return netdutils::status::ok;
+}
+
+void TrafficController::setPermissionForUids(int permission, const std::vector<uid_t>& uids) {
+    std::lock_guard guard(mMutex);
+    if (permission == INetd::PERMISSION_UNINSTALLED) {
+        for (uid_t uid : uids) {
+            // Clean up all permission information for the related uid if all the
+            // packages related to it are uninstalled.
+            mPrivilegedUser.erase(uid);
+            Status ret = mUidPermissionMap.deleteValue(uid);
+            if (!isOk(ret) && ret.code() != ENOENT) {
+                ALOGE("Failed to clean up the permission for %u: %s", uid, strerror(ret.code()));
+            }
+        }
+        return;
+    }
+
+    bool privileged = (permission & INetd::PERMISSION_UPDATE_DEVICE_STATS);
+
+    for (uid_t uid : uids) {
+        if (privileged) {
+            mPrivilegedUser.insert(uid);
+        } else {
+            mPrivilegedUser.erase(uid);
+        }
+
+        // The map stores all the permissions that the UID has, except if the only permission
+        // the UID has is the INTERNET permission, then the UID should not appear in the map.
+        if (permission != INetd::PERMISSION_INTERNET) {
+            Status ret = mUidPermissionMap.writeValue(uid, permission, BPF_ANY);
+            if (!isOk(ret)) {
+                ALOGE("Failed to set permission: %s of uid(%u) to permission map: %s",
+                      UidPermissionTypeToString(permission).c_str(), uid, strerror(ret.code()));
+            }
+        } else {
+            Status ret = mUidPermissionMap.deleteValue(uid);
+            if (!isOk(ret) && ret.code() != ENOENT) {
+                ALOGE("Failed to remove uid %u from permission map: %s", uid, strerror(ret.code()));
+            }
+        }
+    }
+}
+
+std::string getProgramStatus(const char *path) {
+    int ret = access(path, R_OK);
+    if (ret == 0) {
+        return StringPrintf("OK");
+    }
+    if (ret != 0 && errno == ENOENT) {
+        return StringPrintf("program is missing at: %s", path);
+    }
+    return StringPrintf("check Program %s error: %s", path, strerror(errno));
+}
+
+std::string getMapStatus(const base::unique_fd& map_fd, const char* path) {
+    if (map_fd.get() < 0) {
+        return StringPrintf("map fd lost");
+    }
+    if (access(path, F_OK) != 0) {
+        return StringPrintf("map not pinned to location: %s", path);
+    }
+    return StringPrintf("OK");
+}
+
+// NOLINTNEXTLINE(google-runtime-references): grandfathered pass by non-const reference
+void dumpBpfMap(const std::string& mapName, DumpWriter& dw, const std::string& header) {
+    dw.blankline();
+    dw.println("%s:", mapName.c_str());
+    if (!header.empty()) {
+        dw.println(header);
+    }
+}
+
+const String16 TrafficController::DUMP_KEYWORD = String16("trafficcontroller");
+
+void TrafficController::dump(DumpWriter& dw, bool verbose) {
+    std::lock_guard guard(mMutex);
+    ScopedIndent indentTop(dw);
+    dw.println("TrafficController");
+
+    ScopedIndent indentPreBpfModule(dw);
+
+    dw.blankline();
+    dw.println("mCookieTagMap status: %s",
+               getMapStatus(mCookieTagMap.getMap(), COOKIE_TAG_MAP_PATH).c_str());
+    dw.println("mUidCounterSetMap status: %s",
+               getMapStatus(mUidCounterSetMap.getMap(), UID_COUNTERSET_MAP_PATH).c_str());
+    dw.println("mAppUidStatsMap status: %s",
+               getMapStatus(mAppUidStatsMap.getMap(), APP_UID_STATS_MAP_PATH).c_str());
+    dw.println("mStatsMapA status: %s",
+               getMapStatus(mStatsMapA.getMap(), STATS_MAP_A_PATH).c_str());
+    dw.println("mStatsMapB status: %s",
+               getMapStatus(mStatsMapB.getMap(), STATS_MAP_B_PATH).c_str());
+    dw.println("mIfaceIndexNameMap status: %s",
+               getMapStatus(mIfaceIndexNameMap.getMap(), IFACE_INDEX_NAME_MAP_PATH).c_str());
+    dw.println("mIfaceStatsMap status: %s",
+               getMapStatus(mIfaceStatsMap.getMap(), IFACE_STATS_MAP_PATH).c_str());
+    dw.println("mConfigurationMap status: %s",
+               getMapStatus(mConfigurationMap.getMap(), CONFIGURATION_MAP_PATH).c_str());
+    dw.println("mUidOwnerMap status: %s",
+               getMapStatus(mUidOwnerMap.getMap(), UID_OWNER_MAP_PATH).c_str());
+
+    dw.blankline();
+    dw.println("Cgroup ingress program status: %s",
+               getProgramStatus(BPF_INGRESS_PROG_PATH).c_str());
+    dw.println("Cgroup egress program status: %s", getProgramStatus(BPF_EGRESS_PROG_PATH).c_str());
+    dw.println("xt_bpf ingress program status: %s",
+               getProgramStatus(XT_BPF_INGRESS_PROG_PATH).c_str());
+    dw.println("xt_bpf egress program status: %s",
+               getProgramStatus(XT_BPF_EGRESS_PROG_PATH).c_str());
+    dw.println("xt_bpf bandwidth allowlist program status: %s",
+               getProgramStatus(XT_BPF_ALLOWLIST_PROG_PATH).c_str());
+    dw.println("xt_bpf bandwidth denylist program status: %s",
+               getProgramStatus(XT_BPF_DENYLIST_PROG_PATH).c_str());
+
+    if (!verbose) {
+        return;
+    }
+
+    dw.blankline();
+    dw.println("BPF map content:");
+
+    ScopedIndent indentForMapContent(dw);
+
+    // Print CookieTagMap content.
+    dumpBpfMap("mCookieTagMap", dw, "");
+    const auto printCookieTagInfo = [&dw](const uint64_t& key, const UidTagValue& value,
+                                          const BpfMap<uint64_t, UidTagValue>&) {
+        dw.println("cookie=%" PRIu64 " tag=0x%x uid=%u", key, value.tag, value.uid);
+        return base::Result<void>();
+    };
+    base::Result<void> res = mCookieTagMap.iterateWithValue(printCookieTagInfo);
+    if (!res.ok()) {
+        dw.println("mCookieTagMap print end with error: %s", res.error().message().c_str());
+    }
+
+    // Print UidCounterSetMap Content
+    dumpBpfMap("mUidCounterSetMap", dw, "");
+    const auto printUidInfo = [&dw](const uint32_t& key, const uint8_t& value,
+                                    const BpfMap<uint32_t, uint8_t>&) {
+        dw.println("%u %u", key, value);
+        return base::Result<void>();
+    };
+    res = mUidCounterSetMap.iterateWithValue(printUidInfo);
+    if (!res.ok()) {
+        dw.println("mUidCounterSetMap print end with error: %s", res.error().message().c_str());
+    }
+
+    // Print AppUidStatsMap content
+    std::string appUidStatsHeader = StringPrintf("uid rxBytes rxPackets txBytes txPackets");
+    dumpBpfMap("mAppUidStatsMap:", dw, appUidStatsHeader);
+    auto printAppUidStatsInfo = [&dw](const uint32_t& key, const StatsValue& value,
+                                      const BpfMap<uint32_t, StatsValue>&) {
+        dw.println("%u %" PRIu64 " %" PRIu64 " %" PRIu64 " %" PRIu64, key, value.rxBytes,
+                   value.rxPackets, value.txBytes, value.txPackets);
+        return base::Result<void>();
+    };
+    res = mAppUidStatsMap.iterateWithValue(printAppUidStatsInfo);
+    if (!res.ok()) {
+        dw.println("mAppUidStatsMap print end with error: %s", res.error().message().c_str());
+    }
+
+    // Print uidStatsMap content
+    std::string statsHeader = StringPrintf("ifaceIndex ifaceName tag_hex uid_int cnt_set rxBytes"
+                                           " rxPackets txBytes txPackets");
+    dumpBpfMap("mStatsMapA", dw, statsHeader);
+    const auto printStatsInfo = [&dw, this](const StatsKey& key, const StatsValue& value,
+                                            const BpfMap<StatsKey, StatsValue>&) {
+        uint32_t ifIndex = key.ifaceIndex;
+        auto ifname = mIfaceIndexNameMap.readValue(ifIndex);
+        if (!ifname.ok()) {
+            ifname = IfaceValue{"unknown"};
+        }
+        dw.println("%u %s 0x%x %u %u %" PRIu64 " %" PRIu64 " %" PRIu64 " %" PRIu64, ifIndex,
+                   ifname.value().name, key.tag, key.uid, key.counterSet, value.rxBytes,
+                   value.rxPackets, value.txBytes, value.txPackets);
+        return base::Result<void>();
+    };
+    res = mStatsMapA.iterateWithValue(printStatsInfo);
+    if (!res.ok()) {
+        dw.println("mStatsMapA print end with error: %s", res.error().message().c_str());
+    }
+
+    // Print TagStatsMap content.
+    dumpBpfMap("mStatsMapB", dw, statsHeader);
+    res = mStatsMapB.iterateWithValue(printStatsInfo);
+    if (!res.ok()) {
+        dw.println("mStatsMapB print end with error: %s", res.error().message().c_str());
+    }
+
+    // Print ifaceIndexToNameMap content.
+    dumpBpfMap("mIfaceIndexNameMap", dw, "");
+    const auto printIfaceNameInfo = [&dw](const uint32_t& key, const IfaceValue& value,
+                                          const BpfMap<uint32_t, IfaceValue>&) {
+        const char* ifname = value.name;
+        dw.println("ifaceIndex=%u ifaceName=%s", key, ifname);
+        return base::Result<void>();
+    };
+    res = mIfaceIndexNameMap.iterateWithValue(printIfaceNameInfo);
+    if (!res.ok()) {
+        dw.println("mIfaceIndexNameMap print end with error: %s", res.error().message().c_str());
+    }
+
+    // Print ifaceStatsMap content
+    std::string ifaceStatsHeader = StringPrintf("ifaceIndex ifaceName rxBytes rxPackets txBytes"
+                                                " txPackets");
+    dumpBpfMap("mIfaceStatsMap:", dw, ifaceStatsHeader);
+    const auto printIfaceStatsInfo = [&dw, this](const uint32_t& key, const StatsValue& value,
+                                                 const BpfMap<uint32_t, StatsValue>&) {
+        auto ifname = mIfaceIndexNameMap.readValue(key);
+        if (!ifname.ok()) {
+            ifname = IfaceValue{"unknown"};
+        }
+        dw.println("%u %s %" PRIu64 " %" PRIu64 " %" PRIu64 " %" PRIu64, key, ifname.value().name,
+                   value.rxBytes, value.rxPackets, value.txBytes, value.txPackets);
+        return base::Result<void>();
+    };
+    res = mIfaceStatsMap.iterateWithValue(printIfaceStatsInfo);
+    if (!res.ok()) {
+        dw.println("mIfaceStatsMap print end with error: %s", res.error().message().c_str());
+    }
+
+    dw.blankline();
+
+    uint32_t key = UID_RULES_CONFIGURATION_KEY;
+    auto configuration = mConfigurationMap.readValue(key);
+    if (configuration.ok()) {
+        dw.println("current ownerMatch configuration: %d%s", configuration.value(),
+                   uidMatchTypeToString(configuration.value()).c_str());
+    } else {
+        dw.println("mConfigurationMap read ownerMatch configure failed with error: %s",
+                   configuration.error().message().c_str());
+    }
+
+    key = CURRENT_STATS_MAP_CONFIGURATION_KEY;
+    configuration = mConfigurationMap.readValue(key);
+    if (configuration.ok()) {
+        const char* statsMapDescription = "???";
+        switch (configuration.value()) {
+            case SELECT_MAP_A:
+                statsMapDescription = "SELECT_MAP_A";
+                break;
+            case SELECT_MAP_B:
+                statsMapDescription = "SELECT_MAP_B";
+                break;
+                // No default clause, so if we ever add a third map, this code will fail to build.
+        }
+        dw.println("current statsMap configuration: %d %s", configuration.value(),
+                   statsMapDescription);
+    } else {
+        dw.println("mConfigurationMap read stats map configure failed with error: %s",
+                   configuration.error().message().c_str());
+    }
+    dumpBpfMap("mUidOwnerMap", dw, "");
+    const auto printUidMatchInfo = [&dw, this](const uint32_t& key, const UidOwnerValue& value,
+                                               const BpfMap<uint32_t, UidOwnerValue>&) {
+        if (value.rule & IIF_MATCH) {
+            auto ifname = mIfaceIndexNameMap.readValue(value.iif);
+            if (ifname.ok()) {
+                dw.println("%u %s %s", key, uidMatchTypeToString(value.rule).c_str(),
+                           ifname.value().name);
+            } else {
+                dw.println("%u %s %u", key, uidMatchTypeToString(value.rule).c_str(), value.iif);
+            }
+        } else {
+            dw.println("%u %s", key, uidMatchTypeToString(value.rule).c_str());
+        }
+        return base::Result<void>();
+    };
+    res = mUidOwnerMap.iterateWithValue(printUidMatchInfo);
+    if (!res.ok()) {
+        dw.println("mUidOwnerMap print end with error: %s", res.error().message().c_str());
+    }
+    dumpBpfMap("mUidPermissionMap", dw, "");
+    const auto printUidPermissionInfo = [&dw](const uint32_t& key, const int& value,
+                                              const BpfMap<uint32_t, uint8_t>&) {
+        dw.println("%u %s", key, UidPermissionTypeToString(value).c_str());
+        return base::Result<void>();
+    };
+    res = mUidPermissionMap.iterateWithValue(printUidPermissionInfo);
+    if (!res.ok()) {
+        dw.println("mUidPermissionMap print end with error: %s", res.error().message().c_str());
+    }
+
+    dumpBpfMap("mPrivilegedUser", dw, "");
+    for (uid_t uid : mPrivilegedUser) {
+        dw.println("%u ALLOW_UPDATE_DEVICE_STATS", (uint32_t)uid);
+    }
+}
+
+}  // namespace net
+}  // namespace android
diff --git a/service/native/TrafficControllerTest.cpp b/service/native/TrafficControllerTest.cpp
new file mode 100644
index 0000000..f5d5911
--- /dev/null
+++ b/service/native/TrafficControllerTest.cpp
@@ -0,0 +1,741 @@
+/*
+ * Copyright 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * TrafficControllerTest.cpp - unit tests for TrafficController.cpp
+ */
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <fcntl.h>
+#include <inttypes.h>
+#include <linux/inet_diag.h>
+#include <linux/sock_diag.h>
+#include <sys/socket.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <gtest/gtest.h>
+
+#include <android-base/stringprintf.h>
+#include <android-base/strings.h>
+
+#include <netdutils/MockSyscalls.h>
+
+#include "TrafficController.h"
+#include "bpf/BpfUtils.h"
+
+using namespace android::bpf;  // NOLINT(google-build-using-namespace): grandfathered
+
+namespace android {
+namespace net {
+
+using base::Result;
+using netdutils::isOk;
+
+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 uint32_t TEST_COUNTERSET = 1;
+constexpr uint32_t DEFAULT_COUNTERSET = 0;
+
+#define ASSERT_VALID(x) ASSERT_TRUE((x).isValid())
+
+class TrafficControllerTest : public ::testing::Test {
+  protected:
+    TrafficControllerTest() {}
+    TrafficController mTc;
+    BpfMap<uint64_t, UidTagValue> mFakeCookieTagMap;
+    BpfMap<uint32_t, uint8_t> mFakeUidCounterSetMap;
+    BpfMap<uint32_t, StatsValue> mFakeAppUidStatsMap;
+    BpfMap<StatsKey, StatsValue> mFakeStatsMapA;
+    BpfMap<uint32_t, uint8_t> mFakeConfigurationMap;
+    BpfMap<uint32_t, UidOwnerValue> mFakeUidOwnerMap;
+    BpfMap<uint32_t, uint8_t> mFakeUidPermissionMap;
+
+    void SetUp() {
+        std::lock_guard guard(mTc.mMutex);
+        ASSERT_EQ(0, setrlimitForTest());
+
+        mFakeCookieTagMap.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(uint64_t), sizeof(UidTagValue),
+                                          TEST_MAP_SIZE, 0));
+        ASSERT_VALID(mFakeCookieTagMap);
+
+        mFakeUidCounterSetMap.reset(
+            createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), TEST_MAP_SIZE, 0));
+        ASSERT_VALID(mFakeUidCounterSetMap);
+
+        mFakeAppUidStatsMap.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(StatsValue),
+                                            TEST_MAP_SIZE, 0));
+        ASSERT_VALID(mFakeAppUidStatsMap);
+
+        mFakeStatsMapA.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(StatsKey), sizeof(StatsValue),
+                                       TEST_MAP_SIZE, 0));
+        ASSERT_VALID(mFakeStatsMapA);
+
+        mFakeConfigurationMap.reset(
+                createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), 1, 0));
+        ASSERT_VALID(mFakeConfigurationMap);
+
+        mFakeUidOwnerMap.reset(createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(UidOwnerValue),
+                                         TEST_MAP_SIZE, 0));
+        ASSERT_VALID(mFakeUidOwnerMap);
+        mFakeUidPermissionMap.reset(
+                createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint8_t), TEST_MAP_SIZE, 0));
+        ASSERT_VALID(mFakeUidPermissionMap);
+
+        mTc.mCookieTagMap.reset(dupFd(mFakeCookieTagMap.getMap()));
+        ASSERT_VALID(mTc.mCookieTagMap);
+        mTc.mUidCounterSetMap.reset(dupFd(mFakeUidCounterSetMap.getMap()));
+        ASSERT_VALID(mTc.mUidCounterSetMap);
+        mTc.mAppUidStatsMap.reset(dupFd(mFakeAppUidStatsMap.getMap()));
+        ASSERT_VALID(mTc.mAppUidStatsMap);
+        mTc.mStatsMapA.reset(dupFd(mFakeStatsMapA.getMap()));
+        ASSERT_VALID(mTc.mStatsMapA);
+        mTc.mConfigurationMap.reset(dupFd(mFakeConfigurationMap.getMap()));
+        ASSERT_VALID(mTc.mConfigurationMap);
+
+        // Always write to stats map A by default.
+        ASSERT_RESULT_OK(mTc.mConfigurationMap.writeValue(CURRENT_STATS_MAP_CONFIGURATION_KEY,
+                                                          SELECT_MAP_A, BPF_ANY));
+        mTc.mUidOwnerMap.reset(dupFd(mFakeUidOwnerMap.getMap()));
+        ASSERT_VALID(mTc.mUidOwnerMap);
+        mTc.mUidPermissionMap.reset(dupFd(mFakeUidPermissionMap.getMap()));
+        ASSERT_VALID(mTc.mUidPermissionMap);
+        mTc.mPrivilegedUser.clear();
+    }
+
+    int dupFd(const android::base::unique_fd& mapFd) {
+        return fcntl(mapFd.get(), F_DUPFD_CLOEXEC, 0);
+    }
+
+    void populateFakeStats(uint64_t cookie, uint32_t uid, uint32_t tag, StatsKey* key) {
+        UidTagValue cookieMapkey = {.uid = (uint32_t)uid, .tag = tag};
+        EXPECT_RESULT_OK(mFakeCookieTagMap.writeValue(cookie, cookieMapkey, BPF_ANY));
+        *key = {.uid = uid, .tag = tag, .counterSet = TEST_COUNTERSET, .ifaceIndex = 1};
+        StatsValue statsMapValue = {.rxPackets = 1, .rxBytes = 100};
+        uint8_t counterSet = TEST_COUNTERSET;
+        EXPECT_RESULT_OK(mFakeUidCounterSetMap.writeValue(uid, counterSet, BPF_ANY));
+        EXPECT_RESULT_OK(mFakeStatsMapA.writeValue(*key, statsMapValue, BPF_ANY));
+        key->tag = 0;
+        EXPECT_RESULT_OK(mFakeStatsMapA.writeValue(*key, statsMapValue, BPF_ANY));
+        EXPECT_RESULT_OK(mFakeAppUidStatsMap.writeValue(uid, statsMapValue, BPF_ANY));
+        // put tag information back to statsKey
+        key->tag = tag;
+    }
+
+    void checkUidOwnerRuleForChain(ChildChain chain, UidOwnerMatchType match) {
+        uint32_t uid = TEST_UID;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, DENY, DENYLIST));
+        Result<UidOwnerValue> value = mFakeUidOwnerMap.readValue(uid);
+        EXPECT_RESULT_OK(value);
+        EXPECT_TRUE(value.value().rule & match);
+
+        uid = TEST_UID2;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, ALLOW, ALLOWLIST));
+        value = mFakeUidOwnerMap.readValue(uid);
+        EXPECT_RESULT_OK(value);
+        EXPECT_TRUE(value.value().rule & match);
+
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, DENY, ALLOWLIST));
+        value = mFakeUidOwnerMap.readValue(uid);
+        EXPECT_FALSE(value.ok());
+        EXPECT_EQ(ENOENT, value.error().code());
+
+        uid = TEST_UID;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, ALLOW, DENYLIST));
+        value = mFakeUidOwnerMap.readValue(uid);
+        EXPECT_FALSE(value.ok());
+        EXPECT_EQ(ENOENT, value.error().code());
+
+        uid = TEST_UID3;
+        EXPECT_EQ(-ENOENT, mTc.changeUidOwnerRule(chain, uid, ALLOW, DENYLIST));
+        value = mFakeUidOwnerMap.readValue(uid);
+        EXPECT_FALSE(value.ok());
+        EXPECT_EQ(ENOENT, value.error().code());
+    }
+
+    void checkEachUidValue(const std::vector<int32_t>& uids, UidOwnerMatchType match) {
+        for (uint32_t uid : uids) {
+            Result<UidOwnerValue> value = mFakeUidOwnerMap.readValue(uid);
+            EXPECT_RESULT_OK(value);
+            EXPECT_TRUE(value.value().rule & match);
+        }
+        std::set<uint32_t> uidSet(uids.begin(), uids.end());
+        const auto checkNoOtherUid = [&uidSet](const int32_t& key,
+                                               const BpfMap<uint32_t, UidOwnerValue>&) {
+            EXPECT_NE(uidSet.end(), uidSet.find(key));
+            return Result<void>();
+        };
+        EXPECT_RESULT_OK(mFakeUidOwnerMap.iterate(checkNoOtherUid));
+    }
+
+    void checkUidMapReplace(const std::string& name, const std::vector<int32_t>& uids,
+                            UidOwnerMatchType match) {
+        bool isAllowlist = true;
+        EXPECT_EQ(0, mTc.replaceUidOwnerMap(name, isAllowlist, uids));
+        checkEachUidValue(uids, match);
+
+        isAllowlist = false;
+        EXPECT_EQ(0, mTc.replaceUidOwnerMap(name, isAllowlist, uids));
+        checkEachUidValue(uids, match);
+    }
+
+    void expectUidOwnerMapValues(const std::vector<uint32_t>& appUids, uint8_t expectedRule,
+                                 uint32_t expectedIif) {
+        for (uint32_t uid : appUids) {
+            Result<UidOwnerValue> value = mFakeUidOwnerMap.readValue(uid);
+            EXPECT_RESULT_OK(value);
+            EXPECT_EQ(expectedRule, value.value().rule)
+                    << "Expected rule for UID " << uid << " to be " << expectedRule << ", but was "
+                    << value.value().rule;
+            EXPECT_EQ(expectedIif, value.value().iif)
+                    << "Expected iif for UID " << uid << " to be " << expectedIif << ", but was "
+                    << value.value().iif;
+        }
+    }
+
+    template <class Key, class Value>
+    void expectMapEmpty(BpfMap<Key, Value>& map) {
+        auto isEmpty = map.isEmpty();
+        EXPECT_RESULT_OK(isEmpty);
+        EXPECT_TRUE(isEmpty.value());
+    }
+
+    void expectUidPermissionMapValues(const std::vector<uid_t>& appUids, uint8_t expectedValue) {
+        for (uid_t uid : appUids) {
+            Result<uint8_t> value = mFakeUidPermissionMap.readValue(uid);
+            EXPECT_RESULT_OK(value);
+            EXPECT_EQ(expectedValue, value.value())
+                    << "Expected value for UID " << uid << " to be " << expectedValue
+                    << ", but was " << value.value();
+        }
+    }
+
+    void expectPrivilegedUserSet(const std::vector<uid_t>& appUids) {
+        std::lock_guard guard(mTc.mMutex);
+        EXPECT_EQ(appUids.size(), mTc.mPrivilegedUser.size());
+        for (uid_t uid : appUids) {
+            EXPECT_NE(mTc.mPrivilegedUser.end(), mTc.mPrivilegedUser.find(uid));
+        }
+    }
+
+    void expectPrivilegedUserSetEmpty() {
+        std::lock_guard guard(mTc.mMutex);
+        EXPECT_TRUE(mTc.mPrivilegedUser.empty());
+    }
+
+    void addPrivilegedUid(uid_t uid) {
+        std::vector privilegedUid = {uid};
+        mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, privilegedUid);
+    }
+
+    void removePrivilegedUid(uid_t uid) {
+        std::vector privilegedUid = {uid};
+        mTc.setPermissionForUids(INetd::PERMISSION_NONE, privilegedUid);
+    }
+
+    void expectFakeStatsUnchanged(uint64_t cookie, uint32_t tag, uint32_t uid,
+                                  StatsKey tagStatsMapKey) {
+        Result<UidTagValue> cookieMapResult = mFakeCookieTagMap.readValue(cookie);
+        EXPECT_RESULT_OK(cookieMapResult);
+        EXPECT_EQ(uid, cookieMapResult.value().uid);
+        EXPECT_EQ(tag, cookieMapResult.value().tag);
+        Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid);
+        EXPECT_RESULT_OK(counterSetResult);
+        EXPECT_EQ(TEST_COUNTERSET, counterSetResult.value());
+        Result<StatsValue> statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey);
+        EXPECT_RESULT_OK(statsMapResult);
+        EXPECT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
+        EXPECT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
+        tagStatsMapKey.tag = 0;
+        statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey);
+        EXPECT_RESULT_OK(statsMapResult);
+        EXPECT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
+        EXPECT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
+        auto appStatsResult = mFakeAppUidStatsMap.readValue(uid);
+        EXPECT_RESULT_OK(appStatsResult);
+        EXPECT_EQ((uint64_t)1, appStatsResult.value().rxPackets);
+        EXPECT_EQ((uint64_t)100, appStatsResult.value().rxBytes);
+    }
+};
+
+TEST_F(TrafficControllerTest, TestSetCounterSet) {
+    uid_t callingUid = TEST_UID2;
+    addPrivilegedUid(callingUid);
+    ASSERT_EQ(0, mTc.setCounterSet(TEST_COUNTERSET, TEST_UID, callingUid));
+    uid_t uid = TEST_UID;
+    Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid);
+    ASSERT_RESULT_OK(counterSetResult);
+    ASSERT_EQ(TEST_COUNTERSET, counterSetResult.value());
+    ASSERT_EQ(0, mTc.setCounterSet(DEFAULT_COUNTERSET, TEST_UID, callingUid));
+    ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid).ok());
+    expectMapEmpty(mFakeUidCounterSetMap);
+}
+
+TEST_F(TrafficControllerTest, TestSetCounterSetWithoutPermission) {
+    ASSERT_EQ(-EPERM, mTc.setCounterSet(TEST_COUNTERSET, TEST_UID, TEST_UID2));
+    uid_t uid = TEST_UID;
+    ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid).ok());
+    expectMapEmpty(mFakeUidCounterSetMap);
+}
+
+TEST_F(TrafficControllerTest, TestSetInvalidCounterSet) {
+    uid_t callingUid = TEST_UID2;
+    addPrivilegedUid(callingUid);
+    ASSERT_GT(0, mTc.setCounterSet(OVERFLOW_COUNTERSET, TEST_UID, callingUid));
+    uid_t uid = TEST_UID;
+    ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid).ok());
+    expectMapEmpty(mFakeUidCounterSetMap);
+}
+
+TEST_F(TrafficControllerTest, TestDeleteTagDataWithoutPermission) {
+    uint64_t cookie = 1;
+    uid_t uid = TEST_UID;
+    uint32_t tag = TEST_TAG;
+    StatsKey tagStatsMapKey;
+    populateFakeStats(cookie, uid, tag, &tagStatsMapKey);
+    ASSERT_EQ(-EPERM, mTc.deleteTagData(0, TEST_UID, TEST_UID2));
+
+    expectFakeStatsUnchanged(cookie, tag, uid, tagStatsMapKey);
+}
+
+TEST_F(TrafficControllerTest, TestDeleteTagData) {
+    uid_t callingUid = TEST_UID2;
+    addPrivilegedUid(callingUid);
+    uint64_t cookie = 1;
+    uid_t uid = TEST_UID;
+    uint32_t tag = TEST_TAG;
+    StatsKey tagStatsMapKey;
+    populateFakeStats(cookie, uid, tag, &tagStatsMapKey);
+    ASSERT_EQ(0, mTc.deleteTagData(TEST_TAG, TEST_UID, callingUid));
+    ASSERT_FALSE(mFakeCookieTagMap.readValue(cookie).ok());
+    Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid);
+    ASSERT_RESULT_OK(counterSetResult);
+    ASSERT_EQ(TEST_COUNTERSET, counterSetResult.value());
+    ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey).ok());
+    tagStatsMapKey.tag = 0;
+    Result<StatsValue> statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey);
+    ASSERT_RESULT_OK(statsMapResult);
+    ASSERT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
+    ASSERT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
+    auto appStatsResult = mFakeAppUidStatsMap.readValue(TEST_UID);
+    ASSERT_RESULT_OK(appStatsResult);
+    ASSERT_EQ((uint64_t)1, appStatsResult.value().rxPackets);
+    ASSERT_EQ((uint64_t)100, appStatsResult.value().rxBytes);
+}
+
+TEST_F(TrafficControllerTest, TestDeleteAllUidData) {
+    uid_t callingUid = TEST_UID2;
+    addPrivilegedUid(callingUid);
+    uint64_t cookie = 1;
+    uid_t uid = TEST_UID;
+    uint32_t tag = TEST_TAG;
+    StatsKey tagStatsMapKey;
+    populateFakeStats(cookie, uid, tag, &tagStatsMapKey);
+    ASSERT_EQ(0, mTc.deleteTagData(0, TEST_UID, callingUid));
+    ASSERT_FALSE(mFakeCookieTagMap.readValue(cookie).ok());
+    ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid).ok());
+    ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey).ok());
+    tagStatsMapKey.tag = 0;
+    ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey).ok());
+    ASSERT_FALSE(mFakeAppUidStatsMap.readValue(TEST_UID).ok());
+}
+
+TEST_F(TrafficControllerTest, TestDeleteDataWithTwoTags) {
+    uid_t callingUid = TEST_UID2;
+    addPrivilegedUid(callingUid);
+    uint64_t cookie1 = 1;
+    uint64_t cookie2 = 2;
+    uid_t uid = TEST_UID;
+    uint32_t tag1 = TEST_TAG;
+    uint32_t tag2 = TEST_TAG + 1;
+    StatsKey tagStatsMapKey1;
+    StatsKey tagStatsMapKey2;
+    populateFakeStats(cookie1, uid, tag1, &tagStatsMapKey1);
+    populateFakeStats(cookie2, uid, tag2, &tagStatsMapKey2);
+    ASSERT_EQ(0, mTc.deleteTagData(TEST_TAG, TEST_UID, callingUid));
+    ASSERT_FALSE(mFakeCookieTagMap.readValue(cookie1).ok());
+    Result<UidTagValue> cookieMapResult = mFakeCookieTagMap.readValue(cookie2);
+    ASSERT_RESULT_OK(cookieMapResult);
+    ASSERT_EQ(TEST_UID, cookieMapResult.value().uid);
+    ASSERT_EQ(TEST_TAG + 1, cookieMapResult.value().tag);
+    Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid);
+    ASSERT_RESULT_OK(counterSetResult);
+    ASSERT_EQ(TEST_COUNTERSET, counterSetResult.value());
+    ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey1).ok());
+    Result<StatsValue> statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey2);
+    ASSERT_RESULT_OK(statsMapResult);
+    ASSERT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
+    ASSERT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
+}
+
+TEST_F(TrafficControllerTest, TestDeleteDataWithTwoUids) {
+    uid_t callingUid = TEST_UID2;
+    addPrivilegedUid(callingUid);
+    uint64_t cookie1 = 1;
+    uint64_t cookie2 = 2;
+    uid_t uid1 = TEST_UID;
+    uid_t uid2 = TEST_UID + 1;
+    uint32_t tag = TEST_TAG;
+    StatsKey tagStatsMapKey1;
+    StatsKey tagStatsMapKey2;
+    populateFakeStats(cookie1, uid1, tag, &tagStatsMapKey1);
+    populateFakeStats(cookie2, uid2, tag, &tagStatsMapKey2);
+
+    // Delete the stats of one of the uid. Check if it is properly collected by
+    // removedStats.
+    ASSERT_EQ(0, mTc.deleteTagData(0, uid2, callingUid));
+    ASSERT_FALSE(mFakeCookieTagMap.readValue(cookie2).ok());
+    Result<uint8_t> counterSetResult = mFakeUidCounterSetMap.readValue(uid1);
+    ASSERT_RESULT_OK(counterSetResult);
+    ASSERT_EQ(TEST_COUNTERSET, counterSetResult.value());
+    ASSERT_FALSE(mFakeUidCounterSetMap.readValue(uid2).ok());
+    ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey2).ok());
+    tagStatsMapKey2.tag = 0;
+    ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey2).ok());
+    ASSERT_FALSE(mFakeAppUidStatsMap.readValue(uid2).ok());
+    tagStatsMapKey1.tag = 0;
+    Result<StatsValue> statsMapResult = mFakeStatsMapA.readValue(tagStatsMapKey1);
+    ASSERT_RESULT_OK(statsMapResult);
+    ASSERT_EQ((uint64_t)1, statsMapResult.value().rxPackets);
+    ASSERT_EQ((uint64_t)100, statsMapResult.value().rxBytes);
+    auto appStatsResult = mFakeAppUidStatsMap.readValue(uid1);
+    ASSERT_RESULT_OK(appStatsResult);
+    ASSERT_EQ((uint64_t)1, appStatsResult.value().rxPackets);
+    ASSERT_EQ((uint64_t)100, appStatsResult.value().rxBytes);
+
+    // Delete the stats of the other uid.
+    ASSERT_EQ(0, mTc.deleteTagData(0, uid1, callingUid));
+    ASSERT_FALSE(mFakeStatsMapA.readValue(tagStatsMapKey1).ok());
+    ASSERT_FALSE(mFakeAppUidStatsMap.readValue(uid1).ok());
+}
+
+TEST_F(TrafficControllerTest, TestUpdateOwnerMapEntry) {
+    uint32_t uid = TEST_UID;
+    ASSERT_TRUE(isOk(mTc.updateOwnerMapEntry(STANDBY_MATCH, uid, DENY, DENYLIST)));
+    Result<UidOwnerValue> value = mFakeUidOwnerMap.readValue(uid);
+    ASSERT_RESULT_OK(value);
+    ASSERT_TRUE(value.value().rule & STANDBY_MATCH);
+
+    ASSERT_TRUE(isOk(mTc.updateOwnerMapEntry(DOZABLE_MATCH, uid, ALLOW, ALLOWLIST)));
+    value = mFakeUidOwnerMap.readValue(uid);
+    ASSERT_RESULT_OK(value);
+    ASSERT_TRUE(value.value().rule & DOZABLE_MATCH);
+
+    ASSERT_TRUE(isOk(mTc.updateOwnerMapEntry(DOZABLE_MATCH, uid, DENY, ALLOWLIST)));
+    value = mFakeUidOwnerMap.readValue(uid);
+    ASSERT_RESULT_OK(value);
+    ASSERT_FALSE(value.value().rule & DOZABLE_MATCH);
+
+    ASSERT_TRUE(isOk(mTc.updateOwnerMapEntry(STANDBY_MATCH, uid, ALLOW, DENYLIST)));
+    ASSERT_FALSE(mFakeUidOwnerMap.readValue(uid).ok());
+
+    uid = TEST_UID2;
+    ASSERT_FALSE(isOk(mTc.updateOwnerMapEntry(STANDBY_MATCH, uid, ALLOW, DENYLIST)));
+    ASSERT_FALSE(mFakeUidOwnerMap.readValue(uid).ok());
+}
+
+TEST_F(TrafficControllerTest, TestChangeUidOwnerRule) {
+    checkUidOwnerRuleForChain(DOZABLE, DOZABLE_MATCH);
+    checkUidOwnerRuleForChain(STANDBY, STANDBY_MATCH);
+    checkUidOwnerRuleForChain(POWERSAVE, POWERSAVE_MATCH);
+    checkUidOwnerRuleForChain(RESTRICTED, RESTRICTED_MATCH);
+    ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(NONE, TEST_UID, ALLOW, ALLOWLIST));
+    ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(INVALID_CHAIN, TEST_UID, ALLOW, ALLOWLIST));
+}
+
+TEST_F(TrafficControllerTest, TestReplaceUidOwnerMap) {
+    std::vector<int32_t> uids = {TEST_UID, TEST_UID2, TEST_UID3};
+    checkUidMapReplace("fw_dozable", uids, DOZABLE_MATCH);
+    checkUidMapReplace("fw_standby", uids, STANDBY_MATCH);
+    checkUidMapReplace("fw_powersave", uids, POWERSAVE_MATCH);
+    checkUidMapReplace("fw_restricted", uids, RESTRICTED_MATCH);
+    ASSERT_EQ(-EINVAL, mTc.replaceUidOwnerMap("unknow", true, uids));
+}
+
+TEST_F(TrafficControllerTest, TestReplaceSameChain) {
+    std::vector<int32_t> uids = {TEST_UID, TEST_UID2, TEST_UID3};
+    checkUidMapReplace("fw_dozable", uids, DOZABLE_MATCH);
+    std::vector<int32_t> newUids = {TEST_UID2, TEST_UID3};
+    checkUidMapReplace("fw_dozable", newUids, DOZABLE_MATCH);
+}
+
+TEST_F(TrafficControllerTest, TestDenylistUidMatch) {
+    std::vector<uint32_t> appUids = {1000, 1001, 10012};
+    ASSERT_TRUE(isOk(
+            mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, TrafficController::IptOpInsert)));
+    expectUidOwnerMapValues(appUids, PENALTY_BOX_MATCH, 0);
+    ASSERT_TRUE(isOk(
+            mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, TrafficController::IptOpDelete)));
+    expectMapEmpty(mFakeUidOwnerMap);
+}
+
+TEST_F(TrafficControllerTest, TestAllowlistUidMatch) {
+    std::vector<uint32_t> appUids = {1000, 1001, 10012};
+    ASSERT_TRUE(
+            isOk(mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, TrafficController::IptOpInsert)));
+    expectUidOwnerMapValues(appUids, HAPPY_BOX_MATCH, 0);
+    ASSERT_TRUE(
+            isOk(mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, TrafficController::IptOpDelete)));
+    expectMapEmpty(mFakeUidOwnerMap);
+}
+
+TEST_F(TrafficControllerTest, TestReplaceMatchUid) {
+    std::vector<uint32_t> appUids = {1000, 1001, 10012};
+    // Add appUids to the denylist and expect that their values are all PENALTY_BOX_MATCH.
+    ASSERT_TRUE(isOk(
+            mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, TrafficController::IptOpInsert)));
+    expectUidOwnerMapValues(appUids, PENALTY_BOX_MATCH, 0);
+
+    // Add the same UIDs to the allowlist and expect that we get PENALTY_BOX_MATCH |
+    // HAPPY_BOX_MATCH.
+    ASSERT_TRUE(
+            isOk(mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, TrafficController::IptOpInsert)));
+    expectUidOwnerMapValues(appUids, HAPPY_BOX_MATCH | PENALTY_BOX_MATCH, 0);
+
+    // Remove the same UIDs from the allowlist and check the PENALTY_BOX_MATCH is still there.
+    ASSERT_TRUE(
+            isOk(mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, TrafficController::IptOpDelete)));
+    expectUidOwnerMapValues(appUids, PENALTY_BOX_MATCH, 0);
+
+    // Remove the same UIDs from the denylist and check the map is empty.
+    ASSERT_TRUE(isOk(
+            mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, TrafficController::IptOpDelete)));
+    ASSERT_FALSE(mFakeUidOwnerMap.getFirstKey().ok());
+}
+
+TEST_F(TrafficControllerTest, TestDeleteWrongMatchSilentlyFails) {
+    std::vector<uint32_t> appUids = {1000, 1001, 10012};
+    // If the uid does not exist in the map, trying to delete a rule about it will fail.
+    ASSERT_FALSE(
+            isOk(mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, TrafficController::IptOpDelete)));
+    expectMapEmpty(mFakeUidOwnerMap);
+
+    // Add denylist rules for appUids.
+    ASSERT_TRUE(
+            isOk(mTc.updateUidOwnerMap(appUids, HAPPY_BOX_MATCH, TrafficController::IptOpInsert)));
+    expectUidOwnerMapValues(appUids, HAPPY_BOX_MATCH, 0);
+
+    // Delete (non-existent) denylist rules for appUids, and check that this silently does
+    // nothing if the uid is in the map but does not have denylist match. This is required because
+    // NetworkManagementService will try to remove a uid from denylist after adding it to the
+    // allowlist and if the remove fails it will not update the uid status.
+    ASSERT_TRUE(isOk(
+            mTc.updateUidOwnerMap(appUids, PENALTY_BOX_MATCH, TrafficController::IptOpDelete)));
+    expectUidOwnerMapValues(appUids, HAPPY_BOX_MATCH, 0);
+}
+
+TEST_F(TrafficControllerTest, TestAddUidInterfaceFilteringRules) {
+    int iif0 = 15;
+    ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif0, {1000, 1001})));
+    expectUidOwnerMapValues({1000, 1001}, IIF_MATCH, iif0);
+
+    // Add some non-overlapping new uids. They should coexist with existing rules
+    int iif1 = 16;
+    ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif1, {2000, 2001})));
+    expectUidOwnerMapValues({1000, 1001}, IIF_MATCH, iif0);
+    expectUidOwnerMapValues({2000, 2001}, IIF_MATCH, iif1);
+
+    // Overwrite some existing uids
+    int iif2 = 17;
+    ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif2, {1000, 2000})));
+    expectUidOwnerMapValues({1001}, IIF_MATCH, iif0);
+    expectUidOwnerMapValues({2001}, IIF_MATCH, iif1);
+    expectUidOwnerMapValues({1000, 2000}, IIF_MATCH, iif2);
+}
+
+TEST_F(TrafficControllerTest, TestRemoveUidInterfaceFilteringRules) {
+    int iif0 = 15;
+    int iif1 = 16;
+    ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif0, {1000, 1001})));
+    ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif1, {2000, 2001})));
+    expectUidOwnerMapValues({1000, 1001}, IIF_MATCH, iif0);
+    expectUidOwnerMapValues({2000, 2001}, IIF_MATCH, iif1);
+
+    // Rmove some uids
+    ASSERT_TRUE(isOk(mTc.removeUidInterfaceRules({1001, 2001})));
+    expectUidOwnerMapValues({1000}, IIF_MATCH, iif0);
+    expectUidOwnerMapValues({2000}, IIF_MATCH, iif1);
+    checkEachUidValue({1000, 2000}, IIF_MATCH);  // Make sure there are only two uids remaining
+
+    // Remove non-existent uids shouldn't fail
+    ASSERT_TRUE(isOk(mTc.removeUidInterfaceRules({2000, 3000})));
+    expectUidOwnerMapValues({1000}, IIF_MATCH, iif0);
+    checkEachUidValue({1000}, IIF_MATCH);  // Make sure there are only one uid remaining
+
+    // Remove everything
+    ASSERT_TRUE(isOk(mTc.removeUidInterfaceRules({1000})));
+    expectMapEmpty(mFakeUidOwnerMap);
+}
+
+TEST_F(TrafficControllerTest, TestUidInterfaceFilteringRulesCoexistWithExistingMatches) {
+    // Set up existing PENALTY_BOX_MATCH rules
+    ASSERT_TRUE(isOk(mTc.updateUidOwnerMap({1000, 1001, 10012}, PENALTY_BOX_MATCH,
+                                           TrafficController::IptOpInsert)));
+    expectUidOwnerMapValues({1000, 1001, 10012}, PENALTY_BOX_MATCH, 0);
+
+    // Add some partially-overlapping uid owner rules and check result
+    int iif1 = 32;
+    ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif1, {10012, 10013, 10014})));
+    expectUidOwnerMapValues({1000, 1001}, PENALTY_BOX_MATCH, 0);
+    expectUidOwnerMapValues({10012}, PENALTY_BOX_MATCH | IIF_MATCH, iif1);
+    expectUidOwnerMapValues({10013, 10014}, IIF_MATCH, iif1);
+
+    // Removing some PENALTY_BOX_MATCH rules should not change uid interface rule
+    ASSERT_TRUE(isOk(mTc.updateUidOwnerMap({1001, 10012}, PENALTY_BOX_MATCH,
+                                           TrafficController::IptOpDelete)));
+    expectUidOwnerMapValues({1000}, PENALTY_BOX_MATCH, 0);
+    expectUidOwnerMapValues({10012, 10013, 10014}, IIF_MATCH, iif1);
+
+    // Remove all uid interface rules
+    ASSERT_TRUE(isOk(mTc.removeUidInterfaceRules({10012, 10013, 10014})));
+    expectUidOwnerMapValues({1000}, PENALTY_BOX_MATCH, 0);
+    // Make sure these are the only uids left
+    checkEachUidValue({1000}, PENALTY_BOX_MATCH);
+}
+
+TEST_F(TrafficControllerTest, TestUidInterfaceFilteringRulesCoexistWithNewMatches) {
+    int iif1 = 56;
+    // Set up existing uid interface rules
+    ASSERT_TRUE(isOk(mTc.addUidInterfaceRules(iif1, {10001, 10002})));
+    expectUidOwnerMapValues({10001, 10002}, IIF_MATCH, iif1);
+
+    // Add some partially-overlapping doze rules
+    EXPECT_EQ(0, mTc.replaceUidOwnerMap("fw_dozable", true, {10002, 10003}));
+    expectUidOwnerMapValues({10001}, IIF_MATCH, iif1);
+    expectUidOwnerMapValues({10002}, DOZABLE_MATCH | IIF_MATCH, iif1);
+    expectUidOwnerMapValues({10003}, DOZABLE_MATCH, 0);
+
+    // Introduce a third rule type (powersave) on various existing UIDs
+    EXPECT_EQ(0, mTc.replaceUidOwnerMap("fw_powersave", true, {10000, 10001, 10002, 10003}));
+    expectUidOwnerMapValues({10000}, POWERSAVE_MATCH, 0);
+    expectUidOwnerMapValues({10001}, POWERSAVE_MATCH | IIF_MATCH, iif1);
+    expectUidOwnerMapValues({10002}, POWERSAVE_MATCH | DOZABLE_MATCH | IIF_MATCH, iif1);
+    expectUidOwnerMapValues({10003}, POWERSAVE_MATCH | DOZABLE_MATCH, 0);
+
+    // Remove all doze rules
+    EXPECT_EQ(0, mTc.replaceUidOwnerMap("fw_dozable", true, {}));
+    expectUidOwnerMapValues({10000}, POWERSAVE_MATCH, 0);
+    expectUidOwnerMapValues({10001}, POWERSAVE_MATCH | IIF_MATCH, iif1);
+    expectUidOwnerMapValues({10002}, POWERSAVE_MATCH | IIF_MATCH, iif1);
+    expectUidOwnerMapValues({10003}, POWERSAVE_MATCH, 0);
+
+    // Remove all powersave rules, expect ownerMap to only have uid interface rules left
+    EXPECT_EQ(0, mTc.replaceUidOwnerMap("fw_powersave", true, {}));
+    expectUidOwnerMapValues({10001, 10002}, IIF_MATCH, iif1);
+    // Make sure these are the only uids left
+    checkEachUidValue({10001, 10002}, IIF_MATCH);
+}
+
+TEST_F(TrafficControllerTest, TestGrantInternetPermission) {
+    std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
+
+    mTc.setPermissionForUids(INetd::PERMISSION_INTERNET, appUids);
+    expectMapEmpty(mFakeUidPermissionMap);
+    expectPrivilegedUserSetEmpty();
+}
+
+TEST_F(TrafficControllerTest, TestRevokeInternetPermission) {
+    std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
+
+    mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
+    expectUidPermissionMapValues(appUids, INetd::PERMISSION_NONE);
+}
+
+TEST_F(TrafficControllerTest, TestPermissionUninstalled) {
+    std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
+
+    mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, appUids);
+    expectUidPermissionMapValues(appUids, INetd::PERMISSION_UPDATE_DEVICE_STATS);
+    expectPrivilegedUserSet(appUids);
+
+    std::vector<uid_t> uidToRemove = {TEST_UID};
+    mTc.setPermissionForUids(INetd::PERMISSION_UNINSTALLED, uidToRemove);
+
+    std::vector<uid_t> uidRemain = {TEST_UID3, TEST_UID2};
+    expectUidPermissionMapValues(uidRemain, INetd::PERMISSION_UPDATE_DEVICE_STATS);
+    expectPrivilegedUserSet(uidRemain);
+
+    mTc.setPermissionForUids(INetd::PERMISSION_UNINSTALLED, uidRemain);
+    expectMapEmpty(mFakeUidPermissionMap);
+    expectPrivilegedUserSetEmpty();
+}
+
+TEST_F(TrafficControllerTest, TestGrantUpdateStatsPermission) {
+    std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
+
+    mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, appUids);
+    expectUidPermissionMapValues(appUids, INetd::PERMISSION_UPDATE_DEVICE_STATS);
+    expectPrivilegedUserSet(appUids);
+
+    mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
+    expectPrivilegedUserSetEmpty();
+    expectUidPermissionMapValues(appUids, INetd::PERMISSION_NONE);
+}
+
+TEST_F(TrafficControllerTest, TestRevokeUpdateStatsPermission) {
+    std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
+
+    mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, appUids);
+    expectPrivilegedUserSet(appUids);
+
+    std::vector<uid_t> uidToRemove = {TEST_UID};
+    mTc.setPermissionForUids(INetd::PERMISSION_NONE, uidToRemove);
+
+    std::vector<uid_t> uidRemain = {TEST_UID3, TEST_UID2};
+    expectPrivilegedUserSet(uidRemain);
+
+    mTc.setPermissionForUids(INetd::PERMISSION_NONE, uidRemain);
+    expectPrivilegedUserSetEmpty();
+}
+
+TEST_F(TrafficControllerTest, TestGrantWrongPermission) {
+    std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
+
+    mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
+    expectPrivilegedUserSetEmpty();
+    expectUidPermissionMapValues(appUids, INetd::PERMISSION_NONE);
+}
+
+TEST_F(TrafficControllerTest, TestGrantDuplicatePermissionSlientlyFail) {
+    std::vector<uid_t> appUids = {TEST_UID, TEST_UID2, TEST_UID3};
+
+    mTc.setPermissionForUids(INetd::PERMISSION_INTERNET, appUids);
+    expectMapEmpty(mFakeUidPermissionMap);
+
+    std::vector<uid_t> uidToAdd = {TEST_UID};
+    mTc.setPermissionForUids(INetd::PERMISSION_INTERNET, uidToAdd);
+
+    expectPrivilegedUserSetEmpty();
+
+    mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
+    expectUidPermissionMapValues(appUids, INetd::PERMISSION_NONE);
+
+    mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, appUids);
+    expectPrivilegedUserSet(appUids);
+
+    mTc.setPermissionForUids(INetd::PERMISSION_UPDATE_DEVICE_STATS, uidToAdd);
+    expectPrivilegedUserSet(appUids);
+
+    mTc.setPermissionForUids(INetd::PERMISSION_NONE, appUids);
+    expectPrivilegedUserSetEmpty();
+}
+
+}  // namespace net
+}  // namespace android
diff --git a/service/native/include/TrafficController.h b/service/native/include/TrafficController.h
new file mode 100644
index 0000000..3e98b68
--- /dev/null
+++ b/service/native/include/TrafficController.h
@@ -0,0 +1,222 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef NETD_SERVER_TRAFFIC_CONTROLLER_H
+#define NETD_SERVER_TRAFFIC_CONTROLLER_H
+
+#include <linux/bpf.h>
+
+#include "NetlinkListener.h"
+#include "android-base/thread_annotations.h"
+#include "bpf/BpfMap.h"
+#include "bpf_shared.h"
+#include "netdutils/DumpWriter.h"
+#include "netdutils/NetlinkListener.h"
+#include "netdutils/StatusOr.h"
+#include "utils/String16.h"
+
+namespace android {
+namespace net {
+
+class TrafficController {
+  public:
+    /*
+     * Initialize the whole controller
+     */
+    netdutils::Status start();
+
+    /*
+     * Similiar as above, no external lock required.
+     */
+    int setCounterSet(int counterSetNum, uid_t uid, uid_t callingUid) EXCLUDES(mMutex);
+
+    /*
+     * When deleting a tag data, the qtaguid module will grab the spinlock of each
+     * related rb_tree one by one and delete the tag information, counterSet
+     * information, iface stats information and uid stats information one by one.
+     * The new eBPF implementation is done similiarly by removing the entry on
+     * each map one by one. And deleting processes are also protected by the
+     * spinlock of the map. So no additional lock is required.
+     */
+    int deleteTagData(uint32_t tag, uid_t uid, uid_t callingUid) EXCLUDES(mMutex);
+
+    /*
+     * Swap the stats map config from current active stats map to the idle one.
+     */
+    netdutils::Status swapActiveStatsMap() EXCLUDES(mMutex);
+
+    /*
+     * Add the interface name and index pair into the eBPF map.
+     */
+    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 isAllowlist,
+                           const std::vector<int32_t>& uids);
+
+    enum IptOp { IptOpInsert, IptOpDelete };
+
+    netdutils::Status updateOwnerMapEntry(UidOwnerMatchType match, uid_t uid, FirewallRule rule,
+                                          FirewallType type) EXCLUDES(mMutex);
+
+    void dump(netdutils::DumpWriter& dw, bool verbose) EXCLUDES(mMutex);
+
+    netdutils::Status replaceRulesInMap(UidOwnerMatchType match, const std::vector<int32_t>& uids)
+            EXCLUDES(mMutex);
+
+    netdutils::Status addUidInterfaceRules(const int ifIndex, const std::vector<int32_t>& uids)
+            EXCLUDES(mMutex);
+    netdutils::Status removeUidInterfaceRules(const std::vector<int32_t>& uids) EXCLUDES(mMutex);
+
+    netdutils::Status updateUidOwnerMap(const std::vector<uint32_t>& appStrUids,
+                                        UidOwnerMatchType matchType, IptOp op) EXCLUDES(mMutex);
+    static const String16 DUMP_KEYWORD;
+
+    int toggleUidOwnerMap(ChildChain chain, bool enable) EXCLUDES(mMutex);
+
+    static netdutils::StatusOr<std::unique_ptr<netdutils::NetlinkListenerInterface>>
+    makeSkDestroyListener();
+
+    void setPermissionForUids(int permission, const std::vector<uid_t>& uids) EXCLUDES(mMutex);
+
+    FirewallType getFirewallType(ChildChain);
+
+    static const char* LOCAL_DOZABLE;
+    static const char* LOCAL_STANDBY;
+    static const char* LOCAL_POWERSAVE;
+    static const char* LOCAL_RESTRICTED;
+
+  private:
+    /*
+     * mCookieTagMap: Store the corresponding tag and uid for a specific socket.
+     * DO NOT hold any locks when modifying this map, otherwise when the untag
+     * operation is waiting for a lock hold by other process and there are more
+     * sockets being closed than can fit in the socket buffer of the netlink socket
+     * that receives them, then the kernel will drop some of these sockets and we
+     * won't delete their tags.
+     * Map Key: uint64_t socket cookie
+     * Map Value: UidTagValue, contains a uint32 uid and a uint32 tag.
+     */
+    bpf::BpfMap<uint64_t, UidTagValue> mCookieTagMap GUARDED_BY(mMutex);
+
+    /*
+     * mUidCounterSetMap: Store the counterSet of a specific uid.
+     * Map Key: uint32 uid.
+     * Map Value: uint32 counterSet specifies if the traffic is a background
+     * or foreground traffic.
+     */
+    bpf::BpfMap<uint32_t, uint8_t> mUidCounterSetMap GUARDED_BY(mMutex);
+
+    /*
+     * mAppUidStatsMap: Store the total traffic stats for a uid regardless of
+     * tag, counterSet and iface. The stats is used by TrafficStats.getUidStats
+     * API to return persistent stats for a specific uid since device boot.
+     */
+    bpf::BpfMap<uint32_t, StatsValue> mAppUidStatsMap;
+
+    /*
+     * mStatsMapA/mStatsMapB: Store the traffic statistics for a specific
+     * combination of uid, tag, iface and counterSet. These two maps contain
+     * both tagged and untagged traffic.
+     * Map Key: StatsKey contains the uid, tag, counterSet and ifaceIndex
+     * information.
+     * Map Value: Stats, contains packet count and byte count of each
+     * transport protocol on egress and ingress direction.
+     */
+    bpf::BpfMap<StatsKey, StatsValue> mStatsMapA GUARDED_BY(mMutex);
+
+    bpf::BpfMap<StatsKey, StatsValue> mStatsMapB GUARDED_BY(mMutex);
+
+    /*
+     * mIfaceIndexNameMap: Store the index name pair of each interface show up
+     * on the device since boot. The interface index is used by the eBPF program
+     * to correctly match the iface name when receiving a packet.
+     */
+    bpf::BpfMap<uint32_t, IfaceValue> mIfaceIndexNameMap;
+
+    /*
+     * mIfaceStataMap: Store per iface traffic stats gathered from xt_bpf
+     * filter.
+     */
+    bpf::BpfMap<uint32_t, StatsValue> mIfaceStatsMap;
+
+    /*
+     * mConfigurationMap: Store the current network policy about uid filtering
+     * and the current stats map in use. There are two configuration entries in
+     * the map right now:
+     * - Entry with UID_RULES_CONFIGURATION_KEY:
+     *    Store the configuration for the current uid rules. It indicates the device
+     *    is in doze/powersave/standby/restricted mode.
+     * - Entry with CURRENT_STATS_MAP_CONFIGURATION_KEY:
+     *    Stores the current live stats map that kernel program is writing to.
+     *    Userspace can do scraping and cleaning job on the other one depending on the
+     *    current configs.
+     */
+    bpf::BpfMap<uint32_t, uint8_t> mConfigurationMap GUARDED_BY(mMutex);
+
+    /*
+     * mUidOwnerMap: Store uids that are used for bandwidth control uid match.
+     */
+    bpf::BpfMap<uint32_t, UidOwnerValue> mUidOwnerMap GUARDED_BY(mMutex);
+
+    /*
+     * mUidOwnerMap: Store uids that are used for INTERNET permission check.
+     */
+    bpf::BpfMap<uint32_t, uint8_t> mUidPermissionMap GUARDED_BY(mMutex);
+
+    std::unique_ptr<netdutils::NetlinkListenerInterface> mSkDestroyListener;
+
+    netdutils::Status removeRule(uint32_t uid, UidOwnerMatchType match) REQUIRES(mMutex);
+
+    netdutils::Status addRule(uint32_t uid, UidOwnerMatchType match, uint32_t iif = 0)
+            REQUIRES(mMutex);
+
+    // mMutex guards all accesses to mConfigurationMap, mUidOwnerMap, mUidPermissionMap,
+    // mStatsMapA, mStatsMapB and mPrivilegedUser. It is designed to solve the following
+    // problems:
+    // 1. Prevent concurrent access and modification to mConfigurationMap, mUidOwnerMap,
+    //    mUidPermissionMap, and mPrivilegedUser. These data members are controlled by netd but can
+    //    be modified from different threads. TrafficController provides several APIs directly
+    //    called by the binder RPC, and different binder threads can concurrently access these data
+    //    members mentioned above. Some of the data members such as mUidPermissionMap and
+    //    mPrivilegedUsers are also accessed from a different thread when tagging sockets or
+    //    setting the counterSet through FwmarkServer
+    // 2. Coordinate the deletion of uid stats in mStatsMapA and mStatsMapB. The system server
+    //    always call into netd to ask for a live stats map change before it pull and clean up the
+    //    stats from the inactive map. The mMutex will block netd from accessing the stats map when
+    //    the mConfigurationMap is updating the current stats map so netd will not accidentally
+    //    read the map that system_server is cleaning up.
+    std::mutex mMutex;
+
+    netdutils::Status initMaps() EXCLUDES(mMutex);
+
+    // Keep track of uids that have permission UPDATE_DEVICE_STATS so we don't
+    // need to call back to system server for permission check.
+    std::set<uid_t> mPrivilegedUser GUARDED_BY(mMutex);
+
+    bool hasUpdateDeviceStatsPermission(uid_t uid) REQUIRES(mMutex);
+
+    // For testing
+    friend class TrafficControllerTest;
+};
+
+}  // namespace net
+}  // namespace android
+
+#endif  // NETD_SERVER_TRAFFIC_CONTROLLER_H