[Patch bpf v3] skmsg: check sk_rcvbuf limit before queuing to ingress_skb

[Cong Wang <xiyou.wangcong@gmail.com>]

From: Cong Wang <cong.wang@bytedance.com>

Jiang observed OOM frequently when testing our AF_UNIX/UDP
proxy. This is due to the fact that we do not actually limit
the socket memory before queueing skb to ingress_skb. We
charge the skb memory later when handling the psock backlog,
and it is not limited either.

This patch adds checks for sk->sk_rcvbuf right before queuing
to ingress_skb and drops or retries the packets if this limit
exceeds. This is very similar to UDP receive path. Ideally we
should set the skb owner before this check too, but it is hard
to make TCP happy with sk_forward_alloc.

For TCP, we can not just drop the packets on errors. TCP ACKs
are already sent for those packet before reaching
->sk_data_ready(). Instead, we use best effort to retry, this
works because TCP does not remove the skb from receive queue
at that point and exceeding sk_rcvbuf limit is a temporary
situation.

Reported-by: Jiang Wang <jiang.wang@bytedance.com>
Cc: Daniel Borkmann <daniel@iogearbox.net>
Cc: John Fastabend <john.fastabend@gmail.com>
Cc: Lorenz Bauer <lmb@cloudflare.com>
Cc: Jakub Sitnicki <jakub@cloudflare.com>
Signed-off-by: Cong Wang <cong.wang@bytedance.com>
---
v3: add retry logic for TCP
v2: add READ_ONCE()

 net/core/skmsg.c | 15 +++++++++------
 net/ipv4/tcp.c   |  2 ++
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/net/core/skmsg.c b/net/core/skmsg.c
index 2d6249b28928..356c314cd60c 100644
--- a/net/core/skmsg.c
+++ b/net/core/skmsg.c
@@ -877,11 +877,12 @@ static int sk_psock_skb_redirect(struct sk_psock *from, struct sk_buff *skb)
 		return -EIO;
 	}
 	spin_lock_bh(&psock_other->ingress_lock);
-	if (!sk_psock_test_state(psock_other, SK_PSOCK_TX_ENABLED)) {
+	if (!sk_psock_test_state(psock_other, SK_PSOCK_TX_ENABLED) ||
+	    atomic_read(&sk_other->sk_rmem_alloc) > READ_ONCE(sk_other->sk_rcvbuf)) {
 		spin_unlock_bh(&psock_other->ingress_lock);
 		skb_bpf_redirect_clear(skb);
 		sock_drop(from->sk, skb);
-		return -EIO;
+		return -EAGAIN;
 	}
 
 	skb_queue_tail(&psock_other->ingress_skb, skb);
@@ -941,7 +942,7 @@ static int sk_psock_verdict_apply(struct sk_psock *psock, struct sk_buff *skb,
 		}
 
 		skb_bpf_set_ingress(skb);
-
+		err = -EAGAIN;
 		/* If the queue is empty then we can submit directly
 		 * into the msg queue. If its not empty we have to
 		 * queue work otherwise we may get OOO data. Otherwise,
@@ -953,7 +954,8 @@ static int sk_psock_verdict_apply(struct sk_psock *psock, struct sk_buff *skb,
 		}
 		if (err < 0) {
 			spin_lock_bh(&psock->ingress_lock);
-			if (sk_psock_test_state(psock, SK_PSOCK_TX_ENABLED)) {
+			if (sk_psock_test_state(psock, SK_PSOCK_TX_ENABLED) &&
+			    atomic_read(&sk_other->sk_rmem_alloc) <= READ_ONCE(sk_other->sk_rcvbuf)) {
 				skb_queue_tail(&psock->ingress_skb, skb);
 				schedule_work(&psock->work);
 				err = 0;
@@ -1141,8 +1143,9 @@ static int sk_psock_verdict_recv(read_descriptor_t *desc, struct sk_buff *skb,
 		ret = sk_psock_map_verd(ret, skb_bpf_redirect_fetch(skb));
 		skb->sk = NULL;
 	}
-	if (sk_psock_verdict_apply(psock, skb, ret) < 0)
-		len = 0;
+	ret = sk_psock_verdict_apply(psock, skb, ret);
+	if (ret < 0)
+		len = ret;
 out:
 	rcu_read_unlock();
 	return len;
diff --git a/net/ipv4/tcp.c b/net/ipv4/tcp.c
index e8b48df73c85..8b243fcdbb8f 100644
--- a/net/ipv4/tcp.c
+++ b/net/ipv4/tcp.c
@@ -1665,6 +1665,8 @@ int tcp_read_sock(struct sock *sk, read_descriptor_t *desc,
 			if (used <= 0) {
 				if (!copied)
 					copied = used;
+				if (used == -EAGAIN)
+					continue;
 				break;
 			} else if (used <= len) {
 				seq += used;
-- 
2.30.2