[PATCH] mm/memcg: fix UAF in drain_all_stock() async work during offline

[Rik van Riel]

drain_all_stock() queues drain work on remote CPUs via
schedule_drain_work() -> queue_work_on(memcg_wq) and returns
immediately without waiting. The worker, drain_local_memcg_stock()
/ drain_local_obj_stock(), dereferences per-CPU stock caches with
READ_ONCE(stock->cached[i]) and does css_put() / obj_cgroup_put().

mem_cgroup_css_offline() calls drain_all_stock(memcg) to
optimize reclamation latency, but never flushes memcg_wq. If
that races with cgroup removal, free can happen while workers
are still pending, causing UAF. The drain work could also have
been queued by somebody else before offline started (e.g. high
throttling), not just by the offline path itself.

Timeline illustrating the race:

  CPU0 (rmdir + offline)                CPU1 (charge cache holder)
  -------------------------             ----------------------------
  cgroup_rmdir()
    cgroup_destroy_locked()
      kill_css_sync()
      ...

                                      refill_stock(victim)
                                        css_get(victim)
                                        WRITE_ONCE(cached[i]=victim)

  percpu_ref kill confirmed, css_killed_ref_fn() called

  css_killed_work_fn() [offline_wq]
    mem_cgroup_css_offline(victim)
      drain_all_stock(victim)
        is_memcg_drain_needed()
          READ_ONCE(cached) -> victim
        queue_work_on(CPU1, memcg_wq, work)
        // no flush!
      mem_cgroup_private_id_put()
    css_put() -> refcnt may hit 0

  [RCU GP]
  css_free_rwork_fn()
    mem_cgroup_free(victim)
    // victim struct freed

                                        // worker delayed by scheduler/
                                        // WQ concurrency
                                        drain_local_memcg_stock()
                                          old = READ_ONCE(cached[i])
                                          // UAF: old == freed victim
                                          memcg_uncharge(old)
                                          css_put(&old->css)

Fix by having the offline path wait for the workqueue to be
done with the memcg, before freeing the memcg.

Found through a code audit with kres.

Fixes: 591edfb10a94 ("mm: drain memcg stocks on css offlining")
Cc: stable@xxxxxxxxxxxxxxx
Assisted-by: Hermes:muse-spark-1.2 kres
Signed-off-by: Rik van Riel <riel@xxxxxxxxxxx>
---
 mm/memcontrol.c | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 6dc4888a90f3..c95a1f6ec799 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2273,12 +2273,18 @@ static void schedule_drain_work(int cpu, struct work_struct *work)
  * Drains all per-CPU charge caches for given root_memcg resp. subtree
  * of the hierarchy under it.
  */
-void drain_all_stock(struct mem_cgroup *root_memcg)
+static void __drain_all_stock(struct mem_cgroup *root_memcg, bool sync)
 {
 	int cpu, curcpu;
 
-	/* If someone's already draining, avoid adding running more workers. */
-	if (!mutex_trylock(&percpu_charge_mutex))
+	/*
+	 *  If someone's already draining, avoid starting more workers.
+	 *  Synchronous callers need to guarantee all the last things
+	 *  are flushed, e.g. before a memcg is removed.
+	 */
+	if (sync)
+		mutex_lock(&percpu_charge_mutex);
+	else if (!mutex_trylock(&percpu_charge_mutex))
 		return;
 	/*
 	 * Notify other cpus that system-wide "drain" is running
@@ -2316,6 +2322,21 @@ void drain_all_stock(struct mem_cgroup *root_memcg)
 	mutex_unlock(&percpu_charge_mutex);
 }
 
+void drain_all_stock(struct mem_cgroup *root_memcg)
+{
+	__drain_all_stock(root_memcg, false);
+}
+
+void drain_all_stock_sync(struct mem_cgroup *root_memcg)
+{
+	/*
+	 * Make sure the workqueue is done with this memcg
+	 * before freeing it.
+	 */
+	__drain_all_stock(root_memcg, true);
+	flush_workqueue(memcg_wq);
+}
+
 static int memcg_hotplug_cpu_dead(unsigned int cpu)
 {
 	/* no need for the local lock */
@@ -4305,7 +4326,7 @@ static void mem_cgroup_css_offline(struct cgroup_subsys_state *css)
 	wb_memcg_offline(memcg);
 	lru_gen_offline_memcg(memcg);
 
-	drain_all_stock(memcg);
+	drain_all_stock_sync(memcg);
 
 	mem_cgroup_private_id_put(memcg, 1);
 }
-- 
2.55.0