[PATCH bpf v2 2/2] selftests/bpf: Add raw_tp tests for PTR_MAYBE_NULL marking

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

Ensure that pointers with off != 0 are never unmarked as PTR_MAYBE_NULL
when doing NULL checks, while pointers that have off == 0 continue
getting unmarked, and also unmark associated copies with same id but
possibly non-zero offset.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../selftests/bpf/prog_tests/raw_tp_null.c    |  6 ++
 .../selftests/bpf/progs/raw_tp_null_fail.c    | 80 +++++++++++++++++++
 2 files changed, 86 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/raw_tp_null_fail.c

diff --git a/tools/testing/selftests/bpf/prog_tests/raw_tp_null.c b/tools/testing/selftests/bpf/prog_tests/raw_tp_null.c
index 6fa19449297e..13fcd4c31034 100644
--- a/tools/testing/selftests/bpf/prog_tests/raw_tp_null.c
+++ b/tools/testing/selftests/bpf/prog_tests/raw_tp_null.c
@@ -3,6 +3,12 @@
 
 #include <test_progs.h>
 #include "raw_tp_null.skel.h"
+#include "raw_tp_null_fail.skel.h"
+
+void test_raw_tp_null_fail(void)
+{
+	RUN_TESTS(raw_tp_null_fail);
+}
 
 void test_raw_tp_null(void)
 {
diff --git a/tools/testing/selftests/bpf/progs/raw_tp_null_fail.c b/tools/testing/selftests/bpf/progs/raw_tp_null_fail.c
new file mode 100644
index 000000000000..68de752cfe53
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/raw_tp_null_fail.c
@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */
+
+#include <vmlinux.h>
+#include <bpf/bpf_tracing.h>
+#include "bpf_misc.h"
+
+/* r1 with off=0 is checked, which marks r0 with off=8 as non-null */
+SEC("tp_btf/bpf_testmod_test_raw_tp_null")
+__success
+__log_level(2)
+__msg("3: (07) r0 += 8                       ; R0_w=trusted_ptr_or_null_sk_buff(id=1,off=8)")
+__msg("4: (15) if r1 == 0x0 goto pc+1        ; R1_w=trusted_ptr_sk_buff()")
+__msg("5: (bf) r2 = r0                       ; R0_w=trusted_ptr_sk_buff(off=8)")
+int BPF_PROG(test_raw_tp_null_check_zero_off, struct sk_buff *skb)
+{
+	asm volatile (
+		"r1 = *(u64 *)(r1 +0);			\
+		 r0 = r1;				\
+		 r2 = 0;				\
+		 r0 += 8;				\
+		 if r1 == 0 goto jmp;			\
+		 r2 = r0;				\
+		 jmp:					"
+		::
+		: __clobber_all
+	);
+	return 0;
+}
+
+/* r2 with offset is checked, which won't mark r1 with off=0 as non-NULL */
+SEC("tp_btf/bpf_testmod_test_raw_tp_null")
+__success
+__log_level(2)
+__msg("3: (07) r2 += 8                       ; R2_w=trusted_ptr_or_null_sk_buff(id=1,off=8)")
+__msg("4: (15) if r2 == 0x0 goto pc+1        ; R2_w=trusted_ptr_or_null_sk_buff(id=1,off=8)")
+__msg("5: (bf) r2 = r1                       ; R1_w=trusted_ptr_or_null_sk_buff(id=1)")
+int BPF_PROG(test_raw_tp_null_copy_check_with_off, struct sk_buff *skb)
+{
+	asm volatile (
+		"r1 = *(u64 *)(r1 +0);			\
+		 r2 = r1;				\
+		 r3 = 0;				\
+		 r2 += 8;				\
+		 if r2 == 0 goto jmp2;			\
+		 r2 = r1;				\
+		 jmp2:					"
+		::
+		: __clobber_all
+	);
+	return 0;
+}
+
+/* Ensure state doesn't change for r0 and r1 when performing repeated checks.. */
+SEC("tp_btf/bpf_testmod_test_raw_tp_null")
+__success
+__log_level(2)
+__msg("2: (07) r0 += 8                       ; R0_w=trusted_ptr_or_null_sk_buff(id=1,off=8)")
+__msg("3: (15) if r0 == 0x0 goto pc+3        ; R0_w=trusted_ptr_or_null_sk_buff(id=1,off=8)")
+__msg("4: (15) if r0 == 0x0 goto pc+2        ; R0_w=trusted_ptr_or_null_sk_buff(id=1,off=8)")
+__msg("5: (15) if r0 == 0x0 goto pc+1        ; R0_w=trusted_ptr_or_null_sk_buff(id=1,off=8)")
+__msg("6: (bf) r2 = r1                       ; R1=trusted_ptr_or_null_sk_buff(id=1)")
+int BPF_PROG(test_raw_tp_check_with_off, struct sk_buff *skb)
+{
+	asm volatile (
+		"r1 = *(u64 *)(r1 +0);			\
+		 r0 = r1;				\
+		 r0 += 8;				\
+		 if r0 == 0 goto jmp3;			\
+		 if r0 == 0 goto jmp3;			\
+		 if r0 == 0 goto jmp3;			\
+		 r2 = r1;				\
+		 jmp3:					"
+		::
+		: __clobber_all
+	);
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.43.5