[RFC 10/12] io_uring/rsrc: add dmabuf-backed buffer registeration

[Pavel Begunkov <asml.silence@gmail.com>]

Add an ability to register a dmabuf backed io_uring buffer. It also
needs know which device to use for attachment, for that it takes
target_fd and extracts the device through the new file op. Unlike normal
buffers, it also retains the target file so that any imports from
ineligible requests can be rejected in next patches.

Suggested-by: Vishal Verma <vishal1.verma@intel.com>
Suggested-by: David Wei <dw@davidwei.uk>
Signed-off-by: Pavel Begunkov <asml.silence@gmail.com>
---
 io_uring/rsrc.c | 118 +++++++++++++++++++++++++++++++++++++++++++++++-
 io_uring/rsrc.h |   1 +
 2 files changed, 118 insertions(+), 1 deletion(-)

diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index 274274b80b96..f44aa2670bc5 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -10,6 +10,8 @@
 #include <linux/compat.h>
 #include <linux/io_uring.h>
 #include <linux/io_uring/cmd.h>
+#include <linux/dma-map-ops.h>
+#include <linux/dma-buf.h>
 
 #include <uapi/linux/io_uring.h>
 
@@ -18,6 +20,7 @@
 #include "rsrc.h"
 #include "memmap.h"
 #include "register.h"
+#include "dmabuf.h"
 
 struct io_rsrc_update {
 	struct file			*file;
@@ -793,6 +796,117 @@ bool io_check_coalesce_buffer(struct page **page_array, int nr_pages,
 	return true;
 }
 
+struct io_regbuf_dma {
+	struct io_dmabuf		dmabuf;
+	struct dmavec 			*dmav;
+	struct file			*target_file;
+};
+
+static void io_release_reg_dmabuf(struct io_regbuf_dma *db)
+{
+	if (db->dmav)
+		kfree(db->dmav);
+	io_dmabuf_release(&db->dmabuf);
+	if (db->target_file)
+		fput(db->target_file);
+
+	kfree(db);
+}
+
+static void io_release_reg_dmabuf_cb(void *priv)
+{
+	io_release_reg_dmabuf(priv);
+}
+
+static struct io_rsrc_node *io_register_dmabuf(struct io_ring_ctx *ctx,
+						struct io_uring_reg_buffer *rb,
+						struct iovec *iov)
+{
+	struct io_rsrc_node *node = NULL;
+	struct io_mapped_ubuf *imu = NULL;
+	struct io_regbuf_dma *regbuf;
+	struct file *target_file;
+	struct scatterlist *sg;
+	struct device *dev;
+	unsigned int segments;
+	int ret, i;
+
+	if (iov->iov_base || iov->iov_len)
+		return ERR_PTR(-EFAULT);
+
+	regbuf = kzalloc(sizeof(*regbuf), GFP_KERNEL);
+	if (!regbuf) {
+		ret = -ENOMEM;
+		goto err;
+	}
+
+	target_file = fget(rb->target_fd);
+	if (!target_file) {
+		ret = -EBADF;
+		goto err;
+	}
+	regbuf->target_file = target_file;
+
+	if (!target_file->f_op->get_dma_device) {
+		ret = -EOPNOTSUPP;
+		goto err;
+	}
+	dev = target_file->f_op->get_dma_device(target_file);
+	if (IS_ERR(dev)) {
+		ret = PTR_ERR(dev);
+		goto err;
+	}
+
+	ret = io_dmabuf_import(&regbuf->dmabuf, rb->dmabuf_fd, dev,
+				DMA_BIDIRECTIONAL);
+	if (ret)
+		goto err;
+
+	segments = regbuf->dmabuf.sgt->nents;
+	regbuf->dmav = kmalloc_array(segments, sizeof(regbuf->dmav[0]),
+				     GFP_KERNEL_ACCOUNT);
+	if (!regbuf->dmav) {
+		ret = -ENOMEM;
+		goto err;
+	}
+
+	for_each_sgtable_dma_sg(regbuf->dmabuf.sgt, sg, i) {
+		regbuf->dmav[i].addr = sg_dma_address(sg);
+		regbuf->dmav[i].len = sg_dma_len(sg);
+	}
+
+	node = io_rsrc_node_alloc(ctx, IORING_RSRC_BUFFER);
+	if (!node) {
+		ret = -ENOMEM;
+		goto err;
+	}
+	imu = io_alloc_imu(ctx, 0);
+	if (!imu) {
+		ret = -ENOMEM;
+		goto err;
+	}
+
+	imu->nr_bvecs = segments;
+	imu->ubuf = 0;
+	imu->len = regbuf->dmabuf.len;
+	imu->folio_shift = 0;
+	imu->release = io_release_reg_dmabuf_cb;
+	imu->priv = regbuf;
+	imu->flags = IO_IMU_F_DMA;
+	imu->dir = IO_IMU_DEST | IO_IMU_SOURCE;
+	refcount_set(&imu->refs, 1);
+	node->buf = imu;
+	return node;
+err:
+	if (regbuf)
+		io_release_reg_dmabuf(regbuf);
+	if (imu)
+		io_free_imu(ctx, imu);
+	if (node)
+		io_cache_free(&ctx->node_cache, node);
+	return ERR_PTR(ret);
+}
+
 static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
 						   struct io_uring_reg_buffer *rb,
 						   struct iovec *iov,
@@ -808,7 +922,7 @@ static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
 	bool coalesced = false;
 
 	if (rb->dmabuf_fd != -1 || rb->target_fd != -1)
-		return NULL;
+		return io_register_dmabuf(ctx, rb, iov);
 
 	if (!iov->iov_base)
 		return NULL;
@@ -1100,6 +1214,8 @@ static int io_import_fixed(int ddir, struct iov_iter *iter,
 
 	offset = buf_addr - imu->ubuf;
 
+	if (imu->flags & IO_IMU_F_DMA)
+		return -EOPNOTSUPP;
 	if (imu->flags & IO_IMU_F_KBUF)
 		return io_import_kbuf(ddir, iter, imu, len, offset);
 
diff --git a/io_uring/rsrc.h b/io_uring/rsrc.h
index 15ad4a885ae5..f567ad82b76c 100644
--- a/io_uring/rsrc.h
+++ b/io_uring/rsrc.h
@@ -30,6 +30,7 @@ enum {
 
 enum {
 	IO_IMU_F_KBUF			= 1,
+	IO_IMU_F_DMA			= 2,
 };
 
 struct io_mapped_ubuf {
-- 
2.49.0