[mirror_ubuntu-eoan-kernel.git] / tools / testing / selftests / bpf / test_cgroup_storage.c

// SPDX-License-Identifier: GPL-2.0
#include <assert.h>
#include <bpf/bpf.h>
#include <linux/filter.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/sysinfo.h>

#include "bpf_rlimit.h"
#include "cgroup_helpers.h"

char bpf_log_buf[BPF_LOG_BUF_SIZE];

#define TEST_CGROUP "/test-bpf-cgroup-storage-buf/"

int main(int argc, char **argv)
{
	struct bpf_insn prog[] = {
		BPF_LD_MAP_FD(BPF_REG_1, 0), /* percpu map fd */
		BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
			     BPF_FUNC_get_local_storage),
		BPF_LDX_MEM(BPF_W, BPF_REG_3, BPF_REG_0, 0),
		BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, 0x1),
		BPF_STX_MEM(BPF_W, BPF_REG_0, BPF_REG_3, 0),

		BPF_LD_MAP_FD(BPF_REG_1, 0), /* map fd */
		BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
			     BPF_FUNC_get_local_storage),
		BPF_MOV64_IMM(BPF_REG_1, 1),
		BPF_STX_XADD(BPF_DW, BPF_REG_0, BPF_REG_1, 0),
		BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_0, 0),
		BPF_ALU64_IMM(BPF_AND, BPF_REG_1, 0x1),
		BPF_MOV64_REG(BPF_REG_0, BPF_REG_1),
		BPF_EXIT_INSN(),
	};
	size_t insns_cnt = sizeof(prog) / sizeof(struct bpf_insn);
	int error = EXIT_FAILURE;
	int map_fd, percpu_map_fd, prog_fd, cgroup_fd;
	struct bpf_cgroup_storage_key key;
	unsigned long long value;
	unsigned long long *percpu_value;
	int cpu, nproc;

	nproc = get_nprocs_conf();
	percpu_value = malloc(sizeof(*percpu_value) * nproc);
	if (!percpu_value) {
		printf("Not enough memory for per-cpu area (%d cpus)\n", nproc);
		goto err;
	}

	map_fd = bpf_create_map(BPF_MAP_TYPE_CGROUP_STORAGE, sizeof(key),
				sizeof(value), 0, 0);
	if (map_fd < 0) {
		printf("Failed to create map: %s\n", strerror(errno));
		goto out;
	}

	percpu_map_fd = bpf_create_map(BPF_MAP_TYPE_PERCPU_CGROUP_STORAGE,
				       sizeof(key), sizeof(value), 0, 0);
	if (percpu_map_fd < 0) {
		printf("Failed to create map: %s\n", strerror(errno));
		goto out;
	}

	prog[0].imm = percpu_map_fd;
	prog[7].imm = map_fd;
	prog_fd = bpf_load_program(BPF_PROG_TYPE_CGROUP_SKB,
				   prog, insns_cnt, "GPL", 0,
				   bpf_log_buf, BPF_LOG_BUF_SIZE);
	if (prog_fd < 0) {
		printf("Failed to load bpf program: %s\n", bpf_log_buf);
		goto out;
	}

	if (setup_cgroup_environment()) {
		printf("Failed to setup cgroup environment\n");
		goto err;
	}

	/* Create a cgroup, get fd, and join it */
	cgroup_fd = create_and_get_cgroup(TEST_CGROUP);
	if (cgroup_fd < 0) {
		printf("Failed to create test cgroup\n");
		goto err;
	}

	if (join_cgroup(TEST_CGROUP)) {
		printf("Failed to join cgroup\n");
		goto err;
	}

	/* Attach the bpf program */
	if (bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_INET_EGRESS, 0)) {
		printf("Failed to attach bpf program\n");
		goto err;
	}

	if (bpf_map_get_next_key(map_fd, NULL, &key)) {
		printf("Failed to get the first key in cgroup storage\n");
		goto err;
	}

	if (bpf_map_lookup_elem(map_fd, &key, &value)) {
		printf("Failed to lookup cgroup storage 0\n");
		goto err;
	}

	for (cpu = 0; cpu < nproc; cpu++)
		percpu_value[cpu] = 1000;

	if (bpf_map_update_elem(percpu_map_fd, &key, percpu_value, 0)) {
		printf("Failed to update the data in the cgroup storage\n");
		goto err;
	}

	/* Every second packet should be dropped */
	assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
	assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
	assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);

	/* Check the counter in the cgroup local storage */
	if (bpf_map_lookup_elem(map_fd, &key, &value)) {
		printf("Failed to lookup cgroup storage\n");
		goto err;
	}

	if (value != 3) {
		printf("Unexpected data in the cgroup storage: %llu\n", value);
		goto err;
	}

	/* Bump the counter in the cgroup local storage */
	value++;
	if (bpf_map_update_elem(map_fd, &key, &value, 0)) {
		printf("Failed to update the data in the cgroup storage\n");
		goto err;
	}

	/* Every second packet should be dropped */
	assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
	assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
	assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);

	/* Check the final value of the counter in the cgroup local storage */
	if (bpf_map_lookup_elem(map_fd, &key, &value)) {
		printf("Failed to lookup the cgroup storage\n");
		goto err;
	}

	if (value != 7) {
		printf("Unexpected data in the cgroup storage: %llu\n", value);
		goto err;
	}

	/* Check the final value of the counter in the percpu local storage */

	for (cpu = 0; cpu < nproc; cpu++)
		percpu_value[cpu] = 0;

	if (bpf_map_lookup_elem(percpu_map_fd, &key, percpu_value)) {
		printf("Failed to lookup the per-cpu cgroup storage\n");
		goto err;
	}

	value = 0;
	for (cpu = 0; cpu < nproc; cpu++)
		value += percpu_value[cpu];

	if (value != nproc * 1000 + 6) {
		printf("Unexpected data in the per-cpu cgroup storage\n");
		goto err;
	}

	error = 0;
	printf("test_cgroup_storage:PASS\n");

err:
	cleanup_cgroup_environment();
	free(percpu_value);

out:
	return error;
}
Commit	Line	Data
68cfa3ac RG	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>
919646d2	7	#include <sys/sysinfo.h>
68cfa3ac	8
a85da34e	9	#include "bpf_rlimit.h"
68cfa3ac RG	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[] = {
919646d2 RG	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_W, BPF_REG_3, BPF_REG_0, 0),
	24	BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, 0x1),
	25	BPF_STX_MEM(BPF_W, BPF_REG_0, BPF_REG_3, 0),
	26
68cfa3ac RG	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_W, 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;
919646d2	40	int map_fd, percpu_map_fd, prog_fd, cgroup_fd;
68cfa3ac RG	41	struct bpf_cgroup_storage_key key;
68cfa3ac RG	42	unsigned long long value;
919646d2 RG	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	}
68cfa3ac RG	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
919646d2 RG	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;
68cfa3ac RG	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	if (setup_cgroup_environment()) {
	78	printf("Failed to setup cgroup environment\n");
	79	goto err;
	80	}
	81
	82	/* Create a cgroup, get fd, and join it */
	83	cgroup_fd = create_and_get_cgroup(TEST_CGROUP);
a8911d6d	84	if (cgroup_fd < 0) {
68cfa3ac RG	85	printf("Failed to create test cgroup\n");
	86	goto err;
	87	}
	88
	89	if (join_cgroup(TEST_CGROUP)) {
	90	printf("Failed to join cgroup\n");
	91	goto err;
	92	}
	93
	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");
	97	goto err;
	98	}
	99
	100	if (bpf_map_get_next_key(map_fd, NULL, &key)) {
	101	printf("Failed to get the first key in cgroup storage\n");
	102	goto err;
	103	}
	104
	105	if (bpf_map_lookup_elem(map_fd, &key, &value)) {
919646d2 RG	106	printf("Failed to lookup cgroup storage 0\n");
	107	goto err;
	108	}
	109
	110	for (cpu = 0; cpu < nproc; cpu++)
	111	percpu_value[cpu] = 1000;
	112
	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");
68cfa3ac RG	115	goto err;
	116	}
	117
	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);
	122
	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");
	126	goto err;
	127	}
	128
	129	if (value != 3) {
	130	printf("Unexpected data in the cgroup storage: %llu\n", value);
	131	goto err;
	132	}
	133
	134	/* Bump the counter in the cgroup local storage */
	135	value++;
	136	if (bpf_map_update_elem(map_fd, &key, &value, 0)) {
	137	printf("Failed to update the data in the cgroup storage\n");
	138	goto err;
	139	}
	140
	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);
	145
	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");
	149	goto err;
	150	}
	151
	152	if (value != 7) {
	153	printf("Unexpected data in the cgroup storage: %llu\n", value);
	154	goto err;
	155	}
	156
919646d2 RG	157	/* Check the final value of the counter in the percpu local storage */
	158
	159	for (cpu = 0; cpu < nproc; cpu++)
	160	percpu_value[cpu] = 0;
	161
	162	if (bpf_map_lookup_elem(percpu_map_fd, &key, percpu_value)) {
	163	printf("Failed to lookup the per-cpu cgroup storage\n");
	164	goto err;
	165	}
	166
	167	value = 0;
	168	for (cpu = 0; cpu < nproc; cpu++)
	169	value += percpu_value[cpu];
	170
	171	if (value != nproc * 1000 + 6) {
	172	printf("Unexpected data in the per-cpu cgroup storage\n");
	173	goto err;
	174	}
	175
68cfa3ac RG	176	error = 0;
	177	printf("test_cgroup_storage:PASS\n");
	178
	179	err:
	180	cleanup_cgroup_environment();
919646d2	181	free(percpu_value);
68cfa3ac RG	182
	183	out:
	184	return error;
	185	}