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 } 132