1374e74deSXu Liu // SPDX-License-Identifier: GPL-2.0
2374e74deSXu Liu 
3374e74deSXu Liu #include "vmlinux.h"
4374e74deSXu Liu 
5374e74deSXu Liu #include <bpf/bpf_helpers.h>
6374e74deSXu Liu 
7374e74deSXu Liu #define AF_INET6 10
8374e74deSXu Liu 
9374e74deSXu Liu struct {
10374e74deSXu Liu 	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
11374e74deSXu Liu 	__uint(map_flags, BPF_F_NO_PREALLOC);
12374e74deSXu Liu 	__type(key, int);
13374e74deSXu Liu 	__type(value, int);
14*6cbca1eeSXu Liu } sockops_netns_cookies SEC(".maps");
15*6cbca1eeSXu Liu 
16*6cbca1eeSXu Liu struct {
17*6cbca1eeSXu Liu 	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
18*6cbca1eeSXu Liu 	__uint(map_flags, BPF_F_NO_PREALLOC);
19*6cbca1eeSXu Liu 	__type(key, int);
20*6cbca1eeSXu Liu 	__type(value, int);
21*6cbca1eeSXu Liu } sk_msg_netns_cookies SEC(".maps");
22*6cbca1eeSXu Liu 
23*6cbca1eeSXu Liu struct {
24*6cbca1eeSXu Liu 	__uint(type, BPF_MAP_TYPE_SOCKMAP);
25*6cbca1eeSXu Liu 	__uint(max_entries, 2);
26*6cbca1eeSXu Liu 	__type(key, __u32);
27*6cbca1eeSXu Liu 	__type(value, __u64);
28*6cbca1eeSXu Liu } sock_map SEC(".maps");
29374e74deSXu Liu 
30374e74deSXu Liu SEC("sockops")
get_netns_cookie_sockops(struct bpf_sock_ops * ctx)31374e74deSXu Liu int get_netns_cookie_sockops(struct bpf_sock_ops *ctx)
32374e74deSXu Liu {
33374e74deSXu Liu 	struct bpf_sock *sk = ctx->sk;
34374e74deSXu Liu 	int *cookie;
35*6cbca1eeSXu Liu 	__u32 key = 0;
36374e74deSXu Liu 
37374e74deSXu Liu 	if (ctx->family != AF_INET6)
38374e74deSXu Liu 		return 1;
39374e74deSXu Liu 
40374e74deSXu Liu 	if (!sk)
41374e74deSXu Liu 		return 1;
42374e74deSXu Liu 
43*6cbca1eeSXu Liu 	switch (ctx->op) {
44*6cbca1eeSXu Liu 	case BPF_SOCK_OPS_TCP_CONNECT_CB:
45*6cbca1eeSXu Liu 		cookie = bpf_sk_storage_get(&sockops_netns_cookies, sk, 0,
46374e74deSXu Liu 					    BPF_SK_STORAGE_GET_F_CREATE);
47374e74deSXu Liu 		if (!cookie)
48374e74deSXu Liu 			return 1;
49374e74deSXu Liu 
50374e74deSXu Liu 		*cookie = bpf_get_netns_cookie(ctx);
51*6cbca1eeSXu Liu 		break;
52*6cbca1eeSXu Liu 	case BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB:
53*6cbca1eeSXu Liu 		bpf_sock_map_update(ctx, &sock_map, &key, BPF_NOEXIST);
54*6cbca1eeSXu Liu 		break;
55*6cbca1eeSXu Liu 	default:
56*6cbca1eeSXu Liu 		break;
57*6cbca1eeSXu Liu 	}
58374e74deSXu Liu 
59374e74deSXu Liu 	return 1;
60374e74deSXu Liu }
61*6cbca1eeSXu Liu 
62*6cbca1eeSXu Liu SEC("sk_msg")
get_netns_cookie_sk_msg(struct sk_msg_md * msg)63*6cbca1eeSXu Liu int get_netns_cookie_sk_msg(struct sk_msg_md *msg)
64*6cbca1eeSXu Liu {
65*6cbca1eeSXu Liu 	struct bpf_sock *sk = msg->sk;
66*6cbca1eeSXu Liu 	int *cookie;
67*6cbca1eeSXu Liu 
68*6cbca1eeSXu Liu 	if (msg->family != AF_INET6)
69*6cbca1eeSXu Liu 		return 1;
70*6cbca1eeSXu Liu 
71*6cbca1eeSXu Liu 	if (!sk)
72*6cbca1eeSXu Liu 		return 1;
73*6cbca1eeSXu Liu 
74*6cbca1eeSXu Liu 	cookie = bpf_sk_storage_get(&sk_msg_netns_cookies, sk, 0,
75*6cbca1eeSXu Liu 				    BPF_SK_STORAGE_GET_F_CREATE);
76*6cbca1eeSXu Liu 	if (!cookie)
77*6cbca1eeSXu Liu 		return 1;
78*6cbca1eeSXu Liu 
79*6cbca1eeSXu Liu 	*cookie = bpf_get_netns_cookie(msg);
80*6cbca1eeSXu Liu 
81*6cbca1eeSXu Liu 	return 1;
82*6cbca1eeSXu Liu }
83*6cbca1eeSXu Liu 
84*6cbca1eeSXu Liu char _license[] SEC("license") = "GPL";
85