153632e11SMartin KaFai Lau // SPDX-License-Identifier: GPL-2.0
253632e11SMartin KaFai Lau /* Copyright (c) 2020 Facebook */
353632e11SMartin KaFai Lau 
453632e11SMartin KaFai Lau #include <vmlinux.h>
553632e11SMartin KaFai Lau #include <bpf/bpf_tracing.h>
653632e11SMartin KaFai Lau #include <bpf/bpf_core_read.h>
753632e11SMartin KaFai Lau #include <bpf/bpf_helpers.h>
853632e11SMartin KaFai Lau 
953632e11SMartin KaFai Lau struct sk_stg {
1053632e11SMartin KaFai Lau 	__u32 pid;
1153632e11SMartin KaFai Lau 	__u32 last_notclose_state;
1253632e11SMartin KaFai Lau 	char comm[16];
1353632e11SMartin KaFai Lau };
1453632e11SMartin KaFai Lau 
1553632e11SMartin KaFai Lau struct {
1653632e11SMartin KaFai Lau 	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
1753632e11SMartin KaFai Lau 	__uint(map_flags, BPF_F_NO_PREALLOC);
1853632e11SMartin KaFai Lau 	__type(key, int);
1953632e11SMartin KaFai Lau 	__type(value, struct sk_stg);
2053632e11SMartin KaFai Lau } sk_stg_map SEC(".maps");
2153632e11SMartin KaFai Lau 
2253632e11SMartin KaFai Lau /* Testing delete */
2353632e11SMartin KaFai Lau struct {
2453632e11SMartin KaFai Lau 	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
2553632e11SMartin KaFai Lau 	__uint(map_flags, BPF_F_NO_PREALLOC);
2653632e11SMartin KaFai Lau 	__type(key, int);
2753632e11SMartin KaFai Lau 	__type(value, int);
2853632e11SMartin KaFai Lau } del_sk_stg_map SEC(".maps");
2953632e11SMartin KaFai Lau 
3053632e11SMartin KaFai Lau char task_comm[16] = "";
3153632e11SMartin KaFai Lau 
3253632e11SMartin KaFai Lau SEC("tp_btf/inet_sock_set_state")
BPF_PROG(trace_inet_sock_set_state,struct sock * sk,int oldstate,int newstate)3353632e11SMartin KaFai Lau int BPF_PROG(trace_inet_sock_set_state, struct sock *sk, int oldstate,
3453632e11SMartin KaFai Lau 	     int newstate)
3553632e11SMartin KaFai Lau {
3653632e11SMartin KaFai Lau 	struct sk_stg *stg;
3753632e11SMartin KaFai Lau 
3853632e11SMartin KaFai Lau 	if (newstate == BPF_TCP_CLOSE)
3953632e11SMartin KaFai Lau 		return 0;
4053632e11SMartin KaFai Lau 
4153632e11SMartin KaFai Lau 	stg = bpf_sk_storage_get(&sk_stg_map, sk, 0,
4253632e11SMartin KaFai Lau 				 BPF_SK_STORAGE_GET_F_CREATE);
4353632e11SMartin KaFai Lau 	if (!stg)
4453632e11SMartin KaFai Lau 		return 0;
4553632e11SMartin KaFai Lau 
4653632e11SMartin KaFai Lau 	stg->last_notclose_state = newstate;
4753632e11SMartin KaFai Lau 
4853632e11SMartin KaFai Lau 	bpf_sk_storage_delete(&del_sk_stg_map, sk);
4953632e11SMartin KaFai Lau 
5053632e11SMartin KaFai Lau 	return 0;
5153632e11SMartin KaFai Lau }
5253632e11SMartin KaFai Lau 
set_task_info(struct sock * sk)5353632e11SMartin KaFai Lau static void set_task_info(struct sock *sk)
5453632e11SMartin KaFai Lau {
5553632e11SMartin KaFai Lau 	struct task_struct *task;
5653632e11SMartin KaFai Lau 	struct sk_stg *stg;
5753632e11SMartin KaFai Lau 
5853632e11SMartin KaFai Lau 	stg = bpf_sk_storage_get(&sk_stg_map, sk, 0,
5953632e11SMartin KaFai Lau 				 BPF_SK_STORAGE_GET_F_CREATE);
6053632e11SMartin KaFai Lau 	if (!stg)
6153632e11SMartin KaFai Lau 		return;
6253632e11SMartin KaFai Lau 
6353632e11SMartin KaFai Lau 	stg->pid = bpf_get_current_pid_tgid();
6453632e11SMartin KaFai Lau 
6553632e11SMartin KaFai Lau 	task = (struct task_struct *)bpf_get_current_task();
6653632e11SMartin KaFai Lau 	bpf_core_read_str(&stg->comm, sizeof(stg->comm), &task->comm);
6753632e11SMartin KaFai Lau 	bpf_core_read_str(&task_comm, sizeof(task_comm), &task->comm);
6853632e11SMartin KaFai Lau }
6953632e11SMartin KaFai Lau 
7053632e11SMartin KaFai Lau SEC("fentry/inet_csk_listen_start")
BPF_PROG(trace_inet_csk_listen_start,struct sock * sk)71e7049395SKuniyuki Iwashima int BPF_PROG(trace_inet_csk_listen_start, struct sock *sk)
7253632e11SMartin KaFai Lau {
7353632e11SMartin KaFai Lau 	set_task_info(sk);
7453632e11SMartin KaFai Lau 
7553632e11SMartin KaFai Lau 	return 0;
7653632e11SMartin KaFai Lau }
7753632e11SMartin KaFai Lau 
7853632e11SMartin KaFai Lau SEC("fentry/tcp_connect")
BPF_PROG(trace_tcp_connect,struct sock * sk)7953632e11SMartin KaFai Lau int BPF_PROG(trace_tcp_connect, struct sock *sk)
8053632e11SMartin KaFai Lau {
8153632e11SMartin KaFai Lau 	set_task_info(sk);
8253632e11SMartin KaFai Lau 
8353632e11SMartin KaFai Lau 	return 0;
8453632e11SMartin KaFai Lau }
8553632e11SMartin KaFai Lau 
8653632e11SMartin KaFai Lau SEC("fexit/inet_csk_accept")
BPF_PROG(inet_csk_accept,struct sock * sk,int flags,int * err,bool kern,struct sock * accepted_sk)8753632e11SMartin KaFai Lau int BPF_PROG(inet_csk_accept, struct sock *sk, int flags, int *err, bool kern,
8853632e11SMartin KaFai Lau 	     struct sock *accepted_sk)
8953632e11SMartin KaFai Lau {
9053632e11SMartin KaFai Lau 	set_task_info(accepted_sk);
9153632e11SMartin KaFai Lau 
9253632e11SMartin KaFai Lau 	return 0;
9353632e11SMartin KaFai Lau }
9453632e11SMartin KaFai Lau 
95*69f41a78SAlexei Starovoitov SEC("tp_btf/tcp_retransmit_synack")
BPF_PROG(tcp_retransmit_synack,struct sock * sk,struct request_sock * req)96*69f41a78SAlexei Starovoitov int BPF_PROG(tcp_retransmit_synack, struct sock* sk, struct request_sock* req)
97*69f41a78SAlexei Starovoitov {
98*69f41a78SAlexei Starovoitov 	/* load only test */
99*69f41a78SAlexei Starovoitov 	bpf_sk_storage_get(&sk_stg_map, sk, 0, 0);
100*69f41a78SAlexei Starovoitov 	bpf_sk_storage_get(&sk_stg_map, req->sk, 0, 0);
101*69f41a78SAlexei Starovoitov 	return 0;
102*69f41a78SAlexei Starovoitov }
103*69f41a78SAlexei Starovoitov 
104*69f41a78SAlexei Starovoitov SEC("tp_btf/tcp_bad_csum")
BPF_PROG(tcp_bad_csum,struct sk_buff * skb)105*69f41a78SAlexei Starovoitov int BPF_PROG(tcp_bad_csum, struct sk_buff* skb)
106*69f41a78SAlexei Starovoitov {
107*69f41a78SAlexei Starovoitov 	bpf_sk_storage_get(&sk_stg_map, skb->sk, 0, 0);
108*69f41a78SAlexei Starovoitov 	return 0;
109*69f41a78SAlexei Starovoitov }
110*69f41a78SAlexei Starovoitov 
11153632e11SMartin KaFai Lau char _license[] SEC("license") = "GPL";
112