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
ptr_to_u64(const void * ptr)12 static inline __u64 ptr_to_u64(const void *ptr)
13 {
14 return (__u64) (unsigned long) ptr;
15 }
16
sys_bpf(enum bpf_cmd cmd,union bpf_attr * attr,unsigned int size)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
get_root_fd()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
parse_ip(char * p)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
print_ip(unsigned int ip)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
main(int argc,char * argv[])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