bpf/progs/socket_cookie_prog.c

0001 // SPDX-License-Identifier: GPL-2.0
0002 // Copyright (c) 2018 Facebook
0003
0004 #include "vmlinux.h"
0005
0006 #include <bpf/bpf_helpers.h>
0007 #include <bpf/bpf_endian.h>
0008 #include <bpf/bpf_tracing.h>
0009
0010 #define AF_INET6 10
0011
0012 struct socket_cookie {
0013     __u64 cookie_key;
0014     __u32 cookie_value;
0015 };
0016
0017 struct {
0018     __uint(type, BPF_MAP_TYPE_SK_STORAGE);
0019     __uint(map_flags, BPF_F_NO_PREALLOC);
0020     __type(key, int);
0021     __type(value, struct socket_cookie);
0022 } socket_cookies SEC(".maps");
0023
0024 /*
0025  * These three programs get executed in a row on connect() syscalls. The
0026  * userspace side of the test creates a client socket, issues a connect() on it
0027  * and then checks that the local storage associated with this socket has:
0028  * cookie_value == local_port << 8 | 0xFF
0029  * The different parts of this cookie_value are appended by those hooks if they
0030  * all agree on the output of bpf_get_socket_cookie().
0031  */
0032 SEC("cgroup/connect6")
0033 int set_cookie(struct bpf_sock_addr *ctx)
0034 {
0035     struct socket_cookie *p;
0036
0037     if (ctx->family != AF_INET6 || ctx->user_family != AF_INET6)
0038         return 1;
0039
0040     p = bpf_sk_storage_get(&socket_cookies, ctx->sk, 0,
0041                    BPF_SK_STORAGE_GET_F_CREATE);
0042     if (!p)
0043         return 1;
0044
0045     p->cookie_value = 0xF;
0046     p->cookie_key = bpf_get_socket_cookie(ctx);
0047
0048     return 1;
0049 }
0050
0051 SEC("sockops")
0052 int update_cookie_sockops(struct bpf_sock_ops *ctx)
0053 {
0054     struct bpf_sock *sk = ctx->sk;
0055     struct socket_cookie *p;
0056
0057     if (ctx->family != AF_INET6)
0058         return 1;
0059
0060     if (ctx->op != BPF_SOCK_OPS_TCP_CONNECT_CB)
0061         return 1;
0062
0063     if (!sk)
0064         return 1;
0065
0066     p = bpf_sk_storage_get(&socket_cookies, sk, 0, 0);
0067     if (!p)
0068         return 1;
0069
0070     if (p->cookie_key != bpf_get_socket_cookie(ctx))
0071         return 1;
0072
0073     p->cookie_value |= (ctx->local_port << 8);
0074
0075     return 1;
0076 }
0077
0078 SEC("fexit/inet_stream_connect")
0079 int BPF_PROG(update_cookie_tracing, struct socket *sock,
0080          struct sockaddr *uaddr, int addr_len, int flags)
0081 {
0082     struct socket_cookie *p;
0083
0084     if (uaddr->sa_family != AF_INET6)
0085         return 0;
0086
0087     p = bpf_sk_storage_get(&socket_cookies, sock->sk, 0, 0);
0088     if (!p)
0089         return 0;
0090
0091     if (p->cookie_key != bpf_get_socket_cookie(sock->sk))
0092         return 0;
0093
0094     p->cookie_value |= 0xF0;
0095
0096     return 0;
0097 }
0098
0099 char _license[] SEC("license") = "GPL";