libbcc: support BPF_SOCKHASH specify the key type (#3473)
[platform/upstream/bcc.git] / tests / cc / test_sock_table.cc
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(4, 18, 0)
27
28 TEST_CASE("test sock map", "[sockmap]") {
29   {
30     const std::string BPF_PROGRAM = R"(
31 BPF_SOCKMAP(sk_map1, 10);
32 BPF_SOCKMAP(sk_map2, 10);
33 int test(struct bpf_sock_ops *skops)
34 {
35   u32 key = 0, val = 0;
36
37   sk_map2.update(&key, &val);
38   sk_map2.delete(&key);
39   sk_map2.sock_map_update(skops, &key, 0);
40
41   return 0;
42 }
43     )";
44
45     // make sure program is loaded successfully
46     ebpf::BPF bpf;
47     ebpf::StatusTuple res(0);
48     res = bpf.init(BPF_PROGRAM);
49     REQUIRE(res.code() == 0);
50
51     // create a udp socket so we can do some map operations.
52     int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
53     REQUIRE(sockfd >= 0);
54
55     auto sk_map = bpf.get_sockmap_table("sk_map1");
56     int key = 0, val = sockfd;
57
58     res = sk_map.remove_value(key);
59     REQUIRE(res.code() != 0);
60
61     // the socket must be TCP established socket.
62     res = sk_map.update_value(key, val);
63     REQUIRE(res.code() != 0);
64   }
65 }
66
67 TEST_CASE("test sock hash", "[sockhash]") {
68   {
69     const std::string BPF_PROGRAM = R"(
70 BPF_SOCKHASH(sk_hash1, u32, 10);
71 BPF_SOCKHASH(sk_hash2, u32, 10);
72 int test(struct bpf_sock_ops *skops)
73 {
74   u32 key = 0, val = 0;
75
76   sk_hash2.update(&key, &val);
77   sk_hash2.delete(&key);
78   sk_hash2.sock_hash_update(skops, &key, 0);
79
80   return 0;
81 }
82     )";
83
84     // make sure program is loaded successfully
85     ebpf::BPF bpf;
86     ebpf::StatusTuple res(0);
87     res = bpf.init(BPF_PROGRAM);
88     REQUIRE(res.code() == 0);
89
90     // create a udp socket so we can do some map operations.
91     int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
92     REQUIRE(sockfd >= 0);
93
94     auto sk_hash = bpf.get_sockhash_table("sk_hash1");
95     int key = 0, val = sockfd;
96
97     res = sk_hash.remove_value(key);
98     REQUIRE(res.code() != 0);
99
100     // the socket must be TCP established socket.
101     res = sk_hash.update_value(key, val);
102     REQUIRE(res.code() != 0);
103   }
104 }
105
106 #endif