[PATCH 3/4] io_uring: vectored registered buffer import

[Pavel Begunkov <asml.silence@xxxxxxxxx>]

Add a helper that takes a registered buffer and an iovec with addresses
pointing into that registered buffer, and return a new bvec
corresponding to the given iovec. Essentially, each iov entry is
resolved into a bvec array, which gives us an array of arrays of struct
bio_vec, which the function flattens into a single long bvec.

Note, max_segs is overestimated, that can be improved later. The
allocation also can be optimised by doing it inline into the same array.

Signed-off-by: Pavel Begunkov <asml.silence@xxxxxxxxx>
---
 io_uring/rsrc.c | 60 +++++++++++++++++++++++++++++++++++++++++++++++++
 io_uring/rsrc.h |  3 +++
 2 files changed, 63 insertions(+)

diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index fa5f27496aef..6f9f3cb4a2ef 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -1085,6 +1085,66 @@ int io_sqe_buffers_register(struct io_ring_ctx *ctx, void __user *arg,
 	return ret;
 }
 
+struct bio_vec *io_import_fixed_vec(int ddir, struct iov_iter *iter,
+				    struct io_mapped_ubuf *imu,
+				    struct iovec *iovec, int nr_iovs)
+{
+	unsigned long folio_size = (1 << imu->folio_shift);
+	unsigned long folio_mask = folio_size - 1;
+	struct bio_vec *res_bvec;
+	size_t total_len = 0;
+	int max_segs = 0;
+	int bvec_idx = 0;
+	int iov_idx;
+
+	if (WARN_ON_ONCE(!imu))
+		return ERR_PTR(-EFAULT);
+
+	for (iov_idx = 0; iov_idx < nr_iovs; iov_idx++) {
+		size_t iov_len = iovec[iov_idx].iov_len;
+		u64 buf_addr = (u64)iovec[iov_idx].iov_base;
+		u64 buf_end;
+
+		if (unlikely(check_add_overflow(buf_addr, (u64)iov_len, &buf_end)))
+			return ERR_PTR(-EFAULT);
+		/* not inside the mapped region */
+		if (unlikely(buf_addr < imu->ubuf || buf_end > (imu->ubuf + imu->len)))
+			return ERR_PTR(-EFAULT);
+		max_segs += (iov_len >> imu->folio_shift) + 2;
+	}
+
+	res_bvec = kmalloc_array(max_segs, sizeof(*res_bvec), GFP_KERNEL);
+	if (!res_bvec)
+		return ERR_PTR(-ENOMEM);
+
+	for (iov_idx = 0; iov_idx < nr_iovs; iov_idx++) {
+		size_t iov_len = iovec[iov_idx].iov_len;
+		u64 buf_addr = (u64)iovec[iov_idx].iov_base;
+		u64 folio_addr = imu->ubuf & ~folio_mask;
+		struct bio_vec *src_bvec;
+		size_t offset;
+
+		total_len += iov_len;
+		/* by using folio address it also accounts for bvec offset */
+		offset = buf_addr - folio_addr;
+		src_bvec = imu->bvec + (offset >> imu->folio_shift);
+		offset &= folio_mask;
+
+		for (; iov_len; offset = 0, bvec_idx++, src_bvec++) {
+			size_t seg_size = min_t(size_t, iov_len,
+						folio_size - offset);
+
+			res_bvec[bvec_idx].bv_page = src_bvec->bv_page;
+			res_bvec[bvec_idx].bv_offset = offset;
+			res_bvec[bvec_idx].bv_len = seg_size;
+			iov_len -= seg_size;
+		}
+	}
+
+	iov_iter_bvec(iter, ddir, res_bvec, bvec_idx, total_len);
+	return res_bvec;
+}
+
 int io_import_fixed(int ddir, struct iov_iter *iter,
 			   struct io_mapped_ubuf *imu,
 			   u64 buf_addr, size_t len)
diff --git a/io_uring/rsrc.h b/io_uring/rsrc.h
index 8ed588036210..675161cf8b92 100644
--- a/io_uring/rsrc.h
+++ b/io_uring/rsrc.h
@@ -66,6 +66,9 @@ int io_queue_rsrc_removal(struct io_rsrc_data *data, unsigned idx, void *rsrc);
 int io_import_fixed(int ddir, struct iov_iter *iter,
 			   struct io_mapped_ubuf *imu,
 			   u64 buf_addr, size_t len);
+struct bio_vec *io_import_fixed_vec(int ddir, struct iov_iter *iter,
+				    struct io_mapped_ubuf *imu,
+				    struct iovec *iovec, int nr_iovs);
 
 int io_register_clone_buffers(struct io_ring_ctx *ctx, void __user *arg);
 void __io_sqe_buffers_unregister(struct io_ring_ctx *ctx);
-- 
2.46.0