[RFC net-next 5/7] net: tls: keep the TX queue of a device-offloaded socket stable

[Tariq Toukan <tariqt@nvidia.com>]

A device that keeps per-connection TLS TX state (e.g. mlx5e kTLS, where
WQEs update the TIS crypto and record tracking state at execution time)
cannot order work posted on different TX queues. If the flow moves to
another TX queue while work is still in flight on the previous one, the
device state is accessed from both queues with no ordering, corrupting
the record stream on the wire.

The stack moves a socket's flow while packets are in flight: TCP sets
ooo_okay when the rtx queue is empty, a dst reset (e.g. on RTO or PLB
rehash) clears the cached TX queue, and txq_reselection_ms expires it.

Let a driver declare the constraint via a new tlsdev_ops flag,
TLSDEV_F_STICKY_TXQ, and mark sockets offloaded to such a device with a
new sock flag, SOCK_STICKY_TXQ. While such a socket holds sk_wmem_alloc,
keep its cached TX queue: do not clear it on dst changes, do not expire
it, and set ooo_okay from the wmem condition only.

This covers the stack's own queue selection. The rest is the driver's,
as documented: a ndo_select_queue that derives the queue from more than
the cached one must keep its choice stable too, changing it only for an
skb with ooo_okay set, and when TX queues are removed, the device must
be done with their work before the flows move. tc skbedit queue_mapping
bypasses queue selection, and is not supported for such sockets.

Packets in flight are tracked by sk_wmem_alloc, so the flag also
requires the driver not to free an skb, even a dropped one, before the
device is done with the work posted for it.

Found by LLM-assisted code inspection, while answering a question on
kTLS TX TIS reuse for TLS 1.3 rekey on netdev, and reproduced on
ConnectX-6 Dx with forced queue moves.

Reproduced on ConnectX-6 Dx, with 128 channels and 32 kTLS TX offloaded
connections streaming pattern data for 60 seconds to a software kTLS
receiver that checks every byte. Without this patch and the mlx5e one
that follows it, 1 of 33 connections failed decryption with the default
txq_reselection_ms, and 67 of 99 with txq_reselection_ms=1 and the
senders migrating between CPUs every millisecond. With both applied, no
record failed in either test, with no loss of throughput.

Fixes: e8f69799810c ("net/tls: Add generic NIC offload infrastructure")
Fixes: 726e9e8b94b9 ("tcp: refine skb->ooo_okay setting")
Fixes: 4a7708443dec ("net: allow busy connected flows to switch tx queues")
Link: https://lore.kernel.org/netdev/20261022181751.1349984-1-rjethwani@purestorage.com/
Signed-off-by: Tariq Toukan <tariqt@nvidia.com>
---
 Documentation/admin-guide/sysctl/net.rst |  3 +++
 Documentation/networking/tls-offload.rst | 20 ++++++++++++++
 include/net/sock.h                       | 33 +++++++++++++++++++++---
 include/net/tls.h                        | 13 ++++++++++
 net/core/dev.c                           |  3 ++-
 net/core/sock.c                          |  2 +-
 net/ipv4/tcp_output.c                    |  7 ++++-
 net/tls/tls_device.c                     |  3 +++
 8 files changed, 78 insertions(+), 6 deletions(-)

diff --git a/Documentation/admin-guide/sysctl/net.rst b/Documentation/admin-guide/sysctl/net.rst
index fe43e8595958..88bfb4421843 100644
--- a/Documentation/admin-guide/sysctl/net.rst
+++ b/Documentation/admin-guide/sysctl/net.rst
@@ -449,6 +449,9 @@ But switching txq can introduce reorders, especially if the
 old queue is under high pressure. Modern TCP stacks deal
 well with reorders if they happen not too often.
 
+A flow offloaded to a TLS device that keeps per-flow TX state in
+the device switches txq only once it has no packets in flight.
+
 To disable this feature, set the value to 0.
 
 Default : 1000
diff --git a/Documentation/networking/tls-offload.rst b/Documentation/networking/tls-offload.rst
index cf0e599c06a1..f03e48643503 100644
--- a/Documentation/networking/tls-offload.rst
+++ b/Documentation/networking/tls-offload.rst
@@ -190,6 +190,26 @@ TLS handling and confirms the sequence number matches its expectation.
 The device performs encryption and authentication of the record data.
 It replaces the authentication tag and TCP checksum with correct values.
 
+A device may keep the TLS state of a connection in the device, and update it
+as it processes the TX queues, with no ordering between the queues. Its driver
+sets ``TLSDEV_F_STICKY_TXQ`` in ``tlsdev_ops->flags``, and the stack then keeps
+the TX queue selection of an offloaded socket stable while the socket has
+packets in flight (``SOCK_STICKY_TXQ``): the cached queue is neither cleared on
+route changes nor expired by ``txq_reselection_ms``, and TCP sets
+``skb->ooo_okay`` only when no packet of the socket is in flight. Packets in
+flight are tracked by ``sk_wmem_alloc``, so the driver must not free an skb
+before the device is done with all the work it posted for it, even when the skb
+is dropped.
+
+The rest is up to the driver. If its ``ndo_select_queue`` derives the queue from
+more than the socket's cached queue, it must keep its choice stable as well, and
+may change it only for an skb with ``ooo_okay`` set. A queue change driven by
+the skb's priority or traffic class then waits until the socket has no packets
+in flight, which under sustained traffic may never happen. When TX queues are
+removed, the device must be done with their work before the flows that move from
+them start transmitting on other queues. ``tc`` ``skbedit queue_mapping``
+bypasses queue selection and is not supported for such sockets.
+
 RX
 --
 
diff --git a/include/net/sock.h b/include/net/sock.h
index 9bba1e793f1e..ab4b5a5a74f5 100644
--- a/include/net/sock.h
+++ b/include/net/sock.h
@@ -1026,6 +1026,11 @@ enum sock_flags {
 	SOCK_RCVMARK, /* Receive SO_MARK  ancillary data with packet */
 	SOCK_RCVPRIORITY, /* Receive SO_PRIORITY ancillary data with packet */
 	SOCK_TIMESTAMPING_ANY, /* Copy of sk_tsflags & TSFLAGS_ANY */
+	SOCK_STICKY_TXQ, /* Keep the TX queue selection while packets are in
+			  * flight; set for sockets whose TX offload keeps
+			  * per-flow device state updated by queue execution
+			  * (see TLSDEV_F_STICKY_TXQ).
+			  */
 };
 
 #define SK_FLAGS_TIMESTAMP ((1UL << SOCK_TIMESTAMP) | (1UL << SOCK_TIMESTAMPING_RX_SOFTWARE))
@@ -2067,6 +2072,29 @@ static inline void sk_tx_queue_clear(struct sock *sk)
 	WRITE_ONCE(sk->sk_tx_queue_mapping, NO_QUEUE_MAPPING);
 }
 
+#define SK_WMEM_ALLOC_BIAS 1
+
+/* A SOCK_STICKY_TXQ socket keeps its TX queue while it has packets in
+ * flight: a stateful TX offload must not see the flow on a new queue
+ * while its work may still be pending on the previous one.
+ */
+static inline bool sk_tx_queue_pinned(const struct sock *sk)
+{
+	return sock_flag(sk, SOCK_STICKY_TXQ) &&
+	       refcount_read(&sk->sk_wmem_alloc) - SK_WMEM_ALLOC_BIAS >=
+	       SKB_TRUESIZE(1);
+}
+
+/* Like sk_tx_queue_clear(), but keeps a pinned mapping, see
+ * sk_tx_queue_pinned(). Lifecycle resets (allocation, cloning) must use
+ * sk_tx_queue_clear() directly.
+ */
+static inline void sk_tx_queue_maybe_clear(struct sock *sk)
+{
+	if (!sk_tx_queue_pinned(sk))
+		sk_tx_queue_clear(sk);
+}
+
 int sk_tx_queue_get(const struct sock *sk);
 
 static inline void __sk_rx_queue_set(struct sock *sk,
@@ -2237,7 +2265,7 @@ __sk_dst_set(struct sock *sk, struct dst_entry *dst)
 {
 	struct dst_entry *old_dst;
 
-	sk_tx_queue_clear(sk);
+	sk_tx_queue_maybe_clear(sk);
 	WRITE_ONCE(sk->sk_dst_pending_confirm, 0);
 	old_dst = rcu_dereference_protected(sk->sk_dst_cache,
 					    lockdep_sock_is_held(sk));
@@ -2250,7 +2278,7 @@ sk_dst_set(struct sock *sk, struct dst_entry *dst)
 {
 	struct dst_entry *old_dst;
 
-	sk_tx_queue_clear(sk);
+	sk_tx_queue_maybe_clear(sk);
 	WRITE_ONCE(sk->sk_dst_pending_confirm, 0);
 	old_dst = unrcu_pointer(xchg(&sk->sk_dst_cache, RCU_INITIALIZER(dst)));
 	dst_release(old_dst);
@@ -2369,7 +2397,6 @@ static inline int skb_copy_to_page_nocache(struct sock *sk, struct iov_iter *fro
 	return 0;
 }
 
-#define SK_WMEM_ALLOC_BIAS 1
 /**
  * sk_wmem_alloc_get - returns write allocations
  * @sk: socket
diff --git a/include/net/tls.h b/include/net/tls.h
index e57bef58851e..80fd68c9801f 100644
--- a/include/net/tls.h
+++ b/include/net/tls.h
@@ -276,7 +276,20 @@ enum tls_offload_ctx_dir {
 	TLS_OFFLOAD_CTX_DIR_TX,
 };
 
+/* TX offload state (e.g. crypto and record tracking state of a device
+ * connection context) is updated by TX queue execution, and the device
+ * provides no cross-queue ordering. The stack keeps the socket's TX queue
+ * selection stable while the socket has packets in flight
+ * (SOCK_STICKY_TXQ), and the driver keeps its own part of the selection
+ * stable too, see Documentation/networking/tls-offload.rst. Packets in
+ * flight are tracked by sk_wmem_alloc, so the driver must not free an skb
+ * before the device is done with all work it posted for it, even when the
+ * skb is dropped.
+ */
+#define TLSDEV_F_STICKY_TXQ	BIT(0)
+
 struct tlsdev_ops {
+	u32 flags; /* TLSDEV_F_* */
 	int (*tls_dev_add)(struct net_device *netdev, struct sock *sk,
 			   enum tls_offload_ctx_dir direction,
 			   struct tls_crypto_info *crypto_info,
diff --git a/net/core/dev.c b/net/core/dev.c
index 3b60b15e2175..18f7b9159ec5 100644
--- a/net/core/dev.c
+++ b/net/core/dev.c
@@ -4894,7 +4894,8 @@ int sk_tx_queue_get(const struct sock *sk)
 
 	resel = READ_ONCE(sock_net(sk)->core.sysctl_txq_reselection);
 	if (resel && time_is_before_jiffies(
-			READ_ONCE(sk->sk_tx_queue_mapping_jiffies) + resel))
+			READ_ONCE(sk->sk_tx_queue_mapping_jiffies) + resel) &&
+	    !sk_tx_queue_pinned(sk))
 		return -1;
 
 	return val;
diff --git a/net/core/sock.c b/net/core/sock.c
index e2401d735c5a..d90a4f69ddda 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -606,7 +606,7 @@ struct dst_entry *__sk_dst_check(struct sock *sk, u32 cookie)
 	if (dst && READ_ONCE(dst->obsolete) &&
 	    INDIRECT_CALL_INET(dst->ops->check, ip6_dst_check, ipv4_dst_check,
 			       dst, cookie) == NULL) {
-		sk_tx_queue_clear(sk);
+		sk_tx_queue_maybe_clear(sk);
 		WRITE_ONCE(sk->sk_dst_pending_confirm, 0);
 		RCU_INIT_POINTER(sk->sk_dst_cache, NULL);
 		dst_release(dst);
diff --git a/net/ipv4/tcp_output.c b/net/ipv4/tcp_output.c
index 08d0d4a639fe..690f73e9de41 100644
--- a/net/ipv4/tcp_output.c
+++ b/net/ipv4/tcp_output.c
@@ -1661,9 +1661,14 @@ static int __tcp_transmit_skb(struct sock *sk, struct sk_buff *skb,
 	 * 2) Or rtx queue is empty.
 	 *    This mitigates above case if ACK packets for
 	 *    all prior packets were already processed.
+	 * A SOCK_STICKY_TXQ socket (stateful TX offload) relies on 1) only:
+	 * a spurious retransmit clone may still hold wmem and reference
+	 * device state from the old queue when the rtx queue is already
+	 * empty.
 	 */
 	skb->ooo_okay = sk_wmem_alloc_get(sk) < SKB_TRUESIZE(1) ||
-			tcp_rtx_queue_empty(sk);
+			(tcp_rtx_queue_empty(sk) &&
+			 !sock_flag(sk, SOCK_STICKY_TXQ));
 
 	/* If we had to use memory reserve to allocate this skb,
 	 * this might cause drops if packet is looped back :
diff --git a/net/tls/tls_device.c b/net/tls/tls_device.c
index 4cdfbb57ecff..048b6a912d41 100644
--- a/net/tls/tls_device.c
+++ b/net/tls/tls_device.c
@@ -1197,6 +1197,9 @@ int tls_set_device_offload(struct sock *sk)
 	tls_device_attach(ctx, sk, netdev);
 	up_read(&device_offload_lock);
 
+	if (netdev->tlsdev_ops->flags & TLSDEV_F_STICKY_TXQ)
+		sock_set_flag(sk, SOCK_STICKY_TXQ);
+
 	/* following this assignment tls_is_skb_tx_device_offloaded
 	 * will return true and the context might be accessed
 	 * by the netdev's xmit function.
-- 
2.44.0