[PATCH bpf-next v2 2/2] selftests/bpf: Test RCU iterator state pruning

[Ning Ding <dingning04@gmail.com>]

Add a path where RCU protection reaches zero and then starts again.
The iterator is untrusted after this gap and must be rejected.

Signed-off-by: Ning Ding <dingning04@gmail.com>
---
 .../selftests/bpf/progs/iters_task_failure.c  | 46 +++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tools/testing/selftests/bpf/progs/iters_task_failure.c b/tools/testing/selftests/bpf/progs/iters_task_failure.c
index fe3663dedbe14..566a1d3dffeab 100644
--- a/tools/testing/selftests/bpf/progs/iters_task_failure.c
+++ b/tools/testing/selftests/bpf/progs/iters_task_failure.c
@@ -61,6 +61,52 @@ int BPF_PROG(iter_tasks_lock_and_unlock)
 	return 0;
 }
 
+SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
+__failure __msg("expected an RCU CS when using bpf_iter_task_next")
+__flag(BPF_F_TEST_STATE_FREQ)
+int BPF_PROG(iter_tasks_rcu_state_pruning)
+{
+	struct bpf_iter_task it;
+
+	asm volatile (
+		"call %[bpf_rcu_read_lock];"
+		"r1 = %[it];"
+		"r2 = 0;"
+		"r3 = 0;" /* BPF_TASK_ITER_ALL_PROCS */
+		"call %[bpf_iter_task_new];"
+
+		"call %[bpf_get_prandom_u32];"
+		"if w0 == 0 goto unprotected_%=;"
+
+		/* Keep the outer RCU lock active on the straight-line path. */
+		"call %[bpf_rcu_read_lock];"
+		"call %[bpf_rcu_read_unlock];"
+		"goto merge_%=;"
+
+	"unprotected_%=:"
+		/* Create an unprotected gap on the taken path. */
+		"call %[bpf_rcu_read_unlock];"
+		"call %[bpf_rcu_read_lock];"
+
+	"merge_%=: r1 = %[it];"
+		"call %[bpf_iter_task_next];"
+		"r1 = %[it];"
+		"call %[bpf_iter_task_destroy];"
+		"call %[bpf_rcu_read_unlock];"
+		:
+		: __imm_ptr(it),
+		  __imm(bpf_get_prandom_u32),
+		  __imm(bpf_iter_task_new),
+		  __imm(bpf_iter_task_next),
+		  __imm(bpf_iter_task_destroy),
+		  __imm(bpf_rcu_read_lock),
+		  __imm(bpf_rcu_read_unlock)
+		: __clobber_common
+	);
+
+	return 0;
+}
+
 SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
 __failure __msg("expected an RCU CS when using bpf_iter_css_next")
 int BPF_PROG(iter_css_lock_and_unlock)
-- 
2.43.0