selftests/bpf/test_cgroup_storage.c

0001 // SPDX-License-Identifier: GPL-2.0
0002 #include <assert.h>
0003 #include <bpf/bpf.h>
0004 #include <linux/filter.h>
0005 #include <stdio.h>
0006 #include <stdlib.h>
0007 #include <sys/sysinfo.h>
0008
0009 #include "bpf_util.h"
0010 #include "cgroup_helpers.h"
0011 #include "testing_helpers.h"
0012
0013 char bpf_log_buf[BPF_LOG_BUF_SIZE];
0014
0015 #define TEST_CGROUP "/test-bpf-cgroup-storage-buf/"
0016
0017 int main(int argc, char **argv)
0018 {
0019     struct bpf_insn prog[] = {
0020         BPF_LD_MAP_FD(BPF_REG_1, 0), /* percpu map fd */
0021         BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */
0022         BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
0023                  BPF_FUNC_get_local_storage),
0024         BPF_LDX_MEM(BPF_DW, BPF_REG_3, BPF_REG_0, 0),
0025         BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, 0x1),
0026         BPF_STX_MEM(BPF_DW, BPF_REG_0, BPF_REG_3, 0),
0027
0028         BPF_LD_MAP_FD(BPF_REG_1, 0), /* map fd */
0029         BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */
0030         BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
0031                  BPF_FUNC_get_local_storage),
0032         BPF_MOV64_IMM(BPF_REG_1, 1),
0033         BPF_ATOMIC_OP(BPF_DW, BPF_ADD, BPF_REG_0, BPF_REG_1, 0),
0034         BPF_LDX_MEM(BPF_DW, BPF_REG_1, BPF_REG_0, 0),
0035         BPF_ALU64_IMM(BPF_AND, BPF_REG_1, 0x1),
0036         BPF_MOV64_REG(BPF_REG_0, BPF_REG_1),
0037         BPF_EXIT_INSN(),
0038     };
0039     size_t insns_cnt = ARRAY_SIZE(prog);
0040     int error = EXIT_FAILURE;
0041     int map_fd, percpu_map_fd, prog_fd, cgroup_fd;
0042     struct bpf_cgroup_storage_key key;
0043     unsigned long long value;
0044     unsigned long long *percpu_value;
0045     int cpu, nproc;
0046
0047     nproc = bpf_num_possible_cpus();
0048     percpu_value = malloc(sizeof(*percpu_value) * nproc);
0049     if (!percpu_value) {
0050         printf("Not enough memory for per-cpu area (%d cpus)\n", nproc);
0051         goto err;
0052     }
0053
0054     /* Use libbpf 1.0 API mode */
0055     libbpf_set_strict_mode(LIBBPF_STRICT_ALL);
0056
0057     map_fd = bpf_map_create(BPF_MAP_TYPE_CGROUP_STORAGE, NULL, sizeof(key),
0058                 sizeof(value), 0, NULL);
0059     if (map_fd < 0) {
0060         printf("Failed to create map: %s\n", strerror(errno));
0061         goto out;
0062     }
0063
0064     percpu_map_fd = bpf_map_create(BPF_MAP_TYPE_PERCPU_CGROUP_STORAGE, NULL,
0065                        sizeof(key), sizeof(value), 0, NULL);
0066     if (percpu_map_fd < 0) {
0067         printf("Failed to create map: %s\n", strerror(errno));
0068         goto out;
0069     }
0070
0071     prog[0].imm = percpu_map_fd;
0072     prog[7].imm = map_fd;
0073     prog_fd = bpf_test_load_program(BPF_PROG_TYPE_CGROUP_SKB,
0074                    prog, insns_cnt, "GPL", 0,
0075                    bpf_log_buf, BPF_LOG_BUF_SIZE);
0076     if (prog_fd < 0) {
0077         printf("Failed to load bpf program: %s\n", bpf_log_buf);
0078         goto out;
0079     }
0080
0081     cgroup_fd = cgroup_setup_and_join(TEST_CGROUP);
0082
0083     /* Attach the bpf program */
0084     if (bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_INET_EGRESS, 0)) {
0085         printf("Failed to attach bpf program\n");
0086         goto err;
0087     }
0088
0089     if (bpf_map_get_next_key(map_fd, NULL, &key)) {
0090         printf("Failed to get the first key in cgroup storage\n");
0091         goto err;
0092     }
0093
0094     if (bpf_map_lookup_elem(map_fd, &key, &value)) {
0095         printf("Failed to lookup cgroup storage 0\n");
0096         goto err;
0097     }
0098
0099     for (cpu = 0; cpu < nproc; cpu++)
0100         percpu_value[cpu] = 1000;
0101
0102     if (bpf_map_update_elem(percpu_map_fd, &key, percpu_value, 0)) {
0103         printf("Failed to update the data in the cgroup storage\n");
0104         goto err;
0105     }
0106
0107     /* Every second packet should be dropped */
0108     assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
0109     assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
0110     assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
0111
0112     /* Check the counter in the cgroup local storage */
0113     if (bpf_map_lookup_elem(map_fd, &key, &value)) {
0114         printf("Failed to lookup cgroup storage\n");
0115         goto err;
0116     }
0117
0118     if (value != 3) {
0119         printf("Unexpected data in the cgroup storage: %llu\n", value);
0120         goto err;
0121     }
0122
0123     /* Bump the counter in the cgroup local storage */
0124     value++;
0125     if (bpf_map_update_elem(map_fd, &key, &value, 0)) {
0126         printf("Failed to update the data in the cgroup storage\n");
0127         goto err;
0128     }
0129
0130     /* Every second packet should be dropped */
0131     assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
0132     assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
0133     assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
0134
0135     /* Check the final value of the counter in the cgroup local storage */
0136     if (bpf_map_lookup_elem(map_fd, &key, &value)) {
0137         printf("Failed to lookup the cgroup storage\n");
0138         goto err;
0139     }
0140
0141     if (value != 7) {
0142         printf("Unexpected data in the cgroup storage: %llu\n", value);
0143         goto err;
0144     }
0145
0146     /* Check the final value of the counter in the percpu local storage */
0147
0148     for (cpu = 0; cpu < nproc; cpu++)
0149         percpu_value[cpu] = 0;
0150
0151     if (bpf_map_lookup_elem(percpu_map_fd, &key, percpu_value)) {
0152         printf("Failed to lookup the per-cpu cgroup storage\n");
0153         goto err;
0154     }
0155
0156     value = 0;
0157     for (cpu = 0; cpu < nproc; cpu++)
0158         value += percpu_value[cpu];
0159
0160     if (value != nproc * 1000 + 6) {
0161         printf("Unexpected data in the per-cpu cgroup storage\n");
0162         goto err;
0163     }
0164
0165     error = 0;
0166     printf("test_cgroup_storage:PASS\n");
0167
0168 err:
0169     cleanup_cgroup_environment();
0170     free(percpu_value);
0171
0172 out:
0173     return error;
0174 }