[PATCH 7.2.y 1/2] io_uring/memmap: account the pages a compound region really uses

[Sasha Levin <sashal@kernel.org>]

From: Ali Ahmet Memis <ali@iusegentoo.com>

[ Upstream commit f12f0234cc14886bcfd53ffb7c8df4216dad51f1 ]

io_mem_alloc_compound() allocates get_order(size) pages, which rounds a
region size that is not a power of two up to the next order. The pages
past the region are part of the same allocation and cannot be used for
anything else, but io_create_region() accounts reg->size >> PAGE_SHIFT,
so they are never charged against RLIMIT_MEMLOCK. For a ring with 4096
SQ entries and the default CQ size the region is 37 pages while the
allocation is 64.

Account the tail pages together with the region, and fall back to the
exact sized bulk allocation when they do not fit the limit, so a user
close to their limit still gets the region rather than an error.
io_free_region() gives the same amount back, the compound case being the
one that set IO_REGION_F_SINGLE_REF.

Counting how many 4096 entry rings an unprivileged user can create under
a given RLIMIT_MEMLOCK, before and after:

	limit (pages)	before	after
	256		2	2
	300		2	2
	350		3	2
	400		3	3
	512		5	4

Five rings under a 512 page limit really pin 640 pages.

Fixes: dfbbfbf19187 ("io_uring: introduce concept of memory regions")
Link: https://lore.kernel.org/all/87ik5ncj8d.fsf@mailhost.krisman.be/
Signed-off-by: Ali Ahmet Memis <ali@iusegentoo.com>
Link: https://patch.msgid.link/20260806180044.275543-1-ali@iusegentoo.com
Signed-off-by: Jens Axboe <axboe@kernel.dk>
Stable-dep-of: b59cabf9b4ac ("io_uring: do not charge the SQ/CQ rings to RLIMIT_MEMLOCK")
Signed-off-by: Sasha Levin <sashal@kernel.org>
---
 io_uring/memmap.c | 37 ++++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/io_uring/memmap.c b/io_uring/memmap.c
index 23e8a85111bca..48c0eb0124123 100644
--- a/io_uring/memmap.c
+++ b/io_uring/memmap.c
@@ -16,8 +16,10 @@
 #include "zcrx.h"
 
 static bool io_mem_alloc_compound(struct page **pages, int nr_pages,
-				  size_t size, gfp_t gfp)
+				  size_t size, gfp_t gfp,
+				  struct user_struct *user)
 {
+	unsigned long nr_compound, extra;
 	struct page *page;
 	int i, order;
 
@@ -27,9 +29,22 @@ static bool io_mem_alloc_compound(struct page **pages, int nr_pages,
 	else if (order)
 		gfp |= __GFP_COMP;
 
+	/*
+	 * get_order() rounds a non power of two size up, so the allocation
+	 * can hold more pages than the region exposes. Account those too,
+	 * and leave the compound allocation alone if they do not fit.
+	 */
+	nr_compound = 1UL << order;
+	extra = nr_compound - nr_pages;
+	if (extra && user && __io_account_mem(user, extra))
+		return false;
+
 	page = alloc_pages(gfp, order);
-	if (!page)
+	if (!page) {
+		if (extra && user)
+			__io_unaccount_mem(user, extra);
 		return false;
+	}
 
 	for (i = 0; i < nr_pages; i++)
 		pages[i] = page + i;
@@ -105,8 +120,15 @@ void io_free_region(struct user_struct *user, struct io_mapped_region *mr)
 	}
 	if ((mr->flags & IO_REGION_F_VMAP) && mr->ptr)
 		vunmap(mr->ptr);
-	if (mr->nr_pages && user)
-		__io_unaccount_mem(user, mr->nr_pages);
+	if (mr->nr_pages && user) {
+		unsigned long nr_accounted = mr->nr_pages;
+
+		/* a compound region was accounted for the whole allocation */
+		if (mr->flags & IO_REGION_F_SINGLE_REF)
+			nr_accounted = 1UL << get_order(io_region_size(mr));
+
+		__io_unaccount_mem(user, nr_accounted);
+	}
 
 	memset(mr, 0, sizeof(*mr));
 }
@@ -151,7 +173,8 @@ static int io_region_pin_pages(struct io_mapped_region *mr,
 
 static int io_region_allocate_pages(struct io_mapped_region *mr,
 				    struct io_uring_region_desc *reg,
-				    unsigned long mmap_offset)
+				    unsigned long mmap_offset,
+				    struct user_struct *user)
 {
 	gfp_t gfp = GFP_KERNEL_ACCOUNT | __GFP_ZERO | __GFP_NOWARN;
 	size_t size = io_region_size(mr);
@@ -162,7 +185,7 @@ static int io_region_allocate_pages(struct io_mapped_region *mr,
 	if (!pages)
 		return -ENOMEM;
 
-	if (io_mem_alloc_compound(pages, mr->nr_pages, size, gfp)) {
+	if (io_mem_alloc_compound(pages, mr->nr_pages, size, gfp, user)) {
 		mr->flags |= IO_REGION_F_SINGLE_REF;
 		goto done;
 	}
@@ -217,7 +240,7 @@ int io_create_region(struct io_ring_ctx *ctx, struct io_mapped_region *mr,
 	if (reg->flags & IORING_MEM_REGION_TYPE_USER)
 		ret = io_region_pin_pages(mr, reg);
 	else
-		ret = io_region_allocate_pages(mr, reg, mmap_offset);
+		ret = io_region_allocate_pages(mr, reg, mmap_offset, ctx->user);
 	if (ret)
 		goto out_free;
 
-- 
2.53.0