support BPF_MAP_TYPE_SK_STORAGE map

kernel sk local storage is introduced at 5.2. torvalds/linux@6ac99e8 this patch supports BPF_MAP_TYPE_SK_STORAGE in bcc. tested C++ user space APIs and loading of the bpf programs using helpers map.sk_storage_get() and map.sk_storage_delete(). Signed-off-by: Yonghong Song <yhs@fb.com>
iovisor · Jan 3, 2020 · c6b287c · c6b287c
1 parent db82644
commit c6b287c
Show file tree

Hide file tree

Showing 6 changed files with 151 additions and 0 deletions.
diff --git a/src/cc/api/BPF.h b/src/cc/api/BPF.h
@@ -146,6 +146,14 @@ class BPF {
     return BPFPercpuHashTable<KeyType, ValueType>({});
   }
 
+  template <class ValueType>
+  BPFSkStorageTable<ValueType> get_sk_storage_table(const std::string& name) {
+    TableStorage::iterator it;
+    if (bpf_module_->table_storage().Find(Path({bpf_module_->id(), name}), it))
+      return BPFSkStorageTable<ValueType>(it->second);
+    return BPFSkStorageTable<ValueType>({});
+  }
+
   void* get_bsymcache(void) {
     if (bsymcache_ == NULL) {
       bsymcache_ = bcc_buildsymcache_new();

diff --git a/src/cc/api/BPFTable.h b/src/cc/api/BPFTable.h
@@ -409,4 +409,33 @@ class BPFMapInMapTable : public BPFTableBase<int, int> {
   StatusTuple remove_value(const int& index);
 };
 
+template <class ValueType>
+class BPFSkStorageTable : public BPFTableBase<int, ValueType> {
+ public:
+  BPFSkStorageTable(const TableDesc& desc) : BPFTableBase<int, ValueType>(desc) {
+    if (desc.type != BPF_MAP_TYPE_SK_STORAGE)
+      throw std::invalid_argument("Table '" + desc.name +
+                                  "' is not a sk_storage table");
+  }
+
+  virtual StatusTuple get_value(const int& sock_fd, ValueType& value) {
+    if (!this->lookup(const_cast<int*>(&sock_fd), get_value_addr(value)))
+      return StatusTuple(-1, "Error getting value: %s", std::strerror(errno));
+    return StatusTuple(0);
+  }
+
+  virtual StatusTuple update_value(const int& sock_fd, const ValueType& value) {
+    if (!this->update(const_cast<int*>(&sock_fd),
+                      get_value_addr(const_cast<ValueType&>(value))))
+      return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
+    return StatusTuple(0);
+  }
+
+  virtual StatusTuple remove_value(const int& sock_fd) {
+    if (!this->remove(const_cast<int*>(&sock_fd)))
+      return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
+    return StatusTuple(0);
+  }
+};
+
 }  // namespace ebpf
diff --git a/src/cc/export/helpers.h b/src/cc/export/helpers.h
@@ -275,6 +275,18 @@ struct _name##_table_t _name = { .max_entries = (_max_entries) }
 #define BPF_HASH_OF_MAPS(_name, _inner_map_name, _max_entries) \
   BPF_TABLE("hash_of_maps$" _inner_map_name, int, int, _name, _max_entries)
 
+#define BPF_SK_STORAGE(_name, _leaf_type) \
+struct _name##_table_t { \
+  int key; \
+  _leaf_type leaf; \
+  void * (*sk_storage_get) (void *, void *, int); \
+  int (*sk_storage_delete) (void *); \
+  u32 flags; \
+}; \
+__attribute__((section("maps/sk_storage"))) \
+struct _name##_table_t _name = { .flags = BPF_F_NO_PREALLOC }; \
+BPF_ANNOTATE_KV_PAIR(_name, int, _leaf_type)
+
 // packet parsing state machine helpers
 #define cursor_advance(_cursor, _len) \
   ({ void *_tmp = _cursor; _cursor += _len; _tmp; })

diff --git a/src/cc/frontends/clang/b_frontend_action.cc b/src/cc/frontends/clang/b_frontend_action.cc
@@ -887,6 +887,12 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
           } else if (memb_name == "redirect_map") {
             prefix = "bpf_redirect_map";
             suffix = ")";
+          } else if (memb_name == "sk_storage_get") {
+            prefix = "bpf_sk_storage_get";
+            suffix = ")";
+          } else if (memb_name == "sk_storage_delete") {
+            prefix = "bpf_sk_storage_delete";
+            suffix = ")";
           } else {
             error(GET_BEGINLOC(Call), "invalid bpf_table operation %0") << memb_name;
             return false;
@@ -1248,6 +1254,8 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       map_type = BPF_MAP_TYPE_HASH_OF_MAPS;
     } else if (section_attr == "maps/array_of_maps") {
       map_type = BPF_MAP_TYPE_ARRAY_OF_MAPS;
+    } else if (section_attr == "maps/sk_storage") {
+      map_type = BPF_MAP_TYPE_SK_STORAGE;
     } else if (section_attr == "maps/extern") {
       if (!fe_.table_storage().Find(maps_ns_path, table_it)) {
         if (!fe_.table_storage().Find(global_path, table_it)) {

diff --git a/tests/cc/CMakeLists.txt b/tests/cc/CMakeLists.txt
@@ -26,6 +26,7 @@ set(TEST_LIBBCC_SOURCES
 	test_pinned_table.cc
 	test_prog_table.cc
 	test_shared_table.cc
+	test_sk_storage.cc
 	test_usdt_args.cc
 	test_usdt_probes.cc
 	utils.cc)

diff --git a/tests/cc/test_sk_storage.cc b/tests/cc/test_sk_storage.cc
@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2020 Facebook, Inc.
+ *
+ * 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 <linux/version.h>
+#include <unistd.h>
+#include <sys/types.h>
+#include <sys/socket.h>
+#include <string>
+
+#include "BPF.h"
+#include "catch.hpp"
+
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 2, 0)
+
+TEST_CASE("test sk_storage map", "[sk_storage]") {
+  {
+    const std::string BPF_PROGRAM = R"(
+BPF_SK_STORAGE(sk_pkt_cnt, __u64);
+
+int test(struct __sk_buff *skb) {
+  __u64 cnt = 0, *cnt_out;
+  struct bpf_sock *sk;
+
+  sk = skb->sk;
+  if (!sk)
+    return 1;
+
+  sk = bpf_sk_fullsock(sk);
+  if (!sk)
+    return 1;
+
+  cnt_out = sk_pkt_cnt.sk_storage_get(sk, &cnt, BPF_SK_STORAGE_GET_F_CREATE);
+  if (!cnt_out)
+    return 1;
+
+  (*cnt_out)++;
+  return 1;
+}
+    )";
+
+    // make sure program is loaded successfully
+    ebpf::BPF bpf;
+    ebpf::StatusTuple res(0);
+    res = bpf.init(BPF_PROGRAM);
+    REQUIRE(res.code() == 0);
+    int prog_fd;
+    res = bpf.load_func("test", BPF_PROG_TYPE_CGROUP_SKB, prog_fd);
+    REQUIRE(res.code() == 0);
+
+    // create a udp socket so we can do some map operations.
+    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
+    REQUIRE(sockfd >= 0);
+
+    auto sk_table = bpf.get_sk_storage_table<unsigned long long>("sk_pkt_cnt");
+    unsigned long long v = 0, v1 = 10;
+
+    // no sk_storage for the table yet.
+    res = sk_table.get_value(sockfd, v);
+    REQUIRE(res.code() != 0);
+
+    // nothing to remove yet.
+    res = sk_table.remove_value(sockfd);
+    REQUIRE(res.code() != 0);
+
+    // update the table with a certain value.
+    res = sk_table.update_value(sockfd, v1);
+    REQUIRE(res.code() == 0);
+
+    // get_value should be successful now.
+    res = sk_table.get_value(sockfd, v);
+    REQUIRE(res.code() == 0);
+    REQUIRE(v == 10);
+
+    // remove the sk_storage.
+    res = sk_table.remove_value(sockfd);
+    REQUIRE(res.code() == 0);
+  }
+}
+
+#endif