[PATCH bpf-next v1 2/2] selftests/bpf: Add tests for preempt kfuncs

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

Add tests for nested cases, nested count preservation upon different
subprog calls that disable/enable preemption, and test sleepable helper
call in non-preemptible regions.

181/1   preempt_lock/preempt_lock_missing_1:OK
181/2   preempt_lock/preempt_lock_missing_2:OK
181/3   preempt_lock/preempt_lock_missing_3:OK
181/4   preempt_lock/preempt_lock_missing_3_minus_2:OK
181/5   preempt_lock/preempt_lock_missing_1_subprog:OK
181/6   preempt_lock/preempt_lock_missing_2_subprog:OK
181/7   preempt_lock/preempt_lock_missing_2_minus_1_subprog:OK
181/8   preempt_lock/preempt_balance:OK
181/9   preempt_lock/preempt_balance_subprog_test:OK
181/10  preempt_lock/preempt_sleepable_helper:OK
181     preempt_lock:OK
Summary: 1/10 PASSED, 0 SKIPPED, 0 FAILED

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../selftests/bpf/prog_tests/preempt_lock.c   |   9 ++
 .../selftests/bpf/progs/preempt_lock.c        | 119 ++++++++++++++++++
 2 files changed, 128 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/preempt_lock.c
 create mode 100644 tools/testing/selftests/bpf/progs/preempt_lock.c

diff --git a/tools/testing/selftests/bpf/prog_tests/preempt_lock.c b/tools/testing/selftests/bpf/prog_tests/preempt_lock.c
new file mode 100644
index 000000000000..02917c672441
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/preempt_lock.c
@@ -0,0 +1,9 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <test_progs.h>
+#include <network_helpers.h>
+#include <preempt_lock.skel.h>
+
+void test_preempt_lock(void)
+{
+	RUN_TESTS(preempt_lock);
+}
diff --git a/tools/testing/selftests/bpf/progs/preempt_lock.c b/tools/testing/selftests/bpf/progs/preempt_lock.c
new file mode 100644
index 000000000000..53320ea80fa4
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/preempt_lock.c
@@ -0,0 +1,119 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <vmlinux.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include "bpf_misc.h"
+
+void bpf_preempt_disable(void) __ksym;
+void bpf_preempt_enable(void) __ksym;
+
+SEC("?tc")
+__failure __msg("1 bpf_preempt_enable is missing")
+int preempt_lock_missing_1(struct __sk_buff *ctx)
+{
+	bpf_preempt_disable();
+	return 0;
+}
+
+SEC("?tc")
+__failure __msg("2 bpf_preempt_enable(s) are missing")
+int preempt_lock_missing_2(struct __sk_buff *ctx)
+{
+	bpf_preempt_disable();
+	bpf_preempt_disable();
+	return 0;
+}
+
+SEC("?tc")
+__failure __msg("3 bpf_preempt_enable(s) are missing")
+int preempt_lock_missing_3(struct __sk_buff *ctx)
+{
+	bpf_preempt_disable();
+	bpf_preempt_disable();
+	bpf_preempt_disable();
+	return 0;
+}
+
+SEC("?tc")
+__failure __msg("1 bpf_preempt_enable is missing")
+int preempt_lock_missing_3_minus_2(struct __sk_buff *ctx)
+{
+	bpf_preempt_disable();
+	bpf_preempt_disable();
+	bpf_preempt_disable();
+	bpf_preempt_enable();
+	bpf_preempt_enable();
+	return 0;
+}
+
+static __noinline void preempt_disable(void)
+{
+	bpf_preempt_disable();
+}
+
+static __noinline void preempt_enable(void)
+{
+	bpf_preempt_enable();
+}
+
+SEC("?tc")
+__failure __msg("1 bpf_preempt_enable is missing")
+int preempt_lock_missing_1_subprog(struct __sk_buff *ctx)
+{
+	preempt_disable();
+	return 0;
+}
+
+SEC("?tc")
+__failure __msg("2 bpf_preempt_enable(s) are missing")
+int preempt_lock_missing_2_subprog(struct __sk_buff *ctx)
+{
+	preempt_disable();
+	preempt_disable();
+	return 0;
+}
+
+SEC("?tc")
+__failure __msg("1 bpf_preempt_enable is missing")
+int preempt_lock_missing_2_minus_1_subprog(struct __sk_buff *ctx)
+{
+	preempt_disable();
+	preempt_disable();
+	preempt_enable();
+	return 0;
+}
+
+static __noinline void preempt_balance_subprog(void)
+{
+	preempt_disable();
+	preempt_enable();
+}
+
+SEC("?tc")
+__success int preempt_balance(struct __sk_buff *ctx)
+{
+	bpf_preempt_disable();
+	bpf_preempt_enable();
+	return 0;
+}
+
+SEC("?tc")
+__success int preempt_balance_subprog_test(struct __sk_buff *ctx)
+{
+	preempt_balance_subprog();
+	return 0;
+}
+
+SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
+__failure __msg("sleepable helper bpf_copy_from_user#")
+int preempt_sleepable_helper(void *ctx)
+{
+	u32 data;
+
+	bpf_preempt_disable();
+	bpf_copy_from_user(&data, sizeof(data), NULL);
+	bpf_preempt_enable();
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.43.0