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_DW, BPF_REG_3, BPF_REG_0, 0), 24 BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, 0x1), 25 BPF_STX_MEM(BPF_DW, 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_DW, 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 } 186