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
8 #include "bpf_rlimit.h"
9 #include "cgroup_helpers.h"
10
11 char bpf_log_buf[BPF_LOG_BUF_SIZE];
12
13 #define TEST_CGROUP "/test-bpf-cgroup-storage-buf/"
14
main(int argc,char ** argv)15 int main(int argc, char **argv)
16 {
17 struct bpf_insn prog[] = {
18 BPF_LD_MAP_FD(BPF_REG_1, 0), /* map fd */
19 BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */
20 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
21 BPF_FUNC_get_local_storage),
22 BPF_MOV64_IMM(BPF_REG_1, 1),
23 BPF_STX_XADD(BPF_DW, BPF_REG_0, BPF_REG_1, 0),
24 BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_0, 0),
25 BPF_ALU64_IMM(BPF_AND, BPF_REG_1, 0x1),
26 BPF_MOV64_REG(BPF_REG_0, BPF_REG_1),
27 BPF_EXIT_INSN(),
28 };
29 size_t insns_cnt = sizeof(prog) / sizeof(struct bpf_insn);
30 int error = EXIT_FAILURE;
31 int map_fd, prog_fd, cgroup_fd;
32 struct bpf_cgroup_storage_key key;
33 unsigned long long value;
34
35 map_fd = bpf_create_map(BPF_MAP_TYPE_CGROUP_STORAGE, sizeof(key),
36 sizeof(value), 0, 0);
37 if (map_fd < 0) {
38 printf("Failed to create map: %s\n", strerror(errno));
39 goto out;
40 }
41
42 prog[0].imm = map_fd;
43 prog_fd = bpf_load_program(BPF_PROG_TYPE_CGROUP_SKB,
44 prog, insns_cnt, "GPL", 0,
45 bpf_log_buf, BPF_LOG_BUF_SIZE);
46 if (prog_fd < 0) {
47 printf("Failed to load bpf program: %s\n", bpf_log_buf);
48 goto out;
49 }
50
51 if (setup_cgroup_environment()) {
52 printf("Failed to setup cgroup environment\n");
53 goto err;
54 }
55
56 /* Create a cgroup, get fd, and join it */
57 cgroup_fd = create_and_get_cgroup(TEST_CGROUP);
58 if (!cgroup_fd) {
59 printf("Failed to create test cgroup\n");
60 goto err;
61 }
62
63 if (join_cgroup(TEST_CGROUP)) {
64 printf("Failed to join cgroup\n");
65 goto err;
66 }
67
68 /* Attach the bpf program */
69 if (bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_INET_EGRESS, 0)) {
70 printf("Failed to attach bpf program\n");
71 goto err;
72 }
73
74 if (bpf_map_get_next_key(map_fd, NULL, &key)) {
75 printf("Failed to get the first key in cgroup storage\n");
76 goto err;
77 }
78
79 if (bpf_map_lookup_elem(map_fd, &key, &value)) {
80 printf("Failed to lookup cgroup storage\n");
81 goto err;
82 }
83
84 /* Every second packet should be dropped */
85 assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
86 assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
87 assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
88
89 /* Check the counter in the cgroup local storage */
90 if (bpf_map_lookup_elem(map_fd, &key, &value)) {
91 printf("Failed to lookup cgroup storage\n");
92 goto err;
93 }
94
95 if (value != 3) {
96 printf("Unexpected data in the cgroup storage: %llu\n", value);
97 goto err;
98 }
99
100 /* Bump the counter in the cgroup local storage */
101 value++;
102 if (bpf_map_update_elem(map_fd, &key, &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 final value of the counter in the cgroup local storage */
113 if (bpf_map_lookup_elem(map_fd, &key, &value)) {
114 printf("Failed to lookup the cgroup storage\n");
115 goto err;
116 }
117
118 if (value != 7) {
119 printf("Unexpected data in the cgroup storage: %llu\n", value);
120 goto err;
121 }
122
123 error = 0;
124 printf("test_cgroup_storage:PASS\n");
125
126 err:
127 cleanup_cgroup_environment();
128
129 out:
130 return error;
131 }
132