1 | // SPDX-License-Identifier: GPL-2.0 |
2 | #include <assert.h> |
3 | #include <bpf/bpf.h> |
4 | #include <linux/filter.h> |
5 | #include <stdio.h> |
6 | #include <stdlib.h> |
7 | #include <sys/sysinfo.h> |
8 | |
9 | #include "bpf_util.h" |
10 | #include "cgroup_helpers.h" |
11 | #include "testing_helpers.h" |
12 | |
13 | char bpf_log_buf[BPF_LOG_BUF_SIZE]; |
14 | |
15 | #define TEST_CGROUP "/test-bpf-cgroup-storage-buf/" |
16 | |
17 | int main(int argc, char **argv) |
18 | { |
19 | struct bpf_insn prog[] = { |
20 | BPF_LD_MAP_FD(BPF_REG_1, 0), /* percpu map fd */ |
21 | BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */ |
22 | BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, |
23 | BPF_FUNC_get_local_storage), |
24 | BPF_LDX_MEM(BPF_DW, BPF_REG_3, BPF_REG_0, 0), |
25 | BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, 0x1), |
26 | BPF_STX_MEM(BPF_DW, BPF_REG_0, BPF_REG_3, 0), |
27 | |
28 | BPF_LD_MAP_FD(BPF_REG_1, 0), /* map fd */ |
29 | BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */ |
30 | BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, |
31 | BPF_FUNC_get_local_storage), |
32 | BPF_MOV64_IMM(BPF_REG_1, 1), |
33 | BPF_ATOMIC_OP(BPF_DW, BPF_ADD, BPF_REG_0, BPF_REG_1, 0), |
34 | BPF_LDX_MEM(BPF_DW, BPF_REG_1, BPF_REG_0, 0), |
35 | BPF_ALU64_IMM(BPF_AND, BPF_REG_1, 0x1), |
36 | BPF_MOV64_REG(BPF_REG_0, BPF_REG_1), |
37 | BPF_EXIT_INSN(), |
38 | }; |
39 | size_t insns_cnt = ARRAY_SIZE(prog); |
40 | int error = EXIT_FAILURE; |
41 | int map_fd, percpu_map_fd, prog_fd, cgroup_fd; |
42 | struct bpf_cgroup_storage_key key; |
43 | unsigned long long value; |
44 | unsigned long long *percpu_value; |
45 | int cpu, nproc; |
46 | |
47 | nproc = bpf_num_possible_cpus(); |
48 | percpu_value = malloc(sizeof(*percpu_value) * nproc); |
49 | if (!percpu_value) { |
50 | printf("Not enough memory for per-cpu area (%d cpus)\n" , nproc); |
51 | goto err; |
52 | } |
53 | |
54 | /* Use libbpf 1.0 API mode */ |
55 | libbpf_set_strict_mode(LIBBPF_STRICT_ALL); |
56 | |
57 | map_fd = bpf_map_create(BPF_MAP_TYPE_CGROUP_STORAGE, NULL, sizeof(key), |
58 | sizeof(value), 0, NULL); |
59 | if (map_fd < 0) { |
60 | printf("Failed to create map: %s\n" , strerror(errno)); |
61 | goto out; |
62 | } |
63 | |
64 | percpu_map_fd = bpf_map_create(BPF_MAP_TYPE_PERCPU_CGROUP_STORAGE, NULL, |
65 | sizeof(key), sizeof(value), 0, NULL); |
66 | if (percpu_map_fd < 0) { |
67 | printf("Failed to create map: %s\n" , strerror(errno)); |
68 | goto out; |
69 | } |
70 | |
71 | prog[0].imm = percpu_map_fd; |
72 | prog[7].imm = map_fd; |
73 | prog_fd = bpf_test_load_program(type: BPF_PROG_TYPE_CGROUP_SKB, |
74 | insns: prog, insns_cnt, license: "GPL" , kern_version: 0, |
75 | log_buf: bpf_log_buf, log_buf_sz: BPF_LOG_BUF_SIZE); |
76 | if (prog_fd < 0) { |
77 | printf("Failed to load bpf program: %s\n" , bpf_log_buf); |
78 | goto out; |
79 | } |
80 | |
81 | cgroup_fd = cgroup_setup_and_join(TEST_CGROUP); |
82 | |
83 | /* Attach the bpf program */ |
84 | if (bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_INET_EGRESS, 0)) { |
85 | printf("Failed to attach bpf program\n" ); |
86 | goto err; |
87 | } |
88 | |
89 | if (bpf_map_get_next_key(map_fd, NULL, &key)) { |
90 | printf("Failed to get the first key in cgroup storage\n" ); |
91 | goto err; |
92 | } |
93 | |
94 | if (bpf_map_lookup_elem(map_fd, &key, &value)) { |
95 | printf("Failed to lookup cgroup storage 0\n" ); |
96 | goto err; |
97 | } |
98 | |
99 | for (cpu = 0; cpu < nproc; cpu++) |
100 | percpu_value[cpu] = 1000; |
101 | |
102 | if (bpf_map_update_elem(percpu_map_fd, &key, percpu_value, 0)) { |
103 | printf("Failed to update the data in the cgroup storage\n" ); |
104 | goto err; |
105 | } |
106 | |
107 | /* Every second packet should be dropped */ |
108 | assert(system("ping localhost -c 1 -W 1 -q > /dev/null" ) == 0); |
109 | assert(system("ping localhost -c 1 -W 1 -q > /dev/null" )); |
110 | assert(system("ping localhost -c 1 -W 1 -q > /dev/null" ) == 0); |
111 | |
112 | /* Check the counter in the cgroup local storage */ |
113 | if (bpf_map_lookup_elem(map_fd, &key, &value)) { |
114 | printf("Failed to lookup cgroup storage\n" ); |
115 | goto err; |
116 | } |
117 | |
118 | if (value != 3) { |
119 | printf("Unexpected data in the cgroup storage: %llu\n" , value); |
120 | goto err; |
121 | } |
122 | |
123 | /* Bump the counter in the cgroup local storage */ |
124 | value++; |
125 | if (bpf_map_update_elem(map_fd, &key, &value, 0)) { |
126 | printf("Failed to update the data in the cgroup storage\n" ); |
127 | goto err; |
128 | } |
129 | |
130 | /* Every second packet should be dropped */ |
131 | assert(system("ping localhost -c 1 -W 1 -q > /dev/null" ) == 0); |
132 | assert(system("ping localhost -c 1 -W 1 -q > /dev/null" )); |
133 | assert(system("ping localhost -c 1 -W 1 -q > /dev/null" ) == 0); |
134 | |
135 | /* Check the final value of the counter in the cgroup local storage */ |
136 | if (bpf_map_lookup_elem(map_fd, &key, &value)) { |
137 | printf("Failed to lookup the cgroup storage\n" ); |
138 | goto err; |
139 | } |
140 | |
141 | if (value != 7) { |
142 | printf("Unexpected data in the cgroup storage: %llu\n" , value); |
143 | goto err; |
144 | } |
145 | |
146 | /* Check the final value of the counter in the percpu local storage */ |
147 | |
148 | for (cpu = 0; cpu < nproc; cpu++) |
149 | percpu_value[cpu] = 0; |
150 | |
151 | if (bpf_map_lookup_elem(percpu_map_fd, &key, percpu_value)) { |
152 | printf("Failed to lookup the per-cpu cgroup storage\n" ); |
153 | goto err; |
154 | } |
155 | |
156 | value = 0; |
157 | for (cpu = 0; cpu < nproc; cpu++) |
158 | value += percpu_value[cpu]; |
159 | |
160 | if (value != nproc * 1000 + 6) { |
161 | printf("Unexpected data in the per-cpu cgroup storage\n" ); |
162 | goto err; |
163 | } |
164 | |
165 | error = 0; |
166 | printf("test_cgroup_storage:PASS\n" ); |
167 | |
168 | err: |
169 | cleanup_cgroup_environment(); |
170 | free(percpu_value); |
171 | |
172 | out: |
173 | return error; |
174 | } |
175 | |