xref: /linux/tools/testing/selftests/bpf/progs/sock_iter_batch.c (revision 1b98f357dadd6ea613a435fbaef1a5dd7b35fd21)
1 // SPDX-License-Identifier: GPL-2.0
2 // Copyright (c) 2024 Meta
3 
4 #include "vmlinux.h"
5 #include <bpf/bpf_helpers.h>
6 #include <bpf/bpf_core_read.h>
7 #include <bpf/bpf_endian.h>
8 #include "bpf_tracing_net.h"
9 #include "bpf_kfuncs.h"
10 
11 #define ATTR __always_inline
12 #include "test_jhash.h"
13 
14 static bool ipv6_addr_loopback(const struct in6_addr *a)
15 {
16 	return (a->s6_addr32[0] | a->s6_addr32[1] |
17 		a->s6_addr32[2] | (a->s6_addr32[3] ^ bpf_htonl(1))) == 0;
18 }
19 
20 static bool ipv4_addr_loopback(__be32 a)
21 {
22 	return a == bpf_ntohl(0x7f000001);
23 }
24 
25 volatile const unsigned int sf;
26 volatile const __u16 ports[2];
27 unsigned int bucket[2];
28 
29 SEC("iter/tcp")
30 int iter_tcp_soreuse(struct bpf_iter__tcp *ctx)
31 {
32 	struct sock *sk = (struct sock *)ctx->sk_common;
33 	struct inet_hashinfo *hinfo;
34 	unsigned int hash;
35 	__u64 sock_cookie;
36 	struct net *net;
37 	int idx;
38 
39 	if (!sk)
40 		return 0;
41 
42 	sock_cookie = bpf_get_socket_cookie(sk);
43 	sk = bpf_core_cast(sk, struct sock);
44 	if (sk->sk_family != sf ||
45 	    sk->sk_state != TCP_LISTEN ||
46 	    sk->sk_family == AF_INET6 ?
47 	    !ipv6_addr_loopback(&sk->sk_v6_rcv_saddr) :
48 	    !ipv4_addr_loopback(sk->sk_rcv_saddr))
49 		return 0;
50 
51 	if (sk->sk_num == ports[0])
52 		idx = 0;
53 	else if (sk->sk_num == ports[1])
54 		idx = 1;
55 	else
56 		return 0;
57 
58 	/* bucket selection as in inet_lhash2_bucket_sk() */
59 	net = sk->sk_net.net;
60 	hash = jhash2(sk->sk_v6_rcv_saddr.s6_addr32, 4, net->hash_mix);
61 	hash ^= sk->sk_num;
62 	hinfo = net->ipv4.tcp_death_row.hashinfo;
63 	bucket[idx] = hash & hinfo->lhash2_mask;
64 	bpf_seq_write(ctx->meta->seq, &idx, sizeof(idx));
65 	bpf_seq_write(ctx->meta->seq, &sock_cookie, sizeof(sock_cookie));
66 
67 	return 0;
68 }
69 
70 #define udp_sk(ptr) container_of(ptr, struct udp_sock, inet.sk)
71 
72 SEC("iter/udp")
73 int iter_udp_soreuse(struct bpf_iter__udp *ctx)
74 {
75 	struct sock *sk = (struct sock *)ctx->udp_sk;
76 	struct udp_table *udptable;
77 	__u64 sock_cookie;
78 	int idx;
79 
80 	if (!sk)
81 		return 0;
82 
83 	sock_cookie = bpf_get_socket_cookie(sk);
84 	sk = bpf_core_cast(sk, struct sock);
85 	if (sk->sk_family != sf ||
86 	    sk->sk_family == AF_INET6 ?
87 	    !ipv6_addr_loopback(&sk->sk_v6_rcv_saddr) :
88 	    !ipv4_addr_loopback(sk->sk_rcv_saddr))
89 		return 0;
90 
91 	if (sk->sk_num == ports[0])
92 		idx = 0;
93 	else if (sk->sk_num == ports[1])
94 		idx = 1;
95 	else
96 		return 0;
97 
98 	/* bucket selection as in udp_hashslot2() */
99 	udptable = sk->sk_net.net->ipv4.udp_table;
100 	bucket[idx] = udp_sk(sk)->udp_portaddr_hash & udptable->mask;
101 	bpf_seq_write(ctx->meta->seq, &idx, sizeof(idx));
102 	bpf_seq_write(ctx->meta->seq, &sock_cookie, sizeof(sock_cookie));
103 
104 	return 0;
105 }
106 
107 char _license[] SEC("license") = "GPL";
108