[PATCH net-next v2 3/5] tcp: detach page_pool pages before mapping in tcp_zerocopy_receive()

[Mina Almasry]

Once tcp_zerocopy_receive() maps an order-0 RX fragment page into a
userspace VMA via vm_insert_pages(), vm_insert_pages() increments
page->_refcount, which prevents the page from being recycled back into
its page_pool when the RX SKB is freed. However, while the RX SKB still
holds its fragment reference, the page remains marked as a page_pool
page (page_pool_page_is_pp(page) == true) inside the userspace VMA, so
subsequent GUP callers such as sendmsg(MSG_ZEROCOPY) or vmsplice() can
encounter a page_pool page while holding only a raw struct page
reference.

Establish the invariant that pp_ref_count is only incremented by callers
that already hold an active pp_ref_count (>= 1) on that netmem, and a
raw page_ref_count is never converted into a pp_ref_count:
- In can_map_frag(), require !skb_cloned(skb), page_pool_ref_count == 1,
  and page_ref_count == 1 for page_pool pages so the fragment being
  mapped holds the sole reference to the page; cloned SKBs or page_pool
  pages with additional references (including drivers holding an
  undrained fragment bias) fall back to copying.
- In tcp_zerocopy_receive(), convert the fragment's sole pp_ref_count
  into a regular page reference and detach the page from its page_pool
  via get_page(page) + napi_pp_put_page(page_to_netmem(page)) before
  queueing the page for vm_insert_pages(). If vm_insert_pages()
  subsequently fails, the detached page remains referenced by the RX SKB
  via its regular page reference and is freed back to the page allocator
  when the SKB is consumed.

Signed-off-by: Mina Almasry <almasrymina@xxxxxxxxxx>
---
 include/net/page_pool/helpers.h |  5 ++++
 net/core/datagram.c             |  2 ++
 net/ipv4/tcp.c                  | 44 ++++++++++++++++++++++++++++-----
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/include/net/page_pool/helpers.h b/include/net/page_pool/helpers.h
index 28635ce9454e8..c814ce95282eb 100644
--- a/include/net/page_pool/helpers.h
+++ b/include/net/page_pool/helpers.h
@@ -242,6 +242,11 @@ page_pool_get_dma_dir(const struct page_pool *pool)
 	return pool->p.dma_dir;
 }
 
+static inline long page_pool_ref_count(netmem_ref netmem)
+{
+	return atomic_long_read(netmem_get_pp_ref_count_ref(netmem));
+}
+
 static inline void page_pool_fragment_netmem(netmem_ref netmem, long nr)
 {
 	atomic_long_set(netmem_get_pp_ref_count_ref(netmem), nr);
diff --git a/net/core/datagram.c b/net/core/datagram.c
index ecc30a6f83023..808eda4f7efb5 100644
--- a/net/core/datagram.c
+++ b/net/core/datagram.c
@@ -34,6 +34,7 @@
  *
  */
 
+#include <kunit/visibility.h>
 #include <linux/module.h>
 #include <linux/types.h>
 #include <linux/kernel.h>
@@ -702,6 +703,7 @@ int zerocopy_fill_skb_from_iter(struct sk_buff *skb,
 	}
 	return 0;
 }
+EXPORT_SYMBOL_IF_KUNIT(zerocopy_fill_skb_from_iter);
 
 static int
 zerocopy_fill_skb_from_devmem(struct sk_buff *skb, struct iov_iter *from,
diff --git a/net/ipv4/tcp.c b/net/ipv4/tcp.c
index b9de7c302ba83..a984e7e93e4c9 100644
--- a/net/ipv4/tcp.c
+++ b/net/ipv4/tcp.c
@@ -286,6 +286,7 @@
 #include <asm/ioctls.h>
 #include <net/busy_poll.h>
 #include <net/hotdata.h>
+#include <net/page_pool/helpers.h>
 #include <trace/events/tcp.h>
 #include <net/rps.h>
 
@@ -1921,7 +1922,7 @@ static skb_frag_t *skb_advance_to_frag(struct sk_buff *skb, u32 offset_skb,
 	return frag;
 }
 
-static bool can_map_frag(const skb_frag_t *frag)
+static bool can_map_frag(const struct sk_buff *skb, const skb_frag_t *frag)
 {
 	struct page *page;
 
@@ -1935,18 +1936,32 @@ static bool can_map_frag(const skb_frag_t *frag)
 	if (PageCompound(page) || page->mapping)
 		return false;
 
+	/* Before mapping a page_pool page into userspace,
+	 * tcp_zerocopy_receive() converts this fragment's pp_ref_count into a
+	 * regular page reference and detaches the page from its page_pool via
+	 * napi_pp_put_page(). Require this skb_frag_t to hold the sole
+	 * reference to the page so no other fragment or SKB clone expects
+	 * pp_ref_count on it after detachment.
+	 */
+	if (page_pool_page_is_pp(page) &&
+	    (skb_cloned(skb) ||
+	     page_pool_ref_count(page_to_netmem(page)) != 1 ||
+	     page_ref_count(page) != 1))
+		return false;
+
 	return true;
 }
 
-static int find_next_mappable_frag(const skb_frag_t *frag,
+static int find_next_mappable_frag(const struct sk_buff *skb,
+				   const skb_frag_t *frag,
 				   int remaining_in_skb)
 {
 	int offset = 0;
 
-	if (likely(can_map_frag(frag)))
+	if (likely(can_map_frag(skb, frag)))
 		return 0;
 
-	while (offset < remaining_in_skb && !can_map_frag(frag)) {
+	while (offset < remaining_in_skb && !can_map_frag(skb, frag)) {
 		offset += skb_frag_size(frag);
 		++frag;
 	}
@@ -1986,7 +2001,8 @@ static void tcp_zerocopy_set_hint_for_skb(struct sock *sk,
 	 * mappable_offset: Bytes till next mappable frag, *not* counting bytes
 	 * in partial_frag_remainder.
 	 */
-	mappable_offset = find_next_mappable_frag(frag, zc->recv_skip_hint);
+	mappable_offset = find_next_mappable_frag(skb, frag,
+						  zc->recv_skip_hint);
 	zc->recv_skip_hint = mappable_offset + partial_frag_remainder;
 }
 
@@ -2290,7 +2306,7 @@ static int tcp_zerocopy_receive(struct sock *sk,
 				break;
 		}
 
-		mappable_offset = find_next_mappable_frag(frags,
+		mappable_offset = find_next_mappable_frag(skb, frags,
 							  zc->recv_skip_hint);
 		if (mappable_offset) {
 			zc->recv_skip_hint = mappable_offset;
@@ -2300,6 +2316,22 @@ static int tcp_zerocopy_receive(struct sock *sk,
 		if (WARN_ON_ONCE(!page))
 			break;
 
+		/* Convert this fragment's sole pp_ref_count into a regular
+		 * page reference and detach the page from its page_pool before
+		 * mapping it into userspace. Once vm_insert_pages() elevates
+		 * page_ref_count the page cannot be recycled into the
+		 * page_pool anyway (and if vm_insert_pages() fails, the
+		 * detached page is freed to the page allocator when the skb is
+		 * consumed); detaching it up front ensures that a page_pool
+		 * page is never mapped into a userspace VMA (where GUP callers
+		 * such as MSG_ZEROCOPY or vmsplice() would otherwise observe a
+		 * page_pool page while holding only a raw page_ref_count).
+		 */
+		if (page_pool_page_is_pp(page)) {
+			get_page(page);
+			napi_pp_put_page(page_to_netmem(page));
+		}
+
 		prefetchw(page);
 		pages[pages_to_map++] = page;
 		length += PAGE_SIZE;
-- 
2.56.0.455.g6c5a503b38-goog