From d370f95f10fbd04bade48062198cbb2b98b91e4b Mon Sep 17 00:00:00 2001
From: Motomu Utsumi <motomuman@google.com>
Date: Tue, 25 Mar 2025 09:31:16 +0900
Subject: [PATCH] Add macro to declare a sk storage and related accessor
 functions

Test: boot
Bug: 406080305
Change-Id: I03238c72e8024d4424c06803784988570513ff39
---
 bpf/headers/include/bpf_helpers.h | 33 +++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/bpf/headers/include/bpf_helpers.h b/bpf/headers/include/bpf_helpers.h
index 199661de65..e0c524fdff 100644
--- a/bpf/headers/include/bpf_helpers.h
+++ b/bpf/headers/include/bpf_helpers.h
@@ -236,6 +236,11 @@ static void* (*bpf_ringbuf_reserve_unsafe)(const struct bpf_map_def* ringbuf,
         BPF_FUNC_ringbuf_reserve;
 static void (*bpf_ringbuf_submit_unsafe)(const void* data, __u64 flags) = (void*)
         BPF_FUNC_ringbuf_submit;
+static void* (*bpf_sk_storage_get_unsafe) (const struct bpf_map_def* sk_storage, const void* sk,
+                                           const void* value, unsigned long long flags) = (void*)
+        BPF_FUNC_sk_storage_get;
+static int (*bpf_sk_storage_delete_unsafe) (const struct bpf_map_def* sk_storage,
+                                            const void* sk) = (void*) BPF_FUNC_sk_storage_delete;
 
 #define BPF_ANNOTATE_KV_PAIR(name, type_key, type_val)  \
         struct ____btf_map_##name {                     \
@@ -330,6 +335,34 @@ static void (*bpf_ringbuf_submit_unsafe)(const void* data, __u64 flags) = (void*
                            PRIVATE, BPFLOADER_MIN_VER, BPFLOADER_MAX_VER,               \
                            LOAD_ON_ENG, LOAD_ON_USER, LOAD_ON_USERDEBUG)
 
+// Type safe macro to declare a sk storage and related accessor functions.
+// BPF_MAP_TYPE_SK_STORAGE was introduced in kernel 5.2 but this map requires BTF and
+// BTF is enabled on kernel 5.10 or higher.
+#define DEFINE_BPF_SK_STORAGE_EXT(the_map, ValueType, usr, grp, md, selinux, pindir,    \
+                                  share, min_loader, max_loader, ignore_eng,            \
+                                  ignore_user, ignore_userdebug, mapFlags)              \
+    DEFINE_BPF_MAP_BASE(the_map, SK_STORAGE, sizeof(uint32_t), sizeof(ValueType),       \
+                        0, usr, grp, md, selinux, pindir, share,                        \
+                        KVER_5_10, KVER_INF, min_loader, max_loader,                    \
+                        ignore_eng, ignore_user, ignore_userdebug, mapFlags);           \
+    BPF_ANNOTATE_KV_PAIR(the_map, uint32_t, ValueType);                                 \
+                                                                                        \
+    static inline __always_inline __unused ValueType* bpf_##the_map##_get(              \
+            const struct bpf_sock* sk, const ValueType* v, unsigned long long flags) {  \
+        return bpf_sk_storage_get_unsafe(&the_map, sk, v, flags);                       \
+    };                                                                                  \
+                                                                                        \
+    static inline __always_inline __unused int bpf_##the_map##_delete(                  \
+            const struct bpf_sock* sk) {                                                \
+        return bpf_sk_storage_delete_unsafe(&the_map, sk);                              \
+    };
+
+#define DEFINE_BPF_SK_STORAGE(the_map, TypeOfValue)                                      \
+    DEFINE_BPF_SK_STORAGE_EXT(the_map, TypeOfValue,                                      \
+                              AID_ROOT, AID_NET_BW_ACCT, 0060, "fs_bpf_net_shared", "",  \
+                              PRIVATE, BPFLOADER_MIN_VER, BPFLOADER_MAX_VER,             \
+                              LOAD_ON_ENG, LOAD_ON_USER, LOAD_ON_USERDEBUG, 0)
+
 /* There exist buggy kernels with pre-T OS, that due to
  * kernel patch "[ALPS05162612] bpf: fix ubsan error"
  * do not support userspace writes into non-zero index of bpf map arrays.