xref: /linux/io_uring/zcrx.c (revision 15ecd83dc06277385ad71dc7ea26911d9a79acaf)
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 	}
430 
431 	area->free_count = nr_iovs;
432 	area->ifq = ifq;
433 	/* we're only supporting one area per ifq for now */
434 	area->area_id = 0;
435 	area_reg->rq_area_token = (u64)area->area_id << IORING_ZCRX_AREA_SHIFT;
436 	spin_lock_init(&area->freelist_lock);
437 	*res = area;
438 	return 0;
439 err:
440 	if (area)
441 		io_zcrx_free_area(area);
442 	return ret;
443 }
444 
445 static struct io_zcrx_ifq *io_zcrx_ifq_alloc(struct io_ring_ctx *ctx)
446 {
447 	struct io_zcrx_ifq *ifq;
448 
449 	ifq = kzalloc(sizeof(*ifq), GFP_KERNEL);
450 	if (!ifq)
451 		return NULL;
452 
453 	ifq->if_rxq = -1;
454 	ifq->ctx = ctx;
455 	spin_lock_init(&ifq->lock);
456 	spin_lock_init(&ifq->rq_lock);
457 	mutex_init(&ifq->dma_lock);
458 	return ifq;
459 }
460 
461 static void io_zcrx_drop_netdev(struct io_zcrx_ifq *ifq)
462 {
463 	spin_lock(&ifq->lock);
464 	if (ifq->netdev) {
465 		netdev_put(ifq->netdev, &ifq->netdev_tracker);
466 		ifq->netdev = NULL;
467 	}
468 	spin_unlock(&ifq->lock);
469 }
470 
471 static void io_close_queue(struct io_zcrx_ifq *ifq)
472 {
473 	struct net_device *netdev;
474 	netdevice_tracker netdev_tracker;
475 	struct pp_memory_provider_params p = {
476 		.mp_ops = &io_uring_pp_zc_ops,
477 		.mp_priv = ifq,
478 	};
479 
480 	if (ifq->if_rxq == -1)
481 		return;
482 
483 	spin_lock(&ifq->lock);
484 	netdev = ifq->netdev;
485 	netdev_tracker = ifq->netdev_tracker;
486 	ifq->netdev = NULL;
487 	spin_unlock(&ifq->lock);
488 
489 	if (netdev) {
490 		net_mp_close_rxq(netdev, ifq->if_rxq, &p);
491 		netdev_put(netdev, &netdev_tracker);
492 	}
493 	ifq->if_rxq = -1;
494 }
495 
496 static void io_zcrx_ifq_free(struct io_zcrx_ifq *ifq)
497 {
498 	io_close_queue(ifq);
499 	io_zcrx_drop_netdev(ifq);
500 
501 	if (ifq->area)
502 		io_zcrx_free_area(ifq->area);
503 	if (ifq->dev)
504 		put_device(ifq->dev);
505 
506 	io_free_rbuf_ring(ifq);
507 	mutex_destroy(&ifq->dma_lock);
508 	kfree(ifq);
509 }
510 
511 struct io_mapped_region *io_zcrx_get_region(struct io_ring_ctx *ctx,
512 					    unsigned int id)
513 {
514 	struct io_zcrx_ifq *ifq = xa_load(&ctx->zcrx_ctxs, id);
515 
516 	lockdep_assert_held(&ctx->mmap_lock);
517 
518 	return ifq ? &ifq->region : NULL;
519 }
520 
521 int io_register_zcrx_ifq(struct io_ring_ctx *ctx,
522 			  struct io_uring_zcrx_ifq_reg __user *arg)
523 {
524 	struct pp_memory_provider_params mp_param = {};
525 	struct io_uring_zcrx_area_reg area;
526 	struct io_uring_zcrx_ifq_reg reg;
527 	struct io_uring_region_desc rd;
528 	struct io_zcrx_ifq *ifq;
529 	int ret;
530 	u32 id;
531 
532 	/*
533 	 * 1. Interface queue allocation.
534 	 * 2. It can observe data destined for sockets of other tasks.
535 	 */
536 	if (!capable(CAP_NET_ADMIN))
537 		return -EPERM;
538 
539 	/* mandatory io_uring features for zc rx */
540 	if (!(ctx->flags & IORING_SETUP_DEFER_TASKRUN &&
541 	      ctx->flags & IORING_SETUP_CQE32))
542 		return -EINVAL;
543 	if (copy_from_user(&reg, arg, sizeof(reg)))
544 		return -EFAULT;
545 	if (copy_from_user(&rd, u64_to_user_ptr(reg.region_ptr), sizeof(rd)))
546 		return -EFAULT;
547 	if (memchr_inv(&reg.__resv, 0, sizeof(reg.__resv)) ||
548 	    reg.__resv2 || reg.zcrx_id)
549 		return -EINVAL;
550 	if (reg.if_rxq == -1 || !reg.rq_entries || reg.flags)
551 		return -EINVAL;
552 	if (reg.rq_entries > IO_RQ_MAX_ENTRIES) {
553 		if (!(ctx->flags & IORING_SETUP_CLAMP))
554 			return -EINVAL;
555 		reg.rq_entries = IO_RQ_MAX_ENTRIES;
556 	}
557 	reg.rq_entries = roundup_pow_of_two(reg.rq_entries);
558 
559 	if (copy_from_user(&area, u64_to_user_ptr(reg.area_ptr), sizeof(area)))
560 		return -EFAULT;
561 
562 	ifq = io_zcrx_ifq_alloc(ctx);
563 	if (!ifq)
564 		return -ENOMEM;
565 	ifq->rq_entries = reg.rq_entries;
566 
567 	scoped_guard(mutex, &ctx->mmap_lock) {
568 		/* preallocate id */
569 		ret = xa_alloc(&ctx->zcrx_ctxs, &id, NULL, xa_limit_31b, GFP_KERNEL);
570 		if (ret)
571 			goto ifq_free;
572 	}
573 
574 	ret = io_allocate_rbuf_ring(ifq, &reg, &rd, id);
575 	if (ret)
576 		goto err;
577 
578 	ifq->netdev = netdev_get_by_index(current->nsproxy->net_ns, reg.if_idx,
579 					  &ifq->netdev_tracker, GFP_KERNEL);
580 	if (!ifq->netdev) {
581 		ret = -ENODEV;
582 		goto err;
583 	}
584 
585 	ifq->dev = ifq->netdev->dev.parent;
586 	if (!ifq->dev) {
587 		ret = -EOPNOTSUPP;
588 		goto err;
589 	}
590 	get_device(ifq->dev);
591 
592 	ret = io_zcrx_create_area(ifq, &ifq->area, &area);
593 	if (ret)
594 		goto err;
595 
596 	mp_param.mp_ops = &io_uring_pp_zc_ops;
597 	mp_param.mp_priv = ifq;
598 	ret = net_mp_open_rxq(ifq->netdev, reg.if_rxq, &mp_param);
599 	if (ret)
600 		goto err;
601 	ifq->if_rxq = reg.if_rxq;
602 
603 	reg.offsets.rqes = sizeof(struct io_uring);
604 	reg.offsets.head = offsetof(struct io_uring, head);
605 	reg.offsets.tail = offsetof(struct io_uring, tail);
606 	reg.zcrx_id = id;
607 
608 	scoped_guard(mutex, &ctx->mmap_lock) {
609 		/* publish ifq */
610 		ret = -ENOMEM;
611 		if (xa_store(&ctx->zcrx_ctxs, id, ifq, GFP_KERNEL))
612 			goto err;
613 	}
614 
615 	if (copy_to_user(arg, &reg, sizeof(reg)) ||
616 	    copy_to_user(u64_to_user_ptr(reg.region_ptr), &rd, sizeof(rd)) ||
617 	    copy_to_user(u64_to_user_ptr(reg.area_ptr), &area, sizeof(area))) {
618 		ret = -EFAULT;
619 		goto err;
620 	}
621 	return 0;
622 err:
623 	scoped_guard(mutex, &ctx->mmap_lock)
624 		xa_erase(&ctx->zcrx_ctxs, id);
625 ifq_free:
626 	io_zcrx_ifq_free(ifq);
627 	return ret;
628 }
629 
630 void io_unregister_zcrx_ifqs(struct io_ring_ctx *ctx)
631 {
632 	struct io_zcrx_ifq *ifq;
633 	unsigned long id;
634 
635 	lockdep_assert_held(&ctx->uring_lock);
636 
637 	while (1) {
638 		scoped_guard(mutex, &ctx->mmap_lock) {
639 			ifq = xa_find(&ctx->zcrx_ctxs, &id, ULONG_MAX, XA_PRESENT);
640 			if (ifq)
641 				xa_erase(&ctx->zcrx_ctxs, id);
642 		}
643 		if (!ifq)
644 			break;
645 		io_zcrx_ifq_free(ifq);
646 	}
647 
648 	xa_destroy(&ctx->zcrx_ctxs);
649 }
650 
651 static struct net_iov *__io_zcrx_get_free_niov(struct io_zcrx_area *area)
652 {
653 	unsigned niov_idx;
654 
655 	lockdep_assert_held(&area->freelist_lock);
656 
657 	niov_idx = area->freelist[--area->free_count];
658 	return &area->nia.niovs[niov_idx];
659 }
660 
661 static void io_zcrx_return_niov_freelist(struct net_iov *niov)
662 {
663 	struct io_zcrx_area *area = io_zcrx_iov_to_area(niov);
664 
665 	spin_lock_bh(&area->freelist_lock);
666 	area->freelist[area->free_count++] = net_iov_idx(niov);
667 	spin_unlock_bh(&area->freelist_lock);
668 }
669 
670 static void io_zcrx_return_niov(struct net_iov *niov)
671 {
672 	netmem_ref netmem = net_iov_to_netmem(niov);
673 
674 	if (!niov->pp) {
675 		/* copy fallback allocated niovs */
676 		io_zcrx_return_niov_freelist(niov);
677 		return;
678 	}
679 	page_pool_put_unrefed_netmem(niov->pp, netmem, -1, false);
680 }
681 
682 static void io_zcrx_scrub(struct io_zcrx_ifq *ifq)
683 {
684 	struct io_zcrx_area *area = ifq->area;
685 	int i;
686 
687 	if (!area)
688 		return;
689 
690 	/* Reclaim back all buffers given to the user space. */
691 	for (i = 0; i < area->nia.num_niovs; i++) {
692 		struct net_iov *niov = &area->nia.niovs[i];
693 		int nr;
694 
695 		if (!atomic_read(io_get_user_counter(niov)))
696 			continue;
697 		nr = atomic_xchg(io_get_user_counter(niov), 0);
698 		if (nr && !page_pool_unref_netmem(net_iov_to_netmem(niov), nr))
699 			io_zcrx_return_niov(niov);
700 	}
701 }
702 
703 void io_shutdown_zcrx_ifqs(struct io_ring_ctx *ctx)
704 {
705 	struct io_zcrx_ifq *ifq;
706 	unsigned long index;
707 
708 	lockdep_assert_held(&ctx->uring_lock);
709 
710 	xa_for_each(&ctx->zcrx_ctxs, index, ifq) {
711 		io_zcrx_scrub(ifq);
712 		io_close_queue(ifq);
713 	}
714 }
715 
716 static inline u32 io_zcrx_rqring_entries(struct io_zcrx_ifq *ifq)
717 {
718 	u32 entries;
719 
720 	entries = smp_load_acquire(&ifq->rq_ring->tail) - ifq->cached_rq_head;
721 	return min(entries, ifq->rq_entries);
722 }
723 
724 static struct io_uring_zcrx_rqe *io_zcrx_get_rqe(struct io_zcrx_ifq *ifq,
725 						 unsigned mask)
726 {
727 	unsigned int idx = ifq->cached_rq_head++ & mask;
728 
729 	return &ifq->rqes[idx];
730 }
731 
732 static void io_zcrx_ring_refill(struct page_pool *pp,
733 				struct io_zcrx_ifq *ifq)
734 {
735 	unsigned int mask = ifq->rq_entries - 1;
736 	unsigned int entries;
737 	netmem_ref netmem;
738 
739 	spin_lock_bh(&ifq->rq_lock);
740 
741 	entries = io_zcrx_rqring_entries(ifq);
742 	entries = min_t(unsigned, entries, PP_ALLOC_CACHE_REFILL - pp->alloc.count);
743 	if (unlikely(!entries)) {
744 		spin_unlock_bh(&ifq->rq_lock);
745 		return;
746 	}
747 
748 	do {
749 		struct io_uring_zcrx_rqe *rqe = io_zcrx_get_rqe(ifq, mask);
750 		struct io_zcrx_area *area;
751 		struct net_iov *niov;
752 		unsigned niov_idx, area_idx;
753 
754 		area_idx = rqe->off >> IORING_ZCRX_AREA_SHIFT;
755 		niov_idx = (rqe->off & ~IORING_ZCRX_AREA_MASK) >> PAGE_SHIFT;
756 
757 		if (unlikely(rqe->__pad || area_idx))
758 			continue;
759 		area = ifq->area;
760 
761 		if (unlikely(niov_idx >= area->nia.num_niovs))
762 			continue;
763 		niov_idx = array_index_nospec(niov_idx, area->nia.num_niovs);
764 
765 		niov = &area->nia.niovs[niov_idx];
766 		if (!io_zcrx_put_niov_uref(niov))
767 			continue;
768 
769 		netmem = net_iov_to_netmem(niov);
770 		if (page_pool_unref_netmem(netmem, 1) != 0)
771 			continue;
772 
773 		if (unlikely(niov->pp != pp)) {
774 			io_zcrx_return_niov(niov);
775 			continue;
776 		}
777 
778 		io_zcrx_sync_for_device(pp, niov);
779 		net_mp_netmem_place_in_cache(pp, netmem);
780 	} while (--entries);
781 
782 	smp_store_release(&ifq->rq_ring->head, ifq->cached_rq_head);
783 	spin_unlock_bh(&ifq->rq_lock);
784 }
785 
786 static void io_zcrx_refill_slow(struct page_pool *pp, struct io_zcrx_ifq *ifq)
787 {
788 	struct io_zcrx_area *area = ifq->area;
789 
790 	spin_lock_bh(&area->freelist_lock);
791 	while (area->free_count && pp->alloc.count < PP_ALLOC_CACHE_REFILL) {
792 		struct net_iov *niov = __io_zcrx_get_free_niov(area);
793 		netmem_ref netmem = net_iov_to_netmem(niov);
794 
795 		net_mp_niov_set_page_pool(pp, niov);
796 		io_zcrx_sync_for_device(pp, niov);
797 		net_mp_netmem_place_in_cache(pp, netmem);
798 	}
799 	spin_unlock_bh(&area->freelist_lock);
800 }
801 
802 static netmem_ref io_pp_zc_alloc_netmems(struct page_pool *pp, gfp_t gfp)
803 {
804 	struct io_zcrx_ifq *ifq = io_pp_to_ifq(pp);
805 
806 	/* pp should already be ensuring that */
807 	if (unlikely(pp->alloc.count))
808 		goto out_return;
809 
810 	io_zcrx_ring_refill(pp, ifq);
811 	if (likely(pp->alloc.count))
812 		goto out_return;
813 
814 	io_zcrx_refill_slow(pp, ifq);
815 	if (!pp->alloc.count)
816 		return 0;
817 out_return:
818 	return pp->alloc.cache[--pp->alloc.count];
819 }
820 
821 static bool io_pp_zc_release_netmem(struct page_pool *pp, netmem_ref netmem)
822 {
823 	struct net_iov *niov;
824 
825 	if (WARN_ON_ONCE(!netmem_is_net_iov(netmem)))
826 		return false;
827 
828 	niov = netmem_to_net_iov(netmem);
829 	net_mp_niov_clear_page_pool(niov);
830 	io_zcrx_return_niov_freelist(niov);
831 	return false;
832 }
833 
834 static int io_pp_zc_init(struct page_pool *pp)
835 {
836 	struct io_zcrx_ifq *ifq = io_pp_to_ifq(pp);
837 	int ret;
838 
839 	if (WARN_ON_ONCE(!ifq))
840 		return -EINVAL;
841 	if (WARN_ON_ONCE(ifq->dev != pp->p.dev))
842 		return -EINVAL;
843 	if (WARN_ON_ONCE(!pp->dma_map))
844 		return -EOPNOTSUPP;
845 	if (pp->p.order != 0)
846 		return -EOPNOTSUPP;
847 	if (pp->p.dma_dir != DMA_FROM_DEVICE)
848 		return -EOPNOTSUPP;
849 
850 	ret = io_zcrx_map_area(ifq, ifq->area);
851 	if (ret)
852 		return ret;
853 
854 	percpu_ref_get(&ifq->ctx->refs);
855 	return 0;
856 }
857 
858 static void io_pp_zc_destroy(struct page_pool *pp)
859 {
860 	struct io_zcrx_ifq *ifq = io_pp_to_ifq(pp);
861 	struct io_zcrx_area *area = ifq->area;
862 
863 	if (WARN_ON_ONCE(area->free_count != area->nia.num_niovs))
864 		return;
865 	percpu_ref_put(&ifq->ctx->refs);
866 }
867 
868 static int io_pp_nl_fill(void *mp_priv, struct sk_buff *rsp,
869 			 struct netdev_rx_queue *rxq)
870 {
871 	struct nlattr *nest;
872 	int type;
873 
874 	type = rxq ? NETDEV_A_QUEUE_IO_URING : NETDEV_A_PAGE_POOL_IO_URING;
875 	nest = nla_nest_start(rsp, type);
876 	if (!nest)
877 		return -EMSGSIZE;
878 	nla_nest_end(rsp, nest);
879 
880 	return 0;
881 }
882 
883 static void io_pp_uninstall(void *mp_priv, struct netdev_rx_queue *rxq)
884 {
885 	struct pp_memory_provider_params *p = &rxq->mp_params;
886 	struct io_zcrx_ifq *ifq = mp_priv;
887 
888 	io_zcrx_drop_netdev(ifq);
889 	if (ifq->area)
890 		io_zcrx_unmap_area(ifq, ifq->area);
891 
892 	p->mp_ops = NULL;
893 	p->mp_priv = NULL;
894 }
895 
896 static const struct memory_provider_ops io_uring_pp_zc_ops = {
897 	.alloc_netmems		= io_pp_zc_alloc_netmems,
898 	.release_netmem		= io_pp_zc_release_netmem,
899 	.init			= io_pp_zc_init,
900 	.destroy		= io_pp_zc_destroy,
901 	.nl_fill		= io_pp_nl_fill,
902 	.uninstall		= io_pp_uninstall,
903 };
904 
905 static bool io_zcrx_queue_cqe(struct io_kiocb *req, struct net_iov *niov,
906 			      struct io_zcrx_ifq *ifq, int off, int len)
907 {
908 	struct io_uring_zcrx_cqe *rcqe;
909 	struct io_zcrx_area *area;
910 	struct io_uring_cqe *cqe;
911 	u64 offset;
912 
913 	if (!io_defer_get_uncommited_cqe(req->ctx, &cqe))
914 		return false;
915 
916 	cqe->user_data = req->cqe.user_data;
917 	cqe->res = len;
918 	cqe->flags = IORING_CQE_F_MORE;
919 
920 	area = io_zcrx_iov_to_area(niov);
921 	offset = off + (net_iov_idx(niov) << PAGE_SHIFT);
922 	rcqe = (struct io_uring_zcrx_cqe *)(cqe + 1);
923 	rcqe->off = offset + ((u64)area->area_id << IORING_ZCRX_AREA_SHIFT);
924 	rcqe->__pad = 0;
925 	return true;
926 }
927 
928 static struct net_iov *io_zcrx_alloc_fallback(struct io_zcrx_area *area)
929 {
930 	struct net_iov *niov = NULL;
931 
932 	spin_lock_bh(&area->freelist_lock);
933 	if (area->free_count)
934 		niov = __io_zcrx_get_free_niov(area);
935 	spin_unlock_bh(&area->freelist_lock);
936 
937 	if (niov)
938 		page_pool_fragment_netmem(net_iov_to_netmem(niov), 1);
939 	return niov;
940 }
941 
942 static ssize_t io_zcrx_copy_chunk(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
943 				  void *src_base, struct page *src_page,
944 				  unsigned int src_offset, size_t len)
945 {
946 	struct io_zcrx_area *area = ifq->area;
947 	size_t copied = 0;
948 	int ret = 0;
949 
950 	if (area->mem.is_dmabuf)
951 		return -EFAULT;
952 
953 	while (len) {
954 		size_t copy_size = min_t(size_t, PAGE_SIZE, len);
955 		const int dst_off = 0;
956 		struct net_iov *niov;
957 		struct page *dst_page;
958 		void *dst_addr;
959 
960 		niov = io_zcrx_alloc_fallback(area);
961 		if (!niov) {
962 			ret = -ENOMEM;
963 			break;
964 		}
965 
966 		dst_page = io_zcrx_iov_page(niov);
967 		dst_addr = kmap_local_page(dst_page);
968 		if (src_page)
969 			src_base = kmap_local_page(src_page);
970 
971 		memcpy(dst_addr, src_base + src_offset, copy_size);
972 
973 		if (src_page)
974 			kunmap_local(src_base);
975 		kunmap_local(dst_addr);
976 
977 		if (!io_zcrx_queue_cqe(req, niov, ifq, dst_off, copy_size)) {
978 			io_zcrx_return_niov(niov);
979 			ret = -ENOSPC;
980 			break;
981 		}
982 
983 		io_zcrx_get_niov_uref(niov);
984 		src_offset += copy_size;
985 		len -= copy_size;
986 		copied += copy_size;
987 	}
988 
989 	return copied ? copied : ret;
990 }
991 
992 static int io_zcrx_copy_frag(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
993 			     const skb_frag_t *frag, int off, int len)
994 {
995 	struct page *page = skb_frag_page(frag);
996 	u32 p_off, p_len, t, copied = 0;
997 	int ret = 0;
998 
999 	off += skb_frag_off(frag);
1000 
1001 	skb_frag_foreach_page(frag, off, len,
1002 			      page, p_off, p_len, t) {
1003 		ret = io_zcrx_copy_chunk(req, ifq, NULL, page, p_off, p_len);
1004 		if (ret < 0)
1005 			return copied ? copied : ret;
1006 		copied += ret;
1007 	}
1008 	return copied;
1009 }
1010 
1011 static int io_zcrx_recv_frag(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1012 			     const skb_frag_t *frag, int off, int len)
1013 {
1014 	struct net_iov *niov;
1015 
1016 	if (unlikely(!skb_frag_is_net_iov(frag)))
1017 		return io_zcrx_copy_frag(req, ifq, frag, off, len);
1018 
1019 	niov = netmem_to_net_iov(frag->netmem);
1020 	if (niov->pp->mp_ops != &io_uring_pp_zc_ops ||
1021 	    io_pp_to_ifq(niov->pp) != ifq)
1022 		return -EFAULT;
1023 
1024 	if (!io_zcrx_queue_cqe(req, niov, ifq, off + skb_frag_off(frag), len))
1025 		return -ENOSPC;
1026 
1027 	/*
1028 	 * Prevent it from being recycled while user is accessing it.
1029 	 * It has to be done before grabbing a user reference.
1030 	 */
1031 	page_pool_ref_netmem(net_iov_to_netmem(niov));
1032 	io_zcrx_get_niov_uref(niov);
1033 	return len;
1034 }
1035 
1036 static int
1037 io_zcrx_recv_skb(read_descriptor_t *desc, struct sk_buff *skb,
1038 		 unsigned int offset, size_t len)
1039 {
1040 	struct io_zcrx_args *args = desc->arg.data;
1041 	struct io_zcrx_ifq *ifq = args->ifq;
1042 	struct io_kiocb *req = args->req;
1043 	struct sk_buff *frag_iter;
1044 	unsigned start, start_off = offset;
1045 	int i, copy, end, off;
1046 	int ret = 0;
1047 
1048 	len = min_t(size_t, len, desc->count);
1049 	/*
1050 	 * __tcp_read_sock() always calls io_zcrx_recv_skb one last time, even
1051 	 * if desc->count is already 0. This is caused by the if (offset + 1 !=
1052 	 * skb->len) check. Return early in this case to break out of
1053 	 * __tcp_read_sock().
1054 	 */
1055 	if (!len)
1056 		return 0;
1057 	if (unlikely(args->nr_skbs++ > IO_SKBS_PER_CALL_LIMIT))
1058 		return -EAGAIN;
1059 
1060 	if (unlikely(offset < skb_headlen(skb))) {
1061 		ssize_t copied;
1062 		size_t to_copy;
1063 
1064 		to_copy = min_t(size_t, skb_headlen(skb) - offset, len);
1065 		copied = io_zcrx_copy_chunk(req, ifq, skb->data, NULL,
1066 					    offset, to_copy);
1067 		if (copied < 0) {
1068 			ret = copied;
1069 			goto out;
1070 		}
1071 		offset += copied;
1072 		len -= copied;
1073 		if (!len)
1074 			goto out;
1075 		if (offset != skb_headlen(skb))
1076 			goto out;
1077 	}
1078 
1079 	start = skb_headlen(skb);
1080 
1081 	for (i = 0; i < skb_shinfo(skb)->nr_frags; i++) {
1082 		const skb_frag_t *frag;
1083 
1084 		if (WARN_ON(start > offset + len))
1085 			return -EFAULT;
1086 
1087 		frag = &skb_shinfo(skb)->frags[i];
1088 		end = start + skb_frag_size(frag);
1089 
1090 		if (offset < end) {
1091 			copy = end - offset;
1092 			if (copy > len)
1093 				copy = len;
1094 
1095 			off = offset - start;
1096 			ret = io_zcrx_recv_frag(req, ifq, frag, off, copy);
1097 			if (ret < 0)
1098 				goto out;
1099 
1100 			offset += ret;
1101 			len -= ret;
1102 			if (len == 0 || ret != copy)
1103 				goto out;
1104 		}
1105 		start = end;
1106 	}
1107 
1108 	skb_walk_frags(skb, frag_iter) {
1109 		if (WARN_ON(start > offset + len))
1110 			return -EFAULT;
1111 
1112 		end = start + frag_iter->len;
1113 		if (offset < end) {
1114 			copy = end - offset;
1115 			if (copy > len)
1116 				copy = len;
1117 
1118 			off = offset - start;
1119 			ret = io_zcrx_recv_skb(desc, frag_iter, off, copy);
1120 			if (ret < 0)
1121 				goto out;
1122 
1123 			offset += ret;
1124 			len -= ret;
1125 			if (len == 0 || ret != copy)
1126 				goto out;
1127 		}
1128 		start = end;
1129 	}
1130 
1131 out:
1132 	if (offset == start_off)
1133 		return ret;
1134 	desc->count -= (offset - start_off);
1135 	return offset - start_off;
1136 }
1137 
1138 static int io_zcrx_tcp_recvmsg(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1139 				struct sock *sk, int flags,
1140 				unsigned issue_flags, unsigned int *outlen)
1141 {
1142 	unsigned int len = *outlen;
1143 	struct io_zcrx_args args = {
1144 		.req = req,
1145 		.ifq = ifq,
1146 		.sock = sk->sk_socket,
1147 	};
1148 	read_descriptor_t rd_desc = {
1149 		.count = len ? len : UINT_MAX,
1150 		.arg.data = &args,
1151 	};
1152 	int ret;
1153 
1154 	lock_sock(sk);
1155 	ret = tcp_read_sock(sk, &rd_desc, io_zcrx_recv_skb);
1156 	if (len && ret > 0)
1157 		*outlen = len - ret;
1158 	if (ret <= 0) {
1159 		if (ret < 0 || sock_flag(sk, SOCK_DONE))
1160 			goto out;
1161 		if (sk->sk_err)
1162 			ret = sock_error(sk);
1163 		else if (sk->sk_shutdown & RCV_SHUTDOWN)
1164 			goto out;
1165 		else if (sk->sk_state == TCP_CLOSE)
1166 			ret = -ENOTCONN;
1167 		else
1168 			ret = -EAGAIN;
1169 	} else if (unlikely(args.nr_skbs > IO_SKBS_PER_CALL_LIMIT) &&
1170 		   (issue_flags & IO_URING_F_MULTISHOT)) {
1171 		ret = IOU_REQUEUE;
1172 	} else if (sock_flag(sk, SOCK_DONE)) {
1173 		/* Make it to retry until it finally gets 0. */
1174 		if (issue_flags & IO_URING_F_MULTISHOT)
1175 			ret = IOU_REQUEUE;
1176 		else
1177 			ret = -EAGAIN;
1178 	}
1179 out:
1180 	release_sock(sk);
1181 	return ret;
1182 }
1183 
1184 int io_zcrx_recv(struct io_kiocb *req, struct io_zcrx_ifq *ifq,
1185 		 struct socket *sock, unsigned int flags,
1186 		 unsigned issue_flags, unsigned int *len)
1187 {
1188 	struct sock *sk = sock->sk;
1189 	const struct proto *prot = READ_ONCE(sk->sk_prot);
1190 
1191 	if (prot->recvmsg != tcp_recvmsg)
1192 		return -EPROTONOSUPPORT;
1193 
1194 	sock_rps_record_flow(sk);
1195 	return io_zcrx_tcp_recvmsg(req, ifq, sk, flags, issue_flags, len);
1196 }
1197