[PATCH 2/2] io_uring/register: add support for inheriting task restrictions

[Jens Axboe <axboe@kernel.dk>]

By default, the registered task restrictions only apply to the task they
were registered for. Any forked tasks or created threads will not
inherit them.

However, if IORING_REG_RESTRICTIONS_INHERIT is set when registering the
task restrictions, then they will be inherited across process fork or
thread creation.

Signed-off-by: Jens Axboe <axboe@kernel.dk>
---
 include/linux/io_uring_types.h | 2 ++
 include/uapi/linux/io_uring.h  | 7 +++++++
 io_uring/register.c            | 5 ++++-
 io_uring/tctx.c                | 4 +++-
 kernel/fork.c                  | 7 ++++++-
 5 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 54fd30abf2b8..b63b927d8718 100644
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -222,9 +222,11 @@ struct io_rings {
 struct io_restriction {
 	DECLARE_BITMAP(register_op, IORING_REGISTER_LAST);
 	DECLARE_BITMAP(sqe_op, IORING_OP_LAST);
+	refcount_t refs;
 	u8 sqe_flags_allowed;
 	u8 sqe_flags_required;
 	bool registered;
+	bool inherited;
 };
 
 struct io_submit_link {
diff --git a/include/uapi/linux/io_uring.h b/include/uapi/linux/io_uring.h
index 3ecf9c1bfa2d..8d671b5e33e3 100644
--- a/include/uapi/linux/io_uring.h
+++ b/include/uapi/linux/io_uring.h
@@ -807,6 +807,13 @@ struct io_uring_restriction {
 	__u32 resv2[3];
 };
 
+enum {
+	/*
+	 * Registered restrictions are inherited for a fork.
+	 */
+	IORING_REG_RESTRICTIONS_INHERIT	= (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 eac7a6da32b4..36573b362225 100644
--- a/io_uring/register.c
+++ b/io_uring/register.c
@@ -191,7 +191,7 @@ static int io_register_restrictions_task(void __user *arg, unsigned int nr_args)
 	if (copy_from_user(&tres, arg, sizeof(tres)))
 		return -EFAULT;
 
-	if (tres.flags)
+	if (tres.flags & ~IORING_REG_RESTRICTIONS_INHERIT)
 		return -EINVAL;
 	if (!mem_is_zero(tres.resv, sizeof(tres.resv)))
 		return -EINVAL;
@@ -205,6 +205,9 @@ static int io_register_restrictions_task(void __user *arg, unsigned int nr_args)
 		kfree(res);
 		return ret;
 	}
+	if (tres.flags & IORING_REG_RESTRICTIONS_INHERIT)
+		res->inherited = true;
+	refcount_set(&res->refs, 1);
 	current->io_uring_restrict = res;
 	return 0;
 }
diff --git a/io_uring/tctx.c b/io_uring/tctx.c
index c8ad735936dc..f9ad9cbee9be 100644
--- a/io_uring/tctx.c
+++ b/io_uring/tctx.c
@@ -66,7 +66,9 @@ void __io_uring_free(struct task_struct *tsk)
 		kfree(tctx);
 		tsk->io_uring = NULL;
 	}
-	kfree(tsk->io_uring_restrict);
+	if (tsk->io_uring_restrict &&
+	    refcount_dec_and_test(&tsk->io_uring_restrict->refs))
+		kfree(tsk->io_uring_restrict);
 }
 
 __cold int io_uring_alloc_task_context(struct task_struct *task,
diff --git a/kernel/fork.c b/kernel/fork.c
index 6081e1c93e21..505f9397a645 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -97,6 +97,7 @@
 #include <linux/kasan.h>
 #include <linux/scs.h>
 #include <linux/io_uring.h>
+#include <linux/io_uring_types.h>
 #include <linux/bpf.h>
 #include <linux/stackprotector.h>
 #include <linux/user_events.h>
@@ -2129,7 +2130,10 @@ __latent_entropy struct task_struct *copy_process(
 
 #ifdef CONFIG_IO_URING
 	p->io_uring = NULL;
-	p->io_uring_restrict = NULL;
+	if (p->io_uring_restrict && p->io_uring_restrict->inherited)
+		refcount_inc(&p->io_uring_restrict->refs);
+	else
+		p->io_uring_restrict = NULL;
 #endif
 
 	p->default_timer_slack_ns = current->timer_slack_ns;
@@ -2526,6 +2530,7 @@ __latent_entropy struct task_struct *copy_process(
 	mpol_put(p->mempolicy);
 #endif
 bad_fork_cleanup_delayacct:
+	io_uring_free(p);
 	delayacct_tsk_free(p);
 bad_fork_cleanup_count:
 	dec_rlimit_ucounts(task_ucounts(p), UCOUNT_RLIMIT_NPROC, 1);
-- 
2.51.0