[PATCH 2/3] io_uring/register: add MASK support for task filter set

[Jens Axboe <axboe@kernel.dk>]

If IORING_REG_RESTRICTIONS_MASK is set in the flags for an
IORING_REGISTER_RESTRICTIONS_TASK operation, then further restrictions
can be added to the current set. No restrictions may be relaxed this
way. If a current set exists, the passed in set is added to the current
set. If no current set exists, the new set applied will as the current
task io_uring restriction filter.

Signed-off-by: Jens Axboe <axboe@kernel.dk>
---
 include/uapi/linux/io_uring.h |  9 ++++++
 io_uring/register.c           | 54 ++++++++++++++++++++++++++---------
 2 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/include/uapi/linux/io_uring.h b/include/uapi/linux/io_uring.h
index 3ecf9c1bfa2d..e39da481f14c 100644
--- a/include/uapi/linux/io_uring.h
+++ b/include/uapi/linux/io_uring.h
@@ -807,6 +807,15 @@ struct io_uring_restriction {
 	__u32 resv2[3];
 };
 
+enum {
+	/*
+	 * MASK operation to further restrict a filter set. Can clear opcodes
+	 * allowed for SQEs or register operations, clear allowed SQE flags,
+	 * and set further required SQE flags.
+	 */
+	IORING_REG_RESTRICTIONS_MASK	= (1U << 0),
+};
+
 struct io_uring_task_restriction {
 	__u16 flags;
 	__u16 nr_res;
diff --git a/io_uring/register.c b/io_uring/register.c
index 89254d0fbe79..552b22f6b2dc 100644
--- a/io_uring/register.c
+++ b/io_uring/register.c
@@ -104,7 +104,8 @@ static int io_register_personality(struct io_ring_ctx *ctx)
 }
 
 static __cold int io_parse_restrictions(void __user *arg, unsigned int nr_args,
-					struct io_restriction *restrictions)
+					struct io_restriction *restrictions,
+					bool mask_it)
 {
 	struct io_uring_restriction *res;
 	size_t size;
@@ -122,32 +123,41 @@ static __cold int io_parse_restrictions(void __user *arg, unsigned int nr_args,
 		return PTR_ERR(res);
 
 	ret = -EINVAL;
-
 	for (i = 0; i < nr_args; i++) {
 		switch (res[i].opcode) {
 		case IORING_RESTRICTION_REGISTER_OP:
 			if (res[i].register_op >= IORING_REGISTER_LAST)
 				goto err;
-			__set_bit(res[i].register_op, restrictions->register_op);
+			if (mask_it)
+				__clear_bit(res[i].register_op, restrictions->register_op);
+			else
+				__set_bit(res[i].register_op, restrictions->register_op);
 			break;
 		case IORING_RESTRICTION_SQE_OP:
 			if (res[i].sqe_op >= IORING_OP_LAST)
 				goto err;
-			__set_bit(res[i].sqe_op, restrictions->sqe_op);
+			if (mask_it)
+				__clear_bit(res[i].sqe_op, restrictions->sqe_op);
+			else
+				__set_bit(res[i].sqe_op, restrictions->sqe_op);
 			break;
 		case IORING_RESTRICTION_SQE_FLAGS_ALLOWED:
-			restrictions->sqe_flags_allowed = res[i].sqe_flags;
+			if (mask_it)
+				restrictions->sqe_flags_allowed &= res[i].sqe_flags;
+			else
+				restrictions->sqe_flags_allowed = res[i].sqe_flags;
 			break;
 		case IORING_RESTRICTION_SQE_FLAGS_REQUIRED:
-			restrictions->sqe_flags_required = res[i].sqe_flags;
+			if (mask_it)
+				restrictions->sqe_flags_required |= res[i].sqe_flags;
+			else
+				restrictions->sqe_flags_required = res[i].sqe_flags;
 			break;
 		default:
 			goto err;
 		}
 	}
-
 	ret = 0;
-
 err:
 	kfree(res);
 	return ret;
@@ -166,7 +176,7 @@ static __cold int io_register_restrictions(struct io_ring_ctx *ctx,
 	if (ctx->restrictions.registered)
 		return -EBUSY;
 
-	ret = io_parse_restrictions(arg, nr_args, &ctx->restrictions);
+	ret = io_parse_restrictions(arg, nr_args, &ctx->restrictions, false);
 	/* Reset all restrictions if an error happened */
 	if (ret != 0)
 		memset(&ctx->restrictions, 0, sizeof(ctx->restrictions));
@@ -182,29 +192,45 @@ static int io_register_restrictions_task(void __user *arg, unsigned int nr_args)
 	struct io_restriction *res;
 	int ret;
 
-	/* Disallow if task already has registered restrictions */
-	if (current->io_uring_restrict)
-		return -EPERM;
 	if (nr_args != 1)
 		return -EINVAL;
 
 	if (copy_from_user(&tres, arg, sizeof(tres)))
 		return -EFAULT;
 
-	if (tres.flags)
+	if (tres.flags & ~IORING_REG_RESTRICTIONS_MASK)
 		return -EINVAL;
 	if (!mem_is_zero(tres.resv, sizeof(tres.resv)))
 		return -EINVAL;
 
+	/*
+	 * Disallow if task already has registered restrictions, and we're
+	 * not passing in further restrictions to add to an existing set.
+	 */
+	if (current->io_uring_restrict &&
+	    !(tres.flags & IORING_REG_RESTRICTIONS_MASK))
+		return -EPERM;
+
 	res = kzalloc(sizeof(*res), GFP_KERNEL);
 	if (!res)
 		return -ENOMEM;
 
-	ret = io_parse_restrictions(ures->restrictions, tres.nr_res, res);
+	/*
+	 * Can only be set if we're MASK'ing in more restrictions. If so,
+	 * copy existing filters.
+	 */
+	if (current->io_uring_restrict)
+		memcpy(res, current->io_uring_restrict, sizeof(*res));
+
+	ret = io_parse_restrictions(ures->restrictions, tres.nr_res, res,
+				    tres.flags & IORING_REG_RESTRICTIONS_MASK);
 	if (ret) {
 		kfree(res);
 		return ret;
 	}
+	if (current->io_uring_restrict &&
+	    refcount_dec_and_test(&current->io_uring_restrict->refs))
+		kfree(current->io_uring_restrict);
 	refcount_set(&res->refs, 1);
 	current->io_uring_restrict = res;
 	return 0;
-- 
2.51.0