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_rlimit.h"
10 #include "cgroup_helpers.h"
11 
12 char bpf_log_buf[BPF_LOG_BUF_SIZE];
13 
14 #define TEST_CGROUP "/test-bpf-cgroup-storage-buf/"
15 
16 int main(int argc, char **argv)
17 {
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),
26 
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),
36 		BPF_EXIT_INSN(),
37 	};
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;
44 	int cpu, nproc;
45 
46 	nproc = get_nprocs_conf();
47 	percpu_value = malloc(sizeof(*percpu_value) * nproc);
48 	if (!percpu_value) {
49 		printf("Not enough memory for per-cpu area (%d cpus)\n", nproc);
50 		goto err;
51 	}
52 
53 	map_fd = bpf_create_map(BPF_MAP_TYPE_CGROUP_STORAGE, sizeof(key),
54 				sizeof(value), 0, 0);
55 	if (map_fd < 0) {
56 		printf("Failed to create map: %s\n", strerror(errno));
57 		goto out;
58 	}
59 
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));
64 		goto out;
65 	}
66 
67 	prog[0].imm = percpu_map_fd;
68 	prog[7].imm = 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);
72 	if (prog_fd < 0) {
73 		printf("Failed to load bpf program: %s\n", bpf_log_buf);
74 		goto out;
75 	}
76 
77 	cgroup_fd = cgroup_setup_and_join(TEST_CGROUP);
78 
79 	/* Attach the bpf program */
80 	if (bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_INET_EGRESS, 0)) {
81 		printf("Failed to attach bpf program\n");
82 		goto err;
83 	}
84 
85 	if (bpf_map_get_next_key(map_fd, NULL, &key)) {
86 		printf("Failed to get the first key in cgroup storage\n");
87 		goto err;
88 	}
89 
90 	if (bpf_map_lookup_elem(map_fd, &key, &value)) {
91 		printf("Failed to lookup cgroup storage 0\n");
92 		goto err;
93 	}
94 
95 	for (cpu = 0; cpu < nproc; cpu++)
96 		percpu_value[cpu] = 1000;
97 
98 	if (bpf_map_update_elem(percpu_map_fd, &key, percpu_value, 0)) {
99 		printf("Failed to update the data in the cgroup storage\n");
100 		goto err;
101 	}
102 
103 	/* Every second packet should be dropped */
104 	assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
105 	assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
106 	assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
107 
108 	/* Check the counter in the cgroup local storage */
109 	if (bpf_map_lookup_elem(map_fd, &key, &value)) {
110 		printf("Failed to lookup cgroup storage\n");
111 		goto err;
112 	}
113 
114 	if (value != 3) {
115 		printf("Unexpected data in the cgroup storage: %llu\n", value);
116 		goto err;
117 	}
118 
119 	/* Bump the counter in the cgroup local storage */
120 	value++;
121 	if (bpf_map_update_elem(map_fd, &key, &value, 0)) {
122 		printf("Failed to update the data in the cgroup storage\n");
123 		goto err;
124 	}
125 
126 	/* Every second packet should be dropped */
127 	assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
128 	assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
129 	assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
130 
131 	/* Check the final value of the counter in the cgroup local storage */
132 	if (bpf_map_lookup_elem(map_fd, &key, &value)) {
133 		printf("Failed to lookup the cgroup storage\n");
134 		goto err;
135 	}
136 
137 	if (value != 7) {
138 		printf("Unexpected data in the cgroup storage: %llu\n", value);
139 		goto err;
140 	}
141 
142 	/* Check the final value of the counter in the percpu local storage */
143 
144 	for (cpu = 0; cpu < nproc; cpu++)
145 		percpu_value[cpu] = 0;
146 
147 	if (bpf_map_lookup_elem(percpu_map_fd, &key, percpu_value)) {
148 		printf("Failed to lookup the per-cpu cgroup storage\n");
149 		goto err;
150 	}
151 
152 	value = 0;
153 	for (cpu = 0; cpu < nproc; cpu++)
154 		value += percpu_value[cpu];
155 
156 	if (value != nproc * 1000 + 6) {
157 		printf("Unexpected data in the per-cpu cgroup storage\n");
158 		goto err;
159 	}
160 
161 	error = 0;
162 	printf("test_cgroup_storage:PASS\n");
163 
164 err:
165 	cleanup_cgroup_environment();
166 	free(percpu_value);
167 
168 out:
169 	return error;
170 }
171