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_rlimit.h"
  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[] = {
  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
  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;
  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         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);
  84         if (cgroup_fd < 0) {
  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)) {
 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");
 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
 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
 176         error = 0;
 177         printf("test_cgroup_storage:PASS\n");
 178
 179 err:
 180         cleanup_cgroup_environment();
 181         free(percpu_value);
 182
 183 out:
 184         return error;
 185 }