1 #include <stdio.h> 2 #include <unistd.h> 3 #include <linux/bpf.h> 4 #include <asm/unistd.h> 5 #include <fcntl.h> 6 #include <string.h> 7 8 9 #define MAXN 128 10 typedef unsigned int VNode[MAXN]; 11 12 static inline __u64 ptr_to_u64(const void *ptr) 13 { 14 return (__u64) (unsigned long) ptr; 15 } 16 17 static inline int sys_bpf(enum bpf_cmd cmd, union bpf_attr *attr, 18 unsigned int size) 19 { 20 return syscall(__NR_bpf, cmd, attr, size); 21 } 22 23 24 25 int get_root_fd() { 26 union bpf_attr attr_create, attr_pin; 27 int fd, rc; 28 memset(&attr_pin, 0, sizeof(attr_pin)); 29 memset(&attr_create, 0, sizeof(attr_create)); 30 attr_pin.pathname = ptr_to_u64("/sys/fs/bpf/viprules"); 31 fd = sys_bpf(BPF_OBJ_GET, &attr_pin, sizeof(attr_pin)); 32 if (fd<0) { 33 attr_create.map_type = BPF_MAP_TYPE_HASH; 34 attr_create.key_size = sizeof(unsigned int); 35 attr_create.value_size = sizeof(VNode); 36 attr_create.max_entries = 1024; 37 strcpy(attr_create.map_name, "viprules"); 38 printf("root bpf map not initialized, create it now\n"); 39 fd = sys_bpf(BPF_MAP_CREATE, &attr_create, sizeof(attr_create)); 40 if (fd<0) { 41 perror("fail to create bpf map:"); 42 return -1; 43 } 44 attr_pin.bpf_fd = fd; 45 rc = sys_bpf(BPF_OBJ_PIN, &attr_pin, sizeof(attr_pin)); 46 if (rc != 0) { 47 perror("fail to pin root bpf map:"); 48 return -1; 49 } 50 } 51 return fd; 52 } 53 54 unsigned int parse_ip(char *p) { 55 int k, i=0, b; 56 unsigned int ip=0; 57 for (k=0; k<4; k++) { 58 b=0; 59 if (p[i]<'0'||p[i]>'9') return 0; 60 while(p[i]>='0'&&p[i]<='9') { 61 b=b*10+p[i++]-'0'; 62 if (b>256) return 0; 63 } 64 if (p[i]!=0&&p[i]!='.') return 0; 65 ip |= b<<(k*8); 66 if (k==3) break; 67 i++; 68 } 69 if (p[i]!=0) return 0; 70 return ip; 71 } 72 73 void print_ip(unsigned int ip) { 74 int k; 75 int bs[4]; 76 for (k=0; k<4; k++) { 77 bs[k] = ip&0xff; 78 ip>>=8; 79 } 80 printf("%d.%d.%d.%d\n", bs[0], bs[1], bs[2], bs[3]); 81 } 82 83 int main(int argc, char *argv[]) { 84 unsigned int vip, ip; 85 int i, rc, fd, n; 86 union bpf_attr attr_elem; 87 unsigned int key; 88 VNode value; 89 unsigned int rip; 90 int root_fd = get_root_fd(); 91 if (root_fd < 0) return -1; 92 memset(&attr_elem, 0, sizeof(attr_elem)); 93 attr_elem.flags = BPF_ANY; 94 95 96 if (argc==3&&strcmp(argv[1], "list")==0) { 97 vip = parse_ip(argv[2]); 98 if (vip==0) { 99 printf("invalid ip: %s\n", argv[2]); 100 return -1; 101 } 102 attr_elem.map_fd = root_fd; 103 attr_elem.key = ptr_to_u64(&vip); 104 attr_elem.value = ptr_to_u64(&value); 105 rc = sys_bpf(BPF_MAP_LOOKUP_ELEM, &attr_elem, sizeof(attr_elem)); 106 if (rc!=0) { 107 printf("vip not found\n"); 108 return 0; 109 } 110 printf("vip: %s\n", argv[2]); 111 n = value[0]; 112 for (i=0; i<n; i++) print_ip(value[i+1]); 113 } else if (argc==3&&strcmp(argv[1], "pop")==0) { 114 vip = parse_ip(argv[2]); 115 if (vip==0) { 116 printf("invalid ip: %s\n", argv[2]); 117 return -1; 118 } 119 attr_elem.map_fd = root_fd; 120 attr_elem.key = ptr_to_u64(&vip); 121 attr_elem.value = ptr_to_u64(&value); 122 rc = sys_bpf(BPF_MAP_LOOKUP_ELEM, &attr_elem, sizeof(attr_elem)); 123 if (rc!=0) { 124 printf("vip not found\n"); 125 return 0; 126 } 127 if (value[0]==0) { 128 printf("no real ip to pop\n"); 129 return 0; 130 } 131 value[0]-=1; 132 rc = sys_bpf(BPF_MAP_UPDATE_ELEM, &attr_elem, sizeof(attr_elem)); 133 if (rc!=0) { 134 printf("fail to sync count\n"); 135 return -1; 136 } 137 } else if (argc==4&&strcmp(argv[1], "push")==0) { 138 vip = parse_ip(argv[2]); 139 if (vip==0) { 140 printf("invalid ip: %s\n", argv[2]); 141 return -1; 142 } 143 attr_elem.map_fd = root_fd; 144 attr_elem.key = ptr_to_u64(&vip); 145 attr_elem.value = ptr_to_u64(&value); 146 rc = sys_bpf(BPF_MAP_LOOKUP_ELEM, &attr_elem, sizeof(attr_elem)); 147 if (rc!=0) { 148 printf("vip not found, adding new entry now\n"); 149 value[0]=0; 150 rc = sys_bpf(BPF_MAP_UPDATE_ELEM, &attr_elem, sizeof(attr_elem)); 151 if (rc!=0) { 152 perror("fail to register new vip\n"); 153 return -1; 154 } 155 } 156 if (value[0]+1>=MAXN) { 157 printf("rip cap reached, could not append more\n"); 158 return -1; 159 } 160 ip = parse_ip(argv[3]); 161 if (ip==0) { 162 printf("invalid ip: %s\n", argv[3]); 163 return -1; 164 } 165 value[0]+=1; 166 value[value[0]]=ip; 167 rc = sys_bpf(BPF_MAP_UPDATE_ELEM, &attr_elem, sizeof(attr_elem)); 168 if (rc!=0) { 169 printf("fail to sync count\n"); 170 return -1; 171 } 172 } else { 173 return -1; 174 } 175 176 return 0; 177 } 178