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