From 69d96519dbf0bfa1868dc8597d4b9b2cdeb009d7 Mon Sep 17 00:00:00 2001 From: Stanislav Fomichev Date: Wed, 12 Jun 2019 10:30:40 -0700 Subject: selftests/bpf: convert socket_cookie test to sk storage This lets us test that both BPF_PROG_TYPE_CGROUP_SOCK_ADDR and BPF_PROG_TYPE_SOCK_OPS can access underlying bpf_sock. Cc: Martin Lau Signed-off-by: Stanislav Fomichev Signed-off-by: Daniel Borkmann --- .../selftests/bpf/progs/socket_cookie_prog.c | 46 ++++++++++++++-------- 1 file changed, 29 insertions(+), 17 deletions(-) (limited to 'tools/testing/selftests/bpf/progs/socket_cookie_prog.c') diff --git a/tools/testing/selftests/bpf/progs/socket_cookie_prog.c b/tools/testing/selftests/bpf/progs/socket_cookie_prog.c index 9ff8ac4b0bf6..0db15c3210ad 100644 --- a/tools/testing/selftests/bpf/progs/socket_cookie_prog.c +++ b/tools/testing/selftests/bpf/progs/socket_cookie_prog.c @@ -7,25 +7,35 @@ #include "bpf_helpers.h" #include "bpf_endian.h" +struct socket_cookie { + __u64 cookie_key; + __u32 cookie_value; +}; + 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, + .type = BPF_MAP_TYPE_SK_STORAGE, + .key_size = sizeof(int), + .value_size = sizeof(struct socket_cookie), + .map_flags = BPF_F_NO_PREALLOC, }; +BPF_ANNOTATE_KV_PAIR(socket_cookies, int, struct socket_cookie); + 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 +43,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 +52,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; } -- cgit v1.2.3-59-g8ed1b