zcrx.c

   1// SPDX-License-Identifier: GPL-2.0
   2#include <linux/kernel.h>
   3#include <linux/errno.h>
   4#include <linux/dma-map-ops.h>
   5#include <linux/mm.h>
   6#include <linux/nospec.h>
   7#include <linux/io_uring.h>
   8#include <linux/netdevice.h>
   9#include <linux/rtnetlink.h>
  10#include <linux/skbuff_ref.h>
  11#include <linux/anon_inodes.h>
  12
  13#include <net/page_pool/helpers.h>
  14#include <net/page_pool/memory_provider.h>
  15#include <net/netlink.h>
  16#include <net/netdev_queues.h>
  17#include <net/netdev_rx_queue.h>
  18#include <net/tcp.h>
  19#include <net/rps.h>
  20
  21#include <trace/events/page_pool.h>
  22
  23#include <uapi/linux/io_uring.h>
  24
  25#include "io_uring.h"
  26#include "kbuf.h"
  27#include "memmap.h"
  28#include "zcrx.h"
  29#include "rsrc.h"
  30
  31#define IO_ZCRX_AREA_SUPPORTED_FLAGS	(IORING_ZCRX_AREA_DMABUF)
  32
  33#define IO_DMA_ATTR (DMA_ATTR_SKIP_CPU_SYNC | DMA_ATTR_WEAK_ORDERING)
  34
  35static inline struct io_zcrx_ifq *io_pp_to_ifq(struct page_pool *pp)
  36{
  37	return pp->mp_priv;
  38}
  39
  40static inline struct io_zcrx_area *io_zcrx_iov_to_area(const struct net_iov *niov)
  41{
  42	struct net_iov_area *owner = net_iov_owner(niov);
  43
  44	return container_of(owner, struct io_zcrx_area, nia);
  45}
  46
  47static inline struct page *io_zcrx_iov_page(const struct net_iov *niov)
  48{
  49	struct io_zcrx_area *area = io_zcrx_iov_to_area(niov);
  50	unsigned niov_pages_shift;
  51
  52	lockdep_assert(!area->mem.is_dmabuf);
  53
  54	niov_pages_shift = area->ifq->niov_shift - PAGE_SHIFT;
  55	return area->mem.pages[net_iov_idx(niov) << niov_pages_shift];
  56}
  57
  58static int io_populate_area_dma(struct io_zcrx_ifq *ifq,
  59				struct io_zcrx_area *area)
  60{
  61	unsigned niov_size = 1U << ifq->niov_shift;
  62	struct sg_table *sgt = area->mem.sgt;
  63	struct scatterlist *sg;
  64	unsigned i, niov_idx = 0;
  65
  66	for_each_sgtable_dma_sg(sgt, sg, i) {
  67		dma_addr_t dma = sg_dma_address(sg);
  68		unsigned long sg_len = sg_dma_len(sg);
  69
  70		if (WARN_ON_ONCE(sg_len % niov_size))
  71			return -EINVAL;
  72
  73		while (sg_len && niov_idx < area->nia.num_niovs) {
  74			struct net_iov *niov = &area->nia.niovs[niov_idx];
  75
  76			if (net_mp_niov_set_dma_addr(niov, dma))
  77				return -EFAULT;
  78			sg_len -= niov_size;
  79			dma += niov_size;
  80			niov_idx++;
  81		}
  82	}
  83
  84	if (WARN_ON_ONCE(niov_idx != area->nia.num_niovs))
  85		return -EFAULT;
  86	return 0;
  87}
  88
  89static void io_release_dmabuf(struct io_zcrx_mem *mem)
  90{
  91	if (!IS_ENABLED(CONFIG_DMA_SHARED_BUFFER))
  92		return;
  93
  94	if (mem->sgt)
  95		dma_buf_unmap_attachment_unlocked(mem->attach, mem->sgt,
  96						  DMA_FROM_DEVICE);
  97	if (mem->attach)
  98		dma_buf_detach(mem->dmabuf, mem->attach);
  99	if (mem->dmabuf)
 100		dma_buf_put(mem->dmabuf);
 101
 102	mem->sgt = NULL;
 103	mem->attach = NULL;
 104	mem->dmabuf = NULL;
 105}
 106
 107static int io_import_dmabuf(struct io_zcrx_ifq *ifq,
 108			    struct io_zcrx_mem *mem,
 109			    struct io_uring_zcrx_area_reg *area_reg)
 110{
 111	unsigned long off = (unsigned long)area_reg->addr;
 112	unsigned long len = (unsigned long)area_reg->len;
 113	unsigned long total_size = 0;
 114	struct scatterlist *sg;
 115	int dmabuf_fd = area_reg->dmabuf_fd;
 116	int i, ret;
 117
 118	if (off)
 119		return -EINVAL;
 120	if (WARN_ON_ONCE(!ifq->dev))
 121		return -EFAULT;
 122	if (!IS_ENABLED(CONFIG_DMA_SHARED_BUFFER))
 123		return -EINVAL;
 124
 125	mem->is_dmabuf = true;
 126	mem->dmabuf = dma_buf_get(dmabuf_fd);
 127	if (IS_ERR(mem->dmabuf)) {
 128		ret = PTR_ERR(mem->dmabuf);
 129		mem->dmabuf = NULL;
 130		goto err;
 131	}
 132
 133	mem->attach = dma_buf_attach(mem->dmabuf, ifq->dev);
 134	if (IS_ERR(mem->attach)) {
 135		ret = PTR_ERR(mem->attach);
 136		mem->attach = NULL;
 137		goto err;
 138	}
 139
 140	mem->sgt = dma_buf_map_attachment_unlocked(mem->attach, DMA_FROM_DEVICE);
 141	if (IS_ERR(mem->sgt)) {
 142		ret = PTR_ERR(mem->sgt);
 143		mem->sgt = NULL;
 144		goto err;
 145	}
 146
 147	for_each_sgtable_dma_sg(mem->sgt, sg, i)
 148		total_size += sg_dma_len(sg);
 149
 150	if (total_size != len) {
 151		ret = -EINVAL;
 152		goto err;
 153	}
 154
 155	mem->size = len;
 156	return 0;
 157err:
 158	io_release_dmabuf(mem);
 159	return ret;
 160}
 161
 162static unsigned long io_count_account_pages(struct page **pages, unsigned nr_pages)
 163{
 164	struct folio *last_folio = NULL;
 165	unsigned long res = 0;
 166	int i;
 167
 168	for (i = 0; i < nr_pages; i++) {
 169		struct folio *folio = page_folio(pages[i]);
 170
 171		if (folio == last_folio)
 172			continue;
 173		last_folio = folio;
 174		res += folio_nr_pages(folio);
 175	}
 176	return res;
 177}
 178
 179static int io_import_umem(struct io_zcrx_ifq *ifq,
 180			  struct io_zcrx_mem *mem,
 181			  struct io_uring_zcrx_area_reg *area_reg)
 182{
 183	struct page **pages;
 184	int nr_pages, ret;
 185
 186	if (area_reg->dmabuf_fd)
 187		return -EINVAL;
 188	if (!area_reg->addr)
 189		return -EFAULT;
 190	pages = io_pin_pages((unsigned long)area_reg->addr, area_reg->len,
 191				   &nr_pages);
 192	if (IS_ERR(pages))
 193		return PTR_ERR(pages);
 194
 195	ret = sg_alloc_table_from_pages(&mem->page_sg_table, pages, nr_pages,
 196					0, nr_pages << PAGE_SHIFT,
 197					GFP_KERNEL_ACCOUNT);
 198	if (ret) {
 199		unpin_user_pages(pages, nr_pages);
 200		kvfree(pages);
 201		return ret;
 202	}
 203
 204	mem->account_pages = io_count_account_pages(pages, nr_pages);
 205	ret = io_account_mem(ifq->user, ifq->mm_account, mem->account_pages);
 206	if (ret < 0)
 207		mem->account_pages = 0;
 208
 209	mem->sgt = &mem->page_sg_table;
 210	mem->pages = pages;
 211	mem->nr_folios = nr_pages;
 212	mem->size = area_reg->len;
 213	return ret;
 214}
 215
 216static void io_release_area_mem(struct io_zcrx_mem *mem)
 217{
 218	if (mem->is_dmabuf) {
 219		io_release_dmabuf(mem);
 220		return;
 221	}
 222	if (mem->pages) {
 223		unpin_user_pages(mem->pages, mem->nr_folios);
 224		sg_free_table(mem->sgt);
 225		mem->sgt = NULL;
 226		kvfree(mem->pages);
 227	}
 228}
 229
 230static int io_import_area(struct io_zcrx_ifq *ifq,
 231			  struct io_zcrx_mem *mem,
 232			  struct io_uring_zcrx_area_reg *area_reg)
 233{
 234	int ret;
 235
 236	if (area_reg->flags & ~IO_ZCRX_AREA_SUPPORTED_FLAGS)
 237		return -EINVAL;
 238	if (area_reg->rq_area_token)
 239		return -EINVAL;
 240	if (area_reg->__resv2[0] || area_reg->__resv2[1])
 241		return -EINVAL;
 242
 243	ret = io_validate_user_buf_range(area_reg->addr, area_reg->len);
 244	if (ret)
 245		return ret;
 246	if (area_reg->addr & ~PAGE_MASK || area_reg->len & ~PAGE_MASK)
 247		return -EINVAL;
 248
 249	if (area_reg->flags & IORING_ZCRX_AREA_DMABUF)
 250		return io_import_dmabuf(ifq, mem, area_reg);
 251	return io_import_umem(ifq, mem, area_reg);
 252}
 253
 254static void io_zcrx_unmap_area(struct io_zcrx_ifq *ifq,
 255				struct io_zcrx_area *area)
 256{
 257	int i;
 258
 259	guard(mutex)(&ifq->pp_lock);
 260	if (!area->is_mapped)
 261		return;
 262	area->is_mapped = false;
 263
 264	for (i = 0; i < area->nia.num_niovs; i++)
 265		net_mp_niov_set_dma_addr(&area->nia.niovs[i], 0);
 266
 267	if (area->mem.is_dmabuf) {
 268		io_release_dmabuf(&area->mem);
 269	} else {
 270		dma_unmap_sgtable(ifq->dev, &area->mem.page_sg_table,
 271				  DMA_FROM_DEVICE, IO_DMA_ATTR);
 272	}
 273}
 274
 275static int io_zcrx_map_area(struct io_zcrx_ifq *ifq, struct io_zcrx_area *area)
 276{
 277	int ret;
 278
 279	guard(mutex)(&ifq->pp_lock);
 280	if (area->is_mapped)
 281		return 0;
 282
 283	if (!area->mem.is_dmabuf) {
 284		ret = dma_map_sgtable(ifq->dev, &area->mem.page_sg_table,
 285				      DMA_FROM_DEVICE, IO_DMA_ATTR);
 286		if (ret < 0)
 287			return ret;
 288	}
 289
 290	ret = io_populate_area_dma(ifq, area);
 291	if (ret && !area->mem.is_dmabuf)
 292		dma_unmap_sgtable(ifq->dev, &area->mem.page_sg_table,
 293				  DMA_FROM_DEVICE, IO_DMA_ATTR);
 294	if (ret == 0)
 295		area->is_mapped = true;
 296	return ret;
 297}
 298
 299static void io_zcrx_sync_for_device(struct page_pool *pool,
 300				    struct net_iov *niov)
 301{
 302#if defined(CONFIG_HAS_DMA) && defined(CONFIG_DMA_NEED_SYNC)
 303	dma_addr_t dma_addr;
 304
 305	unsigned niov_size;
 306
 307	if (!dma_dev_need_sync(pool->p.dev))
 308		return;
 309
 310	niov_size = 1U << io_pp_to_ifq(pool)->niov_shift;
 311	dma_addr = page_pool_get_dma_addr_netmem(net_iov_to_netmem(niov));
 312	__dma_sync_single_for_device(pool->p.dev, dma_addr + pool->p.offset,
 313				     niov_size, pool->p.dma_dir);
 314#endif
 315}
 316
 317#define IO_RQ_MAX_ENTRIES		32768
 318
 319#define IO_SKBS_PER_CALL_LIMIT	20
 320
 321struct io_zcrx_args {
 322	struct io_kiocb		*req;
 323	struct io_zcrx_ifq	*ifq;
 324	struct socket		*sock;
 325	unsigned		nr_skbs;
 326};
 327
 328static const struct memory_provider_ops io_uring_pp_zc_ops;
 329
 330static inline atomic_t *io_get_user_counter(struct net_iov *niov)
 331{
 332	struct io_zcrx_area *area = io_zcrx_iov_to_area(niov);
 333
 334	return &area->user_refs[net_iov_idx(niov)];
 335}
 336
 337static bool io_zcrx_put_niov_uref(struct net_iov *niov)
 338{
 339	atomic_t *uref = io_get_user_counter(niov);
 340	int old;
 341
 342	old = atomic_read(uref);
 343	do {
 344		if (unlikely(old == 0))
 345			return false;
 346	} while (!atomic_try_cmpxchg(uref, &old, old - 1));
 347
 348	return true;
 349}
 350
 351static void io_zcrx_get_niov_uref(struct net_iov *niov)
 352{
 353	atomic_inc(io_get_user_counter(niov));
 354}
 355
 356static void io_fill_zcrx_offsets(struct io_uring_zcrx_offsets *offsets)
 357{
 358	offsets->head = offsetof(struct io_uring, head);
 359	offsets->tail = offsetof(struct io_uring, tail);
 360	offsets->rqes = ALIGN(sizeof(struct io_uring), L1_CACHE_BYTES);
 361}
 362
 363static int io_allocate_rbuf_ring(struct io_ring_ctx *ctx,
 364				 struct io_zcrx_ifq *ifq,
 365				 struct io_uring_zcrx_ifq_reg *reg,
 366				 struct io_uring_region_desc *rd,
 367				 u32 id)
 368{
 369	u64 mmap_offset;
 370	size_t off, size;
 371	void *ptr;
 372	int ret;
 373
 374	io_fill_zcrx_offsets(&reg->offsets);
 375	off = reg->offsets.rqes;
 376	size = off + sizeof(struct io_uring_zcrx_rqe) * reg->rq_entries;
 377	if (size > rd->size)
 378		return -EINVAL;
 379
 380	mmap_offset = IORING_MAP_OFF_ZCRX_REGION;
 381	mmap_offset += id << IORING_OFF_PBUF_SHIFT;
 382
 383	ret = io_create_region(ctx, &ifq->region, rd, mmap_offset);
 384	if (ret < 0)
 385		return ret;
 386
 387	ptr = io_region_get_ptr(&ifq->region);
 388	ifq->rq_ring = (struct io_uring *)ptr;
 389	ifq->rqes = (struct io_uring_zcrx_rqe *)(ptr + off);
 390
 391	return 0;
 392}
 393
 394static void io_free_rbuf_ring(struct io_zcrx_ifq *ifq)
 395{
 396	io_free_region(ifq->user, &ifq->region);
 397	ifq->rq_ring = NULL;
 398	ifq->rqes = NULL;
 399}
 400
 401static void io_zcrx_free_area(struct io_zcrx_ifq *ifq,
 402			      struct io_zcrx_area *area)
 403{
 404	io_zcrx_unmap_area(ifq, area);
 405	io_release_area_mem(&area->mem);
 406
 407	if (area->mem.account_pages)
 408		io_unaccount_mem(ifq->user, ifq->mm_account,
 409				 area->mem.account_pages);
 410
 411	kvfree(area->freelist);
 412	kvfree(area->nia.niovs);
 413	kvfree(area->user_refs);
 414	kfree(area);
 415}
 416
 417static int io_zcrx_append_area(struct io_zcrx_ifq *ifq,
 418				struct io_zcrx_area *area)
 419{
 420	if (ifq->area)
 421		return -EINVAL;
 422	ifq->area = area;
 423	return 0;
 424}
 425
 426static int io_zcrx_create_area(struct io_zcrx_ifq *ifq,
 427			       struct io_uring_zcrx_area_reg *area_reg)
 428{
 429	struct io_zcrx_area *area;
 430	unsigned nr_iovs;
 431	int i, ret;
 432
 433	ret = -ENOMEM;
 434	area = kzalloc(sizeof(*area), GFP_KERNEL);
 435	if (!area)
 436		goto err;
 437	area->ifq = ifq;
 438
 439	ret = io_import_area(ifq, &area->mem, area_reg);
 440	if (ret)
 441		goto err;
 442
 443	ifq->niov_shift = PAGE_SHIFT;
 444	nr_iovs = area->mem.size >> ifq->niov_shift;
 445	area->nia.num_niovs = nr_iovs;
 446
 447	ret = -ENOMEM;
 448	area->nia.niovs = kvmalloc_array(nr_iovs, sizeof(area->nia.niovs[0]),
 449					 GFP_KERNEL_ACCOUNT | __GFP_ZERO);
 450	if (!area->nia.niovs)
 451		goto err;
 452
 453	area->freelist = kvmalloc_array(nr_iovs, sizeof(area->freelist[0]),
 454					GFP_KERNEL_ACCOUNT | __GFP_ZERO);
 455	if (!area->freelist)
 456		goto err;
 457
 458	area->user_refs = kvmalloc_array(nr_iovs, sizeof(area->user_refs[0]),
 459					GFP_KERNEL_ACCOUNT | __GFP_ZERO);
 460	if (!area->user_refs)
 461		goto err;
 462
 463	for (i = 0; i < nr_iovs; i++) {
 464		struct net_iov *niov = &area->nia.niovs[i];
 465
 466		niov->owner = &area->nia;
 467		area->freelist[i] = i;
 468		atomic_set(&area->user_refs[i], 0);
 469		niov->type = NET_IOV_IOURING;
 470	}
 471
 472	area->free_count = nr_iovs;
 473	/* we're only supporting one area per ifq for now */
 474	area->area_id = 0;
 475	area_reg->rq_area_token = (u64)area->area_id << IORING_ZCRX_AREA_SHIFT;
 476	spin_lock_init(&area->freelist_lock);
 477
 478	ret = io_zcrx_append_area(ifq, area);
 479	if (!ret)
 480		return 0;
 481err:
 482	if (area)
 483		io_zcrx_free_area(ifq, area);
 484	return ret;
 485}
 486
 487static struct io_zcrx_ifq *io_zcrx_ifq_alloc(struct io_ring_ctx *ctx)
 488{
 489	struct io_zcrx_ifq *ifq;
 490
 491	ifq = kzalloc(sizeof(*ifq), GFP_KERNEL);
 492	if (!ifq)
 493		return NULL;
 494
 495	ifq->if_rxq = -1;
 496	spin_lock_init(&ifq->rq_lock);
 497	mutex_init(&ifq->pp_lock);
 498	refcount_set(&ifq->refs, 1);
 499	refcount_set(&ifq->user_refs, 1);
 500	return ifq;
 501}
 502
 503static void io_zcrx_drop_netdev(struct io_zcrx_ifq *ifq)
 504{
 505	guard(mutex)(&ifq->pp_lock);
 506
 507	if (!ifq->netdev)
 508		return;
 509	netdev_put(ifq->netdev, &ifq->netdev_tracker);
 510	ifq->netdev = NULL;
 511}
 512
 513static void io_close_queue(struct io_zcrx_ifq *ifq)
 514{
 515	struct net_device *netdev;
 516	netdevice_tracker netdev_tracker;
 517	struct pp_memory_provider_params p = {
 518		.mp_ops = &io_uring_pp_zc_ops,
 519		.mp_priv = ifq,
 520	};
 521
 522	scoped_guard(mutex, &ifq->pp_lock) {
 523		netdev = ifq->netdev;
 524		netdev_tracker = ifq->netdev_tracker;
 525		ifq->netdev = NULL;
 526	}
 527
 528	if (netdev) {
 529		if (ifq->if_rxq != -1)
 530			net_mp_close_rxq(netdev, ifq->if_rxq, &p);
 531		netdev_put(netdev, &netdev_tracker);
 532	}
 533	ifq->if_rxq = -1;
 534}
 535
 536static void io_zcrx_ifq_free(struct io_zcrx_ifq *ifq)
 537{
 538	io_close_queue(ifq);
 539
 540	if (ifq->area)
 541		io_zcrx_free_area(ifq, ifq->area);
 542	free_uid(ifq->user);
 543	if (ifq->mm_account)
 544		mmdrop(ifq->mm_account);
 545	if (ifq->dev)
 546		put_device(ifq->dev);
 547
 548	io_free_rbuf_ring(ifq);
 549	mutex_destroy(&ifq->pp_lock);
 550	kfree(ifq);
 551}
 552
 553static void io_put_zcrx_ifq(struct io_zcrx_ifq *ifq)
 554{
 555	if (refcount_dec_and_test(&ifq->refs))
 556		io_zcrx_ifq_free(ifq);
 557}
 558
 559static void io_zcrx_return_niov_freelist(struct net_iov *niov)
 560{
 561	struct io_zcrx_area *area = io_zcrx_iov_to_area(niov);
 562
 563	spin_lock_bh(&area->freelist_lock);
 564	area->freelist[area->free_count++] = net_iov_idx(niov);
 565	spin_unlock_bh(&area->freelist_lock);
 566}
 567
 568static void io_zcrx_return_niov(struct net_iov *niov)
 569{
 570	netmem_ref netmem = net_iov_to_netmem(niov);
 571
 572	if (!niov->desc.pp) {
 573		/* copy fallback allocated niovs */
 574		io_zcrx_return_niov_freelist(niov);
 575		return;
 576	}
 577	page_pool_put_unrefed_netmem(niov->desc.pp, netmem, -1, false);
 578}
 579
 580static void io_zcrx_scrub(struct io_zcrx_ifq *ifq)
 581{
 582	struct io_zcrx_area *area = ifq->area;
 583	int i;
 584
 585	if (!area)
 586		return;
 587
 588	/* Reclaim back all buffers given to the user space. */
 589	for (i = 0; i < area->nia.num_niovs; i++) {
 590		struct net_iov *niov = &area->nia.niovs[i];
 591		int nr;
 592
 593		if (!atomic_read(io_get_user_counter(niov)))
 594			continue;
 595		nr = atomic_xchg(io_get_user_counter(niov), 0);
 596		if (nr && !page_pool_unref_netmem(net_iov_to_netmem(niov), nr))
 597			io_zcrx_return_niov(niov);
 598	}
 599}
 600
 601static void zcrx_unregister(struct io_zcrx_ifq *ifq)
 602{
 603	if (refcount_dec_and_test(&ifq->user_refs)) {
 604		io_close_queue(ifq);
 605		io_zcrx_scrub(ifq);
 606	}
 607	io_put_zcrx_ifq(ifq);
 608}
 609
 610struct io_mapped_region *io_zcrx_get_region(struct io_ring_ctx *ctx,
 611					    unsigned int id)
 612{
 613	struct io_zcrx_ifq *ifq = xa_load(&ctx->zcrx_ctxs, id);
 614
 615	lockdep_assert_held(&ctx->mmap_lock);
 616
 617	return ifq ? &ifq->region : NULL;
 618}
 619
 620static int zcrx_box_release(struct inode *inode, struct file *file)
 621{
 622	struct io_zcrx_ifq *ifq = file->private_data;
 623
 624	if (WARN_ON_ONCE(!ifq))
 625		return -EFAULT;
 626	zcrx_unregister(ifq);
 627	return 0;
 628}
 629
 630static const struct file_operations zcrx_box_fops = {
 631	.owner		= THIS_MODULE,
 632	.release	= zcrx_box_release,
 633};
 634
 635static int zcrx_export(struct io_ring_ctx *ctx, struct io_zcrx_ifq *ifq,
 636		       struct zcrx_ctrl *ctrl, void __user *arg)
 637{
 638	struct zcrx_ctrl_export *ce = &ctrl->zc_export;
 639	struct file *file;
 640	int fd = -1;
 641
 642	if (!mem_is_zero(ce, sizeof(*ce)))
 643		return -EINVAL;
 644	fd = get_unused_fd_flags(O_CLOEXEC);
 645	if (fd < 0)
 646		return fd;
 647
 648	ce->zcrx_fd = fd;
 649	if (copy_to_user(arg, ctrl, sizeof(*ctrl))) {
 650		put_unused_fd(fd);
 651		return -EFAULT;
 652	}
 653
 654	refcount_inc(&ifq->refs);
 655	refcount_inc(&ifq->user_refs);
 656
 657	file = anon_inode_create_getfile("[zcrx]", &zcrx_box_fops,
 658					 ifq, O_CLOEXEC, NULL);
 659	if (IS_ERR(file)) {
 660		put_unused_fd(fd);
 661		zcrx_unregister(ifq);
 662		return PTR_ERR(file);
 663	}
 664
 665	fd_install(fd, file);
 666	return 0;
 667}
 668
 669static int import_zcrx(struct io_ring_ctx *ctx,
 670		       struct io_uring_zcrx_ifq_reg __user *arg,
 671		       struct io_uring_zcrx_ifq_reg *reg)
 672{
 673	struct io_zcrx_ifq *ifq;
 674	struct file *file;
 675	int fd, ret;
 676	u32 id;
 677
 678	if (!(ctx->flags & IORING_SETUP_DEFER_TASKRUN))
 679		return -EINVAL;
 680	if (!(ctx->flags & (IORING_SETUP_CQE32|IORING_SETUP_CQE_MIXED)))
 681		return -EINVAL;
 682	if (reg->if_rxq || reg->rq_entries || reg->area_ptr || reg->region_ptr)
 683		return -EINVAL;
 684	if (reg->flags & ~ZCRX_REG_IMPORT)
 685		return -EINVAL;
 686
 687	fd = reg->if_idx;
 688	CLASS(fd, f)(fd);
 689	if (fd_empty(f))
 690		return -EBADF;
 691
 692	file = fd_file(f);
 693	if (file->f_op != &zcrx_box_fops || !file->private_data)
 694		return -EBADF;
 695
 696	ifq = file->private_data;
 697	refcount_inc(&ifq->refs);
 698	refcount_inc(&ifq->user_refs);
 699
 700	scoped_guard(mutex, &ctx->mmap_lock) {
 701		ret = xa_alloc(&ctx->zcrx_ctxs, &id, NULL, xa_limit_31b, GFP_KERNEL);
 702		if (ret)
 703			goto err;
 704	}
 705
 706	reg->zcrx_id = id;
 707	io_fill_zcrx_offsets(&reg->offsets);
 708	if (copy_to_user(arg, reg, sizeof(*reg))) {
 709		ret = -EFAULT;
 710		goto err_xa_erase;
 711	}
 712
 713	scoped_guard(mutex, &ctx->mmap_lock) {
 714		ret = -ENOMEM;
 715		if (xa_store(&ctx->zcrx_ctxs, id, ifq, GFP_KERNEL))
 716			goto err_xa_erase;
 717	}
 718
 719	return 0;
 720err_xa_erase:
 721	scoped_guard(mutex, &ctx->mmap_lock)
 722		xa_erase(&ctx->zcrx_ctxs, id);
 723err:
 724	zcrx_unregister(ifq);
 725	return ret;
 726}
 727
 728int io_register_zcrx_ifq(struct io_ring_ctx *ctx,
 729			  struct io_uring_zcrx_ifq_reg __user *arg)
 730{
 731	struct pp_memory_provider_params mp_param = {};
 732	struct io_uring_zcrx_area_reg area;
 733	struct io_uring_zcrx_ifq_reg reg;
 734	struct io_uring_region_desc rd;
 735	struct io_zcrx_ifq *ifq;
 736	int ret;
 737	u32 id;
 738
 739	/*
 740	 * 1. Interface queue allocation.
 741	 * 2. It can observe data destined for sockets of other tasks.
 742	 */
 743	if (!capable(CAP_NET_ADMIN))
 744		return -EPERM;
 745
 746	/* mandatory io_uring features for zc rx */
 747	if (!(ctx->flags & IORING_SETUP_DEFER_TASKRUN))
 748		return -EINVAL;
 749	if (!(ctx->flags & (IORING_SETUP_CQE32|IORING_SETUP_CQE_MIXED)))
 750		return -EINVAL;
 751	if (copy_from_user(&reg, arg, sizeof(reg)))
 752		return -EFAULT;
 753	if (!mem_is_zero(&reg.__resv, sizeof(reg.__resv)) ||
 754	    reg.__resv2 || reg.zcrx_id)
 755		return -EINVAL;
 756	if (reg.flags & ZCRX_REG_IMPORT)
 757		return import_zcrx(ctx, arg, &reg);
 758	if (copy_from_user(&rd, u64_to_user_ptr(reg.region_ptr), sizeof(rd)))
 759		return -EFAULT;
 760	if (reg.if_rxq == -1 || !reg.rq_entries || reg.flags)
 761		return -EINVAL;
 762	if (reg.rq_entries > IO_RQ_MAX_ENTRIES) {
 763		if (!(ctx->flags & IORING_SETUP_CLAMP))
 764			return -EINVAL;
 765		reg.rq_entries = IO_RQ_MAX_ENTRIES;
 766	}
 767	reg.rq_entries = roundup_pow_of_two(reg.rq_entries);
 768
 769	if (copy_from_user(&area, u64_to_user_ptr(reg.area_ptr), sizeof(area)))
 770		return -EFAULT;
 771
 772	ifq = io_zcrx_ifq_alloc(ctx);
 773	if (!ifq)
 774		return -ENOMEM;
 775
 776	if (ctx->user) {
 777		get_uid(ctx->user);
 778		ifq->user = ctx->user;
 779	}
 780	if (ctx->mm_account) {
 781		mmgrab(ctx->mm_account);
 782		ifq->mm_account = ctx->mm_account;
 783	}
 784	ifq->rq_entries = reg.rq_entries;
 785
 786	scoped_guard(mutex, &ctx->mmap_lock) {
 787		/* preallocate id */
 788		ret = xa_alloc(&ctx->zcrx_ctxs, &id, NULL, xa_limit_31b, GFP_KERNEL);
 789		if (ret)
 790			goto ifq_free;
 791	}
 792
 793	ret = io_allocate_rbuf_ring(ctx, ifq, &reg, &rd, id);
 794	if (ret)
 795		goto err;
 796
 797	ifq->netdev = netdev_get_by_index_lock(current->nsproxy->net_ns, reg.if_idx);
 798	if (!ifq->netdev) {
 799		ret = -ENODEV;
 800		goto err;
 801	}
 802	netdev_hold(ifq->netdev, &ifq->netdev_tracker, GFP_KERNEL);
 803
 804	ifq->dev = netdev_queue_get_dma_dev(ifq->netdev, reg.if_rxq);
 805	if (!ifq->dev) {
 806		ret = -EOPNOTSUPP;
 807		goto netdev_put_unlock;
 808	}
 809	get_device(ifq->dev);
 810
 811	ret = io_zcrx_create_area(ifq, &area);
 812	if (ret)
 813		goto netdev_put_unlock;
 814
 815	mp_param.mp_ops = &io_uring_pp_zc_ops;
 816	mp_param.mp_priv = ifq;
 817	ret = __net_mp_open_rxq(ifq->netdev, reg.if_rxq, &mp_param, NULL);
 818	if (ret)
 819		goto netdev_put_unlock;
 820	netdev_unlock(ifq->netdev);
 821	ifq->if_rxq = reg.if_rxq;
 822
 823	reg.zcrx_id = id;
 824
 825	scoped_guard(mutex, &ctx->mmap_lock) {
 826		/* publish ifq */
 827		ret = -ENOMEM;
 828		if (xa_store(&ctx->zcrx_ctxs, id, ifq, GFP_KERNEL))
 829			goto err;
 830	}
 831
 832	if (copy_to_user(arg, &reg, sizeof(reg)) ||
 833	    copy_to_user(u64_to_user_ptr(reg.region_ptr), &rd, sizeof(rd)) ||
 834	    copy_to_user(u64_to_user_ptr(reg.area_ptr), &area, sizeof(area))) {
 835		ret = -EFAULT;
 836		goto err;
 837	}
 838	return 0;
 839netdev_put_unlock:
 840	netdev_unlock(ifq->netdev);
 841err:
 842	scoped_guard(mutex, &ctx->mmap_lock)
 843		xa_erase(&ctx->zcrx_ctxs, id);
 844ifq_free:
 845	zcrx_unregister(ifq);
 846	return ret;
 847}
 848
 849static struct net_iov *__io_zcrx_get_free_niov(struct io_zcrx_area *area)
 850{
 851	unsigned niov_idx;
 852
 853	lockdep_assert_held(&area->freelist_lock);
 854
 855	niov_idx = area->freelist[--area->free_count];
 856	return &area->nia.niovs[niov_idx];
 857}
 858
 859void io_unregister_zcrx_ifqs(struct io_ring_ctx *ctx)
 860{
 861	struct io_zcrx_ifq *ifq;
 862
 863	lockdep_assert_held(&ctx->uring_lock);
 864
 865	while (1) {
 866		scoped_guard(mutex, &ctx->mmap_lock) {
 867			unsigned long id = 0;
 868
 869			ifq = xa_find(&ctx->zcrx_ctxs, &id, ULONG_MAX, XA_PRESENT);
 870			if (ifq)
 871				xa_erase(&ctx->zcrx_ctxs, id);
 872		}
 873		if (!ifq)
 874			break;
 875		zcrx_unregister(ifq);
 876	}
 877
 878	xa_destroy(&ctx->zcrx_ctxs);
 879}
 880
 881static inline u32 io_zcrx_rqring_entries(struct io_zcrx_ifq *ifq)
 882{
 883	u32 entries;
 884
 885	entries = smp_load_acquire(&ifq->rq_ring->tail) - ifq->cached_rq_head;
 886	return min(entries, ifq->rq_entries);
 887}
 888
 889static struct io_uring_zcrx_rqe *io_zcrx_get_rqe(struct io_zcrx_ifq *ifq,
 890						 unsigned mask)
 891{
 892	unsigned int idx = ifq->cached_rq_head++ & mask;
 893
 894	return &ifq->rqes[idx];
 895}
 896
 897static inline bool io_parse_rqe(struct io_uring_zcrx_rqe *rqe,
 898				struct io_zcrx_ifq *ifq,
 899				struct net_iov **ret_niov)
 900{
 901	__u64 off = READ_ONCE(rqe->off);
 902	unsigned niov_idx, area_idx;
 903	struct io_zcrx_area *area;
 904
 905	area_idx = off >> IORING_ZCRX_AREA_SHIFT;
 906	niov_idx = (off & ~IORING_ZCRX_AREA_MASK) >> ifq->niov_shift;
 907
 908	if (unlikely(rqe->__pad || area_idx))
 909		return false;
 910	area = ifq->area;
 911
 912	if (unlikely(niov_idx >= area->nia.num_niovs))
 913		return false;
 914	niov_idx = array_index_nospec(niov_idx, area->nia.num_niovs);
 915
 916	*ret_niov = &area->nia.niovs[niov_idx];
 917	return true;
 918}
 919
 920static void io_zcrx_ring_refill(struct page_pool *pp,
 921				struct io_zcrx_ifq *ifq)
 922{
 923	unsigned int mask = ifq->rq_entries - 1;
 924	unsigned int entries;
 925
 926	guard(spinlock_bh)(&ifq->rq_lock);
 927
 928	entries = io_zcrx_rqring_entries(ifq);
 929	entries = min_t(unsigned, entries, PP_ALLOC_CACHE_REFILL);
 930	if (unlikely(!entries))
 931		return;
 932
 933	do {
 934		struct io_uring_zcrx_rqe *rqe = io_zcrx_get_rqe(ifq, mask);
 935		struct net_iov *niov;
 936		netmem_ref netmem;
 937
 938		if (!io_parse_rqe(rqe, ifq, &niov))
 939			continue;
 940		if (!io_zcrx_put_niov_uref(niov))
 941			continue;
 942
 943		netmem = net_iov_to_netmem(niov);
 944		if (!page_pool_unref_and_test(netmem))
 945			continue;
 946
 947		if (unlikely(niov->desc.pp != pp)) {
 948			io_zcrx_return_niov(niov);
 949			continue;
 950		}
 951
 952		io_zcrx_sync_for_device(pp, niov);
 953		net_mp_netmem_place_in_cache(pp, netmem);
 954	} while (--entries);
 955
 956	smp_store_release(&ifq->rq_ring->head, ifq->cached_rq_head);
 957}
 958
 959static void io_zcrx_refill_slow(struct page_pool *pp, struct io_zcrx_ifq *ifq)
 960{
 961	struct io_zcrx_area *area = ifq->area;
 962
 963	spin_lock_bh(&area->freelist_lock);
 964	while (area->free_count && pp->alloc.count < PP_ALLOC_CACHE_REFILL) {
 965		struct net_iov *niov = __io_zcrx_get_free_niov(area);
 966		netmem_ref netmem = net_iov_to_netmem(niov);
 967
 968		net_mp_niov_set_page_pool(pp, niov);
 969		io_zcrx_sync_for_device(pp, niov);
 970		net_mp_netmem_place_in_cache(pp, netmem);
 971	}
 972	spin_unlock_bh(&area->freelist_lock);
 973}
 974
 975static netmem_ref io_pp_zc_alloc_netmems(struct page_pool *pp, gfp_t gfp)
 976{
 977	struct io_zcrx_ifq *ifq = io_pp_to_ifq(pp);
 978
 979	/* pp should already be ensuring that */
 980	if (unlikely(pp->alloc.count))
 981		goto out_return;
 982
 983	io_zcrx_ring_refill(pp, ifq);
 984	if (likely(pp->alloc.count))
 985		goto out_return;
 986
 987	io_zcrx_refill_slow(pp, ifq);
 988	if (!pp->alloc.count)
 989		return 0;
 990out_return:
 991	return pp->alloc.cache[--pp->alloc.count];
 992}
 993
 994static bool io_pp_zc_release_netmem(struct page_pool *pp, netmem_ref netmem)
 995{
 996	struct net_iov *niov;
 997
 998	if (WARN_ON_ONCE(!netmem_is_net_iov(netmem)))
 999		return false;
1000
1001	niov = netmem_to_net_iov(netmem);
1002	net_mp_niov_clear_page_pool(niov);
1003	io_zcrx_return_niov_freelist(niov);
1004	return false;
1005}
1006
1007static int io_pp_zc_init(struct page_pool *pp)
1008{
1009	struct io_zcrx_ifq *ifq = io_pp_to_ifq(pp);
1010	int ret;
1011
1012	if (WARN_ON_ONCE(!ifq))
1013		return -EINVAL;
1014	if (WARN_ON_ONCE(ifq->dev != pp->p.dev))
1015		return -EINVAL;
1016	if (WARN_ON_ONCE(!pp->dma_map))
1017		return -EOPNOTSUPP;
1018	if (pp->p.order + PAGE_SHIFT != ifq->niov_shift)
1019		return -EINVAL;
1020	if (pp->p.dma_dir != DMA_FROM_DEVICE)
1021		return -EOPNOTSUPP;
1022
1023	ret = io_zcrx_map_area(ifq, ifq->area);
1024	if (ret)
1025		return ret;
1026
1027	refcount_inc(&ifq->refs);
1028	return 0;
1029}
1030
1031static void io_pp_zc_destroy(struct page_pool *pp)
1032{
1033	io_put_zcrx_ifq(io_pp_to_ifq(pp));
1034}
1035
1036static int io_pp_nl_fill(void *mp_priv, struct sk_buff *rsp,
1037			 struct netdev_rx_queue *rxq)
1038{
1039	struct nlattr *nest;
1040	int type;
1041
1042	type = rxq ? NETDEV_A_QUEUE_IO_URING : NETDEV_A_PAGE_POOL_IO_URING;
1043	nest = nla_nest_start(rsp, type);
1044	if (!nest)
1045		return -EMSGSIZE;
1046	nla_nest_end(rsp, nest);
1047
1048	return 0;
1049}
1050
1051static void io_pp_uninstall(void *mp_priv, struct netdev_rx_queue *rxq)
1052{
1053	struct pp_memory_provider_params *p = &rxq->mp_params;
1054	struct io_zcrx_ifq *ifq = mp_priv;
1055
1056	io_zcrx_drop_netdev(ifq);
1057	if (ifq->area)
1058		io_zcrx_unmap_area(ifq, ifq->area);
1059
1060	p->mp_ops = NULL;
1061	p->mp_priv = NULL;
1062}
1063
1064static const struct memory_provider_ops io_uring_pp_zc_ops = {
1065	.alloc_netmems		= io_pp_zc_alloc_netmems,
1066	.release_netmem		= io_pp_zc_release_netmem,
1067	.init			= io_pp_zc_init,
1068	.destroy		= io_pp_zc_destroy,
1069	.nl_fill		= io_pp_nl_fill,
1070	.uninstall		= io_pp_uninstall,
1071};
1072
1073static unsigned zcrx_parse_rq(netmem_ref *netmem_array, unsigned nr,
1074			      struct io_zcrx_ifq *zcrx)
1075{
1076	unsigned int mask = zcrx->rq_entries - 1;
1077	unsigned int i;
1078
1079	nr = min(nr, io_zcrx_rqring_entries(zcrx));
1080	for (i = 0; i < nr; i++) {
1081		struct io_uring_zcrx_rqe *rqe = io_zcrx_get_rqe(zcrx, mask);
1082		struct net_iov *niov;
1083
1084		if (!io_parse_rqe(rqe, zcrx, &niov))
1085			break;
1086		netmem_array[i] = net_iov_to_netmem(niov);
1087	}
1088
1089	smp_store_release(&zcrx->rq_ring->head, zcrx->cached_rq_head);
1090	return i;
1091}
1092
1093#define ZCRX_FLUSH_BATCH 32
1094
1095static void zcrx_return_buffers(netmem_ref *netmems, unsigned nr)
1096{
1097	unsigned i;
1098
1099	for (i = 0; i < nr; i++) {
1100		netmem_ref netmem = netmems[i];
1101		struct net_iov *niov = netmem_to_net_iov(netmem);
1102
1103		if (!io_zcrx_put_niov_uref(niov))
1104			continue;
1105		if (!page_pool_unref_and_test(netmem))
1106			continue;
1107		io_zcrx_return_niov(niov);
1108	}
1109}
1110
1111static int zcrx_flush_rq(struct io_ring_ctx *ctx, struct io_zcrx_ifq *zcrx,
1112			 struct zcrx_ctrl *ctrl)
1113{
1114	struct zcrx_ctrl_flush_rq *frq = &ctrl->zc_flush;
1115	netmem_ref netmems[ZCRX_FLUSH_BATCH];
1116	unsigned total = 0;
1117	unsigned nr;
1118
1119	if (!mem_is_zero(&frq->__resv, sizeof(frq->__resv)))
1120		return -EINVAL;
1121
1122	do {
1123		scoped_guard(spinlock_bh, &zcrx->rq_lock) {
1124			nr = zcrx_parse_rq(netmems, ZCRX_FLUSH_BATCH, zcrx);
1125			zcrx_return_buffers(netmems, nr);
1126		}
1127
1128		total += nr;
1129
1130		if (fatal_signal_pending(current))
1131			break;
1132		cond_resched();
1133	} while (nr == ZCRX_FLUSH_BATCH && total < zcrx->rq_entries);
1134
1135	return 0;
1136}
1137
1138int io_zcrx_ctrl(struct io_ring_ctx *ctx, void __user *arg, unsigned nr_args)
1139{
1140	struct zcrx_ctrl ctrl;
1141	struct io_zcrx_ifq *zcrx;
1142
1143	if (nr_args)
1144		return -EINVAL;
1145	if (copy_from_user(&ctrl, arg, sizeof(ctrl)))
1146		return -EFAULT;
1147	if (!mem_is_zero(&ctrl.__resv, sizeof(ctrl.__resv)))
1148		return -EFAULT;
1149
1150	zcrx = xa_load(&ctx->zcrx_ctxs, ctrl.zcrx_id);
1151	if (!zcrx)
1152		return -ENXIO;
1153
1154	switch (ctrl.op) {
1155	case ZCRX_CTRL_FLUSH_RQ:
1156		return zcrx_flush_rq(ctx, zcrx, &ctrl);
1157	case ZCRX_CTRL_EXPORT:
1158		return zcrx_export(ctx, zcrx, &ctrl, arg);
1159	}
1160
1161	return -EOPNOTSUPP;
1162}
1163
1164static bool io_zcrx_queue_cqe(struct io_kiocb *req, struct net_iov *niov,
1165			      struct io_zcrx_ifq *ifq, int off, int len)
1166{
1167	struct io_ring_ctx *ctx = req->ctx;
1168	struct io_uring_zcrx_cqe *rcqe;
1169	struct io_zcrx_area *area;
1170	struct io_uring_cqe *cqe;
1171	u64 offset;
1172
1173	if (!io_defer_get_uncommited_cqe(ctx, &cqe))
1174		return false;
1175
1176	cqe->user_data = req->cqe.user_data;
1177	cqe->res = len;
1178	cqe->flags = IORING_CQE_F_MORE;
1179	if (ctx->flags & IORING_SETUP_CQE_MIXED)
1180		cqe->flags |= IORING_CQE_F_32;
1181
1182	area = io_zcrx_iov_to_area(niov);
1183	offset = off + (net_iov_idx(niov) << ifq->niov_shift);
1184	rcqe = (struct io_uring_zcrx_cqe *)(cqe + 1);
1185	rcqe->off = offset + ((u64)area->area_id << IORING_ZCRX_AREA_SHIFT);
1186	rcqe->__pad = 0;
1187	return true;
1188}
1189
1190static struct net_iov *io_alloc_fallback_niov(struct io_zcrx_ifq *ifq)
1191{
1192	struct io_zcrx_area *area = ifq->area;
1193	struct net_iov *niov = NULL;
1194
1195	if (area->mem.is_dmabuf)
1196		return NULL;
1197
1198	spin_lock_bh(&area->freelist_lock);
1199	if (area->free_count)
1200		niov = __io_zcrx_get_free_niov(area);
1201	spin_unlock_bh(&area->freelist_lock);
1202
1203	if (niov)
1204		page_pool_fragment_netmem(net_iov_to_netmem(niov), 1);
1205	return niov;
1206}
1207
1208struct io_copy_cache {
1209	struct page		*page;
1210	unsigned long		offset;
1211	size_t			size;
1212};
1213
1214static ssize_t io_copy_page(struct io_copy_cache *cc, struct page *src_page,
1215			    unsigned int src_offset, size_t len)
1216{
1217	size_t copied = 0;
1218
1219	len = min(len, cc->size);
1220
1221	while (len) {
1222		void *src_addr, *dst_addr;
1223		struct page *dst_page = cc->page;
1224		unsigned dst_offset = cc->offset;
1225		size_t n = len;
1226
1227		if (folio_test_partial_kmap(page_folio(dst_page)) ||
1228		    folio_test_partial_kmap(page_folio(src_page))) {
1229			dst_page += dst_offset / PAGE_SIZE;
1230			dst_offset = offset_in_page(dst_offset);
1231			src_page += src_offset / PAGE_SIZE;
1232			src_offset = offset_in_page(src_offset);
1233			n = min(PAGE_SIZE - src_offset, PAGE_SIZE - dst_offset);
1234			n = min(n, len);
1235		}
1236
1237		dst_addr = kmap_local_page(dst_page) + dst_offset;
1238		src_addr = kmap_local_page(src_page) + src_offset;
1239
1240		memcpy(dst_addr, src_addr, n);
1241
1242		kunmap_local(src_addr);
1243		kunmap_local(dst_addr);
1244
1245		cc->size -= n;
1246		cc->offset += n;
1247		src_offset += n;
1248		len -= n;
1249		copied += n;
1250	}
1251	return copied;
1252}
1253
1254static ssize_t io_zcrx_copy_chunk(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1255				  struct page *src_page, unsigned int src_offset,
1256				  size_t len)
1257{
1258	size_t copied = 0;
1259	int ret = 0;
1260
1261	while (len) {
1262		struct io_copy_cache cc;
1263		struct net_iov *niov;
1264		size_t n;
1265
1266		niov = io_alloc_fallback_niov(ifq);
1267		if (!niov) {
1268			ret = -ENOMEM;
1269			break;
1270		}
1271
1272		cc.page = io_zcrx_iov_page(niov);
1273		cc.offset = 0;
1274		cc.size = PAGE_SIZE;
1275
1276		n = io_copy_page(&cc, src_page, src_offset, len);
1277
1278		if (!io_zcrx_queue_cqe(req, niov, ifq, 0, n)) {
1279			io_zcrx_return_niov(niov);
1280			ret = -ENOSPC;
1281			break;
1282		}
1283
1284		io_zcrx_get_niov_uref(niov);
1285		src_offset += n;
1286		len -= n;
1287		copied += n;
1288	}
1289
1290	return copied ? copied : ret;
1291}
1292
1293static int io_zcrx_copy_frag(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1294			     const skb_frag_t *frag, int off, int len)
1295{
1296	struct page *page = skb_frag_page(frag);
1297
1298	return io_zcrx_copy_chunk(req, ifq, page, off + skb_frag_off(frag), len);
1299}
1300
1301static int io_zcrx_recv_frag(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1302			     const skb_frag_t *frag, int off, int len)
1303{
1304	struct net_iov *niov;
1305	struct page_pool *pp;
1306
1307	if (unlikely(!skb_frag_is_net_iov(frag)))
1308		return io_zcrx_copy_frag(req, ifq, frag, off, len);
1309
1310	niov = netmem_to_net_iov(frag->netmem);
1311	pp = niov->desc.pp;
1312
1313	if (!pp || pp->mp_ops != &io_uring_pp_zc_ops || io_pp_to_ifq(pp) != ifq)
1314		return -EFAULT;
1315
1316	if (!io_zcrx_queue_cqe(req, niov, ifq, off + skb_frag_off(frag), len))
1317		return -ENOSPC;
1318
1319	/*
1320	 * Prevent it from being recycled while user is accessing it.
1321	 * It has to be done before grabbing a user reference.
1322	 */
1323	page_pool_ref_netmem(net_iov_to_netmem(niov));
1324	io_zcrx_get_niov_uref(niov);
1325	return len;
1326}
1327
1328static int
1329io_zcrx_recv_skb(read_descriptor_t *desc, struct sk_buff *skb,
1330		 unsigned int offset, size_t len)
1331{
1332	struct io_zcrx_args *args = desc->arg.data;
1333	struct io_zcrx_ifq *ifq = args->ifq;
1334	struct io_kiocb *req = args->req;
1335	struct sk_buff *frag_iter;
1336	unsigned start, start_off = offset;
1337	int i, copy, end, off;
1338	int ret = 0;
1339
1340	len = min_t(size_t, len, desc->count);
1341	/*
1342	 * __tcp_read_sock() always calls io_zcrx_recv_skb one last time, even
1343	 * if desc->count is already 0. This is caused by the if (offset + 1 !=
1344	 * skb->len) check. Return early in this case to break out of
1345	 * __tcp_read_sock().
1346	 */
1347	if (!len)
1348		return 0;
1349	if (unlikely(args->nr_skbs++ > IO_SKBS_PER_CALL_LIMIT))
1350		return -EAGAIN;
1351
1352	if (unlikely(offset < skb_headlen(skb))) {
1353		ssize_t copied;
1354		size_t to_copy;
1355
1356		to_copy = min_t(size_t, skb_headlen(skb) - offset, len);
1357		copied = io_zcrx_copy_chunk(req, ifq, virt_to_page(skb->data),
1358					    offset_in_page(skb->data) + offset,
1359					    to_copy);
1360		if (copied < 0) {
1361			ret = copied;
1362			goto out;
1363		}
1364		offset += copied;
1365		len -= copied;
1366		if (!len)
1367			goto out;
1368		if (offset != skb_headlen(skb))
1369			goto out;
1370	}
1371
1372	start = skb_headlen(skb);
1373
1374	for (i = 0; i < skb_shinfo(skb)->nr_frags; i++) {
1375		const skb_frag_t *frag;
1376
1377		if (WARN_ON(start > offset + len))
1378			return -EFAULT;
1379
1380		frag = &skb_shinfo(skb)->frags[i];
1381		end = start + skb_frag_size(frag);
1382
1383		if (offset < end) {
1384			copy = end - offset;
1385			if (copy > len)
1386				copy = len;
1387
1388			off = offset - start;
1389			ret = io_zcrx_recv_frag(req, ifq, frag, off, copy);
1390			if (ret < 0)
1391				goto out;
1392
1393			offset += ret;
1394			len -= ret;
1395			if (len == 0 || ret != copy)
1396				goto out;
1397		}
1398		start = end;
1399	}
1400
1401	skb_walk_frags(skb, frag_iter) {
1402		if (WARN_ON(start > offset + len))
1403			return -EFAULT;
1404
1405		end = start + frag_iter->len;
1406		if (offset < end) {
1407			size_t count;
1408
1409			copy = end - offset;
1410			if (copy > len)
1411				copy = len;
1412
1413			off = offset - start;
1414			count = desc->count;
1415			ret = io_zcrx_recv_skb(desc, frag_iter, off, copy);
1416			desc->count = count;
1417			if (ret < 0)
1418				goto out;
1419
1420			offset += ret;
1421			len -= ret;
1422			if (len == 0 || ret != copy)
1423				goto out;
1424		}
1425		start = end;
1426	}
1427
1428out:
1429	if (offset == start_off)
1430		return ret;
1431	desc->count -= (offset - start_off);
1432	return offset - start_off;
1433}
1434
1435static int io_zcrx_tcp_recvmsg(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1436				struct sock *sk, int flags,
1437				unsigned issue_flags, unsigned int *outlen)
1438{
1439	unsigned int len = *outlen;
1440	struct io_zcrx_args args = {
1441		.req = req,
1442		.ifq = ifq,
1443		.sock = sk->sk_socket,
1444	};
1445	read_descriptor_t rd_desc = {
1446		.count = len ? len : UINT_MAX,
1447		.arg.data = &args,
1448	};
1449	int ret;
1450
1451	lock_sock(sk);
1452	ret = tcp_read_sock(sk, &rd_desc, io_zcrx_recv_skb);
1453	if (len && ret > 0)
1454		*outlen = len - ret;
1455	if (ret <= 0) {
1456		if (ret < 0 || sock_flag(sk, SOCK_DONE))
1457			goto out;
1458		if (sk->sk_err)
1459			ret = sock_error(sk);
1460		else if (sk->sk_shutdown & RCV_SHUTDOWN)
1461			goto out;
1462		else if (sk->sk_state == TCP_CLOSE)
1463			ret = -ENOTCONN;
1464		else
1465			ret = -EAGAIN;
1466	} else if (unlikely(args.nr_skbs > IO_SKBS_PER_CALL_LIMIT) &&
1467		   (issue_flags & IO_URING_F_MULTISHOT)) {
1468		ret = IOU_REQUEUE;
1469	} else if (sock_flag(sk, SOCK_DONE)) {
1470		/* Make it to retry until it finally gets 0. */
1471		if (issue_flags & IO_URING_F_MULTISHOT)
1472			ret = IOU_REQUEUE;
1473		else
1474			ret = -EAGAIN;
1475	}
1476out:
1477	release_sock(sk);
1478	return ret;
1479}
1480
1481int io_zcrx_recv(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1482		 struct socket *sock, unsigned int flags,
1483		 unsigned issue_flags, unsigned int *len)
1484{
1485	struct sock *sk = sock->sk;
1486	const struct proto *prot = READ_ONCE(sk->sk_prot);
1487
1488	if (prot->recvmsg != tcp_recvmsg)
1489		return -EPROTONOSUPPORT;
1490
1491	sock_rps_record_flow(sk);
1492	return io_zcrx_tcp_recvmsg(req, ifq, sk, flags, issue_flags, len);
1493}