xref: /linux-tools/ebpf/kernel-samples-bpf/sockconn4lb_admin.c (revision 7889a9fe3f6cd23238c94fad4e1c698d5585c3fe)
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