[PATCH] sched/mmcid: Defer affinity updates until schedule-in

[Weiming Shi <bestswngs@gmail.com>]

set_cpus_allowed_common() passes a remote task's p->mm to
mm_update_cpus_allowed(), but the scheduler locks do not stabilize that
pointer. A same-owner sched_setaffinity() can race the target replacing and
freeing its old mm during exec. Current kernels are only incidentally
serialized by futex_exec_done() taking p->pi_lock when CONFIG_FUTEX=y;
valid CONFIG_FUTEX=n builds remain vulnerable.

  BUG: KASAN: slab-use-after-free in set_cpus_allowed_common (kernel/sched/core.c:10835 kernel/sched/core.c:2829)
  Read of size 4 at addr ffff88810991b04c by task poc/5127
  Call Trace:
   set_cpus_allowed_common (kernel/sched/core.c:10835 kernel/sched/core.c:2829)
   set_cpus_allowed_fair (kernel/sched/fair.c:10046)
   __set_cpus_allowed_ptr_locked (kernel/sched/core.c:2842 kernel/sched/core.c:3224)
   __set_cpus_allowed_ptr (kernel/sched/core.c:3258)
   __sched_setaffinity (kernel/sched/syscalls.c:1161)
   sched_setaffinity (kernel/sched/syscalls.c:1239)
   __x64_sys_sched_setaffinity (kernel/sched/syscalls.c:1275)
   do_syscall_64 (arch/x86/entry/syscall_64.c:84)
   entry_SYSCALL_64_after_hwframe (arch/x86/entry/entry_64.S:121)

  Freed by task 5128:
   kmem_cache_free (mm/slub.c:6757)
   __mmput (kernel/fork.c:1204)
   mmput (kernel/fork.c:1215)
   setup_new_exec (fs/exec.c:1395)
   load_elf_binary (fs/binfmt_elf.c:1024)
   bprm_execve (fs/exec.c:1865)
   do_execveat_common.isra.0 (fs/exec.c:1963)
   __x64_sys_execve (fs/exec.c:2036)
   do_syscall_64 (arch/x86/entry/syscall_64.c:84)
   entry_SYSCALL_64_after_hwframe (arch/x86/entry/entry_64.S:121)

  The buggy address belongs to the cache mm_struct of size 2776
  The buggy address is located 588 bytes inside of a freed object
  Kernel panic - not syncing: KASAN: panic_on_warn set ...

Defer the MM-CID allowed-mask update to the task's next schedule-in, where
its current mm is stable. Synchronize a forked task's shared mm before its
first run so a pending affinity expansion inherited from its parent is not
lost. Protect the fork-time mask snapshot with the task's runqueue lock
because the task is already published before sched_post_fork() runs.

This change was developed with assistance from an LLM.

Fixes: 7e019dcc470f ("sched: Improve cache locality of RSEQ concurrency IDs for intermittent workloads")
Cc: stable@vger.kernel.org
Reported-by: <co+cab19e2ebe4d535a@bugs.sh>
Signed-off-by: Weiming Shi <bestswngs@gmail.com>
---
 include/linux/rseq_types.h |  3 +++
 kernel/fork.c              |  1 +
 kernel/sched/core.c        |  9 ++++-----
 kernel/sched/sched.h       | 12 ++++++++++++
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/include/linux/rseq_types.h b/include/linux/rseq_types.h
index 85739a63e85e..1479ed807add 100644
--- a/include/linux/rseq_types.h
+++ b/include/linux/rseq_types.h
@@ -140,12 +140,15 @@ struct rseq_data { };
 /**
  * struct sched_mm_cid - Storage for per task MM CID data
  * @active:	MM CID is active for the task
+ * @cpus_allowed_changed: Task affinity change needs propagation to the MM;
+ *			 protected by the task's runqueue lock
  * @cid:	The CID associated to the task either permanently or
  *		borrowed from the CPU
  * @node:	Queued in the per MM MMCID list
  */
 struct sched_mm_cid {
 	unsigned int		active;
+	bool			cpus_allowed_changed;
 	unsigned int		cid;
 	struct hlist_node	node;
 };
diff --git a/kernel/fork.c b/kernel/fork.c
index 8b108ed37152..ed5947ae10fd 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -1013,6 +1013,7 @@ static struct task_struct *dup_task_struct(struct task_struct *orig, int node)
 #ifdef CONFIG_SCHED_MM_CID
 	tsk->mm_cid.cid = MM_CID_UNSET;
 	tsk->mm_cid.active = 0;
+	tsk->mm_cid.cpus_allowed_changed = false;
 	INIT_HLIST_NODE(&tsk->mm_cid.node);
 #endif
 
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 04b62c16fa8c..935c7efd349e 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2811,8 +2811,6 @@ int push_cpu_stop(void *arg)
 	return 0;
 }
 
-static inline void mm_update_cpus_allowed(struct mm_struct *mm, const cpumask_t *affmask);
-
 /*
  * sched_class::set_cpus_allowed must do the below, but is not required to
  * actually call this function.
@@ -2826,7 +2824,7 @@ void set_cpus_allowed_common(struct task_struct *p, struct affinity_context *ctx
 
 	cpumask_copy(&p->cpus_mask, ctx->new_mask);
 	p->nr_cpus_allowed = cpumask_weight(ctx->new_mask);
-	mm_update_cpus_allowed(p->mm, ctx->new_mask);
+	mm_cid_notify_affinity_change(p);
 
 	/*
 	 * Swap in a new user_cpus_ptr if SCA_USER flag set
@@ -10826,7 +10824,7 @@ static bool mm_update_max_cids(struct mm_struct *mm)
 	return true;
 }
 
-static inline void mm_update_cpus_allowed(struct mm_struct *mm, const struct cpumask *affmsk)
+void mm_update_cpus_allowed(struct mm_struct *mm, const struct cpumask *affmsk)
 {
 	struct cpumask *mm_allowed;
 	struct mm_mm_cid *mc;
@@ -10986,6 +10984,8 @@ static void sched_mm_cid_fork(struct task_struct *t)
 
 	if (!mm)
 		return;
+	scoped_guard(task_rq_lock, t)
+		mm_update_cpus_allowed(mm, &t->cpus_mask);
 
 	WARN_ON_ONCE(t->mm_cid.cid != MM_CID_UNSET);
 
@@ -11180,7 +11180,6 @@ void mm_init_cid(struct mm_struct *mm, struct task_struct *p)
 	bitmap_zero(mm_cidmask(mm), num_possible_cpus());
 }
 #else /* CONFIG_SCHED_MM_CID */
-static inline void mm_update_cpus_allowed(struct mm_struct *mm, const struct cpumask *affmsk) { }
 static inline void sched_mm_cid_fork(struct task_struct *t) { }
 #endif /* !CONFIG_SCHED_MM_CID */
 
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 6c7c5633524a..6be52141068b 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -3964,6 +3964,13 @@ extern const char *preempt_modes[];
 
 #ifdef CONFIG_SCHED_MM_CID
 
+void mm_update_cpus_allowed(struct mm_struct *mm, const struct cpumask *affmsk);
+
+static inline void mm_cid_notify_affinity_change(struct task_struct *p)
+{
+	p->mm_cid.cpus_allowed_changed = true;
+}
+
 static __always_inline bool cid_on_cpu(unsigned int cid)
 {
 	return cid & MM_CID_ONCPU;
@@ -4152,6 +4159,10 @@ static __always_inline void mm_cid_schedin(struct task_struct *next)
 
 	if (!next->mm_cid.active)
 		return;
+	if (unlikely(next->mm_cid.cpus_allowed_changed)) {
+		mm_update_cpus_allowed(mm, &next->cpus_mask);
+		next->mm_cid.cpus_allowed_changed = false;
+	}
 
 	cpu_cid = __this_cpu_read(mm->mm_cid.pcpu->cid);
 	mode = READ_ONCE(mm->mm_cid.mode);
@@ -4197,6 +4208,7 @@ static inline void mm_cid_switch_to(struct task_struct *prev, struct task_struct
 }
 
 #else /* !CONFIG_SCHED_MM_CID: */
+static inline void mm_cid_notify_affinity_change(struct task_struct *p) { }
 static inline void mm_cid_switch_to(struct task_struct *prev, struct task_struct *next) { }
 #endif /* !CONFIG_SCHED_MM_CID */
 
-- 
2.56.0