xref: /linux/tools/testing/selftests/bpf/progs/test_sockmap_kern.h (revision 9f771739a04919226081a107167596de75108fbb)
1 /* SPDX-License-Identifier: GPL-2.0 */
2 /* Copyright (c) 2017-2018 Covalent IO, Inc. http://covalent.io */
3 #include <stddef.h>
4 #include <string.h>
5 #include <linux/bpf.h>
6 #include <linux/if_ether.h>
7 #include <linux/if_packet.h>
8 #include <linux/ip.h>
9 #include <linux/ipv6.h>
10 #include <linux/in.h>
11 #include <linux/udp.h>
12 #include <linux/tcp.h>
13 #include <linux/pkt_cls.h>
14 #include <sys/socket.h>
15 #include <bpf/bpf_helpers.h>
16 #include <bpf/bpf_endian.h>
17 #include "bpf_misc.h"
18 
19 /* Sockmap sample program connects a client and a backend together
20  * using cgroups.
21  *
22  *    client:X <---> frontend:80 client:X <---> backend:80
23  *
24  * For simplicity we hard code values here and bind 1:1. The hard
25  * coded values are part of the setup in sockmap.sh script that
26  * is associated with this BPF program.
27  *
28  * The bpf_printk is verbose and prints information as connections
29  * are established and verdicts are decided.
30  */
31 
32 struct {
33 	__uint(type, TEST_MAP_TYPE);
34 	__uint(max_entries, 20);
35 	__uint(key_size, sizeof(int));
36 	__uint(value_size, sizeof(int));
37 } sock_map SEC(".maps");
38 
39 struct {
40 	__uint(type, TEST_MAP_TYPE);
41 	__uint(max_entries, 20);
42 	__uint(key_size, sizeof(int));
43 	__uint(value_size, sizeof(int));
44 } sock_map_txmsg SEC(".maps");
45 
46 struct {
47 	__uint(type, TEST_MAP_TYPE);
48 	__uint(max_entries, 20);
49 	__uint(key_size, sizeof(int));
50 	__uint(value_size, sizeof(int));
51 } sock_map_redir SEC(".maps");
52 
53 struct {
54 	__uint(type, BPF_MAP_TYPE_ARRAY);
55 	__uint(max_entries, 1);
56 	__type(key, int);
57 	__type(value, int);
58 } sock_apply_bytes SEC(".maps");
59 
60 struct {
61 	__uint(type, BPF_MAP_TYPE_ARRAY);
62 	__uint(max_entries, 1);
63 	__type(key, int);
64 	__type(value, int);
65 } sock_cork_bytes SEC(".maps");
66 
67 struct {
68 	__uint(type, BPF_MAP_TYPE_ARRAY);
69 	__uint(max_entries, 6);
70 	__type(key, int);
71 	__type(value, int);
72 } sock_bytes SEC(".maps");
73 
74 struct {
75 	__uint(type, BPF_MAP_TYPE_ARRAY);
76 	__uint(max_entries, 1);
77 	__type(key, int);
78 	__type(value, int);
79 } sock_redir_flags SEC(".maps");
80 
81 struct {
82 	__uint(type, BPF_MAP_TYPE_ARRAY);
83 	__uint(max_entries, 3);
84 	__type(key, int);
85 	__type(value, int);
86 } sock_skb_opts SEC(".maps");
87 
88 struct {
89 	__uint(type, TEST_MAP_TYPE);
90 	__uint(max_entries, 20);
91 	__uint(key_size, sizeof(int));
92 	__uint(value_size, sizeof(int));
93 } tls_sock_map SEC(".maps");
94 
95 SEC("sk_skb1")
96 int bpf_prog1(struct __sk_buff *skb)
97 {
98 	int *f, two = 2;
99 
100 	f = bpf_map_lookup_elem(&sock_skb_opts, &two);
101 	if (f && *f) {
102 		return *f;
103 	}
104 	return skb->len;
105 }
106 
107 SEC("sk_skb2")
108 int bpf_prog2(struct __sk_buff *skb)
109 {
110 	__u32 lport = skb->local_port;
111 	__u32 rport = skb->remote_port;
112 	int len, *f, ret, zero = 0;
113 	__u64 flags = 0;
114 
115 	__sink(rport);
116 	if (lport == 10000)
117 		ret = 10;
118 	else
119 		ret = 1;
120 
121 	len = (__u32)skb->data_end - (__u32)skb->data;
122 	__sink(len);
123 
124 	f = bpf_map_lookup_elem(&sock_skb_opts, &zero);
125 	if (f && *f) {
126 		ret = 3;
127 		flags = *f;
128 	}
129 
130 #ifdef SOCKMAP
131 	return bpf_sk_redirect_map(skb, &sock_map, ret, flags);
132 #else
133 	return bpf_sk_redirect_hash(skb, &sock_map, &ret, flags);
134 #endif
135 
136 }
137 
138 static inline void bpf_write_pass(struct __sk_buff *skb, int offset)
139 {
140 	int err = bpf_skb_pull_data(skb, 6 + offset);
141 	void *data_end;
142 	char *c;
143 
144 	if (err)
145 		return;
146 
147 	c = (char *)(long)skb->data;
148 	data_end = (void *)(long)skb->data_end;
149 
150 	if (c + 5 + offset < data_end)
151 		memcpy(c + offset, "PASS", 4);
152 }
153 
154 SEC("sk_skb3")
155 int bpf_prog3(struct __sk_buff *skb)
156 {
157 	int err, *f, ret = SK_PASS;
158 	const int one = 1;
159 
160 	f = bpf_map_lookup_elem(&sock_skb_opts, &one);
161 	if (f && *f) {
162 		__u64 flags = 0;
163 
164 		ret = 0;
165 		flags = *f;
166 
167 		err = bpf_skb_adjust_room(skb, -13, 0, 0);
168 		if (err)
169 			return SK_DROP;
170 		err = bpf_skb_adjust_room(skb, 4, 0, 0);
171 		if (err)
172 			return SK_DROP;
173 		bpf_write_pass(skb, 0);
174 #ifdef SOCKMAP
175 		return bpf_sk_redirect_map(skb, &tls_sock_map, ret, flags);
176 #else
177 		return bpf_sk_redirect_hash(skb, &tls_sock_map, &ret, flags);
178 #endif
179 	}
180 	f = bpf_map_lookup_elem(&sock_skb_opts, &one);
181 	if (f && *f)
182 		ret = SK_DROP;
183 	err = bpf_skb_adjust_room(skb, 4, 0, 0);
184 	if (err)
185 		return SK_DROP;
186 	bpf_write_pass(skb, 13);
187 	return ret;
188 }
189 
190 SEC("sockops")
191 int bpf_sockmap(struct bpf_sock_ops *skops)
192 {
193 	__u32 lport, rport;
194 	int op, ret;
195 
196 	op = (int) skops->op;
197 
198 	switch (op) {
199 	case BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB:
200 		lport = skops->local_port;
201 		rport = skops->remote_port;
202 
203 		if (lport == 10000) {
204 			ret = 1;
205 #ifdef SOCKMAP
206 			bpf_sock_map_update(skops, &sock_map, &ret,
207 						  BPF_NOEXIST);
208 #else
209 			bpf_sock_hash_update(skops, &sock_map, &ret,
210 						   BPF_NOEXIST);
211 #endif
212 		}
213 		break;
214 	case BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB:
215 		lport = skops->local_port;
216 		rport = skops->remote_port;
217 
218 		if (bpf_ntohl(rport) == 10001) {
219 			ret = 10;
220 #ifdef SOCKMAP
221 			bpf_sock_map_update(skops, &sock_map, &ret,
222 						  BPF_NOEXIST);
223 #else
224 			bpf_sock_hash_update(skops, &sock_map, &ret,
225 						   BPF_NOEXIST);
226 #endif
227 		}
228 		break;
229 	default:
230 		break;
231 	}
232 
233 	return 0;
234 }
235 
236 SEC("sk_msg1")
237 int bpf_prog4(struct sk_msg_md *msg)
238 {
239 	int *bytes, zero = 0, one = 1, two = 2, three = 3, four = 4, five = 5;
240 	int *start, *end, *start_push, *end_push, *start_pop, *pop, err = 0;
241 
242 	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
243 	if (bytes)
244 		bpf_msg_apply_bytes(msg, *bytes);
245 	bytes = bpf_map_lookup_elem(&sock_cork_bytes, &zero);
246 	if (bytes)
247 		bpf_msg_cork_bytes(msg, *bytes);
248 	start = bpf_map_lookup_elem(&sock_bytes, &zero);
249 	end = bpf_map_lookup_elem(&sock_bytes, &one);
250 	if (start && end)
251 		bpf_msg_pull_data(msg, *start, *end, 0);
252 	start_push = bpf_map_lookup_elem(&sock_bytes, &two);
253 	end_push = bpf_map_lookup_elem(&sock_bytes, &three);
254 	if (start_push && end_push) {
255 		err = bpf_msg_push_data(msg, *start_push, *end_push, 0);
256 		if (err)
257 			return SK_DROP;
258 	}
259 	start_pop = bpf_map_lookup_elem(&sock_bytes, &four);
260 	pop = bpf_map_lookup_elem(&sock_bytes, &five);
261 	if (start_pop && pop)
262 		bpf_msg_pop_data(msg, *start_pop, *pop, 0);
263 	return SK_PASS;
264 }
265 
266 SEC("sk_msg2")
267 int bpf_prog6(struct sk_msg_md *msg)
268 {
269 	int zero = 0, one = 1, two = 2, three = 3, four = 4, five = 5, key = 0;
270 	int *bytes, *start, *end, *start_push, *end_push, *start_pop, *pop, *f;
271 	int err = 0;
272 	__u64 flags = 0;
273 
274 	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
275 	if (bytes)
276 		bpf_msg_apply_bytes(msg, *bytes);
277 	bytes = bpf_map_lookup_elem(&sock_cork_bytes, &zero);
278 	if (bytes)
279 		bpf_msg_cork_bytes(msg, *bytes);
280 
281 	start = bpf_map_lookup_elem(&sock_bytes, &zero);
282 	end = bpf_map_lookup_elem(&sock_bytes, &one);
283 	if (start && end)
284 		bpf_msg_pull_data(msg, *start, *end, 0);
285 
286 	start_push = bpf_map_lookup_elem(&sock_bytes, &two);
287 	end_push = bpf_map_lookup_elem(&sock_bytes, &three);
288 	if (start_push && end_push) {
289 		err = bpf_msg_push_data(msg, *start_push, *end_push, 0);
290 		if (err)
291 			return SK_DROP;
292 	}
293 
294 	start_pop = bpf_map_lookup_elem(&sock_bytes, &four);
295 	pop = bpf_map_lookup_elem(&sock_bytes, &five);
296 	if (start_pop && pop)
297 		bpf_msg_pop_data(msg, *start_pop, *pop, 0);
298 
299 	f = bpf_map_lookup_elem(&sock_redir_flags, &zero);
300 	if (f && *f) {
301 		key = 2;
302 		flags = *f;
303 	}
304 #ifdef SOCKMAP
305 	return bpf_msg_redirect_map(msg, &sock_map_redir, key, flags);
306 #else
307 	return bpf_msg_redirect_hash(msg, &sock_map_redir, &key, flags);
308 #endif
309 }
310 
311 SEC("sk_msg3")
312 int bpf_prog8(struct sk_msg_md *msg)
313 {
314 	void *data_end = (void *)(long) msg->data_end;
315 	void *data = (void *)(long) msg->data;
316 	int ret = 0, *bytes, zero = 0;
317 
318 	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
319 	if (bytes) {
320 		ret = bpf_msg_apply_bytes(msg, *bytes);
321 		if (ret)
322 			return SK_DROP;
323 	} else {
324 		return SK_DROP;
325 	}
326 
327 	__sink(data_end);
328 	__sink(data);
329 
330 	return SK_PASS;
331 }
332 SEC("sk_msg4")
333 int bpf_prog9(struct sk_msg_md *msg)
334 {
335 	void *data_end = (void *)(long) msg->data_end;
336 	void *data = (void *)(long) msg->data;
337 	int ret = 0, *bytes, zero = 0;
338 
339 	bytes = bpf_map_lookup_elem(&sock_cork_bytes, &zero);
340 	if (bytes) {
341 		if (((__u64)data_end - (__u64)data) >= *bytes)
342 			return SK_PASS;
343 		ret = bpf_msg_cork_bytes(msg, *bytes);
344 		if (ret)
345 			return SK_DROP;
346 	}
347 	return SK_PASS;
348 }
349 
350 SEC("sk_msg5")
351 int bpf_prog10(struct sk_msg_md *msg)
352 {
353 	int *bytes, *start, *end, *start_push, *end_push, *start_pop, *pop;
354 	int zero = 0, one = 1, two = 2, three = 3, four = 4, five = 5, err = 0;
355 
356 	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
357 	if (bytes)
358 		bpf_msg_apply_bytes(msg, *bytes);
359 	bytes = bpf_map_lookup_elem(&sock_cork_bytes, &zero);
360 	if (bytes)
361 		bpf_msg_cork_bytes(msg, *bytes);
362 	start = bpf_map_lookup_elem(&sock_bytes, &zero);
363 	end = bpf_map_lookup_elem(&sock_bytes, &one);
364 	if (start && end)
365 		bpf_msg_pull_data(msg, *start, *end, 0);
366 	start_push = bpf_map_lookup_elem(&sock_bytes, &two);
367 	end_push = bpf_map_lookup_elem(&sock_bytes, &three);
368 	if (start_push && end_push) {
369 		err = bpf_msg_push_data(msg, *start_push, *end_push, 0);
370 		if (err)
371 			return SK_PASS;
372 	}
373 	start_pop = bpf_map_lookup_elem(&sock_bytes, &four);
374 	pop = bpf_map_lookup_elem(&sock_bytes, &five);
375 	if (start_pop && pop)
376 		bpf_msg_pop_data(msg, *start_pop, *pop, 0);
377 	return SK_DROP;
378 }
379 
380 char _license[] SEC("license") = "GPL";
381