xref: /aosp_15_r20/external/bcc/tests/cc/test_map_in_map.cc (revision 387f9dfdfa2baef462e92476d413c7bc2470293e)
1 /*
2  * Copyright (c) 2019 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 <string>
20 
21 #include "BPF.h"
22 #include "catch.hpp"
23 
24 #if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 11, 0)
25 
26 TEST_CASE("test hash of maps", "[hash_of_maps]") {
27   {
28     const std::string BPF_PROGRAM = R"(
29       BPF_ARRAY(cntl, int, 1);
30       BPF_ARRAY(ex1, int, 1024);
31       BPF_ARRAY(ex2, int, 1024);
32       BPF_ARRAY(ex3, u64, 1024);
33       BPF_HASH_OF_MAPS(maps_hash, int, "ex1", 10);
34 
35       int syscall__getuid(void *ctx) {
36          int key = 0, data, *val, cntl_val;
37          void *inner_map;
38 
39          val = cntl.lookup(&key);
40          if (!val || *val == 0)
41            return 0;
42 
43          // cntl_val == 1 : lookup and update
44          cntl_val = *val;
45          inner_map = maps_hash.lookup(&key);
46          if (!inner_map)
47            return 0;
48 
49          if (cntl_val == 1) {
50            val = bpf_map_lookup_elem(inner_map, &key);
51            if (val) {
52              data = 1;
53              bpf_map_update_elem(inner_map, &key, &data, 0);
54            }
55          }
56 
57          return 0;
58       }
59     )";
60 
61     ebpf::BPF bpf;
62     ebpf::StatusTuple res(0);
63     res = bpf.init(BPF_PROGRAM);
64     REQUIRE(res.ok());
65 
66     auto t = bpf.get_map_in_map_table<int>("maps_hash");
67     auto ex1_table = bpf.get_array_table<int>("ex1");
68     auto ex2_table = bpf.get_array_table<int>("ex2");
69     auto ex3_table = bpf.get_array_table<unsigned long long>("ex3");
70     int ex1_fd = ex1_table.get_fd();
71     int ex2_fd = ex2_table.get_fd();
72     int ex3_fd = ex3_table.get_fd();
73 
74     int key = 0, value = 0;
75     res = t.update_value(key, ex1_fd);
76     REQUIRE(res.ok());
77 
78     // updating already-occupied slot will succeed.
79     res = t.update_value(key, ex2_fd);
80     REQUIRE(res.ok());
81     res = t.update_value(key, ex1_fd);
82     REQUIRE(res.ok());
83 
84     // an in-compatible map
85     key = 1;
86     res = t.update_value(key, ex3_fd);
87     REQUIRE(res.code() == -1);
88 
89     // hash table, any valid key should work as long
90     // as hash table is not full.
91     key = 10;
92     res = t.update_value(key, ex2_fd);
93     REQUIRE(res.ok());
94     res = t.remove_value(key);
95     REQUIRE(res.ok());
96 
97     // test effectiveness of map-in-map
98     key = 0;
99     std::string getuid_fnname = bpf.get_syscall_fnname("getuid");
100     res = bpf.attach_kprobe(getuid_fnname, "syscall__getuid");
101     REQUIRE(res.ok());
102 
103     auto cntl_table = bpf.get_array_table<int>("cntl");
104     cntl_table.update_value(0, 1);
105     REQUIRE(getuid() >= 0);
106     res = ex1_table.get_value(key, value);
107     REQUIRE(res.ok());
108     REQUIRE(value > 0);
109 
110     res = bpf.detach_kprobe(getuid_fnname);
111     REQUIRE(res.ok());
112 
113     res = t.remove_value(key);
114     REQUIRE(res.ok());
115   }
116 }
117 
118 TEST_CASE("test hash of maps using custom key", "[hash_of_maps_custom_key]") {
119   {
120     const std::string BPF_PROGRAM = R"(
121         struct custom_key {
122           int value_1;
123           int value_2;
124         };
125 
126         BPF_ARRAY(cntl, int, 1);
127         BPF_TABLE("hash", int, int, ex1, 1024);
128         BPF_TABLE("hash", int, int, ex2, 1024);
129         BPF_HASH_OF_MAPS(maps_hash, struct custom_key, "ex1", 10);
130 
131         int syscall__getuid(void *ctx) {
132           struct custom_key hash_key = {1, 0};
133           int key = 0, data, *val, cntl_val;
134           void *inner_map;
135 
136           val = cntl.lookup(&key);
137           if (!val || *val == 0)
138             return 0;
139 
140           hash_key.value_2 = *val;
141           inner_map = maps_hash.lookup(&hash_key);
142           if (!inner_map)
143             return 0;
144 
145           val = bpf_map_lookup_elem(inner_map, &key);
146           if (!val) {
147             data = 1;
148             bpf_map_update_elem(inner_map, &key, &data, 0);
149           } else {
150             data = 1 + *val;
151             bpf_map_update_elem(inner_map, &key, &data, 0);
152           }
153 
154           return 0;
155         }
156     )";
157 
158     struct custom_key {
159       int value_1;
160       int value_2;
161     };
162 
163     ebpf::BPF bpf;
164     ebpf::StatusTuple res(0);
165     res = bpf.init(BPF_PROGRAM);
166     REQUIRE(res.ok());
167 
168     auto t = bpf.get_map_in_map_table<struct custom_key>("maps_hash");
169     auto ex1_table = bpf.get_hash_table<int, int>("ex1");
170     auto ex2_table = bpf.get_hash_table<int, int>("ex2");
171     auto cntl_table = bpf.get_array_table<int>("cntl");
172     int ex1_fd = ex1_table.get_fd();
173     int ex2_fd = ex2_table.get_fd();
174 
175     // test effectiveness of map-in-map
176     std::string getuid_fnname = bpf.get_syscall_fnname("getuid");
177     res = bpf.attach_kprobe(getuid_fnname, "syscall__getuid");
178     REQUIRE(res.ok());
179 
180     struct custom_key hash_key = {1, 1};
181 
182     res = t.update_value(hash_key, ex1_fd);
183     REQUIRE(res.ok());
184 
185     struct custom_key hash_key2 = {1, 2};
186     res = t.update_value(hash_key2, ex2_fd);
187     REQUIRE(res.ok());
188 
189     int key = 0, value = 0, value2 = 0;
190 
191     // Can't get value when value didn't set.
192     res = ex1_table.get_value(key, value);
193     REQUIRE(!res.ok());
194     REQUIRE(value == 0);
195 
196     // Call syscall__getuid, then set value to ex1_table
197     res = cntl_table.update_value(key, 1);
198     REQUIRE(res.ok());
199     REQUIRE(getuid() >= 0);
200 
201     // Now we can get value from ex1_table
202     res = ex1_table.get_value(key, value);
203     REQUIRE(res.ok());
204     REQUIRE(value >= 1);
205 
206     // Can't get value when value didn't set.
207     res = ex2_table.get_value(key, value2);
208     REQUIRE(!res.ok());
209     REQUIRE(value2 == 0);
210 
211     // Call syscall__getuid, then set value to ex2_table
212     res = cntl_table.update_value(key, 2);
213     REQUIRE(res.ok());
214     REQUIRE(getuid() >= 0);
215 
216     // Now we can get value from ex2_table
217     res = ex2_table.get_value(key, value2);
218     REQUIRE(res.ok());
219     REQUIRE(value > 0);
220 
221     res = bpf.detach_kprobe(getuid_fnname);
222     REQUIRE(res.ok());
223 
224     res = t.remove_value(hash_key);
225     REQUIRE(res.ok());
226     res = t.remove_value(hash_key2);
227     REQUIRE(res.ok());
228   }
229 }
230 
231 TEST_CASE("test array of maps", "[array_of_maps]") {
232   {
233     const std::string BPF_PROGRAM = R"(
234       BPF_ARRAY(cntl, int, 1);
235       BPF_TABLE("hash", int, int, ex1, 1024);
236       BPF_TABLE("hash", int, int, ex2, 1024);
237       BPF_TABLE("hash", u64, u64, ex3, 1024);
238       BPF_ARRAY_OF_MAPS(maps_array, "ex1", 10);
239 
240       int syscall__getuid(void *ctx) {
241          int key = 0, data, *val, cntl_val;
242          void *inner_map;
243 
244          val = cntl.lookup(&key);
245          if (!val || *val == 0)
246            return 0;
247 
248          // cntl_val == 1 : lookup and update
249          // cntl_val == 2 : delete
250          cntl_val = *val;
251          inner_map = maps_array.lookup(&key);
252          if (!inner_map)
253            return 0;
254 
255          if (cntl_val == 1) {
256            val = bpf_map_lookup_elem(inner_map, &key);
257            if (!val) {
258              data = 1;
259              bpf_map_update_elem(inner_map, &key, &data, 0);
260            }
261          } else if (cntl_val == 2) {
262            bpf_map_delete_elem(inner_map, &key);
263          }
264 
265          return 0;
266       }
267     )";
268 
269     ebpf::BPF bpf;
270     ebpf::StatusTuple res(0);
271     res = bpf.init(BPF_PROGRAM);
272     REQUIRE(res.ok());
273 
274     auto t = bpf.get_map_in_map_table<int>("maps_array");
275     auto ex1_table = bpf.get_hash_table<int, int>("ex1");
276     auto ex2_table = bpf.get_hash_table<int, int>("ex2");
277     auto ex3_table =
278         bpf.get_hash_table<unsigned long long, unsigned long long>("ex3");
279     int ex1_fd = ex1_table.get_fd();
280     int ex2_fd = ex2_table.get_fd();
281     int ex3_fd = ex3_table.get_fd();
282 
283     int key = 0, value = 0;
284     res = t.update_value(key, ex1_fd);
285     REQUIRE(res.ok());
286 
287     // updating already-occupied slot will succeed.
288     res = t.update_value(key, ex2_fd);
289     REQUIRE(res.ok());
290     res = t.update_value(key, ex1_fd);
291     REQUIRE(res.ok());
292 
293     // an in-compatible map
294     key = 1;
295     res = t.update_value(key, ex3_fd);
296     REQUIRE(res.code() == -1);
297 
298     // array table, out of bound access
299     key = 10;
300     res = t.update_value(key, ex2_fd);
301     REQUIRE(res.code() == -1);
302 
303     // test effectiveness of map-in-map
304     key = 0;
305     std::string getuid_fnname = bpf.get_syscall_fnname("getuid");
306     res = bpf.attach_kprobe(getuid_fnname, "syscall__getuid");
307     REQUIRE(res.ok());
308 
309     auto cntl_table = bpf.get_array_table<int>("cntl");
310     cntl_table.update_value(0, 1);
311 
312     REQUIRE(getuid() >= 0);
313     res = ex1_table.get_value(key, value);
314     REQUIRE(res.ok());
315     REQUIRE(value == 1);
316 
317     cntl_table.update_value(0, 2);
318     REQUIRE(getuid() >= 0);
319     res = ex1_table.get_value(key, value);
320     REQUIRE(res.code() == -1);
321 
322     res = bpf.detach_kprobe(getuid_fnname);
323     REQUIRE(res.ok());
324 
325     res = t.remove_value(key);
326     REQUIRE(res.ok());
327   }
328 }
329 #endif
330