github.com/cilium/ebpf@v0.15.1-0.20240517100537-8079b37aa138/examples/tcprtt_sockops/tcprtt_sockops.c (about)

     1  //go:build ignore
     2  
     3  #include "common.h"
     4  
     5  #include "bpf_endian.h"
     6  #include "bpf_sockops.h"
     7  #include "bpf_tracing.h"
     8  
     9  #define AF_INET 2
    10  #define SOCKOPS_MAP_SIZE 65535
    11  
    12  char __license[] SEC("license") = "Dual MIT/GPL";
    13  
    14  enum {
    15  	SOCK_TYPE_ACTIVE  = 0,
    16  	SOCK_TYPE_PASSIVE = 1,
    17  };
    18  
    19  struct {
    20  	__uint(type, BPF_MAP_TYPE_HASH);
    21  	__uint(max_entries, SOCKOPS_MAP_SIZE);
    22  	__type(key, struct sk_key);
    23  	__type(value, struct sk_info);
    24  } map_estab_sk SEC(".maps");
    25  
    26  struct sk_key {
    27  	u32 local_ip4;
    28  	u32 remote_ip4;
    29  	u32 local_port;
    30  	u32 remote_port;
    31  };
    32  
    33  struct sk_info {
    34  	struct sk_key sk_key;
    35  	u8 sk_type;
    36  };
    37  
    38  struct {
    39  	__uint(type, BPF_MAP_TYPE_RINGBUF);
    40  	__uint(max_entries, 1 << 24);
    41  } rtt_events SEC(".maps");
    42  
    43  struct rtt_event {
    44  	u16 sport;
    45  	u16 dport;
    46  	u32 saddr;
    47  	u32 daddr;
    48  	u32 srtt;
    49  };
    50  struct rtt_event *unused_event __attribute__((unused));
    51  
    52  static inline void init_sk_key(struct bpf_sock_ops *skops, struct sk_key *sk_key) {
    53  	sk_key->local_ip4   = bpf_ntohl(skops->local_ip4);
    54  	sk_key->remote_ip4  = bpf_ntohl(skops->remote_ip4);
    55  	sk_key->local_port  = skops->local_port;
    56  	sk_key->remote_port = bpf_ntohl(skops->remote_port);
    57  }
    58  
    59  static inline void bpf_sock_ops_establish_cb(struct bpf_sock_ops *skops, u8 sock_type) {
    60  	int err;
    61  	struct sk_info sk_info = {};
    62  	// Only process IPv4 sockets
    63  	if (skops == NULL || skops->family != AF_INET)
    64  		return;
    65  
    66  	// Initialize the 4-tuple key
    67  	init_sk_key(skops, &sk_info.sk_key);
    68  	sk_info.sk_type = sock_type;
    69  
    70  	// Store the socket info in map using the 4-tuple as key
    71  	// We keep track of TCP connections in 'established' state
    72  	err = bpf_map_update_elem(&map_estab_sk, &sk_info.sk_key, &sk_info, BPF_NOEXIST);
    73  	if (err != 0) {
    74  		// Storing the 4-tuple in map has failed, return early.
    75  		// This can happen in case the 4-tuple already exists in the map (i.e. BPF_NOEXIST flag)
    76  		return;
    77  	}
    78  
    79  	// Enable sockops callbacks for RTT and TCP state change
    80  	bpf_sock_ops_cb_flags_set(skops, BPF_SOCK_OPS_RTT_CB_FLAG | BPF_SOCK_OPS_STATE_CB_FLAG);
    81  }
    82  
    83  static inline void bpf_sock_ops_rtt_cb(struct bpf_sock_ops *skops) {
    84  	struct sk_key sk_key = {};
    85  	struct sk_info *sk_info;
    86  	struct rtt_event *rtt_event;
    87  
    88  	// Initialize the 4-tuple key
    89  	init_sk_key(skops, &sk_key);
    90  
    91  	// Retrieve the socket info from map of established connections
    92  	sk_info = bpf_map_lookup_elem(&map_estab_sk, &sk_key);
    93  	if (!sk_info)
    94  		return;
    95  
    96  	rtt_event = bpf_ringbuf_reserve(&rtt_events, sizeof(struct rtt_event), 0);
    97  	if (!rtt_event) {
    98  		return;
    99  	}
   100  
   101  	switch (sk_info->sk_type) {
   102  	case SOCK_TYPE_ACTIVE:
   103  		// If socket is 'active', 'local' means 'source'
   104  		// and 'remote' means 'destination'
   105  		rtt_event->saddr = sk_info->sk_key.local_ip4;
   106  		rtt_event->daddr = sk_info->sk_key.remote_ip4;
   107  		rtt_event->sport = sk_info->sk_key.local_port;
   108  		rtt_event->dport = sk_info->sk_key.remote_port;
   109  		break;
   110  	case SOCK_TYPE_PASSIVE:
   111  		// If socket is 'passive', 'local' means 'destination'
   112  		// and 'remote' means 'source'
   113  		rtt_event->saddr = sk_info->sk_key.remote_ip4;
   114  		rtt_event->daddr = sk_info->sk_key.local_ip4;
   115  		rtt_event->sport = sk_info->sk_key.remote_port;
   116  		rtt_event->dport = sk_info->sk_key.local_port;
   117  		break;
   118  	}
   119  
   120  	// Extract smoothed RTT
   121  	rtt_event->srtt = skops->srtt_us >> 3;
   122  	rtt_event->srtt /= 1000;
   123  
   124  	// Send RTT event data to userspace app via ring buffer
   125  	bpf_ringbuf_submit(rtt_event, 0);
   126  }
   127  
   128  static inline void bpf_sock_ops_state_cb(struct bpf_sock_ops *skops) {
   129  	struct sk_key sk_key = {};
   130  
   131  	// Socket changed state. args[0] stores the previous state.
   132  	// Perform cleanup of map entry if socket is exiting
   133  	// the 'established' state,
   134  	if (skops->args[0] == TCP_ESTABLISHED) {
   135  		init_sk_key(skops, &sk_key);
   136  		bpf_map_delete_elem(&map_estab_sk, &sk_key);
   137  	}
   138  }
   139  
   140  SEC("sockops")
   141  int bpf_sockops_cb(struct bpf_sock_ops *skops) {
   142  	u32 op;
   143  	op = skops->op;
   144  
   145  	switch (op) {
   146  	case BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB:
   147  		bpf_sock_ops_establish_cb(skops, SOCK_TYPE_ACTIVE);
   148  		break;
   149  	case BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB:
   150  		bpf_sock_ops_establish_cb(skops, SOCK_TYPE_PASSIVE);
   151  		break;
   152  	case BPF_SOCK_OPS_RTT_CB:
   153  		bpf_sock_ops_rtt_cb(skops);
   154  		break;
   155  	case BPF_SOCK_OPS_STATE_CB:
   156  		bpf_sock_ops_state_cb(skops);
   157  		break;
   158  	}
   159  
   160  	return 0;
   161  }