xref: /aosp_15_r20/external/bcc/examples/cpp/SkLocalStorageIterator.cc (revision 387f9dfdfa2baef462e92476d413c7bc2470293e)
1 /*
2  * Copyright (c) Facebook, Inc.
3  * Licensed under the Apache License, Version 2.0 (the "License")
4  *
5  * Usage:
6  *   ./SkLocalStorageIterator
7  *
8  * BPF socket local storage map iterator supported is added in 5.9.
9  * But since it takes locks during iterating, it may have performance
10  * implication if in parallel some other bpf program or user space
11  * is doing map update/delete for sockets in the same bucket. The issue
12  * is fixed in 5.10 with the following patch which uses rcu lock instead:
13  *   https://lore.kernel.org/bpf/[email protected]
14  *
15  * This example shows how to dump local storage data from all sockets
16  * associated with one socket local storage map.
17  * An example output likes below:
18  *   family  prot    val
19  *   2       17      20
20  *   2       17      10
21  */
22 
23 #include <unistd.h>
24 #include <fstream>
25 #include <iostream>
26 #include <string>
27 #include <net/if.h>
28 
29 #include "bcc_version.h"
30 #include "BPF.h"
31 
32 const std::string BPF_PROGRAM = R"(
33 
34 #include <linux/bpf.h>
35 #include <linux/seq_file.h>
36 #include <net/sock.h>
37 
38 /* the structure is defined in .c file, so explicitly define
39  * the structure here.
40  */
41 struct bpf_iter__bpf_sk_storage_map {
42   union {
43     struct bpf_iter_meta *meta;
44   };
45   union {
46     struct bpf_map *map;
47   };
48   union {
49     struct sock *sk;
50   };
51   union {
52     void *value;
53   };
54 };
55 
56 BPF_SK_STORAGE(sk_data_map, __u64);
57 
58 struct info_t {
59   __u32 family;
60   __u32 protocol;
61   __u64 val;
62 };
63 
64 BPF_ITER(bpf_sk_storage_map) {
65   struct seq_file *seq = ctx->meta->seq;
66   struct sock *sk = ctx->sk;
67   __u64 *val = ctx->value;
68   struct info_t info = {};
69 
70   if (sk == (void *)0 || val == (void *)0)
71     return 0;
72 
73   info.family = sk->sk_family;
74   info.protocol = sk->sk_protocol;
75   info.val = *val;
76   bpf_seq_write(seq, &info, sizeof(info));
77 
78   return 0;
79 }
80 )";
81 
82 struct info_t {
83   unsigned family;
84   unsigned protocol;
85   unsigned long long val;
86 };
87 
main()88 int main() {
89   ebpf::BPF bpf;
90   auto res = bpf.init(BPF_PROGRAM);
91   if (!res.ok()) {
92     std::cerr << res.msg() << std::endl;
93     return 1;
94   }
95 
96   // create two sockets
97   int sockfd1 = socket(AF_INET, SOCK_DGRAM, 0);
98   if (sockfd1 < 0) {
99     std::cerr << "socket1 create failure: " << sockfd1 << std::endl;
100     return 1;
101   }
102 
103   int sockfd2 = socket(AF_INET, SOCK_DGRAM, 0);
104   if (sockfd2 < 0) {
105     std::cerr << "socket2 create failure: " << sockfd2 << std::endl;
106     close(sockfd1);
107     return 1;
108   }
109 
110   unsigned long long v1 = 10, v2 = 20;
111   auto sk_table = bpf.get_sk_storage_table<unsigned long long>("sk_data_map");
112 
113   res = sk_table.update_value(sockfd1, v1);
114   if (!res.ok()) {
115     std::cerr << "sk_data_map sockfd1 update failure: " << res.msg() << std::endl;
116     close(sockfd2);
117     close(sockfd1);
118     return 1;
119   }
120 
121   res = sk_table.update_value(sockfd2, v2);
122   if (!res.ok()) {
123     std::cerr << "sk_data_map sockfd2 update failure: " << res.msg() << std::endl;
124     close(sockfd2);
125     close(sockfd1);
126     return 1;
127   }
128 
129   int prog_fd;
130   res = bpf.load_func("bpf_iter__bpf_sk_storage_map", BPF_PROG_TYPE_TRACING, prog_fd);
131   if (!res.ok()) {
132     std::cerr << res.msg() << std::endl;
133     return 1;
134   }
135 
136   union bpf_iter_link_info link_info = {};
137   link_info.map.map_fd = sk_table.get_fd();
138   int link_fd = bcc_iter_attach(prog_fd, &link_info, sizeof(union bpf_iter_link_info));
139   if (link_fd < 0) {
140     std::cerr << "bcc_iter_attach failed: " << link_fd << std::endl;
141     close(sockfd2);
142     close(sockfd1);
143     return 1;
144   }
145 
146   int iter_fd = bcc_iter_create(link_fd);
147   if (iter_fd < 0) {
148     std::cerr << "bcc_iter_create failed: " << iter_fd << std::endl;
149     close(link_fd);
150     close(sockfd2);
151     close(sockfd1);
152     return 1;
153   }
154 
155   // Header.
156   printf("family\tprot\tval\n");
157 
158   struct info_t info[20];
159   int len, leftover = 0, info_size = 20 * sizeof(struct info_t);
160   while ((len = read(iter_fd, (char *)info + leftover, info_size - leftover))) {
161     if (len < 0) {
162       if (len == -EAGAIN)
163         continue;
164       std::cerr << "read failed: " << len << std::endl;
165       break;
166     }
167 
168     int num_info = len / sizeof(struct info_t);
169     for (int i = 0; i < num_info; i++) {
170       printf("%d\t%d\t%lld\n", info[i].family, info[i].protocol, info[i].val);
171     }
172 
173     leftover = len % sizeof(struct info_t);
174     if (num_info > 0)
175       memcpy(info, (void *)&info[num_info], leftover);
176   }
177 
178   close(iter_fd);
179   close(link_fd);
180   close(sockfd2);
181   close(sockfd1);
182   return 0;
183 }
184