1 /* This test is a demo of using get_socket_uid and get_socket_cookie 2 * helper function to do per socket based network traffic monitoring. 3 * It requires iptables version higher then 1.6.1. to load pinned eBPF 4 * program into the xt_bpf match. 5 * 6 * TEST: 7 * ./run_cookie_uid_helper_example.sh 8 * Then generate some traffic in variate ways. ping 0 -c 10 would work 9 * but the cookie and uid in this case could both be 0. A sample output 10 * with some traffic generated by web browser is shown below: 11 * 12 * cookie: 877, uid: 0x3e8, Pakcet Count: 20, Bytes Count: 11058 13 * cookie: 132, uid: 0x0, Pakcet Count: 2, Bytes Count: 286 14 * cookie: 812, uid: 0x3e8, Pakcet Count: 3, Bytes Count: 1726 15 * cookie: 802, uid: 0x3e8, Pakcet Count: 2, Bytes Count: 104 16 * cookie: 877, uid: 0x3e8, Pakcet Count: 20, Bytes Count: 11058 17 * cookie: 831, uid: 0x3e8, Pakcet Count: 2, Bytes Count: 104 18 * cookie: 0, uid: 0x0, Pakcet Count: 6, Bytes Count: 712 19 * cookie: 880, uid: 0xfffe, Pakcet Count: 1, Bytes Count: 70 20 * 21 * Clean up: if using shell script, the script file will delete the iptables 22 * rule and unmount the bpf program when exit. Else the iptables rule need 23 * to be deleted by hand, see run_cookie_uid_helper_example.sh for detail. 24 */ 25 26 #define _GNU_SOURCE 27 28 #define offsetof(type, member) __builtin_offsetof(type, member) 29 #define ARRAY_SIZE(x) (sizeof(x) / sizeof(*(x))) 30 31 #include <arpa/inet.h> 32 #include <errno.h> 33 #include <error.h> 34 #include <limits.h> 35 #include <linux/bpf.h> 36 #include <linux/if_ether.h> 37 #include <stdbool.h> 38 #include <stdint.h> 39 #include <stdio.h> 40 #include <stdlib.h> 41 #include <string.h> 42 #include <sys/socket.h> 43 #include <sys/stat.h> 44 #include <sys/types.h> 45 #include <unistd.h> 46 #include <bpf/bpf.h> 47 #include "libbpf.h" 48 49 struct stats { 50 uint32_t uid; 51 uint64_t packets; 52 uint64_t bytes; 53 }; 54 55 static int map_fd, prog_fd; 56 57 static void maps_create(void) 58 { 59 map_fd = bpf_create_map(BPF_MAP_TYPE_HASH, sizeof(uint32_t), 60 sizeof(struct stats), 100, 0); 61 if (map_fd < 0) 62 error(1, errno, "map create failed!\n"); 63 } 64 65 static void prog_load(void) 66 { 67 static char log_buf[1 << 16]; 68 69 struct bpf_insn prog[] = { 70 /* 71 * Save sk_buff for future usage. value stored in R6 to R10 will 72 * not be reset after a bpf helper function call. 73 */ 74 BPF_MOV64_REG(BPF_REG_6, BPF_REG_1), 75 /* 76 * pc1: BPF_FUNC_get_socket_cookie takes one parameter, 77 * R1: sk_buff 78 */ 79 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, 80 BPF_FUNC_get_socket_cookie), 81 /* pc2-4: save &socketCookie to r7 for future usage*/ 82 BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -8), 83 BPF_MOV64_REG(BPF_REG_7, BPF_REG_10), 84 BPF_ALU64_IMM(BPF_ADD, BPF_REG_7, -8), 85 /* 86 * pc5-8: set up the registers for BPF_FUNC_map_lookup_elem, 87 * it takes two parameters (R1: map_fd, R2: &socket_cookie) 88 */ 89 BPF_LD_MAP_FD(BPF_REG_1, map_fd), 90 BPF_MOV64_REG(BPF_REG_2, BPF_REG_7), 91 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, 92 BPF_FUNC_map_lookup_elem), 93 /* 94 * pc9. if r0 != 0x0, go to pc+14, since we have the cookie 95 * stored already 96 * Otherwise do pc10-22 to setup a new data entry. 97 */ 98 BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 14), 99 BPF_MOV64_REG(BPF_REG_1, BPF_REG_6), 100 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, 101 BPF_FUNC_get_socket_uid), 102 /* 103 * Place a struct stats in the R10 stack and sequentially 104 * place the member value into the memory. Packets value 105 * is set by directly place a IMM value 1 into the stack. 106 */ 107 BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, 108 -32 + offsetof(struct stats, uid)), 109 BPF_ST_MEM(BPF_DW, BPF_REG_10, 110 -32 + offsetof(struct stats, packets), 1), 111 /* 112 * __sk_buff is a special struct used for eBPF program to 113 * directly access some sk_buff field. 114 */ 115 BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_6, 116 offsetof(struct __sk_buff, len)), 117 BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_1, 118 -32 + offsetof(struct stats, bytes)), 119 /* 120 * add new map entry using BPF_FUNC_map_update_elem, it takes 121 * 4 parameters (R1: map_fd, R2: &socket_cookie, R3: &stats, 122 * R4: flags) 123 */ 124 BPF_LD_MAP_FD(BPF_REG_1, map_fd), 125 BPF_MOV64_REG(BPF_REG_2, BPF_REG_7), 126 BPF_MOV64_REG(BPF_REG_3, BPF_REG_10), 127 BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, -32), 128 BPF_MOV64_IMM(BPF_REG_4, 0), 129 BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, 130 BPF_FUNC_map_update_elem), 131 BPF_JMP_IMM(BPF_JA, 0, 0, 5), 132 /* 133 * pc24-30 update the packet info to a exist data entry, it can 134 * be done by directly write to pointers instead of using 135 * BPF_FUNC_map_update_elem helper function 136 */ 137 BPF_MOV64_REG(BPF_REG_9, BPF_REG_0), 138 BPF_MOV64_IMM(BPF_REG_1, 1), 139 BPF_STX_XADD(BPF_DW, BPF_REG_9, BPF_REG_1, 140 offsetof(struct stats, packets)), 141 BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_6, 142 offsetof(struct __sk_buff, len)), 143 BPF_STX_XADD(BPF_DW, BPF_REG_9, BPF_REG_1, 144 offsetof(struct stats, bytes)), 145 BPF_LDX_MEM(BPF_W, BPF_REG_0, BPF_REG_6, 146 offsetof(struct __sk_buff, len)), 147 BPF_EXIT_INSN(), 148 }; 149 prog_fd = bpf_load_program(BPF_PROG_TYPE_SOCKET_FILTER, prog, 150 ARRAY_SIZE(prog), "GPL", 0, 151 log_buf, sizeof(log_buf)); 152 if (prog_fd < 0) 153 error(1, errno, "failed to load prog\n%s\n", log_buf); 154 } 155 156 static void prog_attach_iptables(char *file) 157 { 158 int ret; 159 char rules[100]; 160 161 if (bpf_obj_pin(prog_fd, file)) 162 error(1, errno, "bpf_obj_pin"); 163 if (strlen(file) > 50) { 164 printf("file path too long: %s\n", file); 165 exit(1); 166 } 167 sprintf(rules, "iptables -A INPUT -m bpf --object-pinned %s -j ACCEPT", 168 file); 169 ret = system(rules); 170 if (ret < 0) { 171 printf("iptables rule update failed: %d/n", WEXITSTATUS(ret)); 172 exit(1); 173 } 174 } 175 176 static void print_table(void) 177 { 178 struct stats curEntry; 179 uint32_t curN = UINT32_MAX; 180 uint32_t nextN, res; 181 182 while (bpf_map_get_next_key(map_fd, &curN, &nextN) > -1) { 183 curN = nextN; 184 res = bpf_map_lookup_elem(map_fd, &curN, &curEntry); 185 if (res < 0) { 186 error(1, errno, "fail to get entry value of Key: %u\n", 187 curN); 188 } else { 189 printf("cookie: %u, uid: 0x%x, Packet Count: %lu," 190 " Bytes Count: %lu\n", curN, curEntry.uid, 191 curEntry.packets, curEntry.bytes); 192 } 193 } 194 } 195 196 int main(int argc, char *argv[]) 197 { 198 if (argc > 2) { 199 printf("Too many argument provided\n"); 200 return 1; 201 } else if (argc < 2) { 202 printf("Usage: %s bpfObjName\n", argv[0]); 203 return 1; 204 } 205 206 maps_create(); 207 prog_load(); 208 prog_attach_iptables(argv[1]); 209 210 while (true) { 211 print_table(); 212 printf("\n"); 213 sleep(1); 214 }; 215 216 return 0; 217 } 218