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