1 // SPDX-License-Identifier: GPL-2.0
4 #include <linux/filter.h>
7 #include <sys/sysinfo.h>
9 #include "bpf_rlimit.h"
10 #include "cgroup_helpers.h"
12 char bpf_log_buf[BPF_LOG_BUF_SIZE];
14 #define TEST_CGROUP "/test-bpf-cgroup-storage-buf/"
16 int main(int argc, char **argv)
18 struct bpf_insn prog[] = {
19 BPF_LD_MAP_FD(BPF_REG_1, 0), /* percpu map fd */
20 BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */
21 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
22 BPF_FUNC_get_local_storage),
23 BPF_LDX_MEM(BPF_DW, BPF_REG_3, BPF_REG_0, 0),
24 BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, 0x1),
25 BPF_STX_MEM(BPF_DW, BPF_REG_0, BPF_REG_3, 0),
27 BPF_LD_MAP_FD(BPF_REG_1, 0), /* map fd */
28 BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */
29 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
30 BPF_FUNC_get_local_storage),
31 BPF_MOV64_IMM(BPF_REG_1, 1),
32 BPF_STX_XADD(BPF_DW, BPF_REG_0, BPF_REG_1, 0),
33 BPF_LDX_MEM(BPF_DW, BPF_REG_1, BPF_REG_0, 0),
34 BPF_ALU64_IMM(BPF_AND, BPF_REG_1, 0x1),
35 BPF_MOV64_REG(BPF_REG_0, BPF_REG_1),
38 size_t insns_cnt = sizeof(prog) / sizeof(struct bpf_insn);
39 int error = EXIT_FAILURE;
40 int map_fd, percpu_map_fd, prog_fd, cgroup_fd;
41 struct bpf_cgroup_storage_key key;
42 unsigned long long value;
43 unsigned long long *percpu_value;
46 nproc = get_nprocs_conf();
47 percpu_value = malloc(sizeof(*percpu_value) * nproc);
49 printf("Not enough memory for per-cpu area (%d cpus)\n", nproc);
53 map_fd = bpf_create_map(BPF_MAP_TYPE_CGROUP_STORAGE, sizeof(key),
56 printf("Failed to create map: %s\n", strerror(errno));
60 percpu_map_fd = bpf_create_map(BPF_MAP_TYPE_PERCPU_CGROUP_STORAGE,
61 sizeof(key), sizeof(value), 0, 0);
62 if (percpu_map_fd < 0) {
63 printf("Failed to create map: %s\n", strerror(errno));
67 prog[0].imm = percpu_map_fd;
69 prog_fd = bpf_load_program(BPF_PROG_TYPE_CGROUP_SKB,
70 prog, insns_cnt, "GPL", 0,
71 bpf_log_buf, BPF_LOG_BUF_SIZE);
73 printf("Failed to load bpf program: %s\n", bpf_log_buf);
77 if (setup_cgroup_environment()) {
78 printf("Failed to setup cgroup environment\n");
82 /* Create a cgroup, get fd, and join it */
83 cgroup_fd = create_and_get_cgroup(TEST_CGROUP);
85 printf("Failed to create test cgroup\n");
89 if (join_cgroup(TEST_CGROUP)) {
90 printf("Failed to join cgroup\n");
94 /* Attach the bpf program */
95 if (bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_INET_EGRESS, 0)) {
96 printf("Failed to attach bpf program\n");
100 if (bpf_map_get_next_key(map_fd, NULL, &key)) {
101 printf("Failed to get the first key in cgroup storage\n");
105 if (bpf_map_lookup_elem(map_fd, &key, &value)) {
106 printf("Failed to lookup cgroup storage 0\n");
110 for (cpu = 0; cpu < nproc; cpu++)
111 percpu_value[cpu] = 1000;
113 if (bpf_map_update_elem(percpu_map_fd, &key, percpu_value, 0)) {
114 printf("Failed to update the data in the cgroup storage\n");
118 /* Every second packet should be dropped */
119 assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
120 assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
121 assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
123 /* Check the counter in the cgroup local storage */
124 if (bpf_map_lookup_elem(map_fd, &key, &value)) {
125 printf("Failed to lookup cgroup storage\n");
130 printf("Unexpected data in the cgroup storage: %llu\n", value);
134 /* Bump the counter in the cgroup local storage */
136 if (bpf_map_update_elem(map_fd, &key, &value, 0)) {
137 printf("Failed to update the data in the cgroup storage\n");
141 /* Every second packet should be dropped */
142 assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
143 assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
144 assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
146 /* Check the final value of the counter in the cgroup local storage */
147 if (bpf_map_lookup_elem(map_fd, &key, &value)) {
148 printf("Failed to lookup the cgroup storage\n");
153 printf("Unexpected data in the cgroup storage: %llu\n", value);
157 /* Check the final value of the counter in the percpu local storage */
159 for (cpu = 0; cpu < nproc; cpu++)
160 percpu_value[cpu] = 0;
162 if (bpf_map_lookup_elem(percpu_map_fd, &key, percpu_value)) {
163 printf("Failed to lookup the per-cpu cgroup storage\n");
168 for (cpu = 0; cpu < nproc; cpu++)
169 value += percpu_value[cpu];
171 if (value != nproc * 1000 + 6) {
172 printf("Unexpected data in the per-cpu cgroup storage\n");
177 printf("test_cgroup_storage:PASS\n");
180 cleanup_cgroup_environment();