1*a96d1cfbSMartin KaFai Lau // SPDX-License-Identifier: GPL-2.0
2*a96d1cfbSMartin KaFai Lau /* Copyright (c) 2023 Facebook */
3*a96d1cfbSMartin KaFai Lau #include "vmlinux.h"
4*a96d1cfbSMartin KaFai Lau #include "bpf_tracing_net.h"
5*a96d1cfbSMartin KaFai Lau #include <bpf/bpf_helpers.h>
6*a96d1cfbSMartin KaFai Lau #include <bpf/bpf_tracing.h>
7*a96d1cfbSMartin KaFai Lau #include <bpf/bpf_core_read.h>
8*a96d1cfbSMartin KaFai Lau 
9*a96d1cfbSMartin KaFai Lau void *local_storage_ptr = NULL;
10*a96d1cfbSMartin KaFai Lau void *sk_ptr = NULL;
11*a96d1cfbSMartin KaFai Lau int cookie_found = 0;
12*a96d1cfbSMartin KaFai Lau __u64 cookie = 0;
13*a96d1cfbSMartin KaFai Lau __u32 omem = 0;
14*a96d1cfbSMartin KaFai Lau 
15*a96d1cfbSMartin KaFai Lau void *bpf_rdonly_cast(void *, __u32) __ksym;
16*a96d1cfbSMartin KaFai Lau 
17*a96d1cfbSMartin KaFai Lau struct {
18*a96d1cfbSMartin KaFai Lau 	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
19*a96d1cfbSMartin KaFai Lau 	__uint(map_flags, BPF_F_NO_PREALLOC);
20*a96d1cfbSMartin KaFai Lau 	__type(key, int);
21*a96d1cfbSMartin KaFai Lau 	__type(value, int);
22*a96d1cfbSMartin KaFai Lau } sk_storage SEC(".maps");
23*a96d1cfbSMartin KaFai Lau 
24*a96d1cfbSMartin KaFai Lau SEC("fexit/bpf_local_storage_destroy")
BPF_PROG(bpf_local_storage_destroy,struct bpf_local_storage * local_storage)25*a96d1cfbSMartin KaFai Lau int BPF_PROG(bpf_local_storage_destroy, struct bpf_local_storage *local_storage)
26*a96d1cfbSMartin KaFai Lau {
27*a96d1cfbSMartin KaFai Lau 	struct sock *sk;
28*a96d1cfbSMartin KaFai Lau 
29*a96d1cfbSMartin KaFai Lau 	if (local_storage_ptr != local_storage)
30*a96d1cfbSMartin KaFai Lau 		return 0;
31*a96d1cfbSMartin KaFai Lau 
32*a96d1cfbSMartin KaFai Lau 	sk = bpf_rdonly_cast(sk_ptr, bpf_core_type_id_kernel(struct sock));
33*a96d1cfbSMartin KaFai Lau 	if (sk->sk_cookie.counter != cookie)
34*a96d1cfbSMartin KaFai Lau 		return 0;
35*a96d1cfbSMartin KaFai Lau 
36*a96d1cfbSMartin KaFai Lau 	cookie_found++;
37*a96d1cfbSMartin KaFai Lau 	omem = sk->sk_omem_alloc.counter;
38*a96d1cfbSMartin KaFai Lau 	local_storage_ptr = NULL;
39*a96d1cfbSMartin KaFai Lau 
40*a96d1cfbSMartin KaFai Lau 	return 0;
41*a96d1cfbSMartin KaFai Lau }
42*a96d1cfbSMartin KaFai Lau 
43*a96d1cfbSMartin KaFai Lau SEC("fentry/inet6_sock_destruct")
BPF_PROG(inet6_sock_destruct,struct sock * sk)44*a96d1cfbSMartin KaFai Lau int BPF_PROG(inet6_sock_destruct, struct sock *sk)
45*a96d1cfbSMartin KaFai Lau {
46*a96d1cfbSMartin KaFai Lau 	int *value;
47*a96d1cfbSMartin KaFai Lau 
48*a96d1cfbSMartin KaFai Lau 	if (!cookie || sk->sk_cookie.counter != cookie)
49*a96d1cfbSMartin KaFai Lau 		return 0;
50*a96d1cfbSMartin KaFai Lau 
51*a96d1cfbSMartin KaFai Lau 	value = bpf_sk_storage_get(&sk_storage, sk, 0, 0);
52*a96d1cfbSMartin KaFai Lau 	if (value && *value == 0xdeadbeef) {
53*a96d1cfbSMartin KaFai Lau 		cookie_found++;
54*a96d1cfbSMartin KaFai Lau 		sk_ptr = sk;
55*a96d1cfbSMartin KaFai Lau 		local_storage_ptr = sk->sk_bpf_storage;
56*a96d1cfbSMartin KaFai Lau 	}
57*a96d1cfbSMartin KaFai Lau 
58*a96d1cfbSMartin KaFai Lau 	return 0;
59*a96d1cfbSMartin KaFai Lau }
60*a96d1cfbSMartin KaFai Lau 
61*a96d1cfbSMartin KaFai Lau char _license[] SEC("license") = "GPL";
62