[PATCH net 1/1] net: tcp: account zero-copy receive VMA memory

[Ren Wei <weir@nebusec.ai>]

From: Zixuan Chai <petalzu987@gmail.com>

TCP_ZEROCOPY_RECEIVE installs page references in a user VMA and then
consumes the corresponding skb. The socket charge is released at that
point, so a process can retain an unbounded number of receive pages,
including page-table memory, by advancing through a large VMA.

Reserve the VMA size in TCP socket memory accounting while the mapping
exists. This charges the reservation to the socket memory cgroup and
TCP protocol budget, and releases it when the last VMA fragment is
unmapped. Keep the reservation across same-mm VMA splits and moves,
disallow expansion and fork inheritance, and reject inherited VMAs from
the zero-copy path. Use a strict accounting kind so ordinary receive-
buffer minimum allowances cannot bypass the reservation limit.

Fixes: 93ab6cc69162 ("tcp: implement mmap() for zero copy receive")
Cc: stable@vger.kernel.org
Reported-by: VEGA <vega@nebusec.ai>
Assisted-by: LLM
Signed-off-by: Zixuan Chai <petalzu987@gmail.com>
Signed-off-by: Ren Wei <weir@nebusec.ai>
---
 include/net/sock.h          |   1 +
 include/trace/events/sock.h |   3 +-
 net/core/sock.c             |   2 +-
 net/ipv4/tcp.c              | 102 ++++++++++++++++++++++++++++++++++--
 4 files changed, 103 insertions(+), 5 deletions(-)

diff --git a/include/net/sock.h b/include/net/sock.h
index 51185222aac2..cf19055a9c40 100644
--- a/include/net/sock.h
+++ b/include/net/sock.h
@@ -1556,6 +1556,7 @@ void __sk_mem_reclaim(struct sock *sk, int amount);
 
 #define SK_MEM_SEND	0
 #define SK_MEM_RECV	1
+#define SK_MEM_RECV_ZEROCOPY	0x100
 
 /* sysctl_mem values are in pages */
 static inline long sk_prot_mem_limits(const struct sock *sk, int index)
diff --git a/include/trace/events/sock.h b/include/trace/events/sock.h
index b5310439536e..0f10b92c118c 100644
--- a/include/trace/events/sock.h
+++ b/include/trace/events/sock.h
@@ -38,7 +38,8 @@
 
 #define skmem_kind_names			\
 		EM(SK_MEM_SEND)			\
-		EMe(SK_MEM_RECV)
+		EM(SK_MEM_RECV)			\
+		EMe(SK_MEM_RECV_ZEROCOPY)
 
 /* enums need to be exported to user space */
 #undef EM
diff --git a/net/core/sock.c b/net/core/sock.c
index fa60b7494c58..3d85c138e1c7 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -3380,7 +3380,7 @@ int __sk_mem_raise_allocated(struct sock *sk, int size, int amt, int kind)
 		if (atomic_read(&sk->sk_rmem_alloc) < sk_get_rmem0(sk, prot))
 			return 1;
 
-	} else { /* SK_MEM_SEND */
+	} else if (kind == SK_MEM_SEND) {
 		int wmem0 = sk_get_wmem0(sk, prot);
 
 		if (sk->sk_type == SOCK_STREAM) {
diff --git a/net/ipv4/tcp.c b/net/ipv4/tcp.c
index 562752352afe..5cc2bfc8a90d 100644
--- a/net/ipv4/tcp.c
+++ b/net/ipv4/tcp.c
@@ -267,6 +267,7 @@
 #include <linux/errqueue.h>
 #include <linux/static_key.h>
 #include <linux/btf.h>
+#include <linux/sched/mm.h>
 
 #include <net/icmp.h>
 #include <net/inet_common.h>
@@ -1860,23 +1861,118 @@ void tcp_set_rcvbuf(struct sock *sk, int val)
 }
 
 #ifdef CONFIG_MMU
+struct tcp_zc_vma {
+	unsigned long nr_pages;
+	struct mm_struct *mm;
+	refcount_t refcnt;
+	struct sock *sk;
+};
+
+/* Keep the size argument to __sk_mem_raise_allocated() within int. */
+#define TCP_ZEROCOPY_MEM_CHUNK (INT_MAX >> PAGE_SHIFT)
+
+static void tcp_zc_mem_uncharge(struct sock *sk, unsigned long nr_pages)
+{
+	unsigned int chunk;
+
+	while (nr_pages) {
+		chunk = min_t(unsigned long, nr_pages,
+			      TCP_ZEROCOPY_MEM_CHUNK);
+		__sk_mem_reduce_allocated(sk, chunk);
+		nr_pages -= chunk;
+	}
+}
+
+static int tcp_zc_mem_charge(struct sock *sk, unsigned long nr_pages)
+{
+	unsigned long charged = 0;
+	unsigned int chunk;
+	int ret;
+
+	while (nr_pages) {
+		chunk = min_t(unsigned long, nr_pages,
+			      TCP_ZEROCOPY_MEM_CHUNK);
+		ret = __sk_mem_raise_allocated(sk, chunk << PAGE_SHIFT,
+					       chunk, SK_MEM_RECV_ZEROCOPY);
+		if (!ret) {
+			tcp_zc_mem_uncharge(sk, charged);
+			return -ENOMEM;
+		}
+		charged += chunk;
+		nr_pages -= chunk;
+	}
+
+	return 0;
+}
+
+static void tcp_zc_vma_open(struct vm_area_struct *vma)
+{
+	struct tcp_zc_vma *zc_vma = vma->vm_private_data;
+
+	refcount_inc(&zc_vma->refcnt);
+}
+
+static void tcp_zc_vma_close(struct vm_area_struct *vma)
+{
+	struct tcp_zc_vma *zc_vma = vma->vm_private_data;
+
+	if (refcount_dec_and_test(&zc_vma->refcnt)) {
+		tcp_zc_mem_uncharge(zc_vma->sk, zc_vma->nr_pages);
+		mmdrop(zc_vma->mm);
+		sock_put(zc_vma->sk);
+		kfree(zc_vma);
+	}
+}
+
 static const struct vm_operations_struct tcp_vm_ops = {
+	.open = tcp_zc_vma_open,
+	.close = tcp_zc_vma_close,
 };
 
 int tcp_mmap(struct file *file, struct socket *sock,
 	     struct vm_area_struct *vma)
 {
+	struct tcp_zc_vma *zc_vma;
+	unsigned long nr_pages;
+
 	if (vma->vm_flags & (VM_WRITE | VM_EXEC))
 		return -EPERM;
 	vm_flags_clear(vma, VM_MAYWRITE | VM_MAYEXEC);
 
 	/* Instruct vm_insert_page() to not mmap_read_lock(mm) */
-	vm_flags_set(vma, VM_MIXEDMAP);
+	vm_flags_set(vma, VM_DONTCOPY | VM_DONTEXPAND | VM_MIXEDMAP);
+
+	nr_pages = (vma->vm_end - vma->vm_start) >> PAGE_SHIFT;
+	zc_vma = kmalloc_obj(*zc_vma, GFP_KERNEL_ACCOUNT);
+	if (!zc_vma)
+		return -ENOMEM;
+
+	zc_vma->nr_pages = nr_pages;
+	zc_vma->mm = vma->vm_mm;
+	zc_vma->sk = sock->sk;
+	refcount_set(&zc_vma->refcnt, 1);
+	mmgrab(zc_vma->mm);
+	sock_hold(zc_vma->sk);
+	if (tcp_zc_mem_charge(zc_vma->sk, nr_pages)) {
+		mmdrop(zc_vma->mm);
+		sock_put(zc_vma->sk);
+		kfree(zc_vma);
+		return -ENOMEM;
+	}
 
 	vma->vm_ops = &tcp_vm_ops;
+	vma->vm_private_data = zc_vma;
 	return 0;
 }
 
+static bool tcp_zc_vma_valid(const struct vm_area_struct *vma,
+			     const struct mm_struct *mm)
+{
+	const struct tcp_zc_vma *zc_vma = vma->vm_private_data;
+
+	return vma->vm_ops == &tcp_vm_ops && zc_vma && zc_vma->mm == mm;
+}
+
 static skb_frag_t *skb_advance_to_frag(struct sk_buff *skb, u32 offset_skb,
 				       u32 *offset_frag)
 {
@@ -2173,7 +2269,7 @@ static struct vm_area_struct *find_tcp_vma(struct mm_struct *mm,
 	struct vm_area_struct *vma = lock_vma_under_rcu(mm, address);
 
 	if (vma) {
-		if (vma->vm_ops != &tcp_vm_ops) {
+		if (!tcp_zc_vma_valid(vma, mm)) {
 			vma_end_read(vma);
 			return NULL;
 		}
@@ -2183,7 +2279,7 @@ static struct vm_area_struct *find_tcp_vma(struct mm_struct *mm,
 
 	mmap_read_lock(mm);
 	vma = vma_lookup(mm, address);
-	if (!vma || vma->vm_ops != &tcp_vm_ops) {
+	if (!vma || !tcp_zc_vma_valid(vma, mm)) {
 		mmap_read_unlock(mm);
 		return NULL;
 	}