DO NOT MERGE - Merge pi-dev@5234907 into stage-aosp-master

Bug: 120848293
Change-Id: I4f96f8351b170c9b438079080c7a3885237f1cf8
diff --git a/.clang-format b/.clang-format
deleted file mode 100644
index f5c9843..0000000
--- a/.clang-format
+++ /dev/null
@@ -1,15 +0,0 @@
-# https://clang.llvm.org/docs/ClangFormatStyleOptions.html
-# Please keep this file in sync with system/core/.clang-format-4
-#
-BasedOnStyle: Google
-AccessModifierOffset: -2
-AllowShortFunctionsOnASingleLine: Inline
-BreakBeforeTernaryOperators: true
-ColumnLimit: 100
-CommentPragmas: NOLINT:.*
-DerivePointerAlignment: false
-IndentWidth: 4
-ContinuationIndentWidth: 4
-PointerAlignment: Left
-TabWidth: 4
-UseTab: Never
diff --git a/.clang-format b/.clang-format
new file mode 120000
index 0000000..ddcf5a2
--- /dev/null
+++ b/.clang-format
@@ -0,0 +1 @@
+../../build/soong/scripts/system-clang-format
\ No newline at end of file
diff --git a/bpfloader/Android.bp b/bpfloader/Android.bp
index 81cdbe1..261a824 100644
--- a/bpfloader/Android.bp
+++ b/bpfloader/Android.bp
@@ -47,8 +47,7 @@
 
     required: [
         "netd.o",
-// Uncomment once security related patches ready
-//      "time_in_state.o",
+        "time_in_state.o",
     ],
 
 }
diff --git a/libbpf_android/Android.bp b/libbpf_android/Android.bp
index 721792f..da5c6a6 100644
--- a/libbpf_android/Android.bp
+++ b/libbpf_android/Android.bp
@@ -45,6 +45,7 @@
     shared_libs: [
         "libbase",
         "libutils",
+        "libprocessgroup",
         "liblog",
         "libnetdutils",
         "libbpf",
@@ -83,3 +84,29 @@
         "libutils",
     ],
 }
+
+cc_test {
+    name: "libbpf_load_test",
+    srcs: [
+        "BpfLoadTest.cpp",
+    ],
+    defaults: ["bpf_defaults"],
+    cflags: [
+        "-Wall",
+        "-Werror",
+        "-Wno-error=unused-variable",
+    ],
+    static_libs: ["libgmock"],
+    shared_libs: [
+        "libbpf_android",
+        "libbpf",
+        "libbase",
+        "liblog",
+        "libnetdutils",
+        "libutils",
+    ],
+
+    required: [
+        "bpf_load_tp_prog.o",
+    ],
+}
diff --git a/libbpf_android/BpfLoadTest.cpp b/libbpf_android/BpfLoadTest.cpp
new file mode 100644
index 0000000..cae60f0
--- /dev/null
+++ b/libbpf_android/BpfLoadTest.cpp
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2018 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.
+ */
+
+#include <android-base/macros.h>
+#include <gtest/gtest.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <iostream>
+#include "include/bpf/BpfMap.h"
+#include "include/bpf/BpfUtils.h"
+#include "include/libbpf_android.h"
+
+using ::testing::Test;
+
+constexpr const char tp_prog_path[] =
+        "/sys/fs/bpf/prog_bpf_load_tp_prog_tracepoint_sched_sched_switch";
+constexpr const char tp_map_path[] = "/sys/fs/bpf/map_bpf_load_tp_prog_cpu_pid";
+
+namespace android {
+namespace bpf {
+
+class BpfLoadTest : public testing::Test {
+  protected:
+    BpfLoadTest() {}
+    int mProgFd, mMapFd;
+
+    void SetUp() {
+        SKIP_IF_BPF_NOT_SUPPORTED;
+
+        unlink(tp_prog_path);
+        unlink(tp_map_path);
+
+        EXPECT_EQ(android::bpf::loadProg("/system/etc/bpf/bpf_load_tp_prog.o"), 0);
+
+        mProgFd = bpf_obj_get(tp_prog_path);
+        EXPECT_GT(mProgFd, 0);
+
+        mMapFd = bpf_obj_get(tp_map_path);
+        EXPECT_GT(mMapFd, 0);
+
+        int ret = bpf_attach_tracepoint(mProgFd, "sched", "sched_switch");
+        EXPECT_NE(ret, 0);
+    }
+
+    void TearDown() {
+        SKIP_IF_BPF_NOT_SUPPORTED;
+
+        close(mProgFd);
+        close(mMapFd);
+        unlink(tp_prog_path);
+        unlink(tp_map_path);
+    }
+
+    void checkMapNonZero() {
+        // The test program installs a tracepoint on sched:sched_switch
+        // and expects the kernel to populate a PID corresponding to CPU
+        android::bpf::BpfMap<uint32_t, uint32_t> m(mMapFd);
+
+        // Wait for program to run a little
+        sleep(1);
+
+        int non_zero = 0;
+        const auto iterFunc = [&non_zero](const uint32_t& key, const uint32_t& val,
+                                          BpfMap<uint32_t, uint32_t>& map) {
+            if (val && !non_zero) {
+                non_zero = 1;
+            }
+
+            UNUSED(key);
+            UNUSED(map);
+            return android::netdutils::status::ok;
+        };
+
+        EXPECT_OK(m.iterateWithValue(iterFunc));
+        EXPECT_EQ(non_zero, 1);
+    }
+};
+
+TEST_F(BpfLoadTest, bpfCheckMap) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    checkMapNonZero();
+}
+
+}  // namespace bpf
+}  // namespace android
diff --git a/libbpf_android/BpfUtils.cpp b/libbpf_android/BpfUtils.cpp
index 9f482d6..109715b 100644
--- a/libbpf_android/BpfUtils.cpp
+++ b/libbpf_android/BpfUtils.cpp
@@ -23,6 +23,7 @@
 #include <linux/bpf.h>
 #include <linux/if_ether.h>
 #include <linux/in.h>
+#include <linux/pfkeyv2.h>
 #include <stdlib.h>
 #include <string.h>
 #include <sys/mman.h>
@@ -37,13 +38,16 @@
 #include <log/log.h>
 #include <netdutils/MemBlock.h>
 #include <netdutils/Slice.h>
+#include <processgroup/processgroup.h>
 
 using android::base::GetUintProperty;
 using android::base::unique_fd;
 using android::netdutils::MemBlock;
 using android::netdutils::Slice;
 
-constexpr size_t LOG_BUF_SIZE = 65536;
+// The buffer size for the buffer that records program loading logs, needs to be large enough for
+// the largest kernel program.
+constexpr size_t LOG_BUF_SIZE = 0x20000;
 
 namespace android {
 namespace bpf {
@@ -212,7 +216,38 @@
     return sock_cookie;
 }
 
-bool hasBpfSupport() {
+int synchronizeKernelRCU() {
+    // This is a temporary hack for network stats map swap on devices running
+    // 4.9 kernels. The kernel code of socket release on pf_key socket will
+    // explicitly call synchronize_rcu() which is exactly what we need.
+    int pfSocket = socket(AF_KEY, SOCK_RAW | SOCK_CLOEXEC, PF_KEY_V2);
+
+    if (pfSocket < 0) {
+        int ret = -errno;
+        ALOGE("create PF_KEY socket failed: %s", strerror(errno));
+        return ret;
+    }
+
+    // When closing socket, synchronize_rcu() gets called in sock_release().
+    if (close(pfSocket)) {
+        int ret = -errno;
+        ALOGE("failed to close the PF_KEY socket: %s", strerror(errno));
+        return ret;
+    }
+    return 0;
+}
+
+std::string BpfLevelToString(BpfLevel bpfLevel) {
+    switch (bpfLevel) {
+        case BpfLevel::NONE:      return "NONE_SUPPORT";
+        case BpfLevel::BASIC:     return "BPF_LEVEL_BASIC";
+        case BpfLevel::EXTENDED:  return "BPF_LEVEL_EXTENDED";
+        // No default statement. We want to see errors of the form:
+        // "enumeration value 'BPF_LEVEL_xxx' not handled in switch [-Werror,-Wswitch]".
+    }
+}
+
+BpfLevel getBpfSupportLevel() {
     struct utsname buf;
     int kernel_version_major;
     int kernel_version_minor;
@@ -223,18 +258,22 @@
         api_level = GetUintProperty<uint64_t>("ro.build.version.sdk", 0);
     }
 
+    // Check if the device is shipped originally with android P.
+    if (api_level < MINIMUM_API_REQUIRED) return BpfLevel::NONE;
+
     int ret = uname(&buf);
     if (ret) {
-        return false;
+        return BpfLevel::NONE;
     }
     char dummy;
     ret = sscanf(buf.release, "%d.%d%c", &kernel_version_major, &kernel_version_minor, &dummy);
-    if (ret >= 2 &&
-        ((kernel_version_major > 4) || (kernel_version_major == 4 && kernel_version_minor >= 9))) {
-        // Check if the device is shipped originally with android P.
-        return api_level >= MINIMUM_API_REQUIRED;
-    }
-    return false;
+    // Check the device kernel version
+    if (ret < 2) return BpfLevel::NONE;
+    if (kernel_version_major > 4 || (kernel_version_major == 4 && kernel_version_minor >= 14))
+        return BpfLevel::EXTENDED;
+    if (kernel_version_major == 4 && kernel_version_minor >= 9) return BpfLevel::BASIC;
+
+    return BpfLevel::NONE;
 }
 
 int loadAndPinProgram(BpfProgInfo* prog, Slice progBlock) {
@@ -248,7 +287,13 @@
         return ret;
     }
     if (prog->attachType == BPF_CGROUP_INET_EGRESS || prog->attachType == BPF_CGROUP_INET_INGRESS) {
-        unique_fd cg_fd(open(CGROUP_ROOT_PATH, O_DIRECTORY | O_RDONLY | O_CLOEXEC));
+        std::string cg2_path;
+        if (!CgroupGetControllerPath(CGROUPV2_CONTROLLER_NAME, &cg2_path)) {
+            int ret = -errno;
+            ALOGE("Failed to find cgroup v2 root");
+            return ret;
+        }
+        unique_fd cg_fd(open(cg2_path.c_str(), O_DIRECTORY | O_RDONLY | O_CLOEXEC));
         if (cg_fd < 0) {
             int ret = -errno;
             ALOGE("Failed to open the cgroup directory");
diff --git a/libbpf_android/Loader.cpp b/libbpf_android/Loader.cpp
index 332c0b4..d02ed87 100644
--- a/libbpf_android/Loader.cpp
+++ b/libbpf_android/Loader.cpp
@@ -34,6 +34,7 @@
 #include <fstream>
 #include <iostream>
 #include <string>
+#include <vector>
 
 #include <android-base/strings.h>
 
@@ -45,6 +46,7 @@
 using android::base::StartsWith;
 using std::ifstream;
 using std::ios;
+using std::string;
 using std::vector;
 
 namespace android {
@@ -63,13 +65,15 @@
  * is the name of the program, and tracepoint is the type.
  */
 sectionType sectionNameTypes[] = {
-    { "kprobe", BPF_PROG_TYPE_KPROBE },
-    { "tracepoint", BPF_PROG_TYPE_TRACEPOINT },
-    { "skfilter", BPF_PROG_TYPE_SOCKET_FILTER },
-    { "cgroupskb", BPF_PROG_TYPE_CGROUP_SKB },
+    {"kprobe", BPF_PROG_TYPE_KPROBE},
+    {"tracepoint", BPF_PROG_TYPE_TRACEPOINT},
+    {"skfilter", BPF_PROG_TYPE_SOCKET_FILTER},
+    {"cgroupskb", BPF_PROG_TYPE_CGROUP_SKB},
+    {"schedcls", BPF_PROG_TYPE_SCHED_CLS},
+    {"cgroupsock", BPF_PROG_TYPE_CGROUP_SOCK},
 
     /* End of table */
-    { "END", BPF_PROG_TYPE_UNSPEC },
+    {"END", BPF_PROG_TYPE_UNSPEC},
 };
 
 typedef struct {
@@ -496,7 +500,7 @@
     }
 }
 
-static int loadCodeSections(const char* elfPath, vector<codeSection>& cs, string license) {
+static int loadCodeSections(const char* elfPath, vector<codeSection>& cs, const string& license) {
     int ret, fd, kvers;
 
     if ((kvers = getMachineKvers()) < 0) return -1;
diff --git a/libbpf_android/LoaderUtils.h b/libbpf_android/LoaderUtils.h
index f4912e5..9eda833 100644
--- a/libbpf_android/LoaderUtils.h
+++ b/libbpf_android/LoaderUtils.h
@@ -16,20 +16,17 @@
 
 #include <stdio.h>
 #include <iostream>
+#include <string>
 #include <vector>
 
 #include <android-base/strings.h>
 
-using android::base::Split;
-using std::string;
-using std::vector;
-
-static string pathToFilename(string path, bool noext = false) {
-    vector<string> spath = android::base::Split(path, "/");
-    string ret = spath.back();
+static std::string pathToFilename(const std::string& path, bool noext = false) {
+    std::vector<std::string> spath = android::base::Split(path, "/");
+    std::string ret = spath.back();
 
     if (noext) {
-        size_t lastindex = ret.find_last_of(".");
+        size_t lastindex = ret.find_last_of('.');
         return ret.substr(0, lastindex);
     }
     return ret;
@@ -43,21 +40,21 @@
     if (uname(&un)) return -1;
     unameOut = un.release;
 
-    string s = unameOut;
-    string token;
+    std::string s = unameOut;
+    std::string token;
     size_t pos = 0;
     int cur_num = 0;
 
-    while ((pos = s.find(".")) != string::npos && cur_num < 3) {
+    while ((pos = s.find('.')) != std::string::npos && cur_num < 3) {
         token = s.substr(0, pos);
         s.erase(0, pos + 1);
 
-        if ((pos = token.find("-")) != string::npos) token = token.substr(0, pos);
+        if ((pos = token.find('-')) != std::string::npos) token = token.substr(0, pos);
 
         nums[cur_num++] = stoi(token);
     }
 
-    if ((pos = s.find("-")) != string::npos)
+    if ((pos = s.find('-')) != std::string::npos)
         token = s.substr(0, pos);
     else
         token = s;
@@ -72,6 +69,6 @@
         return (65536 * nums[0] + 256 * nums[1] + nums[2]);
 }
 
-static void deslash(string& s) {
+static void deslash(std::string& s) {
     std::replace(s.begin(), s.end(), '/', '_');
 }
diff --git a/libbpf_android/include/bpf/BpfUtils.h b/libbpf_android/include/bpf/BpfUtils.h
index 6bb1f0f..77a18fa 100644
--- a/libbpf_android/include/bpf/BpfUtils.h
+++ b/libbpf_android/include/bpf/BpfUtils.h
@@ -121,12 +121,22 @@
     }
 };
 
+enum class BpfLevel {
+    // Devices shipped before P or kernel version is lower than 4.9 do not
+    // have eBPF enabled.
+    NONE,
+    // Devices shipped in P with android 4.9 kernel only have the basic eBPF
+    // functionalities such as xt_bpf and cgroup skb filter.
+    BASIC,
+    // For devices that have 4.14 kernel. It supports advanced features like
+    // map_in_map and cgroup socket filter.
+    EXTENDED,
+};
+
 #ifndef DEFAULT_OVERFLOWUID
 #define DEFAULT_OVERFLOWUID 65534
 #endif
 
-constexpr const char* CGROUP_ROOT_PATH = "/dev/cg2_bpf";
-
 constexpr const int OVERFLOW_COUNTERSET = 2;
 
 constexpr const uint64_t NONEXISTENT_COOKIE = 0;
@@ -147,13 +157,23 @@
 int attachProgram(bpf_attach_type type, uint32_t prog_fd, uint32_t cg_fd);
 int detachProgram(bpf_attach_type type, uint32_t cg_fd);
 uint64_t getSocketCookie(int sockFd);
-bool hasBpfSupport();
+std::string BpfLevelToString(BpfLevel BpfLevel);
+BpfLevel getBpfSupportLevel();
 int parseProgramsFromFile(const char* path, BpfProgInfo* programs, size_t size,
                           const std::vector<BpfMapInfo>& mapPatterns);
+int synchronizeKernelRCU();
 
-#define SKIP_IF_BPF_NOT_SUPPORTED     \
-    do {                              \
-        if (!hasBpfSupport()) return; \
+#define SKIP_IF_BPF_NOT_SUPPORTED                                                    \
+    do {                                                                             \
+        if (android::bpf::getBpfSupportLevel() == android::bpf::BpfLevel::NONE) {    \
+            GTEST_LOG_(INFO) << "This test is skipped since bpf is not available\n"; \
+            return;                                                                  \
+        }                                                                            \
+    } while (0)
+
+#define SKIP_IF_BPF_SUPPORTED                                                           \
+    do {                                                                                \
+        if (android::bpf::getBpfSupportLevel() != android::bpf::BpfLevel::NONE) return; \
     } while (0)
 
 constexpr int BPF_CONTINUE = 0;
diff --git a/progs/Android.bp b/progs/Android.bp
index 35ba797..4302129 100644
--- a/progs/Android.bp
+++ b/progs/Android.bp
@@ -24,5 +24,8 @@
         "-Wall",
         "-Werror",
     ],
-    include_dirs: ["system/netd/libnetdbpf/include"],
+    include_dirs: [
+        "system/netd/libnetdbpf/include",
+        "system/netd/libnetdutils/include",
+    ],
 }
diff --git a/progs/include/bpf_helpers.h b/progs/include/bpf_helpers.h
new file mode 100644
index 0000000..9fca3cd
--- /dev/null
+++ b/progs/include/bpf_helpers.h
@@ -0,0 +1,49 @@
+/* Common BPF helpers to be used by all BPF programs loaded by Android */
+
+#include <linux/bpf.h>
+#include <stdbool.h>
+#include <stdint.h>
+
+/* place things in different elf sections */
+#define SEC(NAME) __attribute__((section(NAME), used))
+
+/*
+ * Helper functions called from eBPF programs written in C. These are
+ * implemented in the kernel sources.
+ */
+/* generic functions */
+static void* (*bpf_map_lookup_elem)(void* map, void* key) = (void*) BPF_FUNC_map_lookup_elem;
+static int (*bpf_map_update_elem)(void* map, void* key, void* value,
+                                  unsigned long long flags) = (void*) BPF_FUNC_map_update_elem;
+static int (*bpf_map_delete_elem)(void* map, void* key) = (void*) BPF_FUNC_map_delete_elem;
+static int (*bpf_probe_read)(void* dst, int size, void* unsafe_ptr) = (void*) BPF_FUNC_probe_read;
+static unsigned long long (*bpf_ktime_get_ns)(void) = (void*) BPF_FUNC_ktime_get_ns;
+static int (*bpf_trace_printk)(const char* fmt, int fmt_size, ...) = (void*) BPF_FUNC_trace_printk;
+static unsigned long long (*bpf_get_current_pid_tgid)(void) = (void*) BPF_FUNC_get_current_pid_tgid;
+static unsigned long long (*bpf_get_current_uid_gid)(void) = (void*) BPF_FUNC_get_current_uid_gid;
+static unsigned long long (*bpf_get_smp_processor_id)(void) = (void*) BPF_FUNC_get_smp_processor_id;
+/* networking  */
+static uint64_t (*bpf_get_socket_cookie)(struct __sk_buff* skb) = (void*)BPF_FUNC_get_socket_cookie;
+static uint32_t (*bpf_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;
+
+/*
+ * Map structure to be used by Android eBPF C programs. The Android eBPF loader
+ * uses this structure from eBPF object to create maps at boot time.
+ *
+ * The eBPF C program should define structure in the maps section using
+ * SEC("maps") otherwise it will be ignored by the eBPF loader.
+ *
+ * For example:
+ * struct bpf_map_def SEC("maps") mymap { .type=... , .key_size=... }
+ */
+struct bpf_map_def {
+    unsigned int type;
+    unsigned int key_size;
+    unsigned int value_size;
+    unsigned int max_entries;
+    unsigned int map_flags;
+    unsigned int pad1;
+    unsigned int pad2;
+};
diff --git a/progs/netd.c b/progs/netd.c
index 0ea51a9..1d47c73 100644
--- a/progs/netd.c
+++ b/progs/netd.c
@@ -43,19 +43,41 @@
 
 SEC("skfilter/whitelist/xtbpf")
 int xt_bpf_whitelist_prog(struct __sk_buff* skb) {
-    uint32_t sock_uid = get_socket_uid(skb);
+    uint32_t sock_uid = bpf_get_socket_uid(skb);
     if (is_system_uid(sock_uid)) return BPF_MATCH;
-    uint8_t* whitelistMatch = find_map_entry(&uid_owner_map, &sock_uid);
+    uint8_t* whitelistMatch = bpf_map_lookup_elem(&uid_owner_map, &sock_uid);
     if (whitelistMatch) return *whitelistMatch & HAPPY_BOX_MATCH;
     return BPF_NOMATCH;
 }
 
 SEC("skfilter/blacklist/xtbpf")
 int xt_bpf_blacklist_prog(struct __sk_buff* skb) {
-    uint32_t sock_uid = get_socket_uid(skb);
-    uint8_t* blacklistMatch = find_map_entry(&uid_owner_map, &sock_uid);
+    uint32_t sock_uid = bpf_get_socket_uid(skb);
+    uint8_t* blacklistMatch = bpf_map_lookup_elem(&uid_owner_map, &sock_uid);
     if (blacklistMatch) return *blacklistMatch & PENALTY_BOX_MATCH;
     return BPF_NOMATCH;
 }
 
+struct bpf_map_def SEC("maps") uid_permission_map = {
+    .type = BPF_MAP_TYPE_HASH,
+    .key_size = sizeof(uint32_t),
+    .value_size = sizeof(uint8_t),
+    .max_entries = UID_OWNER_MAP_SIZE,
+};
+
+SEC("cgroupsock/inet/creat")
+int inet_socket_create(struct bpf_sock* sk) {
+    uint64_t gid_uid = bpf_get_current_uid_gid();
+    /*
+     * A given app is guaranteed to have the same app ID in all the profiles in
+     * which it is installed, and install permission is granted to app for all
+     * user at install time so we only check the appId part of a request uid at
+     * run time. See UserHandle#isSameApp for detail.
+     */
+    uint32_t appId = (gid_uid & 0xffffffff) % PER_USER_RANGE;
+    uint8_t* internetPermission = bpf_map_lookup_elem(&uid_permission_map, &appId);
+    if (internetPermission) return *internetPermission & ALLOW_SOCK_CREATE;
+    return NO_PERMISSION;
+}
+
 char _license[] SEC("license") = "Apache 2.0";
diff --git a/progs/netd.h b/progs/netd.h
index 6b2d103..fcf3a72 100644
--- a/progs/netd.h
+++ b/progs/netd.h
@@ -19,6 +19,7 @@
  * program.
  */
 
+#include <bpf_helpers.h>
 #include <linux/bpf.h>
 #include <linux/if.h>
 #include <linux/if_ether.h>
@@ -30,8 +31,6 @@
 #include <stdint.h>
 #include "netdbpf/bpf_shared.h"
 
-#define SEC(NAME) __attribute__((section(NAME), used))
-
 struct uid_tag {
     uint32_t uid;
     uint32_t tag;
@@ -55,16 +54,6 @@
     char name[IFNAMSIZ];
 };
 
-/* helper functions called from eBPF programs written in C */
-static void* (*find_map_entry)(void* map, void* key) = (void*)BPF_FUNC_map_lookup_elem;
-static int (*write_to_map_entry)(void* map, void* key, void* value,
-                                 uint64_t flags) = (void*)BPF_FUNC_map_update_elem;
-static int (*delete_map_entry)(void* map, void* key) = (void*)BPF_FUNC_map_delete_elem;
-static uint64_t (*get_socket_cookie)(struct __sk_buff* skb) = (void*)BPF_FUNC_get_socket_cookie;
-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;
-
 // This is defined for cgroup bpf filter only.
 #define BPF_PASS 1
 #define BPF_DROP 0
@@ -82,16 +71,6 @@
 #define TCP_FLAG_OFF 13
 #define RST_OFFSET 2
 
-/* loader usage */
-struct bpf_map_def {
-    unsigned int type;
-    unsigned int key_size;
-    unsigned int value_size;
-    unsigned int max_entries;
-    unsigned int map_flags;
-    unsigned int pad[2];
-};
-
 struct bpf_map_def SEC("maps") cookie_tag_map = {
     .type = BPF_MAP_TYPE_HASH,
     .key_size = sizeof(uint64_t),
@@ -110,21 +89,21 @@
     .type = BPF_MAP_TYPE_HASH,
     .key_size = sizeof(uint32_t),
     .value_size = sizeof(struct stats_value),
-    .max_entries = UID_STATS_MAP_SIZE,
+    .max_entries = APP_STATS_MAP_SIZE,
 };
 
-struct bpf_map_def SEC("maps") uid_stats_map = {
+struct bpf_map_def SEC("maps") stats_map_A = {
     .type = BPF_MAP_TYPE_HASH,
     .key_size = sizeof(struct stats_key),
     .value_size = sizeof(struct stats_value),
-    .max_entries = UID_STATS_MAP_SIZE,
+    .max_entries = STATS_MAP_SIZE,
 };
 
-struct bpf_map_def SEC("maps") tag_stats_map = {
+struct bpf_map_def SEC("maps") stats_map_B = {
     .type = BPF_MAP_TYPE_HASH,
     .key_size = sizeof(struct stats_key),
     .value_size = sizeof(struct stats_value),
-    .max_entries = TAG_STATS_MAP_SIZE,
+    .max_entries = STATS_MAP_SIZE,
 };
 
 struct bpf_map_def SEC("maps") iface_stats_map = {
@@ -162,11 +141,11 @@
 static __always_inline inline void bpf_update_stats(struct __sk_buff* skb, struct bpf_map_def* map,
                                                     int direction, void* key) {
     struct stats_value* value;
-    value = find_map_entry(map, key);
+    value = bpf_map_lookup_elem(map, key);
     if (!value) {
         struct stats_value newValue = {};
-        write_to_map_entry(map, key, &newValue, BPF_NOEXIST);
-        value = find_map_entry(map, key);
+        bpf_map_update_elem(map, key, &newValue, BPF_NOEXIST);
+        value = bpf_map_lookup_elem(map, key);
     }
     if (value) {
         if (direction == BPF_EGRESS) {
@@ -218,9 +197,9 @@
     return false;
 }
 
-static __always_inline BpfConfig getConfig() {
-    uint32_t mapSettingKey = CONFIGURATION_KEY;
-    BpfConfig* config = find_map_entry(&configuration_map, &mapSettingKey);
+static __always_inline BpfConfig getConfig(uint32_t configKey) {
+    uint32_t mapSettingKey = configKey;
+    BpfConfig* config = bpf_map_lookup_elem(&configuration_map, &mapSettingKey);
     if (!config) {
         // Couldn't read configuration entry. Assume everything is disabled.
         return DEFAULT_CONFIG;
@@ -233,12 +212,12 @@
 
     if ((uid <= MAX_SYSTEM_UID) && (uid >= MIN_SYSTEM_UID)) return BPF_PASS;
 
-    BpfConfig enabledRules = getConfig();
+    BpfConfig enabledRules = getConfig(UID_RULES_CONFIGURATION_KEY);
     if (!enabledRules) {
         return BPF_PASS;
     }
 
-    uint8_t* uidEntry = find_map_entry(&uid_owner_map, &uid);
+    uint8_t* uidEntry = bpf_map_lookup_elem(&uid_owner_map, &uid);
     uint8_t uidRules = uidEntry ? *uidEntry : 0;
     if ((enabledRules & DOZABLE_MATCH) && !(uidRules & DOZABLE_MATCH)) {
         return BPF_DROP;
@@ -252,8 +231,17 @@
     return BPF_PASS;
 }
 
+static __always_inline inline void update_stats_with_config(struct __sk_buff* skb, int direction,
+                                                            void* key, uint8_t selectedMap) {
+    if (selectedMap == SELECT_MAP_A) {
+        bpf_update_stats(skb, &stats_map_A, direction, key);
+    } else if (selectedMap == SELECT_MAP_B) {
+        bpf_update_stats(skb, &stats_map_B, direction, key);
+    }
+}
+
 static __always_inline inline int bpf_traffic_account(struct __sk_buff* skb, int direction) {
-    uint32_t sock_uid = get_socket_uid(skb);
+    uint32_t sock_uid = bpf_get_socket_uid(skb);
     int match = bpf_owner_match(skb, sock_uid);
     if ((direction == BPF_EGRESS) && (match == BPF_DROP)) {
         // If an outbound packet is going to be dropped, we do not count that
@@ -261,8 +249,8 @@
         return match;
     }
 
-    uint64_t cookie = get_socket_cookie(skb);
-    struct uid_tag* utag = find_map_entry(&cookie_tag_map, &cookie);
+    uint64_t cookie = bpf_get_socket_cookie(skb);
+    struct uid_tag* utag = bpf_map_lookup_elem(&cookie_tag_map, &cookie);
     uint32_t uid, tag;
     if (utag) {
         uid = utag->uid;
@@ -274,15 +262,21 @@
 
     struct stats_key key = {.uid = uid, .tag = tag, .counterSet = 0, .ifaceIndex = skb->ifindex};
 
-    uint8_t* counterSet = find_map_entry(&uid_counterset_map, &uid);
+    uint8_t* counterSet = bpf_map_lookup_elem(&uid_counterset_map, &uid);
     if (counterSet) key.counterSet = (uint32_t)*counterSet;
 
+    uint32_t mapSettingKey = CURRENT_STATS_MAP_CONFIGURATION_KEY;
+    uint8_t* selectedMap = bpf_map_lookup_elem(&configuration_map, &mapSettingKey);
+    if (!selectedMap) {
+        return match;
+    }
+
     if (tag) {
-        bpf_update_stats(skb, &tag_stats_map, direction, &key);
+        update_stats_with_config(skb, direction, &key, *selectedMap);
     }
 
     key.tag = 0;
-    bpf_update_stats(skb, &uid_stats_map, direction, &key);
+    update_stats_with_config(skb, direction, &key, *selectedMap);
     bpf_update_stats(skb, &app_uid_stats_map, direction, &uid);
     return match;
 }