bpf/progs/test_migrate_reuseport.c

0001 // SPDX-License-Identifier: GPL-2.0
0002 /*
0003  * Check if we can migrate child sockets.
0004  *
0005  *   1. If reuse_md->migrating_sk is NULL (SYN packet),
0006  *        return SK_PASS without selecting a listener.
0007  *   2. If reuse_md->migrating_sk is not NULL (socket migration),
0008  *        select a listener (reuseport_map[migrate_map[cookie]])
0009  *
0010  * Author: Kuniyuki Iwashima <kuniyu@amazon.co.jp>
0011  */
0012
0013 #include <stddef.h>
0014 #include <string.h>
0015 #include <linux/bpf.h>
0016 #include <linux/if_ether.h>
0017 #include <linux/ip.h>
0018 #include <linux/ipv6.h>
0019 #include <linux/tcp.h>
0020 #include <linux/in.h>
0021 #include <bpf/bpf_endian.h>
0022 #include <bpf/bpf_helpers.h>
0023
0024 struct {
0025     __uint(type, BPF_MAP_TYPE_REUSEPORT_SOCKARRAY);
0026     __uint(max_entries, 256);
0027     __type(key, int);
0028     __type(value, __u64);
0029 } reuseport_map SEC(".maps");
0030
0031 struct {
0032     __uint(type, BPF_MAP_TYPE_HASH);
0033     __uint(max_entries, 256);
0034     __type(key, __u64);
0035     __type(value, int);
0036 } migrate_map SEC(".maps");
0037
0038 int migrated_at_close = 0;
0039 int migrated_at_close_fastopen = 0;
0040 int migrated_at_send_synack = 0;
0041 int migrated_at_recv_ack = 0;
0042 __be16 server_port;
0043
0044 SEC("xdp")
0045 int drop_ack(struct xdp_md *xdp)
0046 {
0047     void *data_end = (void *)(long)xdp->data_end;
0048     void *data = (void *)(long)xdp->data;
0049     struct ethhdr *eth = data;
0050     struct tcphdr *tcp = NULL;
0051
0052     if (eth + 1 > data_end)
0053         goto pass;
0054
0055     switch (bpf_ntohs(eth->h_proto)) {
0056     case ETH_P_IP: {
0057         struct iphdr *ip = (struct iphdr *)(eth + 1);
0058
0059         if (ip + 1 > data_end)
0060             goto pass;
0061
0062         if (ip->protocol != IPPROTO_TCP)
0063             goto pass;
0064
0065         tcp = (struct tcphdr *)((void *)ip + ip->ihl * 4);
0066         break;
0067     }
0068     case ETH_P_IPV6: {
0069         struct ipv6hdr *ipv6 = (struct ipv6hdr *)(eth + 1);
0070
0071         if (ipv6 + 1 > data_end)
0072             goto pass;
0073
0074         if (ipv6->nexthdr != IPPROTO_TCP)
0075             goto pass;
0076
0077         tcp = (struct tcphdr *)(ipv6 + 1);
0078         break;
0079     }
0080     default:
0081         goto pass;
0082     }
0083
0084     if (tcp + 1 > data_end)
0085         goto pass;
0086
0087     if (tcp->dest != server_port)
0088         goto pass;
0089
0090     if (!tcp->syn && tcp->ack)
0091         return XDP_DROP;
0092
0093 pass:
0094     return XDP_PASS;
0095 }
0096
0097 SEC("sk_reuseport/migrate")
0098 int migrate_reuseport(struct sk_reuseport_md *reuse_md)
0099 {
0100     int *key, flags = 0, state, err;
0101     __u64 cookie;
0102
0103     if (!reuse_md->migrating_sk)
0104         return SK_PASS;
0105
0106     state = reuse_md->migrating_sk->state;
0107     cookie = bpf_get_socket_cookie(reuse_md->sk);
0108
0109     key = bpf_map_lookup_elem(&migrate_map, &cookie);
0110     if (!key)
0111         return SK_DROP;
0112
0113     err = bpf_sk_select_reuseport(reuse_md, &reuseport_map, key, flags);
0114     if (err)
0115         return SK_PASS;
0116
0117     switch (state) {
0118     case BPF_TCP_ESTABLISHED:
0119         __sync_fetch_and_add(&migrated_at_close, 1);
0120         break;
0121     case BPF_TCP_SYN_RECV:
0122         __sync_fetch_and_add(&migrated_at_close_fastopen, 1);
0123         break;
0124     case BPF_TCP_NEW_SYN_RECV:
0125         if (!reuse_md->len)
0126             __sync_fetch_and_add(&migrated_at_send_synack, 1);
0127         else
0128             __sync_fetch_and_add(&migrated_at_recv_ack, 1);
0129         break;
0130     }
0131
0132     return SK_PASS;
0133 }
0134
0135 char _license[] SEC("license") = "GPL";