summaryrefslogtreecommitdiffstats
path: root/tools/testing/selftests/bpf/progs/socket_cookie_prog.c
diff options
context:
space:
mode:
Diffstat (limited to 'tools/testing/selftests/bpf/progs/socket_cookie_prog.c')
-rw-r--r--tools/testing/selftests/bpf/progs/socket_cookie_prog.c49
1 files changed, 31 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..6aabb681fb9a 100644
--- a/tools/testing/selftests/bpf/progs/socket_cookie_prog.c
+++ b/tools/testing/selftests/bpf/progs/socket_cookie_prog.c
@@ -7,25 +7,36 @@
#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 {
+ __u32 type;
+ __u32 map_flags;
+ int *key;
+ struct socket_cookie *value;
+} socket_cookies SEC(".maps") = {
+ .type = BPF_MAP_TYPE_SK_STORAGE,
+ .map_flags = BPF_F_NO_PREALLOC,
};
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 +44,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 +53,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;
}