xref: /linux/io_uring/zcrx.c (revision 4f372263ef92ed2af55a8c226750b72021ff8d0f)
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 
12 #include <net/page_pool/helpers.h>
13 #include <net/page_pool/memory_provider.h>
14 #include <net/netlink.h>
15 #include <net/netdev_rx_queue.h>
16 #include <net/tcp.h>
17 #include <net/rps.h>
18 
19 #include <trace/events/page_pool.h>
20 
21 #include <uapi/linux/io_uring.h>
22 
23 #include "io_uring.h"
24 #include "kbuf.h"
25 #include "memmap.h"
26 #include "zcrx.h"
27 #include "rsrc.h"
28 
29 #define IO_DMA_ATTR (DMA_ATTR_SKIP_CPU_SYNC | DMA_ATTR_WEAK_ORDERING)
30 
31 static inline struct io_zcrx_ifq *io_pp_to_ifq(struct page_pool *pp)
32 {
33 	return pp->mp_priv;
34 }
35 
36 static inline struct io_zcrx_area *io_zcrx_iov_to_area(const struct net_iov *niov)
37 {
38 	struct net_iov_area *owner = net_iov_owner(niov);
39 
40 	return container_of(owner, struct io_zcrx_area, nia);
41 }
42 
43 static inline struct page *io_zcrx_iov_page(const struct net_iov *niov)
44 {
45 	struct io_zcrx_area *area = io_zcrx_iov_to_area(niov);
46 
47 	return area->mem.pages[net_iov_idx(niov)];
48 }
49 
50 static void io_release_dmabuf(struct io_zcrx_mem *mem)
51 {
52 	if (!IS_ENABLED(CONFIG_DMA_SHARED_BUFFER))
53 		return;
54 
55 	if (mem->sgt)
56 		dma_buf_unmap_attachment_unlocked(mem->attach, mem->sgt,
57 						  DMA_FROM_DEVICE);
58 	if (mem->attach)
59 		dma_buf_detach(mem->dmabuf, mem->attach);
60 	if (mem->dmabuf)
61 		dma_buf_put(mem->dmabuf);
62 
63 	mem->sgt = NULL;
64 	mem->attach = NULL;
65 	mem->dmabuf = NULL;
66 }
67 
68 static int io_import_dmabuf(struct io_zcrx_ifq *ifq,
69 			    struct io_zcrx_mem *mem,
70 			    struct io_uring_zcrx_area_reg *area_reg)
71 {
72 	unsigned long off = (unsigned long)area_reg->addr;
73 	unsigned long len = (unsigned long)area_reg->len;
74 	unsigned long total_size = 0;
75 	struct scatterlist *sg;
76 	int dmabuf_fd = area_reg->dmabuf_fd;
77 	int i, ret;
78 
79 	if (WARN_ON_ONCE(!ifq->dev))
80 		return -EFAULT;
81 	if (!IS_ENABLED(CONFIG_DMA_SHARED_BUFFER))
82 		return -EINVAL;
83 
84 	mem->is_dmabuf = true;
85 	mem->dmabuf = dma_buf_get(dmabuf_fd);
86 	if (IS_ERR(mem->dmabuf)) {
87 		ret = PTR_ERR(mem->dmabuf);
88 		mem->dmabuf = NULL;
89 		goto err;
90 	}
91 
92 	mem->attach = dma_buf_attach(mem->dmabuf, ifq->dev);
93 	if (IS_ERR(mem->attach)) {
94 		ret = PTR_ERR(mem->attach);
95 		mem->attach = NULL;
96 		goto err;
97 	}
98 
99 	mem->sgt = dma_buf_map_attachment_unlocked(mem->attach, DMA_FROM_DEVICE);
100 	if (IS_ERR(mem->sgt)) {
101 		ret = PTR_ERR(mem->sgt);
102 		mem->sgt = NULL;
103 		goto err;
104 	}
105 
106 	for_each_sgtable_dma_sg(mem->sgt, sg, i)
107 		total_size += sg_dma_len(sg);
108 
109 	if (total_size < off + len)
110 		return -EINVAL;
111 
112 	mem->dmabuf_offset = off;
113 	mem->size = len;
114 	return 0;
115 err:
116 	io_release_dmabuf(mem);
117 	return ret;
118 }
119 
120 static int io_zcrx_map_area_dmabuf(struct io_zcrx_ifq *ifq, struct io_zcrx_area *area)
121 {
122 	unsigned long off = area->mem.dmabuf_offset;
123 	struct scatterlist *sg;
124 	unsigned i, niov_idx = 0;
125 
126 	if (!IS_ENABLED(CONFIG_DMA_SHARED_BUFFER))
127 		return -EINVAL;
128 
129 	for_each_sgtable_dma_sg(area->mem.sgt, sg, i) {
130 		dma_addr_t dma = sg_dma_address(sg);
131 		unsigned long sg_len = sg_dma_len(sg);
132 		unsigned long sg_off = min(sg_len, off);
133 
134 		off -= sg_off;
135 		sg_len -= sg_off;
136 		dma += sg_off;
137 
138 		while (sg_len && niov_idx < area->nia.num_niovs) {
139 			struct net_iov *niov = &area->nia.niovs[niov_idx];
140 
141 			if (net_mp_niov_set_dma_addr(niov, dma))
142 				return 0;
143 			sg_len -= PAGE_SIZE;
144 			dma += PAGE_SIZE;
145 			niov_idx++;
146 		}
147 	}
148 	return niov_idx;
149 }
150 
151 static int io_import_umem(struct io_zcrx_ifq *ifq,
152 			  struct io_zcrx_mem *mem,
153 			  struct io_uring_zcrx_area_reg *area_reg)
154 {
155 	struct page **pages;
156 	int nr_pages;
157 
158 	if (area_reg->dmabuf_fd)
159 		return -EINVAL;
160 	if (!area_reg->addr)
161 		return -EFAULT;
162 	pages = io_pin_pages((unsigned long)area_reg->addr, area_reg->len,
163 				   &nr_pages);
164 	if (IS_ERR(pages))
165 		return PTR_ERR(pages);
166 
167 	mem->pages = pages;
168 	mem->nr_folios = nr_pages;
169 	mem->size = area_reg->len;
170 	return 0;
171 }
172 
173 static void io_release_area_mem(struct io_zcrx_mem *mem)
174 {
175 	if (mem->is_dmabuf) {
176 		io_release_dmabuf(mem);
177 		return;
178 	}
179 	if (mem->pages) {
180 		unpin_user_pages(mem->pages, mem->nr_folios);
181 		kvfree(mem->pages);
182 	}
183 }
184 
185 static int io_import_area(struct io_zcrx_ifq *ifq,
186 			  struct io_zcrx_mem *mem,
187 			  struct io_uring_zcrx_area_reg *area_reg)
188 {
189 	int ret;
190 
191 	ret = io_validate_user_buf_range(area_reg->addr, area_reg->len);
192 	if (ret)
193 		return ret;
194 	if (area_reg->addr & ~PAGE_MASK || area_reg->len & ~PAGE_MASK)
195 		return -EINVAL;
196 
197 	if (area_reg->flags & IORING_ZCRX_AREA_DMABUF)
198 		return io_import_dmabuf(ifq, mem, area_reg);
199 	return io_import_umem(ifq, mem, area_reg);
200 }
201 
202 static void io_zcrx_unmap_umem(struct io_zcrx_ifq *ifq,
203 				struct io_zcrx_area *area, int nr_mapped)
204 {
205 	int i;
206 
207 	for (i = 0; i < nr_mapped; i++) {
208 		netmem_ref netmem = net_iov_to_netmem(&area->nia.niovs[i]);
209 		dma_addr_t dma = page_pool_get_dma_addr_netmem(netmem);
210 
211 		dma_unmap_page_attrs(ifq->dev, dma, PAGE_SIZE,
212 				     DMA_FROM_DEVICE, IO_DMA_ATTR);
213 	}
214 }
215 
216 static void __io_zcrx_unmap_area(struct io_zcrx_ifq *ifq,
217 				 struct io_zcrx_area *area, int nr_mapped)
218 {
219 	int i;
220 
221 	if (area->mem.is_dmabuf)
222 		io_release_dmabuf(&area->mem);
223 	else
224 		io_zcrx_unmap_umem(ifq, area, nr_mapped);
225 
226 	for (i = 0; i < area->nia.num_niovs; i++)
227 		net_mp_niov_set_dma_addr(&area->nia.niovs[i], 0);
228 }
229 
230 static void io_zcrx_unmap_area(struct io_zcrx_ifq *ifq, struct io_zcrx_area *area)
231 {
232 	guard(mutex)(&ifq->dma_lock);
233 
234 	if (area->is_mapped)
235 		__io_zcrx_unmap_area(ifq, area, area->nia.num_niovs);
236 	area->is_mapped = false;
237 }
238 
239 static int io_zcrx_map_area_umem(struct io_zcrx_ifq *ifq, struct io_zcrx_area *area)
240 {
241 	int i;
242 
243 	for (i = 0; i < area->nia.num_niovs; i++) {
244 		struct net_iov *niov = &area->nia.niovs[i];
245 		dma_addr_t dma;
246 
247 		dma = dma_map_page_attrs(ifq->dev, area->mem.pages[i], 0,
248 					 PAGE_SIZE, DMA_FROM_DEVICE, IO_DMA_ATTR);
249 		if (dma_mapping_error(ifq->dev, dma))
250 			break;
251 		if (net_mp_niov_set_dma_addr(niov, dma)) {
252 			dma_unmap_page_attrs(ifq->dev, dma, PAGE_SIZE,
253 					     DMA_FROM_DEVICE, IO_DMA_ATTR);
254 			break;
255 		}
256 	}
257 	return i;
258 }
259 
260 static int io_zcrx_map_area(struct io_zcrx_ifq *ifq, struct io_zcrx_area *area)
261 {
262 	unsigned nr;
263 
264 	guard(mutex)(&ifq->dma_lock);
265 	if (area->is_mapped)
266 		return 0;
267 
268 	if (area->mem.is_dmabuf)
269 		nr = io_zcrx_map_area_dmabuf(ifq, area);
270 	else
271 		nr = io_zcrx_map_area_umem(ifq, area);
272 
273 	if (nr != area->nia.num_niovs) {
274 		__io_zcrx_unmap_area(ifq, area, nr);
275 		return -EINVAL;
276 	}
277 
278 	area->is_mapped = true;
279 	return 0;
280 }
281 
282 static void io_zcrx_sync_for_device(const struct page_pool *pool,
283 				    struct net_iov *niov)
284 {
285 #if defined(CONFIG_HAS_DMA) && defined(CONFIG_DMA_NEED_SYNC)
286 	dma_addr_t dma_addr;
287 
288 	if (!dma_dev_need_sync(pool->p.dev))
289 		return;
290 
291 	dma_addr = page_pool_get_dma_addr_netmem(net_iov_to_netmem(niov));
292 	__dma_sync_single_for_device(pool->p.dev, dma_addr + pool->p.offset,
293 				     PAGE_SIZE, pool->p.dma_dir);
294 #endif
295 }
296 
297 #define IO_RQ_MAX_ENTRIES		32768
298 
299 #define IO_SKBS_PER_CALL_LIMIT	20
300 
301 struct io_zcrx_args {
302 	struct io_kiocb		*req;
303 	struct io_zcrx_ifq	*ifq;
304 	struct socket		*sock;
305 	unsigned		nr_skbs;
306 };
307 
308 static const struct memory_provider_ops io_uring_pp_zc_ops;
309 
310 static inline atomic_t *io_get_user_counter(struct net_iov *niov)
311 {
312 	struct io_zcrx_area *area = io_zcrx_iov_to_area(niov);
313 
314 	return &area->user_refs[net_iov_idx(niov)];
315 }
316 
317 static bool io_zcrx_put_niov_uref(struct net_iov *niov)
318 {
319 	atomic_t *uref = io_get_user_counter(niov);
320 
321 	if (unlikely(!atomic_read(uref)))
322 		return false;
323 	atomic_dec(uref);
324 	return true;
325 }
326 
327 static void io_zcrx_get_niov_uref(struct net_iov *niov)
328 {
329 	atomic_inc(io_get_user_counter(niov));
330 }
331 
332 static int io_allocate_rbuf_ring(struct io_zcrx_ifq *ifq,
333 				 struct io_uring_zcrx_ifq_reg *reg,
334 				 struct io_uring_region_desc *rd,
335 				 u32 id)
336 {
337 	u64 mmap_offset;
338 	size_t off, size;
339 	void *ptr;
340 	int ret;
341 
342 	off = sizeof(struct io_uring);
343 	size = off + sizeof(struct io_uring_zcrx_rqe) * reg->rq_entries;
344 	if (size > rd->size)
345 		return -EINVAL;
346 
347 	mmap_offset = IORING_MAP_OFF_ZCRX_REGION;
348 	mmap_offset += id << IORING_OFF_PBUF_SHIFT;
349 
350 	ret = io_create_region(ifq->ctx, &ifq->region, rd, mmap_offset);
351 	if (ret < 0)
352 		return ret;
353 
354 	ptr = io_region_get_ptr(&ifq->region);
355 	ifq->rq_ring = (struct io_uring *)ptr;
356 	ifq->rqes = (struct io_uring_zcrx_rqe *)(ptr + off);
357 	return 0;
358 }
359 
360 static void io_free_rbuf_ring(struct io_zcrx_ifq *ifq)
361 {
362 	io_free_region(ifq->ctx, &ifq->region);
363 	ifq->rq_ring = NULL;
364 	ifq->rqes = NULL;
365 }
366 
367 static void io_zcrx_free_area(struct io_zcrx_area *area)
368 {
369 	io_zcrx_unmap_area(area->ifq, area);
370 	io_release_area_mem(&area->mem);
371 
372 	kvfree(area->freelist);
373 	kvfree(area->nia.niovs);
374 	kvfree(area->user_refs);
375 	kfree(area);
376 }
377 
378 #define IO_ZCRX_AREA_SUPPORTED_FLAGS	(IORING_ZCRX_AREA_DMABUF)
379 
380 static int io_zcrx_create_area(struct io_zcrx_ifq *ifq,
381 			       struct io_zcrx_area **res,
382 			       struct io_uring_zcrx_area_reg *area_reg)
383 {
384 	struct io_zcrx_area *area;
385 	unsigned nr_iovs;
386 	int i, ret;
387 
388 	if (area_reg->flags & ~IO_ZCRX_AREA_SUPPORTED_FLAGS)
389 		return -EINVAL;
390 	if (area_reg->rq_area_token)
391 		return -EINVAL;
392 	if (area_reg->__resv2[0] || area_reg->__resv2[1])
393 		return -EINVAL;
394 
395 	ret = -ENOMEM;
396 	area = kzalloc(sizeof(*area), GFP_KERNEL);
397 	if (!area)
398 		goto err;
399 
400 	ret = io_import_area(ifq, &area->mem, area_reg);
401 	if (ret)
402 		goto err;
403 
404 	nr_iovs = area->mem.size >> PAGE_SHIFT;
405 	area->nia.num_niovs = nr_iovs;
406 
407 	ret = -ENOMEM;
408 	area->nia.niovs = kvmalloc_array(nr_iovs, sizeof(area->nia.niovs[0]),
409 					 GFP_KERNEL | __GFP_ZERO);
410 	if (!area->nia.niovs)
411 		goto err;
412 
413 	area->freelist = kvmalloc_array(nr_iovs, sizeof(area->freelist[0]),
414 					GFP_KERNEL | __GFP_ZERO);
415 	if (!area->freelist)
416 		goto err;
417 
418 	area->user_refs = kvmalloc_array(nr_iovs, sizeof(area->user_refs[0]),
419 					GFP_KERNEL | __GFP_ZERO);
420 	if (!area->user_refs)
421 		goto err;
422 
423 	for (i = 0; i < nr_iovs; i++) {
424 		struct net_iov *niov = &area->nia.niovs[i];
425 
426 		niov->owner = &area->nia;
427 		area->freelist[i] = i;
428 		atomic_set(&area->user_refs[i], 0);
429 		niov->type = NET_IOV_IOURING;
430 	}
431 
432 	area->free_count = nr_iovs;
433 	area->ifq = ifq;
434 	/* we're only supporting one area per ifq for now */
435 	area->area_id = 0;
436 	area_reg->rq_area_token = (u64)area->area_id << IORING_ZCRX_AREA_SHIFT;
437 	spin_lock_init(&area->freelist_lock);
438 	*res = area;
439 	return 0;
440 err:
441 	if (area)
442 		io_zcrx_free_area(area);
443 	return ret;
444 }
445 
446 static struct io_zcrx_ifq *io_zcrx_ifq_alloc(struct io_ring_ctx *ctx)
447 {
448 	struct io_zcrx_ifq *ifq;
449 
450 	ifq = kzalloc(sizeof(*ifq), GFP_KERNEL);
451 	if (!ifq)
452 		return NULL;
453 
454 	ifq->if_rxq = -1;
455 	ifq->ctx = ctx;
456 	spin_lock_init(&ifq->lock);
457 	spin_lock_init(&ifq->rq_lock);
458 	mutex_init(&ifq->dma_lock);
459 	return ifq;
460 }
461 
462 static void io_zcrx_drop_netdev(struct io_zcrx_ifq *ifq)
463 {
464 	spin_lock(&ifq->lock);
465 	if (ifq->netdev) {
466 		netdev_put(ifq->netdev, &ifq->netdev_tracker);
467 		ifq->netdev = NULL;
468 	}
469 	spin_unlock(&ifq->lock);
470 }
471 
472 static void io_close_queue(struct io_zcrx_ifq *ifq)
473 {
474 	struct net_device *netdev;
475 	netdevice_tracker netdev_tracker;
476 	struct pp_memory_provider_params p = {
477 		.mp_ops = &io_uring_pp_zc_ops,
478 		.mp_priv = ifq,
479 	};
480 
481 	if (ifq->if_rxq == -1)
482 		return;
483 
484 	spin_lock(&ifq->lock);
485 	netdev = ifq->netdev;
486 	netdev_tracker = ifq->netdev_tracker;
487 	ifq->netdev = NULL;
488 	spin_unlock(&ifq->lock);
489 
490 	if (netdev) {
491 		net_mp_close_rxq(netdev, ifq->if_rxq, &p);
492 		netdev_put(netdev, &netdev_tracker);
493 	}
494 	ifq->if_rxq = -1;
495 }
496 
497 static void io_zcrx_ifq_free(struct io_zcrx_ifq *ifq)
498 {
499 	io_close_queue(ifq);
500 	io_zcrx_drop_netdev(ifq);
501 
502 	if (ifq->area)
503 		io_zcrx_free_area(ifq->area);
504 	if (ifq->dev)
505 		put_device(ifq->dev);
506 
507 	io_free_rbuf_ring(ifq);
508 	mutex_destroy(&ifq->dma_lock);
509 	kfree(ifq);
510 }
511 
512 struct io_mapped_region *io_zcrx_get_region(struct io_ring_ctx *ctx,
513 					    unsigned int id)
514 {
515 	struct io_zcrx_ifq *ifq = xa_load(&ctx->zcrx_ctxs, id);
516 
517 	lockdep_assert_held(&ctx->mmap_lock);
518 
519 	return ifq ? &ifq->region : NULL;
520 }
521 
522 int io_register_zcrx_ifq(struct io_ring_ctx *ctx,
523 			  struct io_uring_zcrx_ifq_reg __user *arg)
524 {
525 	struct pp_memory_provider_params mp_param = {};
526 	struct io_uring_zcrx_area_reg area;
527 	struct io_uring_zcrx_ifq_reg reg;
528 	struct io_uring_region_desc rd;
529 	struct io_zcrx_ifq *ifq;
530 	int ret;
531 	u32 id;
532 
533 	/*
534 	 * 1. Interface queue allocation.
535 	 * 2. It can observe data destined for sockets of other tasks.
536 	 */
537 	if (!capable(CAP_NET_ADMIN))
538 		return -EPERM;
539 
540 	/* mandatory io_uring features for zc rx */
541 	if (!(ctx->flags & IORING_SETUP_DEFER_TASKRUN &&
542 	      ctx->flags & IORING_SETUP_CQE32))
543 		return -EINVAL;
544 	if (copy_from_user(&reg, arg, sizeof(reg)))
545 		return -EFAULT;
546 	if (copy_from_user(&rd, u64_to_user_ptr(reg.region_ptr), sizeof(rd)))
547 		return -EFAULT;
548 	if (memchr_inv(&reg.__resv, 0, sizeof(reg.__resv)) ||
549 	    reg.__resv2 || reg.zcrx_id)
550 		return -EINVAL;
551 	if (reg.if_rxq == -1 || !reg.rq_entries || reg.flags)
552 		return -EINVAL;
553 	if (reg.rq_entries > IO_RQ_MAX_ENTRIES) {
554 		if (!(ctx->flags & IORING_SETUP_CLAMP))
555 			return -EINVAL;
556 		reg.rq_entries = IO_RQ_MAX_ENTRIES;
557 	}
558 	reg.rq_entries = roundup_pow_of_two(reg.rq_entries);
559 
560 	if (copy_from_user(&area, u64_to_user_ptr(reg.area_ptr), sizeof(area)))
561 		return -EFAULT;
562 
563 	ifq = io_zcrx_ifq_alloc(ctx);
564 	if (!ifq)
565 		return -ENOMEM;
566 	ifq->rq_entries = reg.rq_entries;
567 
568 	scoped_guard(mutex, &ctx->mmap_lock) {
569 		/* preallocate id */
570 		ret = xa_alloc(&ctx->zcrx_ctxs, &id, NULL, xa_limit_31b, GFP_KERNEL);
571 		if (ret)
572 			goto ifq_free;
573 	}
574 
575 	ret = io_allocate_rbuf_ring(ifq, &reg, &rd, id);
576 	if (ret)
577 		goto err;
578 
579 	ifq->netdev = netdev_get_by_index(current->nsproxy->net_ns, reg.if_idx,
580 					  &ifq->netdev_tracker, GFP_KERNEL);
581 	if (!ifq->netdev) {
582 		ret = -ENODEV;
583 		goto err;
584 	}
585 
586 	ifq->dev = ifq->netdev->dev.parent;
587 	if (!ifq->dev) {
588 		ret = -EOPNOTSUPP;
589 		goto err;
590 	}
591 	get_device(ifq->dev);
592 
593 	ret = io_zcrx_create_area(ifq, &ifq->area, &area);
594 	if (ret)
595 		goto err;
596 
597 	mp_param.mp_ops = &io_uring_pp_zc_ops;
598 	mp_param.mp_priv = ifq;
599 	ret = net_mp_open_rxq(ifq->netdev, reg.if_rxq, &mp_param);
600 	if (ret)
601 		goto err;
602 	ifq->if_rxq = reg.if_rxq;
603 
604 	reg.offsets.rqes = sizeof(struct io_uring);
605 	reg.offsets.head = offsetof(struct io_uring, head);
606 	reg.offsets.tail = offsetof(struct io_uring, tail);
607 	reg.zcrx_id = id;
608 
609 	scoped_guard(mutex, &ctx->mmap_lock) {
610 		/* publish ifq */
611 		ret = -ENOMEM;
612 		if (xa_store(&ctx->zcrx_ctxs, id, ifq, GFP_KERNEL))
613 			goto err;
614 	}
615 
616 	if (copy_to_user(arg, &reg, sizeof(reg)) ||
617 	    copy_to_user(u64_to_user_ptr(reg.region_ptr), &rd, sizeof(rd)) ||
618 	    copy_to_user(u64_to_user_ptr(reg.area_ptr), &area, sizeof(area))) {
619 		ret = -EFAULT;
620 		goto err;
621 	}
622 	return 0;
623 err:
624 	scoped_guard(mutex, &ctx->mmap_lock)
625 		xa_erase(&ctx->zcrx_ctxs, id);
626 ifq_free:
627 	io_zcrx_ifq_free(ifq);
628 	return ret;
629 }
630 
631 void io_unregister_zcrx_ifqs(struct io_ring_ctx *ctx)
632 {
633 	struct io_zcrx_ifq *ifq;
634 	unsigned long id;
635 
636 	lockdep_assert_held(&ctx->uring_lock);
637 
638 	while (1) {
639 		scoped_guard(mutex, &ctx->mmap_lock) {
640 			ifq = xa_find(&ctx->zcrx_ctxs, &id, ULONG_MAX, XA_PRESENT);
641 			if (ifq)
642 				xa_erase(&ctx->zcrx_ctxs, id);
643 		}
644 		if (!ifq)
645 			break;
646 		io_zcrx_ifq_free(ifq);
647 	}
648 
649 	xa_destroy(&ctx->zcrx_ctxs);
650 }
651 
652 static struct net_iov *__io_zcrx_get_free_niov(struct io_zcrx_area *area)
653 {
654 	unsigned niov_idx;
655 
656 	lockdep_assert_held(&area->freelist_lock);
657 
658 	niov_idx = area->freelist[--area->free_count];
659 	return &area->nia.niovs[niov_idx];
660 }
661 
662 static void io_zcrx_return_niov_freelist(struct net_iov *niov)
663 {
664 	struct io_zcrx_area *area = io_zcrx_iov_to_area(niov);
665 
666 	spin_lock_bh(&area->freelist_lock);
667 	area->freelist[area->free_count++] = net_iov_idx(niov);
668 	spin_unlock_bh(&area->freelist_lock);
669 }
670 
671 static void io_zcrx_return_niov(struct net_iov *niov)
672 {
673 	netmem_ref netmem = net_iov_to_netmem(niov);
674 
675 	if (!niov->pp) {
676 		/* copy fallback allocated niovs */
677 		io_zcrx_return_niov_freelist(niov);
678 		return;
679 	}
680 	page_pool_put_unrefed_netmem(niov->pp, netmem, -1, false);
681 }
682 
683 static void io_zcrx_scrub(struct io_zcrx_ifq *ifq)
684 {
685 	struct io_zcrx_area *area = ifq->area;
686 	int i;
687 
688 	if (!area)
689 		return;
690 
691 	/* Reclaim back all buffers given to the user space. */
692 	for (i = 0; i < area->nia.num_niovs; i++) {
693 		struct net_iov *niov = &area->nia.niovs[i];
694 		int nr;
695 
696 		if (!atomic_read(io_get_user_counter(niov)))
697 			continue;
698 		nr = atomic_xchg(io_get_user_counter(niov), 0);
699 		if (nr && !page_pool_unref_netmem(net_iov_to_netmem(niov), nr))
700 			io_zcrx_return_niov(niov);
701 	}
702 }
703 
704 void io_shutdown_zcrx_ifqs(struct io_ring_ctx *ctx)
705 {
706 	struct io_zcrx_ifq *ifq;
707 	unsigned long index;
708 
709 	lockdep_assert_held(&ctx->uring_lock);
710 
711 	xa_for_each(&ctx->zcrx_ctxs, index, ifq) {
712 		io_zcrx_scrub(ifq);
713 		io_close_queue(ifq);
714 	}
715 }
716 
717 static inline u32 io_zcrx_rqring_entries(struct io_zcrx_ifq *ifq)
718 {
719 	u32 entries;
720 
721 	entries = smp_load_acquire(&ifq->rq_ring->tail) - ifq->cached_rq_head;
722 	return min(entries, ifq->rq_entries);
723 }
724 
725 static struct io_uring_zcrx_rqe *io_zcrx_get_rqe(struct io_zcrx_ifq *ifq,
726 						 unsigned mask)
727 {
728 	unsigned int idx = ifq->cached_rq_head++ & mask;
729 
730 	return &ifq->rqes[idx];
731 }
732 
733 static void io_zcrx_ring_refill(struct page_pool *pp,
734 				struct io_zcrx_ifq *ifq)
735 {
736 	unsigned int mask = ifq->rq_entries - 1;
737 	unsigned int entries;
738 	netmem_ref netmem;
739 
740 	spin_lock_bh(&ifq->rq_lock);
741 
742 	entries = io_zcrx_rqring_entries(ifq);
743 	entries = min_t(unsigned, entries, PP_ALLOC_CACHE_REFILL - pp->alloc.count);
744 	if (unlikely(!entries)) {
745 		spin_unlock_bh(&ifq->rq_lock);
746 		return;
747 	}
748 
749 	do {
750 		struct io_uring_zcrx_rqe *rqe = io_zcrx_get_rqe(ifq, mask);
751 		struct io_zcrx_area *area;
752 		struct net_iov *niov;
753 		unsigned niov_idx, area_idx;
754 
755 		area_idx = rqe->off >> IORING_ZCRX_AREA_SHIFT;
756 		niov_idx = (rqe->off & ~IORING_ZCRX_AREA_MASK) >> PAGE_SHIFT;
757 
758 		if (unlikely(rqe->__pad || area_idx))
759 			continue;
760 		area = ifq->area;
761 
762 		if (unlikely(niov_idx >= area->nia.num_niovs))
763 			continue;
764 		niov_idx = array_index_nospec(niov_idx, area->nia.num_niovs);
765 
766 		niov = &area->nia.niovs[niov_idx];
767 		if (!io_zcrx_put_niov_uref(niov))
768 			continue;
769 
770 		netmem = net_iov_to_netmem(niov);
771 		if (page_pool_unref_netmem(netmem, 1) != 0)
772 			continue;
773 
774 		if (unlikely(niov->pp != pp)) {
775 			io_zcrx_return_niov(niov);
776 			continue;
777 		}
778 
779 		io_zcrx_sync_for_device(pp, niov);
780 		net_mp_netmem_place_in_cache(pp, netmem);
781 	} while (--entries);
782 
783 	smp_store_release(&ifq->rq_ring->head, ifq->cached_rq_head);
784 	spin_unlock_bh(&ifq->rq_lock);
785 }
786 
787 static void io_zcrx_refill_slow(struct page_pool *pp, struct io_zcrx_ifq *ifq)
788 {
789 	struct io_zcrx_area *area = ifq->area;
790 
791 	spin_lock_bh(&area->freelist_lock);
792 	while (area->free_count && pp->alloc.count < PP_ALLOC_CACHE_REFILL) {
793 		struct net_iov *niov = __io_zcrx_get_free_niov(area);
794 		netmem_ref netmem = net_iov_to_netmem(niov);
795 
796 		net_mp_niov_set_page_pool(pp, niov);
797 		io_zcrx_sync_for_device(pp, niov);
798 		net_mp_netmem_place_in_cache(pp, netmem);
799 	}
800 	spin_unlock_bh(&area->freelist_lock);
801 }
802 
803 static netmem_ref io_pp_zc_alloc_netmems(struct page_pool *pp, gfp_t gfp)
804 {
805 	struct io_zcrx_ifq *ifq = io_pp_to_ifq(pp);
806 
807 	/* pp should already be ensuring that */
808 	if (unlikely(pp->alloc.count))
809 		goto out_return;
810 
811 	io_zcrx_ring_refill(pp, ifq);
812 	if (likely(pp->alloc.count))
813 		goto out_return;
814 
815 	io_zcrx_refill_slow(pp, ifq);
816 	if (!pp->alloc.count)
817 		return 0;
818 out_return:
819 	return pp->alloc.cache[--pp->alloc.count];
820 }
821 
822 static bool io_pp_zc_release_netmem(struct page_pool *pp, netmem_ref netmem)
823 {
824 	struct net_iov *niov;
825 
826 	if (WARN_ON_ONCE(!netmem_is_net_iov(netmem)))
827 		return false;
828 
829 	niov = netmem_to_net_iov(netmem);
830 	net_mp_niov_clear_page_pool(niov);
831 	io_zcrx_return_niov_freelist(niov);
832 	return false;
833 }
834 
835 static int io_pp_zc_init(struct page_pool *pp)
836 {
837 	struct io_zcrx_ifq *ifq = io_pp_to_ifq(pp);
838 	int ret;
839 
840 	if (WARN_ON_ONCE(!ifq))
841 		return -EINVAL;
842 	if (WARN_ON_ONCE(ifq->dev != pp->p.dev))
843 		return -EINVAL;
844 	if (WARN_ON_ONCE(!pp->dma_map))
845 		return -EOPNOTSUPP;
846 	if (pp->p.order != 0)
847 		return -EOPNOTSUPP;
848 	if (pp->p.dma_dir != DMA_FROM_DEVICE)
849 		return -EOPNOTSUPP;
850 
851 	ret = io_zcrx_map_area(ifq, ifq->area);
852 	if (ret)
853 		return ret;
854 
855 	percpu_ref_get(&ifq->ctx->refs);
856 	return 0;
857 }
858 
859 static void io_pp_zc_destroy(struct page_pool *pp)
860 {
861 	struct io_zcrx_ifq *ifq = io_pp_to_ifq(pp);
862 	struct io_zcrx_area *area = ifq->area;
863 
864 	if (WARN_ON_ONCE(area->free_count != area->nia.num_niovs))
865 		return;
866 	percpu_ref_put(&ifq->ctx->refs);
867 }
868 
869 static int io_pp_nl_fill(void *mp_priv, struct sk_buff *rsp,
870 			 struct netdev_rx_queue *rxq)
871 {
872 	struct nlattr *nest;
873 	int type;
874 
875 	type = rxq ? NETDEV_A_QUEUE_IO_URING : NETDEV_A_PAGE_POOL_IO_URING;
876 	nest = nla_nest_start(rsp, type);
877 	if (!nest)
878 		return -EMSGSIZE;
879 	nla_nest_end(rsp, nest);
880 
881 	return 0;
882 }
883 
884 static void io_pp_uninstall(void *mp_priv, struct netdev_rx_queue *rxq)
885 {
886 	struct pp_memory_provider_params *p = &rxq->mp_params;
887 	struct io_zcrx_ifq *ifq = mp_priv;
888 
889 	io_zcrx_drop_netdev(ifq);
890 	if (ifq->area)
891 		io_zcrx_unmap_area(ifq, ifq->area);
892 
893 	p->mp_ops = NULL;
894 	p->mp_priv = NULL;
895 }
896 
897 static const struct memory_provider_ops io_uring_pp_zc_ops = {
898 	.alloc_netmems		= io_pp_zc_alloc_netmems,
899 	.release_netmem		= io_pp_zc_release_netmem,
900 	.init			= io_pp_zc_init,
901 	.destroy		= io_pp_zc_destroy,
902 	.nl_fill		= io_pp_nl_fill,
903 	.uninstall		= io_pp_uninstall,
904 };
905 
906 static bool io_zcrx_queue_cqe(struct io_kiocb *req, struct net_iov *niov,
907 			      struct io_zcrx_ifq *ifq, int off, int len)
908 {
909 	struct io_uring_zcrx_cqe *rcqe;
910 	struct io_zcrx_area *area;
911 	struct io_uring_cqe *cqe;
912 	u64 offset;
913 
914 	if (!io_defer_get_uncommited_cqe(req->ctx, &cqe))
915 		return false;
916 
917 	cqe->user_data = req->cqe.user_data;
918 	cqe->res = len;
919 	cqe->flags = IORING_CQE_F_MORE;
920 
921 	area = io_zcrx_iov_to_area(niov);
922 	offset = off + (net_iov_idx(niov) << PAGE_SHIFT);
923 	rcqe = (struct io_uring_zcrx_cqe *)(cqe + 1);
924 	rcqe->off = offset + ((u64)area->area_id << IORING_ZCRX_AREA_SHIFT);
925 	rcqe->__pad = 0;
926 	return true;
927 }
928 
929 static struct net_iov *io_zcrx_alloc_fallback(struct io_zcrx_area *area)
930 {
931 	struct net_iov *niov = NULL;
932 
933 	spin_lock_bh(&area->freelist_lock);
934 	if (area->free_count)
935 		niov = __io_zcrx_get_free_niov(area);
936 	spin_unlock_bh(&area->freelist_lock);
937 
938 	if (niov)
939 		page_pool_fragment_netmem(net_iov_to_netmem(niov), 1);
940 	return niov;
941 }
942 
943 static ssize_t io_zcrx_copy_chunk(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
944 				  void *src_base, struct page *src_page,
945 				  unsigned int src_offset, size_t len)
946 {
947 	struct io_zcrx_area *area = ifq->area;
948 	size_t copied = 0;
949 	int ret = 0;
950 
951 	if (area->mem.is_dmabuf)
952 		return -EFAULT;
953 
954 	while (len) {
955 		size_t copy_size = min_t(size_t, PAGE_SIZE, len);
956 		const int dst_off = 0;
957 		struct net_iov *niov;
958 		struct page *dst_page;
959 		void *dst_addr;
960 
961 		niov = io_zcrx_alloc_fallback(area);
962 		if (!niov) {
963 			ret = -ENOMEM;
964 			break;
965 		}
966 
967 		dst_page = io_zcrx_iov_page(niov);
968 		dst_addr = kmap_local_page(dst_page);
969 		if (src_page)
970 			src_base = kmap_local_page(src_page);
971 
972 		memcpy(dst_addr, src_base + src_offset, copy_size);
973 
974 		if (src_page)
975 			kunmap_local(src_base);
976 		kunmap_local(dst_addr);
977 
978 		if (!io_zcrx_queue_cqe(req, niov, ifq, dst_off, copy_size)) {
979 			io_zcrx_return_niov(niov);
980 			ret = -ENOSPC;
981 			break;
982 		}
983 
984 		io_zcrx_get_niov_uref(niov);
985 		src_offset += copy_size;
986 		len -= copy_size;
987 		copied += copy_size;
988 	}
989 
990 	return copied ? copied : ret;
991 }
992 
993 static int io_zcrx_copy_frag(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
994 			     const skb_frag_t *frag, int off, int len)
995 {
996 	struct page *page = skb_frag_page(frag);
997 	u32 p_off, p_len, t, copied = 0;
998 	int ret = 0;
999 
1000 	off += skb_frag_off(frag);
1001 
1002 	skb_frag_foreach_page(frag, off, len,
1003 			      page, p_off, p_len, t) {
1004 		ret = io_zcrx_copy_chunk(req, ifq, NULL, page, p_off, p_len);
1005 		if (ret < 0)
1006 			return copied ? copied : ret;
1007 		copied += ret;
1008 	}
1009 	return copied;
1010 }
1011 
1012 static int io_zcrx_recv_frag(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1013 			     const skb_frag_t *frag, int off, int len)
1014 {
1015 	struct net_iov *niov;
1016 
1017 	if (unlikely(!skb_frag_is_net_iov(frag)))
1018 		return io_zcrx_copy_frag(req, ifq, frag, off, len);
1019 
1020 	niov = netmem_to_net_iov(frag->netmem);
1021 	if (!niov->pp || niov->pp->mp_ops != &io_uring_pp_zc_ops ||
1022 	    io_pp_to_ifq(niov->pp) != ifq)
1023 		return -EFAULT;
1024 
1025 	if (!io_zcrx_queue_cqe(req, niov, ifq, off + skb_frag_off(frag), len))
1026 		return -ENOSPC;
1027 
1028 	/*
1029 	 * Prevent it from being recycled while user is accessing it.
1030 	 * It has to be done before grabbing a user reference.
1031 	 */
1032 	page_pool_ref_netmem(net_iov_to_netmem(niov));
1033 	io_zcrx_get_niov_uref(niov);
1034 	return len;
1035 }
1036 
1037 static int
1038 io_zcrx_recv_skb(read_descriptor_t *desc, struct sk_buff *skb,
1039 		 unsigned int offset, size_t len)
1040 {
1041 	struct io_zcrx_args *args = desc->arg.data;
1042 	struct io_zcrx_ifq *ifq = args->ifq;
1043 	struct io_kiocb *req = args->req;
1044 	struct sk_buff *frag_iter;
1045 	unsigned start, start_off = offset;
1046 	int i, copy, end, off;
1047 	int ret = 0;
1048 
1049 	len = min_t(size_t, len, desc->count);
1050 	/*
1051 	 * __tcp_read_sock() always calls io_zcrx_recv_skb one last time, even
1052 	 * if desc->count is already 0. This is caused by the if (offset + 1 !=
1053 	 * skb->len) check. Return early in this case to break out of
1054 	 * __tcp_read_sock().
1055 	 */
1056 	if (!len)
1057 		return 0;
1058 	if (unlikely(args->nr_skbs++ > IO_SKBS_PER_CALL_LIMIT))
1059 		return -EAGAIN;
1060 
1061 	if (unlikely(offset < skb_headlen(skb))) {
1062 		ssize_t copied;
1063 		size_t to_copy;
1064 
1065 		to_copy = min_t(size_t, skb_headlen(skb) - offset, len);
1066 		copied = io_zcrx_copy_chunk(req, ifq, skb->data, NULL,
1067 					    offset, to_copy);
1068 		if (copied < 0) {
1069 			ret = copied;
1070 			goto out;
1071 		}
1072 		offset += copied;
1073 		len -= copied;
1074 		if (!len)
1075 			goto out;
1076 		if (offset != skb_headlen(skb))
1077 			goto out;
1078 	}
1079 
1080 	start = skb_headlen(skb);
1081 
1082 	for (i = 0; i < skb_shinfo(skb)->nr_frags; i++) {
1083 		const skb_frag_t *frag;
1084 
1085 		if (WARN_ON(start > offset + len))
1086 			return -EFAULT;
1087 
1088 		frag = &skb_shinfo(skb)->frags[i];
1089 		end = start + skb_frag_size(frag);
1090 
1091 		if (offset < end) {
1092 			copy = end - offset;
1093 			if (copy > len)
1094 				copy = len;
1095 
1096 			off = offset - start;
1097 			ret = io_zcrx_recv_frag(req, ifq, frag, off, copy);
1098 			if (ret < 0)
1099 				goto out;
1100 
1101 			offset += ret;
1102 			len -= ret;
1103 			if (len == 0 || ret != copy)
1104 				goto out;
1105 		}
1106 		start = end;
1107 	}
1108 
1109 	skb_walk_frags(skb, frag_iter) {
1110 		if (WARN_ON(start > offset + len))
1111 			return -EFAULT;
1112 
1113 		end = start + frag_iter->len;
1114 		if (offset < end) {
1115 			copy = end - offset;
1116 			if (copy > len)
1117 				copy = len;
1118 
1119 			off = offset - start;
1120 			ret = io_zcrx_recv_skb(desc, frag_iter, off, copy);
1121 			if (ret < 0)
1122 				goto out;
1123 
1124 			offset += ret;
1125 			len -= ret;
1126 			if (len == 0 || ret != copy)
1127 				goto out;
1128 		}
1129 		start = end;
1130 	}
1131 
1132 out:
1133 	if (offset == start_off)
1134 		return ret;
1135 	desc->count -= (offset - start_off);
1136 	return offset - start_off;
1137 }
1138 
1139 static int io_zcrx_tcp_recvmsg(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1140 				struct sock *sk, int flags,
1141 				unsigned issue_flags, unsigned int *outlen)
1142 {
1143 	unsigned int len = *outlen;
1144 	struct io_zcrx_args args = {
1145 		.req = req,
1146 		.ifq = ifq,
1147 		.sock = sk->sk_socket,
1148 	};
1149 	read_descriptor_t rd_desc = {
1150 		.count = len ? len : UINT_MAX,
1151 		.arg.data = &args,
1152 	};
1153 	int ret;
1154 
1155 	lock_sock(sk);
1156 	ret = tcp_read_sock(sk, &rd_desc, io_zcrx_recv_skb);
1157 	if (len && ret > 0)
1158 		*outlen = len - ret;
1159 	if (ret <= 0) {
1160 		if (ret < 0 || sock_flag(sk, SOCK_DONE))
1161 			goto out;
1162 		if (sk->sk_err)
1163 			ret = sock_error(sk);
1164 		else if (sk->sk_shutdown & RCV_SHUTDOWN)
1165 			goto out;
1166 		else if (sk->sk_state == TCP_CLOSE)
1167 			ret = -ENOTCONN;
1168 		else
1169 			ret = -EAGAIN;
1170 	} else if (unlikely(args.nr_skbs > IO_SKBS_PER_CALL_LIMIT) &&
1171 		   (issue_flags & IO_URING_F_MULTISHOT)) {
1172 		ret = IOU_REQUEUE;
1173 	} else if (sock_flag(sk, SOCK_DONE)) {
1174 		/* Make it to retry until it finally gets 0. */
1175 		if (issue_flags & IO_URING_F_MULTISHOT)
1176 			ret = IOU_REQUEUE;
1177 		else
1178 			ret = -EAGAIN;
1179 	}
1180 out:
1181 	release_sock(sk);
1182 	return ret;
1183 }
1184 
1185 int io_zcrx_recv(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1186 		 struct socket *sock, unsigned int flags,
1187 		 unsigned issue_flags, unsigned int *len)
1188 {
1189 	struct sock *sk = sock->sk;
1190 	const struct proto *prot = READ_ONCE(sk->sk_prot);
1191 
1192 	if (prot->recvmsg != tcp_recvmsg)
1193 		return -EPROTONOSUPPORT;
1194 
1195 	sock_rps_record_flow(sk);
1196 	return io_zcrx_tcp_recvmsg(req, ifq, sk, flags, issue_flags, len);
1197 }
1198