[RFC PATCH net-next 01/12] net: add an exclusive TCP receive claim

[Alex Markuze <amarkuze@redhat.com>]

The io_uring opaque object collector, added later in this series,
consumes the receive queue of a TCP socket through tcp_read_sock() with
its own actor. It reports stream offsets that count from the first byte
that was unread when it attached. If another reader consumes bytes while
the collector is attached, those offsets no longer match the data, and
the application selects the wrong ranges.

Add a receive owner to struct sock: a pointer to a struct sock_rx_owner
that names a tcp_read_sock() actor and its data. The owner is published
and cleared under the socket lock and sk_callback_lock. While an owner is
set:

- tcp_read_sock() and tcp_read_sock_noack() return -EBUSY unless the
  caller passes the owning actor and data. This also stops splice().
- tcp_recvmsg_locked(), tcp_zerocopy_receive() and tcp_read_skb() return
  -EBUSY.
- sk_wait_data() also returns when an owner appears. A reader that sleeps
  in tcp_recvmsg_locked() then stops with the bytes that it copied, or
  with -EBUSY. The claimant wakes sleeping readers after it publishes the
  owner.
- tcp_disconnect(), connect(AF_UNSPEC), tcp_set_ulp(), sk_psock_init()
  and TCP_REPAIR return -EBUSY, because each of them resets or takes
  over the receive sequence. TCP_REPAIR checks the owner before its
  permission check, so a caller without the permission also gets -EBUSY.
- sk_clone() starts the child socket without an owner.

CONFIG_SOCK_RX_OWNER is a hidden option that the io_uring feature
selects. Without it, struct sock does not change and the checks compile
away. With it, struct sock grows by one pointer. Each affected receive
path reads the pointer once with READ_ONCE(), and tcp_recvmsg_locked()
reads it once per loop iteration.

An optional data_ready callback notifies the owner independently of
ordinary TCP low-water readiness. TCP invokes it under the read side
of sk_callback_lock, which pins the owner and its data; the callback
must not sleep. With receive ownership disabled, the helper is empty.

This patch adds no user of the claim or callback.

Signed-off-by: Alex Markuze <amarkuze@redhat.com>
---
 include/linux/net.h  |  8 ++++++++
 include/net/sock.h   | 46 ++++++++++++++++++++++++++++++++++++++++++++
 net/Kconfig          |  7 +++++++
 net/core/skmsg.c     |  4 ++++
 net/core/sock.c      |  6 +++++-
 net/ipv4/af_inet.c   |  2 ++
 net/ipv4/tcp.c       | 22 ++++++++++++++++++++-
 net/ipv4/tcp_input.c |  1 +
 net/ipv4/tcp_ulp.c   |  2 ++
 9 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/include/linux/net.h b/include/linux/net.h
index e2a866fbcfa5..d8fe1812e02b 100644
--- a/include/linux/net.h
+++ b/include/linux/net.h
@@ -230,6 +230,14 @@ struct sk_buff;
 struct proto_accept_arg;
 typedef int (*sk_read_actor_t)(read_descriptor_t *, struct sk_buff *,
 			       unsigned int, size_t);
+
+/* Published under socket and callback locks; controls the receive sequence. */
+struct sock_rx_owner {
+	sk_read_actor_t actor;
+	void *data;
+	/* TCP receive context, under sk_callback_lock: must not sleep. */
+	void (*data_ready)(struct sock *sk, void *data);
+};
 typedef int (*skb_read_actor_t)(struct sock *, struct sk_buff *);
 
 
diff --git a/include/net/sock.h b/include/net/sock.h
index 173b37b2b8bb..4da1796b3b8c 100644
--- a/include/net/sock.h
+++ b/include/net/sock.h
@@ -326,6 +326,7 @@ struct sk_filter;
   *	@sk_zckey: counter to order MSG_ZEROCOPY notifications
   *	@sk_socket: Identd and reporting IO signals
   *	@sk_user_data: RPC layer private data. Write-protected by @sk_callback_lock.
+  *	@sk_rx_owner: Exclusive receive owner under socket and callback locks.
   *	@sk_frag: cached page frag
   *	@sk_peek_off: current peek_offset value
   *	@sk_send_head: front of stuff to transmit
@@ -575,6 +576,9 @@ struct sock {
 	u8			sk_bpf_cb_flags;
 
 	void			*sk_user_data;
+#ifdef CONFIG_SOCK_RX_OWNER
+	struct sock_rx_owner	*sk_rx_owner;
+#endif
 #ifdef CONFIG_SECURITY
 	void			*sk_security;
 #endif
@@ -617,6 +621,48 @@ enum sk_pacing {
 	SK_PACING_FQ		= 2,
 };
 
+static inline bool sock_rx_owned(const struct sock *sk)
+{
+#ifdef CONFIG_SOCK_RX_OWNER
+	return unlikely(READ_ONCE(sk->sk_rx_owner));
+#else
+	return false;
+#endif
+}
+
+/*
+ * Notify the receive owner independently of ordinary socket low-water wakes.
+ * The callback lock pins both owner and data until the callback returns;
+ * publication and removal take its write side while owning the socket.
+ */
+static inline void sock_rx_owner_data_ready(struct sock *sk)
+{
+#ifdef CONFIG_SOCK_RX_OWNER
+	struct sock_rx_owner *owner;
+
+	if (!sock_rx_owned(sk))
+		return;
+	read_lock_bh(&sk->sk_callback_lock);
+	owner = sk->sk_rx_owner;
+	if (owner && owner->data_ready)
+		owner->data_ready(sk, owner->data);
+	read_unlock_bh(&sk->sk_callback_lock);
+#endif
+}
+
+static inline bool sock_rx_owner_conflict(const struct sock *sk,
+					  read_descriptor_t *desc,
+					  sk_read_actor_t actor)
+{
+#ifdef CONFIG_SOCK_RX_OWNER
+	const struct sock_rx_owner *owner = READ_ONCE(sk->sk_rx_owner);
+
+	return owner && (owner->actor != actor || owner->data != desc->arg.data);
+#else
+	return false;
+#endif
+}
+
 /* flag bits in sk_user_data
  *
  * - SK_USER_DATA_NOCOPY:      Pointer stored in sk_user_data might
diff --git a/net/Kconfig b/net/Kconfig
index 76ab44aa439a..6ece51ab26ba 100644
--- a/net/Kconfig
+++ b/net/Kconfig
@@ -26,6 +26,13 @@ menuconfig NET
 
 if NET
 
+config SOCK_RX_OWNER
+	bool
+	depends on INET
+	help
+	  Allow a TCP stream actor to claim exclusive receive sequence access.
+	  Receive paths reject competing readers while the claim is held.
+
 config WANT_COMPAT_NETLINK_MESSAGES
 	bool
 	help
diff --git a/net/core/skmsg.c b/net/core/skmsg.c
index df385a5a961e..5c753bcb808c 100644
--- a/net/core/skmsg.c
+++ b/net/core/skmsg.c
@@ -754,6 +754,10 @@ struct sk_psock *sk_psock_init(struct sock *sk, int node)
 
 	write_lock_bh(&sk->sk_callback_lock);
 
+	if (sock_rx_owned(sk)) {
+		psock = ERR_PTR(-EBUSY);
+		goto out;
+	}
 	if (sk_is_inet(sk) && inet_csk_has_ulp(sk)) {
 		psock = ERR_PTR(-EINVAL);
 		goto out;
diff --git a/net/core/sock.c b/net/core/sock.c
index e7405be8941f..3fcdc56e624e 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -2614,6 +2614,9 @@ struct sock *sk_clone(const struct sock *sk, const gfp_t priority,
 	 */
 	if (sk_user_data_is_nocopy(newsk))
 		newsk->sk_user_data = NULL;
+#ifdef CONFIG_SOCK_RX_OWNER
+	newsk->sk_rx_owner = NULL;
+#endif
 
 	newsk->sk_err	   = 0;
 	newsk->sk_err_soft = 0;
@@ -3389,7 +3392,8 @@ int sk_wait_data(struct sock *sk, long *timeo, const struct sk_buff *skb)
 
 	add_wait_queue(sk_sleep(sk), &wait);
 	sk_set_bit(SOCKWQ_ASYNC_WAITDATA, sk);
-	rc = sk_wait_event(sk, timeo, skb_peek_tail(&sk->sk_receive_queue) != skb, &wait);
+	rc = sk_wait_event(sk, timeo, sock_rx_owned(sk) ||
+			   skb_peek_tail(&sk->sk_receive_queue) != skb, &wait);
 	sk_clear_bit(SOCKWQ_ASYNC_WAITDATA, sk);
 	remove_wait_queue(sk_sleep(sk), &wait);
 	return rc;
diff --git a/net/ipv4/af_inet.c b/net/ipv4/af_inet.c
index 7677b0343abd..54c30a59bbce 100644
--- a/net/ipv4/af_inet.c
+++ b/net/ipv4/af_inet.c
@@ -649,6 +649,8 @@ int __inet_stream_connect(struct socket *sock, struct sockaddr_unsized *uaddr,
 			return -EINVAL;
 
 		if (uaddr->sa_family == AF_UNSPEC) {
+			if (sock_rx_owned(sk))
+				return -EBUSY;
 			sk->sk_disconnects++;
 			err = sk->sk_prot->disconnect(sk, flags);
 			sock->state = err ? SS_DISCONNECTING : SS_UNCONNECTED;
diff --git a/net/ipv4/tcp.c b/net/ipv4/tcp.c
index b9de7c302ba8..6de6ce153fda 100644
--- a/net/ipv4/tcp.c
+++ b/net/ipv4/tcp.c
@@ -1753,6 +1753,8 @@ out:
 int tcp_read_sock(struct sock *sk, read_descriptor_t *desc,
 		  sk_read_actor_t recv_actor)
 {
+	if (sock_rx_owner_conflict(sk, desc, recv_actor))
+		return -EBUSY;
 	return __tcp_read_sock(sk, desc, recv_actor, false,
 			       &tcp_sk(sk)->copied_seq);
 }
@@ -1762,6 +1764,8 @@ int tcp_read_sock_noack(struct sock *sk, read_descriptor_t *desc,
 			sk_read_actor_t recv_actor, bool noack,
 			u32 *copied_seq)
 {
+	if (sock_rx_owner_conflict(sk, desc, recv_actor))
+		return -EBUSY;
 	return __tcp_read_sock(sk, desc, recv_actor, noack, copied_seq);
 }
 
@@ -1770,6 +1774,8 @@ int tcp_read_skb(struct sock *sk, skb_read_actor_t recv_actor)
 	struct sk_buff *skb;
 	int copied = 0;
 
+	if (sock_rx_owned(sk))
+		return -EBUSY;
 	if (sk->sk_state == TCP_LISTEN)
 		return -ENOTCONN;
 
@@ -2224,6 +2230,8 @@ static int tcp_zerocopy_receive(struct sock *sk,
 
 	zc->copybuf_len = 0;
 	zc->msg_flags = 0;
+	if (sock_rx_owned(sk))
+		return -EBUSY;
 
 	if (address & (PAGE_SIZE - 1) || address != zc->address)
 		return -EINVAL;
@@ -2678,6 +2686,8 @@ static int tcp_recvmsg_locked(struct sock *sk, struct msghdr *msg, size_t len,
 	u32 peek_offset = 0;
 	u32 urg_hole = 0;
 
+	if (sock_rx_owned(sk))
+		return -EBUSY;
 	err = -ENOTCONN;
 	if (sk->sk_state == TCP_LISTEN)
 		goto out;
@@ -2717,6 +2727,11 @@ static int tcp_recvmsg_locked(struct sock *sk, struct msghdr *msg, size_t len,
 	do {
 		u32 offset;
 
+		/* Ownership may have changed while this reader slept for data. */
+		if (sock_rx_owned(sk)) {
+			copied = copied ?: -EBUSY;
+			break;
+		}
 		/* Are we at urgent data? Stop if we have read anything or have SIGURG pending. */
 		if (unlikely(tp->urg_data) && tp->urg_seq == *seq) {
 			if (copied)
@@ -3374,6 +3389,9 @@ int tcp_disconnect(struct sock *sk, int flags)
 	struct request_sock *req;
 	u32 seq;
 
+	if (sock_rx_owned(sk))
+		return -EBUSY;
+
 	if (old_state != TCP_CLOSE)
 		tcp_set_state(sk, TCP_CLOSE);
 
@@ -3988,7 +4006,9 @@ int do_tcp_setsockopt(struct sock *sk, int level, int optname,
 		break;
 
 	case TCP_REPAIR:
-		if (!tcp_can_repair_sock(sk))
+		if (sock_rx_owned(sk))
+			err = -EBUSY;
+		else if (!tcp_can_repair_sock(sk))
 			err = -EPERM;
 		else if (val == TCP_REPAIR_ON) {
 			tp->repair = 1;
diff --git a/net/ipv4/tcp_input.c b/net/ipv4/tcp_input.c
index e1f2d9765f81..2d17fcdf73dd 100644
--- a/net/ipv4/tcp_input.c
+++ b/net/ipv4/tcp_input.c
@@ -5637,6 +5637,7 @@ err:
 
 void tcp_data_ready(struct sock *sk)
 {
+	sock_rx_owner_data_ready(sk);
 	if (tcp_epollin_ready(sk, sk->sk_rcvlowat) || sock_flag(sk, SOCK_DONE))
 		READ_ONCE(sk->sk_data_ready)(sk);
 }
diff --git a/net/ipv4/tcp_ulp.c b/net/ipv4/tcp_ulp.c
index b58045df101e..77abc21bcb63 100644
--- a/net/ipv4/tcp_ulp.c
+++ b/net/ipv4/tcp_ulp.c
@@ -164,6 +164,8 @@ int tcp_set_ulp(struct sock *sk, const char *name)
 
 	sock_owned_by_me(sk);
 
+	if (sock_rx_owned(sk))
+		return -EBUSY;
 	ulp_ops = __tcp_ulp_find_autoload(name);
 	if (!ulp_ops)
 		return -ENOENT;

base-commit: 4dd16e31135ea8437a720222df34e8745f385d6d
-- 
2.34.1