[PATCH] selftests/sched_ext: Validate select_cpu_and mask constraints

[Tianyi Chen <hi@tychen.cc>]

The syscall test checks whether BPF_PROG_TEST_RUN succeeds, but only
prints the CPU selection result returned by the BPF program.

Exercise empty and affinity-disjoint custom masks and require -EBUSY.
Also exercise a legal singleton mask, allowing its CPU to be busy.
Interpret retval as signed and reject unexpected errors or selection
outside the custom mask or the task's affinity.

Choose CPUs from the runner's affinity, restore it after the disjoint
case, and skip only that case when fewer than two CPUs are available.

The test passes in a two-CPU VM with full and single-CPU affinity.
Forcing selection of CPU 0 against an empty mask makes it fail with
-ERANGE; the restored test passes again.

Assisted-by: LLM
Signed-off-by: Tianyi Chen <hi@tychen.cc>
---
 .../selftests/sched_ext/allowed_cpus.bpf.c    | 24 +++++-
 .../selftests/sched_ext/allowed_cpus.c        | 77 ++++++++++++++++---
 2 files changed, 86 insertions(+), 15 deletions(-)

diff --git a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
index 9dd72d0da29..08c2fe0e1c3 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
@@ -147,23 +147,41 @@ void BPF_STRUCT_OPS(allowed_cpus_exit, struct scx_exit_info *ei)
 }
 
 struct task_cpu_arg {
-	pid_t pid;
+	u64 pid;
+	s64 custom_cpu;
 };
 
 SEC("syscall")
 int select_cpu_from_user(struct task_cpu_arg *input)
 {
 	struct task_struct *p;
-	int cpu;
+	struct bpf_cpumask *mask;
+	s32 cpu;
 
 	p = bpf_task_from_pid(input->pid);
 	if (!p)
 		return -EINVAL;
 
+	mask = bpf_cpumask_create();
+	if (!mask) {
+		bpf_task_release(p);
+		return -ENOMEM;
+	}
+
+	/* A negative custom_cpu leaves the custom mask empty. */
+	if (input->custom_cpu >= 0)
+		bpf_cpumask_set_cpu(input->custom_cpu, mask);
+
 	bpf_rcu_read_lock();
-	cpu = scx_bpf_select_cpu_and(p, bpf_get_smp_processor_id(), 0, p->cpus_ptr, 0);
+	cpu = scx_bpf_select_cpu_and(p, bpf_get_smp_processor_id(), 0,
+				     cast_mask(mask), 0);
+	if (cpu >= 0 &&
+	    (!bpf_cpumask_test_cpu(cpu, cast_mask(mask)) ||
+	     !bpf_cpumask_test_cpu(cpu, p->cpus_ptr)))
+		cpu = -ERANGE;
 	bpf_rcu_read_unlock();
 
+	bpf_cpumask_release(mask);
 	bpf_task_release(p);
 
 	return cpu;
diff --git a/tools/testing/selftests/sched_ext/allowed_cpus.c b/tools/testing/selftests/sched_ext/allowed_cpus.c
index 093f285ab4b..cc08c42ad95 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.c
@@ -2,7 +2,9 @@
 /*
  * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>
  */
+#define _GNU_SOURCE
 #include <bpf/bpf.h>
+#include <sched.h>
 #include <scx/common.h>
 #include <sys/wait.h>
 #include <unistd.h>
@@ -23,17 +25,19 @@ static enum scx_test_status setup(void **ctx)
 	return SCX_TEST_PASS;
 }
 
-static int test_select_cpu_from_user(const struct allowed_cpus *skel)
+static int test_select_cpu_from_user(const struct allowed_cpus *skel,
+				     const char *name, int custom_cpu,
+				     bool expect_busy)
 {
 	int fd, ret;
-	__u64 args[1];
+	__s32 cpu;
+	__u64 args[] = { getpid(), (__u64)(__s64)custom_cpu };
 
 	LIBBPF_OPTS(bpf_test_run_opts, attr,
 		.ctx_in = args,
 		.ctx_size_in = sizeof(args),
 	);
 
-	args[0] = getpid();
 	fd = bpf_program__fd(skel->progs.select_cpu_from_user);
 	if (fd < 0)
 		return fd;
@@ -42,7 +46,13 @@ static int test_select_cpu_from_user(const struct allowed_cpus *skel)
 	if (ret < 0)
 		return ret;
 
-	fprintf(stderr, "%s: CPU %d\n", __func__, attr.retval);
+	/* test_run returns the signed BPF result through an unsigned field. */
+	cpu = (__s32)attr.retval;
+	if ((expect_busy && cpu != -EBUSY) ||
+	    (!expect_busy && cpu != -EBUSY && cpu != custom_cpu)) {
+		SCX_ERR("%s: unexpected CPU selection result %d", name, cpu);
+		return -EINVAL;
+	}
 
 	return 0;
 }
@@ -50,21 +60,64 @@ static int test_select_cpu_from_user(const struct allowed_cpus *skel)
 static enum scx_test_status run(void *ctx)
 {
 	struct allowed_cpus *skel = ctx;
+	enum scx_test_status status = SCX_TEST_FAIL;
+	cpu_set_t original, pinned;
+	int first = -1, second = -1, cpu;
 	struct bpf_link *link;
 
+	SCX_FAIL_IF(sched_getaffinity(0, sizeof(original), &original),
+		    "Failed to get affinity (%d)", errno);
+	for (cpu = 0; cpu < CPU_SETSIZE; cpu++) {
+		if (!CPU_ISSET(cpu, &original))
+			continue;
+		if (first < 0)
+			first = cpu;
+		else {
+			second = cpu;
+			break;
+		}
+	}
+	SCX_FAIL_IF(first < 0, "No CPU in affinity mask");
+
 	link = bpf_map__attach_struct_ops(skel->maps.allowed_cpus_ops);
 	SCX_FAIL_IF(!link, "Failed to attach scheduler");
 
-	/* Pick an idle CPU from user-space */
-	SCX_FAIL_IF(test_select_cpu_from_user(skel), "Failed to pick idle CPU");
-
-	/* Just sleeping is fine, plenty of scheduling events happening */
+	if (test_select_cpu_from_user(skel, "empty mask", -1, true))
+		goto out;
+
+	/* A legal candidate may be busy; selection need not succeed. */
+	if (test_select_cpu_from_user(skel, "legal candidate", first, false))
+		goto out;
+
+	if (second >= 0) {
+		CPU_ZERO(&pinned);
+		CPU_SET(first, &pinned);
+		if (sched_setaffinity(0, sizeof(pinned), &pinned)) {
+			SCX_ERR("Failed to pin task (%d)", errno);
+			goto out;
+		}
+		if (test_select_cpu_from_user(skel, "disjoint masks", second, true))
+			goto restore;
+	} else {
+		fprintf(stderr, "Skipping disjoint masks: need two allowed CPUs\n");
+	}
+
+	/* Just sleeping is fine, plenty of scheduling events happening. */
 	sleep(1);
-
-	SCX_EQ(skel->data->uei.kind, EXIT_KIND(SCX_EXIT_NONE));
+	if (skel->data->uei.kind != EXIT_KIND(SCX_EXIT_NONE)) {
+		SCX_ERR("Scheduler exited unexpectedly");
+		goto restore;
+	}
+	status = SCX_TEST_PASS;
+
+restore:
+	if (second >= 0 && sched_setaffinity(0, sizeof(original), &original)) {
+		SCX_ERR("Failed to restore affinity (%d)", errno);
+		status = SCX_TEST_FAIL;
+	}
+out:
 	bpf_link__destroy(link);
-
-	return SCX_TEST_PASS;
+	return status;
 }
 
 static void cleanup(void *ctx)
-- 
2.55.0