[PATCH] io_uring: fix cloned compound buffer accounting and R_DISABLED restriction bypass

[Hui Peng <benquike@gmail.com>]

Fix two issues in io_uring buffer registration and restriction
enforcement:

1. In io_uring/rsrc.c, when registered compound buffers are cloned
   across rings via IORING_REGISTER_BUFFERS2 /
   IORING_RSRC_REGISTER_SPARSE, unaccounting on release can underflow
   mm->pinned_vm and user->locked_vm if head pages are unaccounted
   multiple times or against a different accounting context. Track per-
   imu accounting ownership cleanly.
2. In io_uring/register.c, enforce IO_RING_F_REG_RESTRICTED on rings
   created with IORING_SETUP_R_DISABLED so restricted opcodes cannot be
   invoked before restrictions are registered and enabled.

Fixes: 735729844819 ("io_uring: move rsrc related data, core, and commands")
Fixes: c43203154d8a ("io_uring/register: move io_uring_register(2) related code to register.c")
Assisted-by: LLM
Signed-off-by: Hui Peng <benquike@gmail.com>
---
diff --git a/io_uring/register.c b/io_uring/register.c
index 02bc103bcc9d..ad6f2a3c98a0 100644
--- a/io_uring/register.c
+++ b/io_uring/register.c
@@ -764,7 +764,7 @@ static int __io_uring_register(struct io_ring_ctx *ctx, unsigned opcode,
 	if (ctx->submitter_task && ctx->submitter_task != current)
 		return -EEXIST;
 
-	if ((ctx->int_flags & IO_RING_F_REG_RESTRICTED) && !(ctx->flags & IORING_SETUP_R_DISABLED)) {
+	if (ctx->int_flags & IO_RING_F_REG_RESTRICTED) {
 		opcode = array_index_nospec(opcode, IORING_REGISTER_LAST);
 		if (!test_bit(opcode, ctx->restrictions.register_op))
 			return -EACCES;
diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index 51b46e624ddd..1efaf29514e8 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -174,8 +174,8 @@ static void io_free_imu(struct io_ring_ctx *ctx, struct io_mapped_ubuf *imu)
 		kvfree(imu);
 }
 
-static unsigned long io_buffer_unaccount_pages(struct io_ring_ctx *ctx,
-					       struct io_mapped_ubuf *imu)
+static unsigned long io_imu_unaccount_hpages(struct io_ring_ctx *ctx,
+					     struct io_mapped_ubuf *imu)
 {
 	struct page *seen = NULL;
 	unsigned long acct = 0;
@@ -188,17 +188,14 @@ static unsigned long io_buffer_unaccount_pages(struct io_ring_ctx *ctx,
 		struct page *page = imu->bvec[i].bv_page;
 		struct page *hpage;
 
-		if (!PageCompound(page)) {
-			acct++;
+		if (!PageCompound(page))
 			continue;
-		}
 
 		hpage = compound_head(page);
 		if (hpage == seen)
 			continue;
 		seen = hpage;
 
-		/* Unaccount on last reference */
 		if (hpage_acct_unref(ctx, hpage))
 			acct += page_size(hpage) >> PAGE_SHIFT;
 		cond_resched();
@@ -207,18 +204,38 @@ static unsigned long io_buffer_unaccount_pages(struct io_ring_ctx *ctx,
 	return acct;
 }
 
+static unsigned long io_imu_unaccount_reg_pages(struct io_ring_ctx *ctx,
+						struct io_mapped_ubuf *imu)
+{
+	unsigned long acct = 0;
+	int i;
+
+	if (imu->flags & IO_REGBUF_F_KBUF || !ctx->user)
+		return 0;
+
+	for (i = 0; i < imu->nr_bvecs; i++) {
+		if (!PageCompound(imu->bvec[i].bv_page))
+			acct++;
+	}
+	return acct;
+}
+
 static void io_buffer_unmap(struct io_ring_ctx *ctx, struct io_mapped_ubuf *imu)
 {
-	unsigned long acct_pages = 0;
+	unsigned long acct_pages;
 
-	/* Always decrement, so it works for cloned buffers too */
-	acct_pages = io_buffer_unaccount_pages(ctx, imu);
+	/* Compound hpages are accounted per-ring in ctx->hpage_acct */
+	acct_pages = io_imu_unaccount_hpages(ctx, imu);
 
 	if (unlikely(refcount_read(&imu->refs) > 1)) {
-		if (!refcount_dec_and_test(&imu->refs))
+		if (!refcount_dec_and_test(&imu->refs)) {
+			if (acct_pages)
+				io_unaccount_mem(ctx->user, ctx->mm_account, acct_pages);
 			return;
+		}
 	}
 
+	acct_pages += io_imu_unaccount_reg_pages(ctx, imu);
 	if (acct_pages)
 		io_unaccount_mem(ctx->user, ctx->mm_account, acct_pages);
 	imu->release(imu->priv);
@@ -1280,6 +1297,7 @@ static int io_buffer_acct_cloned_hpages(struct io_ring_ctx *ctx,
 					struct io_mapped_ubuf *imu)
 {
 	struct page *seen = NULL;
+	unsigned long acct = 0;
 	int i, ret = 0;
 
 	if (imu->flags & IO_REGBUF_F_KBUF || !ctx->user)
@@ -1302,10 +1320,14 @@ static int io_buffer_acct_cloned_hpages(struct io_ring_ctx *ctx,
 		ret = hpage_acct_ref(ctx, hpage, &acct_new);
 		if (ret)
 			break;
+		if (acct_new)
+			acct += page_size(hpage) >> PAGE_SHIFT;
 
 		cond_resched();
 	}
 
+	if (!ret && acct)
+		ret = io_account_mem(ctx->user, ctx->mm_account, acct);
 	if (!ret)
 		return 0;