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
   8 #include "bpf_rlimit.h"
   9 #include "cgroup_helpers.h"
  10
  11 char bpf_log_buf[BPF_LOG_BUF_SIZE];
  12
  13 #define TEST_CGROUP "/test-bpf-cgroup-storage-buf/"
  14
  15 int main(int argc, char **argv)
  16 {
  17         struct bpf_insn prog[] = {
  18                 BPF_LD_MAP_FD(BPF_REG_1, 0), /* map fd */
  19                 BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */
  20                 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
  21                              BPF_FUNC_get_local_storage),
  22                 BPF_MOV64_IMM(BPF_REG_1, 1),
  23                 BPF_STX_XADD(BPF_DW, BPF_REG_0, BPF_REG_1, 0),
  24                 BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_0, 0),
  25                 BPF_ALU64_IMM(BPF_AND, BPF_REG_1, 0x1),
  26                 BPF_MOV64_REG(BPF_REG_0, BPF_REG_1),
  27                 BPF_EXIT_INSN(),
  28         };
  29         size_t insns_cnt = sizeof(prog) / sizeof(struct bpf_insn);
  30         int error = EXIT_FAILURE;
  31         int map_fd, prog_fd, cgroup_fd;
  32         struct bpf_cgroup_storage_key key;
  33         unsigned long long value;
  34
  35         map_fd = bpf_create_map(BPF_MAP_TYPE_CGROUP_STORAGE, sizeof(key),
  36                                 sizeof(value), 0, 0);
  37         if (map_fd < 0) {
  38                 printf("Failed to create map: %s\n", strerror(errno));
  39                 goto out;
  40         }
  41
  42         prog[0].imm = map_fd;
  43         prog_fd = bpf_load_program(BPF_PROG_TYPE_CGROUP_SKB,
  44                                    prog, insns_cnt, "GPL", 0,
  45                                    bpf_log_buf, BPF_LOG_BUF_SIZE);
  46         if (prog_fd < 0) {
  47                 printf("Failed to load bpf program: %s\n", bpf_log_buf);
  48                 goto out;
  49         }
  50
  51         if (setup_cgroup_environment()) {
  52                 printf("Failed to setup cgroup environment\n");
  53                 goto err;
  54         }
  55
  56         /* Create a cgroup, get fd, and join it */
  57         cgroup_fd = create_and_get_cgroup(TEST_CGROUP);
  58         if (!cgroup_fd) {
  59                 printf("Failed to create test cgroup\n");
  60                 goto err;
  61         }
  62
  63         if (join_cgroup(TEST_CGROUP)) {
  64                 printf("Failed to join cgroup\n");
  65                 goto err;
  66         }
  67
  68         /* Attach the bpf program */
  69         if (bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_INET_EGRESS, 0)) {
  70                 printf("Failed to attach bpf program\n");
  71                 goto err;
  72         }
  73
  74         if (bpf_map_get_next_key(map_fd, NULL, &key)) {
  75                 printf("Failed to get the first key in cgroup storage\n");
  76                 goto err;
  77         }
  78
  79         if (bpf_map_lookup_elem(map_fd, &key, &value)) {
  80                 printf("Failed to lookup cgroup storage\n");
  81                 goto err;
  82         }
  83
  84         /* Every second packet should be dropped */
  85         assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
  86         assert(system("ping localhost -c 1 -W 1 -q > /dev/null"));
  87         assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0);
  88
  89         /* Check the counter in the cgroup local storage */
  90         if (bpf_map_lookup_elem(map_fd, &key, &value)) {
  91                 printf("Failed to lookup cgroup storage\n");
  92                 goto err;
  93         }
  94
  95         if (value != 3) {
  96                 printf("Unexpected data in the cgroup storage: %llu\n", value);
  97                 goto err;
  98         }
  99
 100         /* Bump the counter in the cgroup local storage */
 101         value++;
 102         if (bpf_map_update_elem(map_fd, &key, &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 final value of the counter in the cgroup local storage */
 113         if (bpf_map_lookup_elem(map_fd, &key, &value)) {
 114                 printf("Failed to lookup the cgroup storage\n");
 115                 goto err;
 116         }
 117
 118         if (value != 7) {
 119                 printf("Unexpected data in the cgroup storage: %llu\n", value);
 120                 goto err;
 121         }
 122
 123         error = 0;
 124         printf("test_cgroup_storage:PASS\n");
 125
 126 err:
 127         cleanup_cgroup_environment();
 128
 129 out:
 130         return error;
 131 }