[PATCH bpf-next v1 14/15] selftests/bpf: Add C tests for PTR_TO_BTF_ID in map

[Kumar Kartikeya Dwivedi <memxor@gmail.com>]

This uses the __kptr* macros as well, and tries to test the stuff that
is supposed to work, since we have negative tests in test_verifier
suite.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../selftests/bpf/prog_tests/map_btf_ptr.c    |  13 +++
 .../testing/selftests/bpf/progs/map_btf_ptr.c | 105 ++++++++++++++++++
 .../testing/selftests/bpf/progs/test_bpf_nf.c |  31 ++++++
 3 files changed, 149 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/map_btf_ptr.c
 create mode 100644 tools/testing/selftests/bpf/progs/map_btf_ptr.c

diff --git a/tools/testing/selftests/bpf/prog_tests/map_btf_ptr.c b/tools/testing/selftests/bpf/prog_tests/map_btf_ptr.c
new file mode 100644
index 000000000000..8fb6acf1b89d
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/map_btf_ptr.c
@@ -0,0 +1,13 @@
+#include <test_progs.h>
+
+#include "map_btf_ptr.skel.h"
+
+void test_map_btf_ptr(void)
+{
+	struct map_btf_ptr *skel;
+
+	skel = map_btf_ptr__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "map_btf_ptr__open_and_load"))
+		return;
+	map_btf_ptr__destroy(skel);
+}
diff --git a/tools/testing/selftests/bpf/progs/map_btf_ptr.c b/tools/testing/selftests/bpf/progs/map_btf_ptr.c
new file mode 100644
index 000000000000..b0c2ba595290
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/map_btf_ptr.c
@@ -0,0 +1,105 @@
+#include <vmlinux.h>
+#include <bpf/bpf_tracing.h>
+#include <bpf/bpf_helpers.h>
+
+#define xchg(dst, src) __sync_lock_test_and_set(&(dst), (src))
+
+struct map_value {
+	struct prog_test_ref_kfunc __kptr *unref_ptr;
+	/* Workarounds for https://lore.kernel.org/bpf/20220220071333.sltv4jrwniool2qy@apollo.legion */
+	struct prog_test_ref_kfunc __kptr __attribute__((btf_type_tag("kernel.bpf.ref"))) *ref_ptr;
+	struct prog_test_ref_kfunc __kptr __attribute__((btf_type_tag("kernel.bpf.percpu"))) *percpu_ptr;
+	struct prog_test_ref_kfunc __kptr __attribute__((btf_type_tag("kernel.bpf.user"))) *user_ptr;
+};
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__type(key, int);
+	__type(value, struct map_value);
+	__uint(max_entries, 1);
+} array_map SEC(".maps");
+
+extern struct prog_test_ref_kfunc *bpf_kfunc_call_test_acquire(unsigned long *sp) __ksym;
+extern void bpf_kfunc_call_test_release(struct prog_test_ref_kfunc *p) __ksym;
+
+SEC("tc")
+int map_btf_ptr(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *p;
+	char buf[sizeof(*p)];
+	struct map_value *v;
+
+	v = bpf_map_lookup_elem(&array_map, &(int){0});
+	if (!v)
+		return 0;
+	p = v->unref_ptr;
+	/* store untrusted_ptr_or_null_ */
+	v->unref_ptr = p;
+	if (!p)
+		return 0;
+	if (p->a + p->b > 100)
+		return 1;
+	/* store untrusted_ptr_ */
+	v->unref_ptr = p;
+	/* store NULL */
+	v->unref_ptr = NULL;
+
+	p = v->ref_ptr;
+	/* store ptr_or_null_ */
+	v->unref_ptr = p;
+	if (!p)
+		return 0;
+	if (p->a + p->b > 100)
+		return 1;
+	/* store NULL */
+	p = xchg(v->ref_ptr, NULL);
+	if (!p)
+		return 0;
+	if (p->a + p->b > 100) {
+		bpf_kfunc_call_test_release(p);
+		return 1;
+	}
+	/* store ptr_ */
+	v->unref_ptr = p;
+	bpf_kfunc_call_test_release(p);
+
+	p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
+	if (!p)
+		return 0;
+	/* store ptr_ */
+	p = xchg(v->ref_ptr, p);
+	if (!p)
+		return 0;
+	if (p->a + p->b > 100) {
+		bpf_kfunc_call_test_release(p);
+		return 1;
+	}
+	bpf_kfunc_call_test_release(p);
+
+	p = v->percpu_ptr;
+	/* store percpu_ptr_or_null_ */
+	v->percpu_ptr = p;
+	if (!p)
+		return 0;
+	p = bpf_this_cpu_ptr(p);
+	if (p->a + p->b > 100)
+		return 1;
+	/* store percpu_ptr_ */
+	v->percpu_ptr = p;
+	/* store NULL */
+	v->percpu_ptr = NULL;
+
+	p = v->user_ptr;
+	/* store user_ptr_or_null_ */
+	v->user_ptr = p;
+	if (!p)
+		return 0;
+	bpf_probe_read_user(buf, sizeof(buf), p);
+	/* store user_ptr_ */
+	v->user_ptr = p;
+	/* store NULL */
+	v->user_ptr = NULL;
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
diff --git a/tools/testing/selftests/bpf/progs/test_bpf_nf.c b/tools/testing/selftests/bpf/progs/test_bpf_nf.c
index f00a9731930e..74e3892be544 100644
--- a/tools/testing/selftests/bpf/progs/test_bpf_nf.c
+++ b/tools/testing/selftests/bpf/progs/test_bpf_nf.c
@@ -30,8 +30,21 @@ struct nf_conn *bpf_xdp_ct_lookup(struct xdp_md *, struct bpf_sock_tuple *, u32,
 				  struct bpf_ct_opts___local *, u32) __ksym;
 struct nf_conn *bpf_skb_ct_lookup(struct __sk_buff *, struct bpf_sock_tuple *, u32,
 				  struct bpf_ct_opts___local *, u32) __ksym;
+struct nf_conn *bpf_ct_kptr_get(struct nf_conn **, struct bpf_sock_tuple *, u32,
+				u8, u8) __ksym;
 void bpf_ct_release(struct nf_conn *) __ksym;
 
+struct nf_map_value {
+	struct nf_conn __kptr_ref *ct;
+};
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__type(key, int);
+	__type(value, struct nf_map_value);
+	__uint(max_entries, 1);
+} array_map SEC(".maps");
+
 static __always_inline void
 nf_ct_test(struct nf_conn *(*func)(void *, struct bpf_sock_tuple *, u32,
 				   struct bpf_ct_opts___local *, u32),
@@ -101,10 +114,27 @@ nf_ct_test(struct nf_conn *(*func)(void *, struct bpf_sock_tuple *, u32,
 		test_eafnosupport = opts_def.error;
 }
 
+static __always_inline void
+nf_ct_test_kptr(void)
+{
+	struct bpf_sock_tuple tuple = {};
+	struct nf_map_value *v;
+	struct nf_conn *ct;
+
+	v = bpf_map_lookup_elem(&array_map, &(int){0});
+	if (!v)
+		return;
+	ct = bpf_ct_kptr_get(&v->ct, &tuple, sizeof(tuple.ipv4), IPPROTO_TCP, IP_CT_DIR_ORIGINAL);
+	if (!ct)
+		return;
+	bpf_ct_release(ct);
+}
+
 SEC("xdp")
 int nf_xdp_ct_test(struct xdp_md *ctx)
 {
 	nf_ct_test((void *)bpf_xdp_ct_lookup, ctx);
+	nf_ct_test_kptr();
 	return 0;
 }
 
@@ -112,6 +142,7 @@ SEC("tc")
 int nf_skb_ct_test(struct __sk_buff *ctx)
 {
 	nf_ct_test((void *)bpf_skb_ct_lookup, ctx);
+	nf_ct_test_kptr();
 	return 0;
 }
 
-- 
2.35.1