[PATCH] sched_ext: Prevent RB-tree corruption in scx_bpf_task_set_dsq_vtime()

[Cheng-Yang Chou <yphbchou0911@gmail.com>]

scx_bpf_task_set_dsq_vtime() allows modifying a task's dsq_vtime without
checking if it is already enqueued on SCX_DSQ_PRIQ. Since dsq_vtime is
the rb-tree sorting key, mutating it in-place violates the BST invariant
and corrupts the tree structure.

In ops.dispatch():
	p = scx_bpf_dsq_peek(PRIO_DSQ); // Get a task already in the DSQ
	if (p) {
		// This illegally returns %true
		scx_bpf_task_set_dsq_vtime(p, 0xFFFFFFFFFFFFFFFF);
	}

Fix this by adding a check for the SCX_TASK_DSQ_ON_PRIQ flag. Disallow
vtime modification and trigger scx_error() if the task is already queued
on a priority DSQ.

Fixes: 3035addfaf28 ("sched_ext: Add scx_bpf_task_set_slice() and scx_bpf_task_set_dsq_vtime()")
Cc: stable@vger.kernel.org # v6.19+
Signed-off-by: Cheng-Yang Chou <yphbchou0911@gmail.com>
---
 kernel/sched/ext.c | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 012ca8bd70fb..7a54f9bc5e7a 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -8540,7 +8540,8 @@ __bpf_kfunc bool scx_bpf_task_set_slice(struct task_struct *p, u64 slice,
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
  *
  * Set @p's virtual time to @vtime. Returns %true on success, %false if the
- * calling scheduler doesn't have authority over @p.
+ * calling scheduler doesn't have authority over @p. If @p is already enqueued
+ * on a priority DSQ, scx_error() is triggered and %false is returned.
  */
 __bpf_kfunc bool scx_bpf_task_set_dsq_vtime(struct task_struct *p, u64 vtime,
 					    const struct bpf_prog_aux *aux)
@@ -8552,6 +8553,11 @@ __bpf_kfunc bool scx_bpf_task_set_dsq_vtime(struct task_struct *p, u64 vtime,
 	if (unlikely(!scx_task_on_sched(sch, p)))
 		return false;
 
+	if (unlikely(READ_ONCE(p->scx.dsq_flags) & SCX_TASK_DSQ_ON_PRIQ)) {
+		scx_error(sch, "vtime modification disallowed while on a priority DSQ");
+		return false;
+	}
+
 	p->scx.dsq_vtime = vtime;
 	return true;
 }
-- 
2.48.1