[PATCH bpf v2 2/2] selftests/bpf: Test that a drained css iterator keeps returning NULL

[Kang Chen]

Add a syscall program that walks the current cgroup's css descendants,
drains the iterator and calls bpf_iter_css_next() once more.  Store the
raw result in a global so the verifier cannot fold the call to zero.

The test holds an RCU read lock and releases the cgroup reference taken by
bpf_cgroup_from_id().  The user-space test checks that the constructor
succeeded, that at least one css was drained, and that the extra next()
returned NULL.

Signed-off-by: Kang Chen <k.chen@xxxxxxxxxxxxxxxx>
---
v2:
- New selftest in v2.

 .../selftests/bpf/prog_tests/css_iter.c       | 35 +++++++++++++
 tools/testing/selftests/bpf/progs/css_iter.c  | 49 +++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/css_iter.c
 create mode 100644 tools/testing/selftests/bpf/progs/css_iter.c

diff --git a/tools/testing/selftests/bpf/prog_tests/css_iter.c b/tools/testing/selftests/bpf/prog_tests/css_iter.c
new file mode 100644
index 000000000000..9bf8f056d3cd
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/css_iter.c
@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <test_progs.h>
+#include <bpf/libbpf.h>
+#include "css_iter.skel.h"
+
+static void test_css_iter_next_after_drain(void)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, topts);
+	struct css_iter *skel;
+	int err, fd;
+
+	skel = css_iter__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "css_iter__open_and_load"))
+		return;
+
+	fd = bpf_program__fd(skel->progs.iter_css_next_after_drain);
+	err = bpf_prog_test_run_opts(fd, &topts);
+	if (!ASSERT_OK(err, "test_run_opts err"))
+		goto out;
+	if (!ASSERT_OK(topts.retval, "test_run_opts retval"))
+		goto out;
+
+	ASSERT_OK(skel->bss->new_ret, "new_ret");
+	ASSERT_GT(skel->bss->css_count, 0, "css_count");
+	ASSERT_EQ(skel->data->next_after_drain, 0, "next_after_drain");
+
+out:
+	css_iter__destroy(skel);
+}
+
+void test_css_iter(void)
+{
+	if (test__start_subtest("next_after_drain"))
+		test_css_iter_next_after_drain();
+}
diff --git a/tools/testing/selftests/bpf/progs/css_iter.c b/tools/testing/selftests/bpf/progs/css_iter.c
new file mode 100644
index 000000000000..2ebac638bf78
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/css_iter.c
@@ -0,0 +1,49 @@
+// SPDX-License-Identifier: GPL-2.0
+#include "vmlinux.h"
+#include "bpf_experimental.h"
+#include <bpf/bpf_helpers.h>
+
+char _license[] SEC("license") = "GPL";
+
+/* Result of a bpf_iter_css_next() call after the iterator was drained */
+__u64 next_after_drain = -1;
+int new_ret;
+__u64 css_count;
+
+struct cgroup *bpf_cgroup_from_id(u64 cgid) __ksym;
+void bpf_cgroup_release(struct cgroup *p) __ksym;
+void bpf_rcu_read_lock(void) __ksym;
+void bpf_rcu_read_unlock(void) __ksym;
+
+SEC("syscall")
+int iter_css_next_after_drain(const void *ctx)
+{
+	u64 cg_id = bpf_get_current_cgroup_id();
+	struct cgroup *cgrp = bpf_cgroup_from_id(cg_id);
+	struct cgroup_subsys_state *pos;
+	struct bpf_iter_css it;
+
+	if (!cgrp) {
+		new_ret = -1;
+		return 0;
+	}
+
+	bpf_rcu_read_lock();
+	new_ret = bpf_iter_css_new(&it, &cgrp->self,
+				   BPF_CGROUP_ITER_DESCENDANTS_PRE);
+	if (new_ret)
+		goto out_destroy;
+
+	while ((pos = bpf_iter_css_next(&it)))
+		css_count++;
+
+	/* Store the raw value; the verifier would remove a branch on it */
+	next_after_drain = (__u64)bpf_iter_css_next(&it);
+
+out_destroy:
+	/* The destructor is called even when the constructor failed. */
+	bpf_iter_css_destroy(&it);
+	bpf_rcu_read_unlock();
+	bpf_cgroup_release(cgrp);
+	return 0;
+}
-- 
2.34.1