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(BPF_PROG_TYPE_CGROUP_SKB,
74 				   prog, insns_cnt, "GPL", 0,
75 				   bpf_log_buf, 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