tools/testing/selftests/bpf/test_cgroup_storage.c

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