[RFC 1/9] net: allow __tcp_read_sock actors to steal skbs

Currently __tcp_read_sock() owns skbs and expects them to be present
when the actor function returns (modulo collapsing). For zcrx
optimisations I want to be able to take ownership of the skb in the
callback, add a helper doing that. It's only implemented for tcp, hence
keep "tcp" in the helper name. It could be later extended to other
protocols but would need some whitelisting mechanism.

Signed-off-by: Pavel Begunkov <[email protected]>
---
 include/linux/net.h |  1 +
 include/net/tcp.h   | 13 +++++++++++++
 net/ipv4/tcp.c      | 11 +++++++++++
 3 files changed, 25 insertions(+)

diff --git a/include/linux/net.h b/include/linux/net.h
index f268f395ce47..ed882aeac4a5 100644
--- a/include/linux/net.h
+++ b/include/linux/net.h
@@ -165,6 +165,7 @@ typedef struct {
 		void *data;
 	} arg;
 	int error;
+	bool stolen;
 } read_descriptor_t;
 
 struct vm_area_struct;
diff --git a/include/net/tcp.h b/include/net/tcp.h
index ecbadcb3a744..3d25707b73c3 100644
--- a/include/net/tcp.h
+++ b/include/net/tcp.h
@@ -3089,6 +3089,19 @@ static inline int tcp_recv_should_stop(struct sock *sk)
 	       signal_pending(current);
 }
 
+static inline bool tcp_read_sock_steal_skb(read_descriptor_t *desc,
+					   struct sk_buff *skb,
+					   struct sock *sk)
+{
+	if (skb_shared(skb))
+		return false;
+
+	desc->stolen = true;
+	__skb_unlink(skb, &sk->sk_receive_queue);
+	skb_orphan(skb);
+	return true;
+}
+
 INDIRECT_CALLABLE_DECLARE(union tcp_seq_and_ts_off
 			  tcp_v4_init_seq_and_ts_off(const struct net *net,
 						     const struct sk_buff *skb));
diff --git a/net/ipv4/tcp.c b/net/ipv4/tcp.c
index 432fa28e47d4..309a0e6b0173 100644
--- a/net/ipv4/tcp.c
+++ b/net/ipv4/tcp.c
@@ -1677,6 +1677,7 @@ static int __tcp_read_sock(struct sock *sk, read_descriptor_t *desc,
 		return -ENOTCONN;
 	while ((skb = tcp_recv_skb(sk, seq, &offset)) != NULL) {
 		if (offset < skb->len) {
+			u8 tcp_flags = TCP_SKB_CB(skb)->tcp_flags;
 			int used;
 			size_t len;
 
@@ -1689,6 +1690,7 @@ static int __tcp_read_sock(struct sock *sk, read_descriptor_t *desc,
 				if (!len)
 					break;
 			}
+			desc->stolen = false;
 			used = recv_actor(desc, skb, offset, len);
 			if (used <= 0) {
 				if (!copied)
@@ -1701,6 +1703,14 @@ static int __tcp_read_sock(struct sock *sk, read_descriptor_t *desc,
 			copied += used;
 			offset += used;
 
+			if (desc->stolen) {
+				if (tcp_flags & TCPHDR_FIN) {
+					++seq;
+					break;
+				}
+				goto next;
+			}
+
 			/* If recv_actor drops the lock (e.g. TCP splice
 			 * receive) the skb pointer might be invalid when
 			 * getting here: tcp_collapse might have deleted it
@@ -1721,6 +1731,7 @@ static int __tcp_read_sock(struct sock *sk, read_descriptor_t *desc,
 			break;
 		}
 		tcp_eat_recv_skb(sk, skb);
+next:
 		if (!desc->count)
 			break;
 		WRITE_ONCE(*copied_seq, seq);
-- 
2.54.0