tools/testing/selftests/bpf/progs/test_sk_storage_tracing.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 /* Copyright (c) 2020 Facebook */
   3
   4 #include <vmlinux.h>
   5 #include <bpf/bpf_tracing.h>
   6 #include <bpf/bpf_core_read.h>
   7 #include <bpf/bpf_helpers.h>
   8
   9 struct sk_stg {
  10         __u32 pid;
  11         __u32 last_notclose_state;
  12         char comm[16];
  13 };
  14
  15 struct {
  16         __uint(type, BPF_MAP_TYPE_SK_STORAGE);
  17         __uint(map_flags, BPF_F_NO_PREALLOC);
  18         __type(key, int);
  19         __type(value, struct sk_stg);
  20 } sk_stg_map SEC(".maps");
  21
  22 /* Testing delete */
  23 struct {
  24         __uint(type, BPF_MAP_TYPE_SK_STORAGE);
  25         __uint(map_flags, BPF_F_NO_PREALLOC);
  26         __type(key, int);
  27         __type(value, int);
  28 } del_sk_stg_map SEC(".maps");
  29
  30 char task_comm[16] = "";
  31
  32 SEC("tp_btf/inet_sock_set_state")
  33 int BPF_PROG(trace_inet_sock_set_state, struct sock *sk, int oldstate,
  34              int newstate)
  35 {
  36         struct sk_stg *stg;
  37
  38         if (newstate == BPF_TCP_CLOSE)
  39                 return 0;
  40
  41         stg = bpf_sk_storage_get(&sk_stg_map, sk, 0,
  42                                  BPF_SK_STORAGE_GET_F_CREATE);
  43         if (!stg)
  44                 return 0;
  45
  46         stg->last_notclose_state = newstate;
  47
  48         bpf_sk_storage_delete(&del_sk_stg_map, sk);
  49
  50         return 0;
  51 }
  52
  53 static void set_task_info(struct sock *sk)
  54 {
  55         struct task_struct *task;
  56         struct sk_stg *stg;
  57
  58         stg = bpf_sk_storage_get(&sk_stg_map, sk, 0,
  59                                  BPF_SK_STORAGE_GET_F_CREATE);
  60         if (!stg)
  61                 return;
  62
  63         stg->pid = bpf_get_current_pid_tgid();
  64
  65         task = (struct task_struct *)bpf_get_current_task();
  66         bpf_core_read_str(&stg->comm, sizeof(stg->comm), &task->comm);
  67         bpf_core_read_str(&task_comm, sizeof(task_comm), &task->comm);
  68 }
  69
  70 SEC("fentry/inet_csk_listen_start")
  71 int BPF_PROG(trace_inet_csk_listen_start, struct sock *sk, int backlog)
  72 {
  73         set_task_info(sk);
  74
  75         return 0;
  76 }
  77
  78 SEC("fentry/tcp_connect")
  79 int BPF_PROG(trace_tcp_connect, struct sock *sk)
  80 {
  81         set_task_info(sk);
  82
  83         return 0;
  84 }
  85
  86 SEC("fexit/inet_csk_accept")
  87 int BPF_PROG(inet_csk_accept, struct sock *sk, int flags, int *err, bool kern,
  88              struct sock *accepted_sk)
  89 {
  90         set_task_info(accepted_sk);
  91
  92         return 0;
  93 }
  94
  95 char _license[] SEC("license") = "GPL";