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_util.h" 10 #include "cgroup_helpers.h" 11 #include "testing_helpers.h" 12 13 char bpf_log_buf[BPF_LOG_BUF_SIZE]; 14 15 #define TEST_CGROUP "/test-bpf-cgroup-storage-buf/" 16 17 int main(int argc, char **argv) 18 { 19 struct bpf_insn prog[] = { 20 BPF_LD_MAP_FD(BPF_REG_1, 0), /* percpu map fd */ 21 BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */ 22 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, 23 BPF_FUNC_get_local_storage), 24 BPF_LDX_MEM(BPF_DW, BPF_REG_3, BPF_REG_0, 0), 25 BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, 0x1), 26 BPF_STX_MEM(BPF_DW, BPF_REG_0, BPF_REG_3, 0), 27 28 BPF_LD_MAP_FD(BPF_REG_1, 0), /* map fd */ 29 BPF_MOV64_IMM(BPF_REG_2, 0), /* flags, not used */ 30 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, 31 BPF_FUNC_get_local_storage), 32 BPF_MOV64_IMM(BPF_REG_1, 1), 33 BPF_ATOMIC_OP(BPF_DW, BPF_ADD, BPF_REG_0, BPF_REG_1, 0), 34 BPF_LDX_MEM(BPF_DW, BPF_REG_1, BPF_REG_0, 0), 35 BPF_ALU64_IMM(BPF_AND, BPF_REG_1, 0x1), 36 BPF_MOV64_REG(BPF_REG_0, BPF_REG_1), 37 BPF_EXIT_INSN(), 38 }; 39 size_t insns_cnt = ARRAY_SIZE(prog); 40 int error = EXIT_FAILURE; 41 int map_fd, percpu_map_fd, prog_fd, cgroup_fd; 42 struct bpf_cgroup_storage_key key; 43 unsigned long long value; 44 unsigned long long *percpu_value; 45 int cpu, nproc; 46 47 nproc = bpf_num_possible_cpus(); 48 percpu_value = malloc(sizeof(*percpu_value) * nproc); 49 if (!percpu_value) { 50 printf("Not enough memory for per-cpu area (%d cpus)\n", nproc); 51 goto err; 52 } 53 54 /* Use libbpf 1.0 API mode */ 55 libbpf_set_strict_mode(LIBBPF_STRICT_ALL); 56 57 map_fd = bpf_map_create(BPF_MAP_TYPE_CGROUP_STORAGE, NULL, sizeof(key), 58 sizeof(value), 0, NULL); 59 if (map_fd < 0) { 60 printf("Failed to create map: %s\n", strerror(errno)); 61 goto out; 62 } 63 64 percpu_map_fd = bpf_map_create(BPF_MAP_TYPE_PERCPU_CGROUP_STORAGE, NULL, 65 sizeof(key), sizeof(value), 0, NULL); 66 if (percpu_map_fd < 0) { 67 printf("Failed to create map: %s\n", strerror(errno)); 68 goto out; 69 } 70 71 prog[0].imm = percpu_map_fd; 72 prog[7].imm = map_fd; 73 prog_fd = bpf_test_load_program(BPF_PROG_TYPE_CGROUP_SKB, 74 prog, insns_cnt, "GPL", 0, 75 bpf_log_buf, BPF_LOG_BUF_SIZE); 76 if (prog_fd < 0) { 77 printf("Failed to load bpf program: %s\n", bpf_log_buf); 78 goto out; 79 } 80 81 cgroup_fd = cgroup_setup_and_join(TEST_CGROUP); 82 83 /* Attach the bpf program */ 84 if (bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_INET_EGRESS, 0)) { 85 printf("Failed to attach bpf program\n"); 86 goto err; 87 } 88 89 if (bpf_map_get_next_key(map_fd, NULL, &key)) { 90 printf("Failed to get the first key in cgroup storage\n"); 91 goto err; 92 } 93 94 if (bpf_map_lookup_elem(map_fd, &key, &value)) { 95 printf("Failed to lookup cgroup storage 0\n"); 96 goto err; 97 } 98 99 for (cpu = 0; cpu < nproc; cpu++) 100 percpu_value[cpu] = 1000; 101 102 if (bpf_map_update_elem(percpu_map_fd, &key, percpu_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 counter in the cgroup local storage */ 113 if (bpf_map_lookup_elem(map_fd, &key, &value)) { 114 printf("Failed to lookup cgroup storage\n"); 115 goto err; 116 } 117 118 if (value != 3) { 119 printf("Unexpected data in the cgroup storage: %llu\n", value); 120 goto err; 121 } 122 123 /* Bump the counter in the cgroup local storage */ 124 value++; 125 if (bpf_map_update_elem(map_fd, &key, &value, 0)) { 126 printf("Failed to update the data in the cgroup storage\n"); 127 goto err; 128 } 129 130 /* Every second packet should be dropped */ 131 assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0); 132 assert(system("ping localhost -c 1 -W 1 -q > /dev/null")); 133 assert(system("ping localhost -c 1 -W 1 -q > /dev/null") == 0); 134 135 /* Check the final value of the counter in the cgroup local storage */ 136 if (bpf_map_lookup_elem(map_fd, &key, &value)) { 137 printf("Failed to lookup the cgroup storage\n"); 138 goto err; 139 } 140 141 if (value != 7) { 142 printf("Unexpected data in the cgroup storage: %llu\n", value); 143 goto err; 144 } 145 146 /* Check the final value of the counter in the percpu local storage */ 147 148 for (cpu = 0; cpu < nproc; cpu++) 149 percpu_value[cpu] = 0; 150 151 if (bpf_map_lookup_elem(percpu_map_fd, &key, percpu_value)) { 152 printf("Failed to lookup the per-cpu cgroup storage\n"); 153 goto err; 154 } 155 156 value = 0; 157 for (cpu = 0; cpu < nproc; cpu++) 158 value += percpu_value[cpu]; 159 160 if (value != nproc * 1000 + 6) { 161 printf("Unexpected data in the per-cpu cgroup storage\n"); 162 goto err; 163 } 164 165 error = 0; 166 printf("test_cgroup_storage:PASS\n"); 167 168 err: 169 cleanup_cgroup_environment(); 170 free(percpu_value); 171 172 out: 173 return error; 174 } 175