[PATCH] io_uring/sqpoll: pin task across task-work publication

[Jérémy Jean]

SQPOLL can consume a request immediately after mpscq_push(). If this
drops the final ring reference, ctx and tctx may be freed before
io_req_normal_work_add() finishes using them. KASAN reports:

  BUG: KASAN: slab-use-after-free in io_req_normal_work_add+0x439/0x510
  Read of size 4 at addr ff11000000ca0000 by task repro/55
  (...)
  BUG: KASAN: slab-use-after-free in queue_work_on+0x25/0x70
  Write of size 8 at addr ff11000000c3bd00 by task repro/55

Handle SQPOLL before publication. Pin tctx->task, publish the request,
then use only the pinned task so the publication tail cannot dereference
freed contexts.

Fixes: af5d68f8892f ("io_uring/sqpoll: manage task_work privately")
Assisted-by: LLM
Signed-off-by: Jérémy Jean <Jeremy.Jean@xxxxxxxxxxxxxxxxx>
---
 io_uring/tw.c | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/io_uring/tw.c b/io_uring/tw.c
index f573bcc3af6a..2e37a3fd7a33 100644
--- a/io_uring/tw.c
+++ b/io_uring/tw.c
@@ -210,6 +210,19 @@ void io_req_normal_work_add(struct io_kiocb *req)
 	struct io_uring_task *tctx = req->tctx;
 	struct io_ring_ctx *ctx = req->ctx;
 
+	/* SQPOLL may consume and retire the request immediately after push. */
+	if (ctx->flags & IORING_SETUP_SQPOLL) {
+		struct task_struct *task = tctx->task;
+		bool first;
+
+		get_task_struct(task);
+		first = mpscq_push(&tctx->task_list, &req->io_task_work.node);
+		if (first)
+			__set_notify_signal(task);
+		put_task_struct(task);
+		return;
+	}
+
 	/* tw run already pending, nothing else to do */
 	if (!mpscq_push(&tctx->task_list, &req->io_task_work.node))
 		return;
@@ -221,12 +234,6 @@ void io_req_normal_work_add(struct io_kiocb *req)
 	if (ctx->flags & IORING_SETUP_TASKRUN_FLAG)
 		atomic_or(IORING_SQ_TASKRUN, &ctx->rings->sq_flags);
 
-	/* SQPOLL doesn't need the task_work added, it'll run it itself */
-	if (ctx->flags & IORING_SETUP_SQPOLL) {
-		__set_notify_signal(tctx->task);
-		return;
-	}
-
 	if (likely(!task_work_add(tctx->task, &tctx->task_work, ctx->notify_method)))
 		return;
 
-- 
2.47.3