xref: /aosp_15_r20/external/bcc/tests/cc/test_sk_storage.cc (revision 387f9dfdfa2baef462e92476d413c7bc2470293e)
1 /*
2  * Copyright (c) 2020 Facebook, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include <linux/version.h>
18 #include <unistd.h>
19 #include <sys/types.h>
20 #include <sys/socket.h>
21 #include <string>
22 
23 #include "BPF.h"
24 #include "catch.hpp"
25 
26 #if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 2, 0)
27 
28 TEST_CASE("test sk_storage map", "[sk_storage][!mayfail]") {
29   {
30     const std::string BPF_PROGRAM = R"(
31 BPF_SK_STORAGE(sk_pkt_cnt, __u64);
32 
33 int test(struct __sk_buff *skb) {
34   __u64 cnt = 0, *cnt_out;
35   struct bpf_sock *sk;
36 
37   sk = skb->sk;
38   if (!sk)
39     return 1;
40 
41   sk = bpf_sk_fullsock(sk);
42   if (!sk)
43     return 1;
44 
45   cnt_out = sk_pkt_cnt.sk_storage_get(sk, &cnt, BPF_SK_STORAGE_GET_F_CREATE);
46   if (!cnt_out)
47     return 1;
48 
49   (*cnt_out)++;
50   return 1;
51 }
52     )";
53 
54     // make sure program is loaded successfully
55     ebpf::BPF bpf;
56     ebpf::StatusTuple res(0);
57     res = bpf.init(BPF_PROGRAM);
58     REQUIRE(res.ok());
59     int prog_fd;
60     res = bpf.load_func("test", BPF_PROG_TYPE_CGROUP_SKB, prog_fd);
61     REQUIRE(res.ok());
62 
63     // create a udp socket so we can do some map operations.
64     int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
65     REQUIRE(sockfd >= 0);
66 
67     auto sk_table = bpf.get_sk_storage_table<unsigned long long>("sk_pkt_cnt");
68     unsigned long long v = 0, v1 = 10;
69 
70     // no sk_storage for the table yet.
71     res = sk_table.get_value(sockfd, v);
72     REQUIRE(!res.ok());
73 
74     // nothing to remove yet.
75     res = sk_table.remove_value(sockfd);
76     REQUIRE(!res.ok());
77 
78     // update the table with a certain value.
79     res = sk_table.update_value(sockfd, v1);
80     REQUIRE(res.ok());
81 
82     // get_value should be successful now.
83     res = sk_table.get_value(sockfd, v);
84     REQUIRE(res.ok());
85     REQUIRE(v == 10);
86 
87     // remove the sk_storage.
88     res = sk_table.remove_value(sockfd);
89     REQUIRE(res.ok());
90   }
91 }
92 
93 #endif
94