io_uring: prototype for device memory tx

index 2a8d7b14f1d1..5ae24847f5c4 100644
@@ -90,7 +90,7 @@ struct msghdr {
 	unsigned int	msg_flags;	/* flags on received message */
 	__kernel_size_t	msg_controllen;	/* ancillary data buffer length */
 	struct ubuf_info *msg_ubuf;
-	int (*sg_from_iter)(struct sk_buff *skb,
+	int (*sg_from_iter)(struct sk_buff *skb, struct ubuf_info *ubuf,
 			    struct iov_iter *from, size_t length);
 };
 

index 00a7df803b99..cf273d6f02b1 100644
@@ -116,9 +116,9 @@ struct io_recvzc {
 	struct io_zcrx_ifq		*ifq;
 };
 
-static int io_sg_from_iter_iovec(struct sk_buff *skb,
+static int io_sg_from_iter_iovec(struct sk_buff *skb, struct ubuf_info *ubuf,
 				 struct iov_iter *from, size_t length);
-static int io_sg_from_iter(struct sk_buff *skb,
+static int io_sg_from_iter(struct sk_buff *skb, struct ubuf_info *ubuf,
 			   struct iov_iter *from, size_t length);
 
 int io_shutdown_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)

@@ -1447,14 +1447,14 @@ int io_send_zc_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 	return 0;
 }
 
-static int io_sg_from_iter_iovec(struct sk_buff *skb,
+static int io_sg_from_iter_iovec(struct sk_buff *skb, struct ubuf_info *ubuf,
 				 struct iov_iter *from, size_t length)
 {
 	skb_zcopy_downgrade_managed(skb);
 	return zerocopy_fill_skb_from_iter(skb, from, length);
 }
 
-static int io_sg_from_iter(struct sk_buff *skb,
+static int io_sg_from_iter(struct sk_buff *skb, struct ubuf_info *ubuf,
 			   struct iov_iter *from, size_t length)
 {
 	struct skb_shared_info *shinfo = skb_shinfo(skb);

index c285c6465923..f15886f40efc 100644
@@ -753,7 +753,7 @@ int __zerocopy_sg_from_iter(struct msghdr *msg, struct sock *sk,
 	int ret;
 
 	if (msg && msg->msg_ubuf && msg->sg_from_iter)
-		ret = msg->sg_from_iter(skb, from, length);
+		ret = msg->sg_from_iter(skb, msg->msg_ubuf, from, length);
 	else if (binding)
 		ret = zerocopy_fill_skb_from_devmem(skb, from, length, binding);
 	else

index bccacd21b6c3..71024c7ce884 100644
@@ -103,6 +103,7 @@ struct net_iov_area {
 	struct net_iov *niovs;
 	size_t num_niovs;
 
+	struct net_device *netdev;
 	/* Offset into the dma-buf where this chunk starts.  */
 	unsigned long base_virtual;
 };

index 7ad52f499f87..ef82e064e796 100644
@@ -516,6 +516,7 @@ static int io_zcrx_create_area(struct io_zcrx_ifq *ifq,
 			goto err;
 	}
 
+	area->nia.netdev = ifq->netdev;
 	area->free_count = nr_iovs;
 	/* we're only supporting one area per ifq for now */
 	area->area_id = 0;

@@ -554,6 +555,7 @@ static void io_zcrx_drop_netdev(struct io_zcrx_ifq *ifq)
 
 	if (!ifq->netdev)
 		return;
+	WRITE_ONCE(ifq->area->nia.netdev, NULL);
 	netdev_put(ifq->netdev, &ifq->netdev_tracker);
 	ifq->netdev = NULL;
 }

@@ -568,6 +570,7 @@ static void io_close_queue(struct io_zcrx_ifq *ifq)
 	};
 
 	scoped_guard(mutex, &ifq->pp_lock) {
+		WRITE_ONCE(ifq->area->nia.netdev, NULL);
 		netdev = ifq->netdev;
 		netdev_tracker = ifq->netdev_tracker;
 		ifq->netdev = NULL;

index 4b3d5cfdf6e0..703b55778c32 100644
@@ -4004,9 +4004,15 @@ static struct sk_buff *validate_xmit_unreadable_skb(struct sk_buff *skb,
 	shinfo = skb_shinfo(skb);
 
 	if (shinfo->nr_frags > 0) {
+		struct net_device *trgt_dev;
+
 		niov = netmem_to_net_iov(skb_frag_netmem(&shinfo->frags[0]));
-		if (net_is_devmem_iov(niov) &&
-		    READ_ONCE(net_devmem_iov_binding(niov)->dev) != dev)
+		if (net_is_devmem_iov(niov))
+			trgt_dev = READ_ONCE(net_devmem_iov_binding(niov)->dev);
+		else
+			trgt_dev = READ_ONCE(net_iov_owner(niov)->netdev);
+
+		if (trgt_dev != dev)
 			goto out_free;
 	}
 

index ef82e064e796..f501fc75d7b6 100644
@@ -32,6 +32,8 @@
 
 #define IO_DMA_ATTR (DMA_ATTR_SKIP_CPU_SYNC | DMA_ATTR_WEAK_ORDERING)
 
+static void ifq_pcpu_release(struct percpu_ref *ref);
+
 static inline struct io_zcrx_ifq *io_pp_to_ifq(struct page_pool *pp)
 {
 	return pp->mp_priv;

@@ -535,16 +537,22 @@ static int io_zcrx_create_area(struct io_zcrx_ifq *ifq,
 static struct io_zcrx_ifq *io_zcrx_ifq_alloc(struct io_ring_ctx *ctx)
 {
 	struct io_zcrx_ifq *ifq;
+	int ret;
 
 	ifq = kzalloc_obj(*ifq);
 	if (!ifq)
 		return NULL;
+	ret = percpu_ref_init(&ifq->refs, ifq_pcpu_release, 0, GFP_KERNEL_ACCOUNT);
+	if (ret) {
+		kfree(ifq);
+		return NULL;
+	}
+	percpu_ref_get(&ifq->refs);
 
 	ifq->if_rxq = -1;
 	spin_lock_init(&ifq->ctx_lock);
 	spin_lock_init(&ifq->rq.lock);
 	mutex_init(&ifq->pp_lock);
-	refcount_set(&ifq->refs, 1);
 	refcount_set(&ifq->user_refs, 1);
 	return ifq;
 }

@@ -606,13 +614,28 @@ static void io_zcrx_ifq_free(struct io_zcrx_ifq *ifq)
 	io_free_rbuf_ring(ifq);
 	free_uid(ifq->user);
 	mutex_destroy(&ifq->pp_lock);
+	percpu_ref_exit(&ifq->refs);
 	kfree(ifq);
 }
 
+static inline void zcrx_release_work(struct work_struct *work)
+{
+	struct io_zcrx_ifq *ifq = container_of(work, struct io_zcrx_ifq, release_work);
+
+	io_zcrx_ifq_free(ifq);
+}
+
+static void ifq_pcpu_release(struct percpu_ref *ref)
+{
+	struct io_zcrx_ifq *ifq = container_of(ref, struct io_zcrx_ifq, refs);
+
+	INIT_WORK(&ifq->release_work, zcrx_release_work);
+	queue_work(system_wq, &ifq->release_work);
+}
+
 static void io_put_zcrx_ifq(struct io_zcrx_ifq *ifq)
 {
-	if (refcount_dec_and_test(&ifq->refs))
-		io_zcrx_ifq_free(ifq);
+	percpu_ref_put(&ifq->refs);
 }
 
 static void io_zcrx_return_niov_freelist(struct net_iov *niov)

@@ -683,6 +706,7 @@ static void zcrx_unregister_user(struct io_zcrx_ifq *ifq, struct io_ring_ctx *ct
 	if (refcount_dec_and_test(&ifq->user_refs)) {
 		io_close_queue(ifq);
 		io_zcrx_scrub(ifq);
+		percpu_ref_kill(&ifq->refs);
 	}
 }
 

@@ -727,7 +751,7 @@ static int zcrx_export(struct io_ring_ctx *ctx, struct io_zcrx_ifq *ifq,
 	if (!mem_is_zero(ce, sizeof(*ce)))
 		return -EINVAL;
 
-	refcount_inc(&ifq->refs);
+	percpu_ref_get(&ifq->refs);
 	refcount_inc(&ifq->user_refs);
 
 	file = anon_inode_create_getfile("[zcrx]", &zcrx_box_fops,

@@ -784,7 +808,7 @@ static int import_zcrx(struct io_ring_ctx *ctx,
 		return -EBADF;
 
 	ifq = file->private_data;
-	refcount_inc(&ifq->refs);
+	percpu_ref_get(&ifq->refs);
 	refcount_inc(&ifq->user_refs);
 
 	scoped_guard(mutex, &ctx->mmap_lock) {

@@ -1285,7 +1309,7 @@ static int io_pp_zc_init(struct page_pool *pp)
 	if (pp->p.dma_dir != DMA_FROM_DEVICE)
 		return -EOPNOTSUPP;
 
-	refcount_inc(&ifq->refs);
+	percpu_ref_get(&ifq->refs);
 	return 0;
 }
 

index fa00900e479e..205b9b40c74d 100644
@@ -7,6 +7,7 @@
 #include <linux/socket.h>
 #include <net/page_pool/types.h>
 #include <net/net_trackers.h>
+#include <linux/percpu-refcount.h>
 
 #define ZCRX_SUPPORTED_REG_FLAGS	(ZCRX_REG_IMPORT | ZCRX_REG_NODEV)
 #define ZCRX_FEATURES			(ZCRX_FEATURE_RX_PAGE_SIZE |\

@@ -64,7 +65,7 @@ struct io_zcrx_ifq {
 	struct device			*dev;
 	struct net_device		*netdev;
 	netdevice_tracker		netdev_tracker;
-	refcount_t			refs;
+	struct percpu_ref		refs;
 	/* counts userspace facing users like io_uring */
 	refcount_t			user_refs;
 

@@ -81,6 +82,7 @@ struct io_zcrx_ifq {
 	u32				fired_notifs;
 	u64				notif_data;
 	struct zcrx_notif_stats		*notif_stats;
+	struct work_struct		release_work;
 };
 
 #if defined(CONFIG_IO_URING_ZCRX)

index b58f39fed4d5..19f31f0d38d5 100644
@@ -2,6 +2,8 @@
 #ifndef _LINUX_IO_URING_NET_H
 #define _LINUX_IO_URING_NET_H
 
+#include <net/netmem.h>
+
 struct io_uring_cmd;
 
 #if defined(CONFIG_IO_URING)

@@ -15,4 +17,12 @@ static inline int io_uring_cmd_sock(struct io_uring_cmd *cmd,
 }
 #endif
 
+#if defined(CONFIG_IO_URING_ZCRX)
+void zcrx_ref_niov(struct net_iov *niov);
+#else
+static inline void zcrx_ref_niov(struct net_iov *niov)
+{
+}
+#endif
+
 #endif

128 lines not on this page · raw message

index 205b9b40c74d..406e0399bd7b 100644
@@ -32,6 +32,7 @@ struct io_zcrx_area {
 	struct net_iov_area	nia;
 	struct io_zcrx_ifq	*ifq;
 	atomic_t		*user_refs;
+	struct net_iov		*tx_niovs;
 
 	bool			is_mapped;
 	u16			area_id;

@@ -96,6 +97,10 @@ int io_zcrx_recv(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
 		 unsigned issue_flags, unsigned int *len);
 struct io_mapped_region *io_zcrx_get_region(struct io_ring_ctx *ctx,
 					    unsigned int id);
+
+int io_zcrx_fill_tx_skb(struct sk_buff *skb, struct io_zcrx_ifq *zcrx,
+			 struct iov_iter *from, size_t length);
+
 #else
 static inline int io_register_zcrx(struct io_ring_ctx *ctx,
 				   struct io_uring_zcrx_ifq_reg __user *arg)

@@ -124,6 +129,11 @@ static inline int io_zcrx_ctrl(struct io_ring_ctx *ctx,
 {
 	return -EOPNOTSUPP;
 }
+static inline int io_zcrx_fill_tx_skb(struct sk_buff *skb, struct io_zcrx_ifq *zcrx,
+					struct iov_iter *from, size_t length)
+{
+	return -EOPNOTSUPP;
+}
 #endif
 
 int io_recvzc(struct io_kiocb *req, unsigned int issue_flags);

index 18dabb4e9cfa..f871e3f2299b 100644
@@ -92,6 +92,7 @@
 #include <linux/user_namespace.h>
 #include <linux/indirect_call_wrapper.h>
 #include <linux/textsearch.h>
+#include <linux/io_uring/net.h>
 
 #include "dev.h"
 #include "devmem.h"

@@ -7476,6 +7477,8 @@ void __get_netmem(netmem_ref netmem)
 
 	if (net_is_devmem_iov(niov))
 		net_devmem_get_net_iov(netmem_to_net_iov(netmem));
+	else if (niov->type == NET_IOV_IOURING)
+		zcrx_ref_niov(niov);
 }
 EXPORT_SYMBOL(__get_netmem);
 

index 8d0f2ee24e0c..8af371ba6c06 100644
@@ -27,8 +27,13 @@ struct io_rsrc_update {
 	u32				offset;
 };
 
+struct io_uring_regbuf_desc {
+	__u64 uaddr;
+	__u64 size;
+};
+
 static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
-						   struct iovec *iov);
+						   struct io_uring_regbuf_desc *desc);
 
 static int hpage_acct_ref(struct io_ring_ctx *ctx, struct page *hpage,
 			  bool *acct_new)

@@ -81,6 +86,15 @@ static bool hpage_acct_unref(struct io_ring_ctx *ctx, struct page *hpage)
 
 #define IO_CACHED_BVECS_SEGS	32
 
+static void io_iov_to_regbuf_desc(const struct iovec *iov,
+				  struct io_uring_regbuf_desc *desc)
+{
+	*desc = (struct io_uring_regbuf_desc) {
+		.uaddr = (u64)(uintptr_t)iov->iov_base,
+		.size = iov->iov_len,
+	};
+}
+
 int __io_account_mem(struct user_struct *user, unsigned long nr_pages)
 {
 	unsigned long page_limit, cur_pages, new_pages;

@@ -381,6 +395,7 @@ static int __io_sqe_buffers_update(struct io_ring_ctx *ctx,
 		return -EINVAL;
 
 	for (done = 0; done < nr_args; done++) {
+		struct io_uring_regbuf_desc desc;
 		struct io_rsrc_node *node;
 		u64 tag = 0;
 

@@ -394,7 +409,9 @@ static int __io_sqe_buffers_update(struct io_ring_ctx *ctx,
 			err = -EFAULT;
 			break;
 		}
-		node = io_sqe_buffer_register(ctx, iov);
+
+		io_iov_to_regbuf_desc(iov, &desc);
+		node = io_sqe_buffer_register(ctx, &desc);
 		if (IS_ERR(node)) {
 			err = PTR_ERR(node);
 			break;

@@ -853,26 +870,26 @@ bool io_check_coalesce_buffer(struct page **page_array, int nr_pages,
 }
 
 static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
-						   struct iovec *iov)
+						   struct io_uring_regbuf_desc *desc)
 {
+	unsigned long uaddr = (unsigned long)desc->uaddr;
+	size_t size = desc->size;
 	struct io_mapped_ubuf *imu = NULL;
 	struct page **pages = NULL;
 	struct io_rsrc_node *node;
 	unsigned long off;
-	size_t size;
 	int ret, nr_pages, i;
 	struct io_imu_folio_data data;
 	bool coalesced = false;
 
-	if (!iov->iov_base) {
-		if (iov->iov_len)
+	if (!uaddr) {
+		if (size)
 			return ERR_PTR(-EFAULT);
 		/* remove the buffer without installing a new one */
 		return NULL;
 	}
 
-	ret = io_validate_user_buf_range((unsigned long)iov->iov_base,
-					 iov->iov_len);
+	ret = io_validate_user_buf_range(uaddr, size);
 	if (ret)
 		return ERR_PTR(ret);
 

@@ -881,8 +898,7 @@ static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
 		return ERR_PTR(-ENOMEM);
 
 	ret = -ENOMEM;
-	pages = io_pin_pages((unsigned long) iov->iov_base, iov->iov_len,
-				&nr_pages);
+	pages = io_pin_pages(uaddr, size, &nr_pages);
 	if (IS_ERR(pages)) {
 		ret = PTR_ERR(pages);
 		pages = NULL;

@@ -904,10 +920,9 @@ static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
 	if (ret)
 		goto done;
 
-	size = iov->iov_len;
 	/* store original address for later verification */
-	imu->ubuf = (unsigned long) iov->iov_base;
-	imu->len = iov->iov_len;
+	imu->ubuf = uaddr;
+	imu->len = size;
 	imu->folio_shift = PAGE_SHIFT;
 	imu->release = io_release_ubuf;
 	imu->priv = imu;

@@ -917,7 +932,7 @@ static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
 		imu->folio_shift = data.folio_shift;
 	refcount_set(&imu->refs, 1);
 
-	off = (unsigned long)iov->iov_base & ~PAGE_MASK;
+	off = uaddr & ~PAGE_MASK;
 	if (coalesced)
 		off += data.first_folio_page_idx << PAGE_SHIFT;
 

@@ -969,6 +984,7 @@ int io_sqe_buffers_register(struct io_ring_ctx *ctx, void __user *arg,
 		memset(iov, 0, sizeof(*iov));
 
 	for (i = 0; i < nr_args; i++) {
+		struct io_uring_regbuf_desc desc;
 		struct io_rsrc_node *node;
 		u64 tag = 0;
 

@@ -992,7 +1008,8 @@ int io_sqe_buffers_register(struct io_ring_ctx *ctx, void __user *arg,
 			}
 		}
 
-		node = io_sqe_buffer_register(ctx, iov);
+		io_iov_to_regbuf_desc(iov, &desc);
+		node = io_sqe_buffer_register(ctx, &desc);
 		if (IS_ERR(node)) {
 			ret = PTR_ERR(node);
 			break;

index 909fb7aea638..98b259901185 100644
@@ -790,13 +790,38 @@ struct io_uring_rsrc_update {
 
 struct io_uring_rsrc_update2 {
 	__u32 offset;
-	__u32 resv;
+	__u32 flags;
 	__aligned_u64 data;
 	__aligned_u64 tags;
 	__u32 nr;
 	__u32 resv2;
 };
 
+/* struct io_uring_rsrc_update2::flags */
+enum io_uring_rsrc_reg_flags {
+	/*
+	 * Use the extended descriptor format for buffer updates,
+	 * see struct io_uring_regbuf_desc
+	 */
+	IORING_RSRC_UPDATE_EXTENDED		= 1U << 1,
+};
+
+/* Buffer registration type, passed in struct io_uring_regbuf_desc::type */
+enum io_uring_regbuf_type {
+	IO_REGBUF_TYPE_EMPTY,
+	IO_REGBUF_TYPE_UADDR,
+
+	__IO_REGBUF_TYPE_MAX,
+};
+
+struct io_uring_regbuf_desc {
+	__u32 type; /* enum io_uring_regbuf_type */
+	__u32 flags;
+	__u64 size;
+	__u64 uaddr;
+	__u64 __resv[7];
+};
+
 /* Skip updating fd indexes set to this value in the fd table */
 #define IORING_REGISTER_FILES_SKIP	(-2)
 

151 lines not on this page · raw message

index 24fc3232a66a..d57e8a0380b5 100644
@@ -1381,6 +1381,11 @@ static int io_clone_buffers(struct io_ring_ctx *ctx, struct io_ring_ctx *src_ctx
 		if (!src_node) {
 			dst_node = NULL;
 		} else {
+			if (src_node->buf->flags & IO_REGBUF_F_UNCLONEABLE) {
+				io_rsrc_data_free(ctx, &data);
+				return -ENOMEM;
+			}
+
 			dst_node = io_rsrc_node_alloc(ctx, IORING_RSRC_BUFFER);
 			if (!dst_node) {
 				io_rsrc_data_free(ctx, &data);

index 98ae8ef51009..83aa86e6f320 100644
@@ -29,7 +29,8 @@ enum {
 };
 
 enum {
-	IO_REGBUF_F_KBUF		= 1,
+	IO_REGBUF_F_KBUF		= 1 << 0,
+	IO_REGBUF_F_UNCLONEABLE		= 1 << 1,
 };
 
 struct io_mapped_ubuf {

index d57e8a0380b5..05877b4c0ee5 100644
@@ -1245,9 +1245,9 @@ inline struct io_rsrc_node *io_find_buf_node(struct io_kiocb *req,
 	return NULL;
 }
 
-int io_import_reg_buf(struct io_kiocb *req, struct iov_iter *iter,
+int __io_import_reg_buf(struct io_kiocb *req, struct iov_iter *iter,
 			u64 buf_addr, size_t len, int ddir,
-			unsigned issue_flags)
+			unsigned issue_flags, unsigned import_flags)
 {
 	struct io_rsrc_node *node;
 

@@ -1656,9 +1656,9 @@ static int io_kern_bvec_size(struct iovec *iov, unsigned nr_iovs,
 	return 0;
 }
 
-int io_import_reg_vec(int ddir, struct iov_iter *iter,
+int __io_import_reg_vec(int ddir, struct iov_iter *iter,
 			struct io_kiocb *req, struct iou_vec *vec,
-			unsigned nr_iovs, unsigned issue_flags)
+			unsigned nr_iovs, unsigned issue_flags, unsigned import_flags)
 {
 	struct io_rsrc_node *node;
 	struct io_mapped_ubuf *imu;

index 83aa86e6f320..f12abaf63270 100644
@@ -65,12 +65,28 @@ int io_rsrc_data_alloc(struct io_rsrc_data *data, unsigned nr);
 
 struct io_rsrc_node *io_find_buf_node(struct io_kiocb *req,
 				      unsigned issue_flags);
-int io_import_reg_buf(struct io_kiocb *req, struct iov_iter *iter,
+int __io_import_reg_buf(struct io_kiocb *req, struct iov_iter *iter,
 			u64 buf_addr, size_t len, int ddir,
-			unsigned issue_flags);
-int io_import_reg_vec(int ddir, struct iov_iter *iter,
+			unsigned issue_flags, unsigned import_flags);
+int __io_import_reg_vec(int ddir, struct iov_iter *iter,
 			struct io_kiocb *req, struct iou_vec *vec,
-			unsigned nr_iovs, unsigned issue_flags);
+			unsigned nr_iovs, unsigned issue_flags,
+			unsigned import_flags);
+
+static inline int io_import_reg_buf(struct io_kiocb *req, struct iov_iter *iter,
+				    u64 buf_addr, size_t len, int ddir,
+				    unsigned issue_flags)
+{
+	return __io_import_reg_buf(req, iter, buf_addr, len, ddir, issue_flags, 0);
+}
+
+static inline int io_import_reg_vec(int ddir, struct iov_iter *iter,
+				    struct io_kiocb *req, struct iou_vec *vec,
+				    unsigned nr_iovs, unsigned issue_flags)
+{
+	return __io_import_reg_vec(ddir, iter, req, vec, nr_iovs, issue_flags, 0);
+}
+
 int io_prep_reg_iovec(struct io_kiocb *req, struct iou_vec *iv,
 			const struct iovec __user *uvec, size_t uvec_segs);
 

index 98b259901185..4e86ad3ee977 100644
@@ -810,6 +810,7 @@ enum io_uring_rsrc_reg_flags {
 enum io_uring_regbuf_type {
 	IO_REGBUF_TYPE_EMPTY,
 	IO_REGBUF_TYPE_UADDR,
+	IO_REGBUF_TYPE_ZCRX,
 
 	__IO_REGBUF_TYPE_MAX,
 };

141 lines not on this page · raw message

index f12abaf63270..a44f7dfcd470 100644
@@ -31,6 +31,11 @@ enum {
 enum {
 	IO_REGBUF_F_KBUF		= 1 << 0,
 	IO_REGBUF_F_UNCLONEABLE		= 1 << 1,
+	IO_REGBUF_F_ZCRX		= 1 << 2,
+};
+
+enum {
+	IO_REGBUF_IMPORT_ALLOW_ZCRX	= 1 << 0,
 };
 
 struct io_mapped_ubuf {

index cf273d6f02b1..2ffbb59ceee4 100644
@@ -1497,6 +1497,14 @@ static int io_sg_from_iter(struct sk_buff *skb, struct ubuf_info *ubuf,
 	return ret;
 }
 
+static int io_sg_from_zcrx_iter(struct sk_buff *skb, struct ubuf_info *ubuf,
+				struct iov_iter *from, size_t length)
+{
+	struct io_notif_data *nd = container_of(ubuf, struct io_notif_data, uarg);
+
+	return io_zcrx_fill_tx_skb(skb, nd->zcrx, from, length);
+}
+
 static int io_send_zc_import(struct io_kiocb *req,
 			     struct io_async_msghdr *kmsg,
 			     unsigned int issue_flags)

@@ -1510,15 +1518,17 @@ static int io_send_zc_import(struct io_kiocb *req,
 	notif->buf_index = req->buf_index;
 
 	if (!(sr->flags & IORING_SEND_VECTORIZED)) {
-		ret = io_import_reg_buf(notif, &kmsg->msg.msg_iter,
-					(u64)(uintptr_t)sr->buf, sr->len,
-					ITER_SOURCE, issue_flags);
+		ret = __io_import_reg_buf(notif, &kmsg->msg.msg_iter,
+					  (u64)(uintptr_t)sr->buf, sr->len,
+					  ITER_SOURCE, issue_flags,
+					  IO_REGBUF_IMPORT_ALLOW_ZCRX);
 	} else {
 		unsigned uvec_segs = kmsg->msg.msg_iter.nr_segs;
 
-		ret = io_import_reg_vec(ITER_SOURCE, &kmsg->msg.msg_iter,
+		ret = __io_import_reg_vec(ITER_SOURCE, &kmsg->msg.msg_iter,
 					notif, &kmsg->vec, uvec_segs,
-					issue_flags);
+					issue_flags,
+					IO_REGBUF_IMPORT_ALLOW_ZCRX);
 	}
 
 	if (unlikely(ret))

@@ -1545,9 +1555,18 @@ int io_sendmsg_zc(struct io_kiocb *req, unsigned int issue_flags)
 		return -EAGAIN;
 
 	if (req->flags & REQ_F_IMPORT_BUFFER) {
+		struct io_mapped_ubuf *buf;
+
 		ret = io_send_zc_import(req, kmsg, issue_flags);
 		if (unlikely(ret))
 			return ret;
+
+		buf = sr->notif->buf_node->buf;
+
+		if (buf->flags & IO_REGBUF_TYPE_ZCRX) {
+			kmsg->msg.sg_from_iter = io_sg_from_zcrx_iter;
+			io_notif_to_data(sr->notif)->zcrx = buf->priv;
+		}
 	}
 
 	msg_flags = sr->msg_flags;

index f3589cfef4a9..2dfd5bf23302 100644
@@ -17,7 +17,10 @@ struct io_notif_data {
 	struct io_notif_data	*next;
 	struct io_notif_data	*head;
 
-	unsigned		account_pages;
+	union {
+		unsigned		account_pages;
+		void			*zcrx;
+	};
 	bool			zc_report;
 	bool			zc_used;
 	bool			zc_copied;