[PATCH 1/3] io_uring: move ctx->restrictions to be dynamically allocated

[Jens Axboe <axboe@kernel.dk>]

In preparation for being able to share restrictions, move them to
allocated data rather than be embedded in the ring. This makes it more
obvious when they will have potentially different lifetimes.

Signed-off-by: Jens Axboe <axboe@kernel.dk>
---
 include/linux/io_uring_types.h |  4 +++-
 io_uring/io_uring.c            | 12 ++++++-----
 io_uring/register.c            | 37 +++++++++++++++++++++++++++-------
 io_uring/register.h            |  2 ++
 4 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 211686ad89fd..c664c84247f1 100644
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -220,6 +220,7 @@ struct io_rings {
 };
 
 struct io_restriction {
+	refcount_t refs;
 	DECLARE_BITMAP(register_op, IORING_REGISTER_LAST);
 	DECLARE_BITMAP(sqe_op, IORING_OP_LAST);
 	u8 sqe_flags_allowed;
@@ -342,6 +343,8 @@ struct io_ring_ctx {
 		struct io_alloc_cache	rw_cache;
 		struct io_alloc_cache	cmd_cache;
 
+		struct io_restriction	*restrictions;
+
 		/*
 		 * Any cancelable uring_cmd is added to this list in
 		 * ->uring_cmd() by io_uring_cmd_insert_cancelable()
@@ -413,7 +416,6 @@ struct io_ring_ctx {
 
 	/* Keep this last, we don't need it for the fast path */
 	struct wait_queue_head		poll_wq;
-	struct io_restriction		restrictions;
 
 	/* Stores zcrx object pointers of type struct io_zcrx_ifq */
 	struct xarray			zcrx_ctxs;
diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index 2cde22af78a3..eec8da38a596 100644
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -2058,15 +2058,15 @@ static inline bool io_check_restriction(struct io_ring_ctx *ctx,
 {
 	if (!ctx->op_restricted)
 		return true;
-	if (!test_bit(req->opcode, ctx->restrictions.sqe_op))
+	if (!test_bit(req->opcode, ctx->restrictions->sqe_op))
 		return false;
 
-	if ((sqe_flags & ctx->restrictions.sqe_flags_required) !=
-	    ctx->restrictions.sqe_flags_required)
+	if ((sqe_flags & ctx->restrictions->sqe_flags_required) !=
+	    ctx->restrictions->sqe_flags_required)
 		return false;
 
-	if (sqe_flags & ~(ctx->restrictions.sqe_flags_allowed |
-			  ctx->restrictions.sqe_flags_required))
+	if (sqe_flags & ~(ctx->restrictions->sqe_flags_allowed |
+			  ctx->restrictions->sqe_flags_required))
 		return false;
 
 	return true;
@@ -2850,6 +2850,8 @@ static __cold void io_ring_ctx_free(struct io_ring_ctx *ctx)
 	percpu_ref_exit(&ctx->refs);
 	free_uid(ctx->user);
 	io_req_caches_free(ctx);
+	if (ctx->restrictions)
+		io_put_restrictions(ctx->restrictions);
 
 	WARN_ON_ONCE(ctx->nr_req_allocated);
 
diff --git a/io_uring/register.c b/io_uring/register.c
index 8551f13920dc..6c99b441d886 100644
--- a/io_uring/register.c
+++ b/io_uring/register.c
@@ -163,9 +163,28 @@ static __cold int io_parse_restrictions(void __user *arg, unsigned int nr_args,
 	return ret;
 }
 
+void io_put_restrictions(struct io_restriction *res)
+{
+	if (refcount_dec_and_test(&res->refs))
+		kfree(res);
+}
+
+static struct io_restriction *io_alloc_restrictions(void)
+{
+	struct io_restriction *res;
+
+	res = kzalloc(sizeof(*res), GFP_KERNEL);
+	if (!res)
+		return ERR_PTR(-ENOMEM);
+
+	refcount_set(&res->refs, 1);
+	return res;
+}
+
 static __cold int io_register_restrictions(struct io_ring_ctx *ctx,
 					   void __user *arg, unsigned int nr_args)
 {
+	struct io_restriction *res;
 	int ret;
 
 	/* Restrictions allowed only if rings started disabled */
@@ -173,19 +192,23 @@ static __cold int io_register_restrictions(struct io_ring_ctx *ctx,
 		return -EBADFD;
 
 	/* We allow only a single restrictions registration */
-	if (ctx->restrictions.op_registered || ctx->restrictions.reg_registered)
+	if (ctx->restrictions)
 		return -EBUSY;
 
-	ret = io_parse_restrictions(arg, nr_args, &ctx->restrictions);
-	/* Reset all restrictions if an error happened */
+	res = io_alloc_restrictions();
+	if (IS_ERR(res))
+		return PTR_ERR(res);
+
+	ret = io_parse_restrictions(arg, nr_args, res);
 	if (ret < 0) {
-		memset(&ctx->restrictions, 0, sizeof(ctx->restrictions));
+		io_put_restrictions(res);
 		return ret;
 	}
-	if (ctx->restrictions.op_registered)
+	if (res->op_registered)
 		ctx->op_restricted = 1;
-	if (ctx->restrictions.reg_registered)
+	if (res->reg_registered)
 		ctx->reg_restricted = 1;
+	ctx->restrictions = res;
 	return 0;
 }
 
@@ -637,7 +660,7 @@ static int __io_uring_register(struct io_ring_ctx *ctx, unsigned opcode,
 
 	if (ctx->reg_restricted && !(ctx->flags & IORING_SETUP_R_DISABLED)) {
 		opcode = array_index_nospec(opcode, IORING_REGISTER_LAST);
-		if (!test_bit(opcode, ctx->restrictions.register_op))
+		if (!test_bit(opcode, ctx->restrictions->register_op))
 			return -EACCES;
 	}
 
diff --git a/io_uring/register.h b/io_uring/register.h
index a5f39d5ef9e0..99c59894d163 100644
--- a/io_uring/register.h
+++ b/io_uring/register.h
@@ -6,4 +6,6 @@ int io_eventfd_unregister(struct io_ring_ctx *ctx);
 int io_unregister_personality(struct io_ring_ctx *ctx, unsigned id);
 struct file *io_uring_register_get_file(unsigned int fd, bool registered);
 
+void io_put_restrictions(struct io_restriction *res);
+
 #endif
-- 
2.51.0