refactor bpf kernel program

Refactor the program into several inline function and resolve the
duplication of ingress program and egress program. The Program compiled
from the new inline function should be no difference compare to the old
ones.

Bug: 30950746
Test: bpf program can be loaded.
      atest android.net.cts.TrafficStatsTest
      atest android.app.usage.cts.NetworkUsageStatsTest
Change-Id: I8091e8bdf7f3fe118c26907028e3f845eb41ab65
diff --git a/bpfloader/bpf_egress.c b/bpfloader/bpf_egress.c
index 5f577d0..9c2cc5f 100644
--- a/bpfloader/bpf_egress.c
+++ b/bpfloader/bpf_egress.c
@@ -23,49 +23,5 @@
 
 ELF_SEC(BPF_PROG_SEC_NAME)
 int bpf_cgroup_egress(struct __sk_buff* skb) {
-    uint64_t cookie = get_socket_cookie(skb);
-    struct uid_tag* utag = find_map_entry(COOKIE_TAG_MAP, &cookie);
-    uint32_t uid, tag;
-    if (utag) {
-        uid = utag->uid;
-        tag = utag->tag;
-    } else {
-        uid = get_socket_uid(skb);
-        tag = 0;
-    }
-
-    struct stats_key key = {.uid = uid, .tag = tag, .counterSet = 0, .ifaceIndex = skb->ifindex};
-
-    uint32_t* counterSet;
-    counterSet = find_map_entry(UID_COUNTERSET_MAP, &uid);
-    if (counterSet) key.counterSet = *counterSet;
-
-    int ret;
-    if (tag) {
-        struct stats_value* tagValue;
-        tagValue = find_map_entry(TAG_STATS_MAP, &key);
-        if (!tagValue) {
-            struct stats_value newValue = {};
-            write_to_map_entry(TAG_STATS_MAP, &key, &newValue, BPF_NOEXIST);
-            tagValue = find_map_entry(TAG_STATS_MAP, &key);
-        }
-        if (tagValue) {
-            __sync_fetch_and_add(&tagValue->txPackets, 1);
-            __sync_fetch_and_add(&tagValue->txBytes, skb->len);
-        }
-    }
-
-    key.tag = 0;
-    struct stats_value* value;
-    value = find_map_entry(UID_STATS_MAP, &key);
-    if (!value) {
-        struct stats_value newValue = {};
-        write_to_map_entry(UID_STATS_MAP, &key, &newValue, BPF_NOEXIST);
-        value = find_map_entry(UID_STATS_MAP, &key);
-    }
-    if (value) {
-        __sync_fetch_and_add(&value->txPackets, 1);
-        __sync_fetch_and_add(&value->txBytes, skb->len);
-    }
-    return BPF_PASS;
+    return bpf_traffic_account(skb, BPF_EGRESS);
 }
diff --git a/bpfloader/bpf_ingress.c b/bpfloader/bpf_ingress.c
index 0a5c5b8..d25877e 100644
--- a/bpfloader/bpf_ingress.c
+++ b/bpfloader/bpf_ingress.c
@@ -23,49 +23,5 @@
 
 ELF_SEC(BPF_PROG_SEC_NAME)
 int bpf_cgroup_ingress(struct __sk_buff* skb) {
-    uint64_t cookie = get_socket_cookie(skb);
-    struct uid_tag* utag = find_map_entry(COOKIE_TAG_MAP, &cookie);
-    uint32_t uid, tag;
-    if (utag) {
-        uid = utag->uid;
-        tag = utag->tag;
-    } else {
-        uid = get_socket_uid(skb);
-        tag = 0;
-    }
-
-    struct stats_key key = {.uid = uid, .tag = tag, .counterSet = 0, .ifaceIndex = skb->ifindex};
-
-    uint32_t* counterSet;
-    counterSet = find_map_entry(UID_COUNTERSET_MAP, &uid);
-    if (counterSet) key.counterSet = *counterSet;
-
-    int ret;
-    if (tag) {
-        struct stats_value* tagValue;
-        tagValue = find_map_entry(TAG_STATS_MAP, &key);
-        if (!tagValue) {
-            struct stats_value newValue = {};
-            write_to_map_entry(TAG_STATS_MAP, &key, &newValue, BPF_NOEXIST);
-            tagValue = find_map_entry(TAG_STATS_MAP, &key);
-        }
-        if (tagValue) {
-            __sync_fetch_and_add(&tagValue->rxPackets, 1);
-            __sync_fetch_and_add(&tagValue->rxBytes, skb->len);
-        }
-    }
-
-    key.tag = 0;
-    struct stats_value* value;
-    value = find_map_entry(UID_STATS_MAP, &key);
-    if (!value) {
-        struct stats_value newValue = {};
-        write_to_map_entry(UID_STATS_MAP, &key, &newValue, BPF_NOEXIST);
-        value = find_map_entry(UID_STATS_MAP, &key);
-    }
-    if (value) {
-        __sync_fetch_and_add(&value->rxPackets, 1);
-        __sync_fetch_and_add(&value->rxBytes, skb->len);
-    }
-    return BPF_PASS;
+    return bpf_traffic_account(skb, BPF_INGRESS);
 }
diff --git a/bpfloader/bpf_kern.h b/bpfloader/bpf_kern.h
index 2a2ef8c..806b45a 100644
--- a/bpfloader/bpf_kern.h
+++ b/bpfloader/bpf_kern.h
@@ -75,3 +75,51 @@
     }
     return BPF_PASS;
 }
+
+static __always_inline inline void bpf_update_stats(struct __sk_buff* skb, uint64_t map,
+                                                    int direction, struct stats_key key) {
+    struct stats_value* value;
+    value = find_map_entry(map, &key);
+    if (!value) {
+        struct stats_value newValue = {};
+        write_to_map_entry(map, &key, &newValue, BPF_NOEXIST);
+        value = find_map_entry(map, &key);
+    }
+    if (value) {
+      if (direction == BPF_INGRESS) {
+        __sync_fetch_and_add(&value->rxPackets, 1);
+        __sync_fetch_and_add(&value->rxBytes, skb->len);
+      } else {
+        __sync_fetch_and_add(&value->txPackets, 1);
+        __sync_fetch_and_add(&value->txBytes, skb->len);
+      }
+    }
+}
+
+static __always_inline inline int bpf_traffic_account(struct __sk_buff* skb, int direction) {
+    uint64_t cookie = get_socket_cookie(skb);
+    struct uid_tag* utag = find_map_entry(COOKIE_TAG_MAP, &cookie);
+    uint32_t uid, tag;
+    if (utag) {
+        uid = utag->uid;
+        tag = utag->tag;
+    } else {
+        uid = get_socket_uid(skb);
+        tag = 0;
+    }
+
+    struct stats_key key = {.uid = uid, .tag = tag, .counterSet = 0, .ifaceIndex = skb->ifindex};
+
+    uint32_t* counterSet;
+    counterSet = find_map_entry(UID_COUNTERSET_MAP, &uid);
+    if (counterSet) key.counterSet = *counterSet;
+
+    int ret;
+    if (tag) {
+        bpf_update_stats(skb, TAG_STATS_MAP, direction, key);
+    }
+
+    key.tag = 0;
+    bpf_update_stats(skb, UID_STATS_MAP, direction, key);
+    return BPF_PASS;
+}
diff --git a/bpfloader/cgroup_bpf_egress_prog.o b/bpfloader/cgroup_bpf_egress_prog.o
index b8a09ee..ab29366 100644
--- a/bpfloader/cgroup_bpf_egress_prog.o
+++ b/bpfloader/cgroup_bpf_egress_prog.o
Binary files differ
diff --git a/bpfloader/cgroup_bpf_ingress_prog.o b/bpfloader/cgroup_bpf_ingress_prog.o
index 4d15cf1..da7ccca 100644
--- a/bpfloader/cgroup_bpf_ingress_prog.o
+++ b/bpfloader/cgroup_bpf_ingress_prog.o
Binary files differ