Add two new helpers, io_uring_fixed_index_get() and
io_uring_fixed_index_put(). io_uring_fixed_index_get() constructs an
iter for a fixed buffer at a given index and acquires a refcount on
the underlying node. io_uring_fixed_index_put() decrements this
refcount. The caller is responsible for ensuring
io_uring_fixed_index_put() is properly called for releasing the refcount
after it is done using the iter it obtained through
io_uring_fixed_index_get().

The struct io_rsrc_node pointer needs to be returned in
io_uring_fixed_index_get() because the buffer at the index may be
unregistered/replaced in the meantime between this and the
io_uring_fixed_index_put() call. io_uring_fixed_index_put() takes in the
struct io_rsrc_node pointer as an arg.

This is a preparatory patch needed for fuse-over-io-uring support, as
the metadata for fuse requests will be stored at the last index, which
will be different from the buf index set on the sqe.

Signed-off-by: Joanne Koong <joannelkoong@gmail.com>
---
 include/linux/io_uring/cmd.h | 20 ++++++++++++
 io_uring/rsrc.c              | 59 ++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/include/linux/io_uring/cmd.h b/include/linux/io_uring/cmd.h
index a488e945f883..de3f550598cf 100644
--- a/include/linux/io_uring/cmd.h
+++ b/include/linux/io_uring/cmd.h
@@ -44,6 +44,14 @@ int io_uring_cmd_import_fixed_vec(struct io_uring_cmd *ioucmd,
 				  size_t uvec_segs,
 				  int ddir, struct iov_iter *iter,
 				  unsigned issue_flags);
+struct io_rsrc_node *io_uring_fixed_index_get(struct io_uring_cmd *cmd,
+					      int buf_index, unsigned int off,
+					      size_t len, int ddir,
+					      struct iov_iter *iter,
+					      unsigned int issue_flags);
+void io_uring_fixed_index_put(struct io_uring_cmd *cmd,
+			      struct io_rsrc_node *node,
+			      unsigned int issue_flags);
 
 /*
  * Completes the request, i.e. posts an io_uring CQE and deallocates @ioucmd
@@ -108,6 +116,18 @@ static inline int io_uring_cmd_import_fixed_vec(struct io_uring_cmd *ioucmd,
 {
 	return -EOPNOTSUPP;
 }
+static inline struct io_rsrc_node *
+io_uring_fixed_index_get(struct io_uring_cmd *cmd, int buf_index,
+			 unsigned int off, size_t len, int ddir,
+			 struct iov_iter *iter, unsigned int issue_flags)
+{
+	return ERR_PTR(-EOPNOTSUPP);
+}
+static inline void io_uring_fixed_index_put(struct io_uring_cmd *cmd,
+					    struct io_rsrc_node *node,
+					    unsigned int issue_flags)
+{
+}
 static inline void __io_uring_cmd_done(struct io_uring_cmd *cmd, s32 ret,
 		u64 ret2, unsigned issue_flags, bool is_cqe32)
 {
diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index 41c89f5c616d..fa41cae5e922 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -1152,6 +1152,65 @@ int io_import_reg_buf(struct io_kiocb *req, struct iov_iter *iter,
 	return io_import_fixed(ddir, iter, node->buf, buf_addr, len);
 }
 
+struct io_rsrc_node *io_uring_fixed_index_get(struct io_uring_cmd *cmd,
+					      int buf_index, unsigned int off,
+					      size_t len, int ddir,
+					      struct iov_iter *iter,
+					      unsigned int issue_flags)
+{
+	struct io_ring_ctx *ctx = cmd_to_io_kiocb(cmd)->ctx;
+	struct io_rsrc_node *node;
+	struct io_mapped_ubuf *imu;
+	u64 addr;
+	int err;
+
+	io_ring_submit_lock(ctx, issue_flags);
+
+	node = io_rsrc_node_lookup(&ctx->buf_table, buf_index);
+	if (!node) {
+		io_ring_submit_unlock(ctx, issue_flags);
+		return ERR_PTR(-EINVAL);
+	}
+
+	node->refs++;
+
+	io_ring_submit_unlock(ctx, issue_flags);
+
+	imu = node->buf;
+	if (!imu) {
+		err = -EFAULT;
+		goto error;
+	}
+
+	if (check_add_overflow(imu->ubuf, off, &addr)) {
+		err = -EINVAL;
+		goto error;
+	}
+
+	err = io_import_fixed(ddir, iter, imu, addr, len);
+	if (err)
+		goto error;
+
+	return node;
+
+error:
+	io_uring_fixed_index_put(cmd, node, issue_flags);
+	return ERR_PTR(err);
+}
+EXPORT_SYMBOL_GPL(io_uring_fixed_index_get);
+
+void io_uring_fixed_index_put(struct io_uring_cmd *cmd,
+			      struct io_rsrc_node *node,
+			      unsigned int issue_flags)
+{
+	struct io_ring_ctx *ctx = cmd_to_io_kiocb(cmd)->ctx;
+
+	io_ring_submit_lock(ctx, issue_flags);
+	io_put_rsrc_node(ctx, node);
+	io_ring_submit_unlock(ctx, issue_flags);
+}
+EXPORT_SYMBOL_GPL(io_uring_fixed_index_put);
+
 /* Lock two rings at once. The rings must be different! */
 static void lock_two_rings(struct io_ring_ctx *ctx1, struct io_ring_ctx *ctx2)
 {
-- 
2.47.3