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
 13char bpf_log_buf[BPF_LOG_BUF_SIZE];
 14
 15#define TEST_CGROUP "/test-bpf-cgroup-storage-buf/"
 16
 17int 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
168err:
169	cleanup_cgroup_environment();
170	free(percpu_value);
171
172out:
173	return error;
174}