diff options
Diffstat (limited to 'tools/testing/selftests/bpf/progs/socket_cookie_prog.c')
| -rw-r--r-- | tools/testing/selftests/bpf/progs/socket_cookie_prog.c | 46 |
1 files changed, 28 insertions, 18 deletions
diff --git a/tools/testing/selftests/bpf/progs/socket_cookie_prog.c b/tools/testing/selftests/bpf/progs/socket_cookie_prog.c index 9ff8ac4b0bf6..e4440fdd94cb 100644 --- a/tools/testing/selftests/bpf/progs/socket_cookie_prog.c +++ b/tools/testing/selftests/bpf/progs/socket_cookie_prog.c @@ -7,25 +7,33 @@ #include "bpf_helpers.h" #include "bpf_endian.h" -struct bpf_map_def SEC("maps") socket_cookies = { - .type = BPF_MAP_TYPE_HASH, - .key_size = sizeof(__u64), - .value_size = sizeof(__u32), - .max_entries = 1 << 8, +struct socket_cookie { + __u64 cookie_key; + __u32 cookie_value; }; +struct { + __uint(type, BPF_MAP_TYPE_SK_STORAGE); + __uint(map_flags, BPF_F_NO_PREALLOC); + __type(key, int); + __type(value, struct socket_cookie); +} socket_cookies SEC(".maps"); + SEC("cgroup/connect6") int set_cookie(struct bpf_sock_addr *ctx) { - __u32 cookie_value = 0xFF; - __u64 cookie_key; + struct socket_cookie *p; if (ctx->family != AF_INET6 || ctx->user_family != AF_INET6) return 1; - cookie_key = bpf_get_socket_cookie(ctx); - if (bpf_map_update_elem(&socket_cookies, &cookie_key, &cookie_value, 0)) - return 0; + p = bpf_sk_storage_get(&socket_cookies, ctx->sk, 0, + BPF_SK_STORAGE_GET_F_CREATE); + if (!p) + return 1; + + p->cookie_value = 0xFF; + p->cookie_key = bpf_get_socket_cookie(ctx); return 1; } @@ -33,9 +41,8 @@ int set_cookie(struct bpf_sock_addr *ctx) SEC("sockops") int update_cookie(struct bpf_sock_ops *ctx) { - __u32 new_cookie_value; - __u32 *cookie_value; - __u64 cookie_key; + struct bpf_sock *sk; + struct socket_cookie *p; if (ctx->family != AF_INET6) return 1; @@ -43,14 +50,17 @@ int update_cookie(struct bpf_sock_ops *ctx) if (ctx->op != BPF_SOCK_OPS_TCP_CONNECT_CB) return 1; - cookie_key = bpf_get_socket_cookie(ctx); + if (!ctx->sk) + return 1; + + p = bpf_sk_storage_get(&socket_cookies, ctx->sk, 0, 0); + if (!p) + return 1; - cookie_value = bpf_map_lookup_elem(&socket_cookies, &cookie_key); - if (!cookie_value) + if (p->cookie_key != bpf_get_socket_cookie(ctx)) return 1; - new_cookie_value = (ctx->local_port << 8) | *cookie_value; - bpf_map_update_elem(&socket_cookies, &cookie_key, &new_cookie_value, 0); + p->cookie_value = (ctx->local_port << 8) | p->cookie_value; return 1; } |