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