[PATCH net-next 0/9] psp: support rekeying psp protected tcp connections

[PATCH net-next 0/9] psp: support rekeying psp protected tcp connections

[Daniel Zahka]

The PSP architecture spec specifies the need for rekeying connections
in the event of a device key rotation. After a device key rotation has
occurred, a new rx derived key needs to be generated and sent out to
the other end of the connection sometime before the next device key
rotation occurs.

Because PSP connections involve two different keys at each endpoint,
one for decrypting ingress traffic, and one for encrypting egress
traffic, there are two types of rekeying events that need to be
supported. From the perspective of one endpoint of the connection:

1. rx rekey: we need to allocate a new spi and decryption key on the
current device key to provide to our peer.

2. tx rekey: our peer has provided us with a new spi + encryption key
pair which we should use for encrypting traffic immediately.

In the case of rx rekeying, there is a period where it makes sense to
accept packets authenticated from either the previous or current
spi. To deal with that we allow a socket to keep a chain of a max of
two psp assocs at any time. If authentication state does not match the
most recent assoc, the previous one will be tried.

In the case of tx rekeying, as soon as we install the new tx key, we
have no use for the previous one, and it can be disposed of
immediately. The only catch, is in the case where hw uses a key handle
in tx descriptor state (as opposed to inlining the key directly). If
this is the case, psp core needs to be sure that any of these
unaccounted for references to key state are gone by the time it tries
to sync a deleted key to hw.

To deal with this race condition, the series includes a driver api for
implementing deferred tx key deletion, where the driver can signal
back to the core when it is safe to dispose of old tx keys.

Lastly, some test cases for rekeying are included that go through key
rotations and rekeying.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
Daniel Zahka (9):
      psp: support rx rekey operation
      psp: move code from psp_sock_assoc_set_tx() into helper functions
      psp: support tx rekey operation
      psp: refactor psp_dev_tx_key_del()
      psp: add driver api for deferred tx key deletion
      psp: add core tracked stats for deferred key deletion
      mlx5: psp: implement deferred tx key deletion
      selftests: drv-net: psp: lift psp connection setup out of _data_basic_send() testcase
      selftests: drv-net: psp: add tests for rekeying connections

 Documentation/netlink/specs/psp.yaml               |  14 ++
 .../net/ethernet/mellanox/mlx5/core/en_accel/psp.c | 101 +++++++++-
 .../net/ethernet/mellanox/mlx5/core/en_accel/psp.h |   7 +
 drivers/net/ethernet/mellanox/mlx5/core/en_stats.h |   1 +
 drivers/net/ethernet/mellanox/mlx5/core/en_tx.c    |   1 +
 include/net/psp/functions.h                        |   6 +
 include/net/psp/types.h                            |  39 ++++
 include/uapi/linux/psp.h                           |   2 +
 net/psp/psp.h                                      |   7 +-
 net/psp/psp_main.c                                 | 101 +++++++++-
 net/psp/psp_nl.c                                   |   9 +-
 net/psp/psp_sock.c                                 | 213 +++++++++++++++------
 tools/testing/selftests/drivers/net/psp.py         | 133 ++++++++++++-
 .../testing/selftests/drivers/net/psp_responder.c  | 131 +++++++++++++
 14 files changed, 685 insertions(+), 80 deletions(-)
---
base-commit: 9a9424c756feee9ee6e717405a9d6fa7bacdef08
change-id: 20260202-psp-3c8e2f65c5c4

Best regards,
-- 
Daniel Zahka <daniel.zahka@gmail.com>

[PATCH net-next 1/9] psp: support rx rekey operation

[Daniel Zahka]

Support rx-assoc netlink operation on an already keyed socket.

Performing an rx-rekey on a socket involves splicing a new psp_assoc
object with newly generated rx state and reused tx state into a
socket's psp_assoc field. Add a 'prev' field to struct psp_assoc,
because we need to be able to accept skb's matching the previous
assoc. For SADB implementations, after copying the tx state of the
assoc, we need to mark the prior assoc as not needing tx key
cleanup. The 'tx_moved' field is introduced to struct psp_assoc for
this purpose.

With this change, skb's matching the current or previous psp_assoc
will be accepted. It might be reasonable to stop accepting skb's on
the previous assoc after seeing skb's with the new assoc and after
waiting a grace period, but that is not implemented for now.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 include/net/psp/functions.h |  6 ++++++
 include/net/psp/types.h     |  3 +++
 net/psp/psp.h               |  6 +++---
 net/psp/psp_nl.c            |  6 +-----
 net/psp/psp_sock.c          | 41 +++++++++++++++++++++++++++--------------
 5 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/include/net/psp/functions.h b/include/net/psp/functions.h
index c5c23a54774e..0f8b16f2f093 100644
--- a/include/net/psp/functions.h
+++ b/include/net/psp/functions.h
@@ -103,6 +103,12 @@ __psp_sk_rx_policy_check(struct sk_buff *skb, struct psp_assoc *pas)
 		return 0;
 	}
 
+	if (pas->prev && psp_pse_matches_pas(pse, pas->prev)) {
+		pas->peer_tx = 1;
+		pas->prev->peer_tx = 1;
+		return 0;
+	}
+
 	if (!pse) {
 		if (!pas->tx.spi ||
 		    (!pas->peer_tx && psp_is_allowed_nondata(skb, pas)))
diff --git a/include/net/psp/types.h b/include/net/psp/types.h
index 25a9096d4e7d..52844ac6f870 100644
--- a/include/net/psp/types.h
+++ b/include/net/psp/types.h
@@ -139,6 +139,9 @@ struct psp_assoc {
 
 	u32 upgrade_seq;
 
+	struct psp_assoc *prev;
+
+	bool tx_moved;
 	struct psp_key_parsed tx;
 	struct psp_key_parsed rx;
 
diff --git a/net/psp/psp.h b/net/psp/psp.h
index 9f19137593a0..b86539d5137a 100644
--- a/net/psp/psp.h
+++ b/net/psp/psp.h
@@ -21,9 +21,9 @@ void psp_nl_notify_dev(struct psp_dev *psd, u32 cmd);
 struct psp_assoc *psp_assoc_create(struct psp_dev *psd);
 struct psp_dev *psp_dev_get_for_sock(struct sock *sk);
 void psp_dev_tx_key_del(struct psp_dev *psd, struct psp_assoc *pas);
-int psp_sock_assoc_set_rx(struct sock *sk, struct psp_assoc *pas,
-			  struct psp_key_parsed *key,
-			  struct netlink_ext_ack *extack);
+void psp_sock_assoc_set_rx(struct sock *sk, struct psp_assoc *pas,
+			   struct psp_key_parsed *key,
+			   struct netlink_ext_ack *extack);
 int psp_sock_assoc_set_tx(struct sock *sk, struct psp_dev *psd,
 			  u32 version, struct psp_key_parsed *key,
 			  struct netlink_ext_ack *extack);
diff --git a/net/psp/psp_nl.c b/net/psp/psp_nl.c
index 6afd7707ec12..a540a1e77694 100644
--- a/net/psp/psp_nl.c
+++ b/net/psp/psp_nl.c
@@ -446,11 +446,7 @@ int psp_nl_rx_assoc_doit(struct sk_buff *skb, struct genl_info *info)
 		goto err_free_pas;
 	}
 
-	err = psp_sock_assoc_set_rx(socket->sk, pas, &key, info->extack);
-	if (err) {
-		NL_SET_BAD_ATTR(info->extack, info->attrs[PSP_A_ASSOC_SOCK_FD]);
-		goto err_free_pas;
-	}
+	psp_sock_assoc_set_rx(socket->sk, pas, &key, info->extack);
 	psp_assoc_put(pas);
 
 	return psp_nl_reply_send(rsp, info);
diff --git a/net/psp/psp_sock.c b/net/psp/psp_sock.c
index f785672b7df6..3a8abd023f99 100644
--- a/net/psp/psp_sock.c
+++ b/net/psp/psp_sock.c
@@ -85,7 +85,7 @@ static int psp_dev_tx_key_add(struct psp_dev *psd, struct psp_assoc *pas,
 
 void psp_dev_tx_key_del(struct psp_dev *psd, struct psp_assoc *pas)
 {
-	if (pas->tx.spi)
+	if (pas->tx.spi && !pas->tx_moved)
 		psd->ops->tx_key_del(psd, pas);
 	list_del(&pas->assocs_list);
 }
@@ -99,6 +99,7 @@ static void psp_assoc_free(struct work_struct *work)
 	if (psd->ops)
 		psp_dev_tx_key_del(psd, pas);
 	mutex_unlock(&psd->lock);
+	psp_assoc_put(pas->prev);
 	psp_dev_put(psd);
 	kfree(pas);
 }
@@ -129,30 +130,42 @@ void psp_sk_assoc_free(struct sock *sk)
 	psp_assoc_put(pas);
 }
 
-int psp_sock_assoc_set_rx(struct sock *sk, struct psp_assoc *pas,
-			  struct psp_key_parsed *key,
-			  struct netlink_ext_ack *extack)
+static void psp_sock_rx_rekey(struct psp_assoc *pas, struct psp_assoc *prev)
 {
-	int err;
+	lockdep_assert_held(&pas->psd->lock);
+
+	pas->peer_tx = prev->peer_tx;
+	pas->upgrade_seq = prev->upgrade_seq;
+
+	/* steal refcount from sk->psp_assoc */
+	pas->prev = prev;
+
+	memcpy(&pas->tx, &prev->tx, sizeof(pas->tx));
+	memcpy(pas->drv_data, prev->drv_data, pas->psd->caps->assoc_drv_spc);
+	prev->tx_moved = true;
+
+	psp_assoc_put(prev->prev);
+	prev->prev = NULL;
+}
+
+void psp_sock_assoc_set_rx(struct sock *sk, struct psp_assoc *pas,
+			   struct psp_key_parsed *key,
+			   struct netlink_ext_ack *extack)
+{
+	struct psp_assoc *prev;
 
 	memcpy(&pas->rx, key, sizeof(*key));
 
 	lock_sock(sk);
 
-	if (psp_sk_assoc(sk)) {
-		NL_SET_ERR_MSG(extack, "Socket already has PSP state");
-		err = -EBUSY;
-		goto exit_unlock;
-	}
+	prev = psp_sk_assoc(sk);
+	if (prev)
+		psp_sock_rx_rekey(pas, prev);
 
 	refcount_inc(&pas->refcnt);
 	rcu_assign_pointer(sk->psp_assoc, pas);
-	err = 0;
 
-exit_unlock:
 	release_sock(sk);
-
-	return err;
 }
 
 static int psp_sock_recv_queue_check(struct sock *sk, struct psp_assoc *pas)

-- 
2.47.3

[PATCH net-next 2/9] psp: move code from psp_sock_assoc_set_tx() into helper functions

[Daniel Zahka]

This commit is a pure refactor. Its purpose is to lift code that needs
to be called from both initial tx establishment and tx rekeying into
dedicated functions.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 net/psp/psp_sock.c | 94 +++++++++++++++++++++++++++++++++---------------------
 1 file changed, 58 insertions(+), 36 deletions(-)

diff --git a/net/psp/psp_sock.c b/net/psp/psp_sock.c
index 3a8abd023f99..9b0ecce8350f 100644
--- a/net/psp/psp_sock.c
+++ b/net/psp/psp_sock.c
@@ -187,12 +187,67 @@ static int psp_sock_recv_queue_check(struct sock *sk, struct psp_assoc *pas)
 	return 0;
 }
 
+static int
+psp_pas_set_tx_key(struct psp_dev *psd, struct psp_assoc *pas,
+		   struct psp_key_parsed *key, struct netlink_ext_ack *extack)
+{
+	struct psp_assoc *dummy;
+	int rc;
+
+	/* Pass a fake association to drivers to make sure they don't
+	 * try to store pointers to it. For re-keying we'll need to
+	 * re-allocate the assoc structures.
+	 */
+	dummy = psp_assoc_dummy(pas);
+	if (!dummy)
+		return -ENOMEM;
+
+	memcpy(&dummy->tx, key, sizeof(*key));
+	rc = psp_dev_tx_key_add(psd, dummy, extack);
+	if (rc)
+		goto exit_free_dummy;
+
+	memcpy(pas->drv_data, dummy->drv_data, psd->caps->assoc_drv_spc);
+	memcpy(&pas->tx, key, sizeof(*key));
+
+exit_free_dummy:
+	kfree(dummy);
+	return rc;
+}
+
+static int
+psp_sock_set_tx_key(struct sock *sk, struct psp_dev *psd, struct psp_assoc *pas,
+		    struct psp_key_parsed *key, struct netlink_ext_ack *extack)
+{
+	struct inet_connection_sock *icsk;
+	int err;
+
+	err = psp_sock_recv_queue_check(sk, pas);
+	if (err) {
+		NL_SET_ERR_MSG(extack, "Socket has incompatible segments already in the recv queue");
+		return err;
+	}
+
+	err = psp_pas_set_tx_key(psd, pas, key, extack);
+	if (err)
+		return err;
+
+	WRITE_ONCE(sk->sk_validate_xmit_skb, psp_validate_xmit);
+	tcp_write_collapse_fence(sk);
+	pas->upgrade_seq = tcp_sk(sk)->rcv_nxt;
+
+	icsk = inet_csk(sk);
+	icsk->icsk_ext_hdr_len += psp_sk_overhead(sk);
+	icsk->icsk_sync_mss(sk, icsk->icsk_pmtu_cookie);
+
+	return err;
+}
+
 int psp_sock_assoc_set_tx(struct sock *sk, struct psp_dev *psd,
 			  u32 version, struct psp_key_parsed *key,
 			  struct netlink_ext_ack *extack)
 {
-	struct inet_connection_sock *icsk;
-	struct psp_assoc *pas, *dummy;
+	struct psp_assoc *pas;
 	int err;
 
 	lock_sock(sk);
@@ -220,40 +275,7 @@ int psp_sock_assoc_set_tx(struct sock *sk, struct psp_dev *psd,
 		goto exit_unlock;
 	}
 
-	err = psp_sock_recv_queue_check(sk, pas);
-	if (err) {
-		NL_SET_ERR_MSG(extack, "Socket has incompatible segments already in the recv queue");
-		goto exit_unlock;
-	}
-
-	/* Pass a fake association to drivers to make sure they don't
-	 * try to store pointers to it. For re-keying we'll need to
-	 * re-allocate the assoc structures.
-	 */
-	dummy = psp_assoc_dummy(pas);
-	if (!dummy) {
-		err = -ENOMEM;
-		goto exit_unlock;
-	}
-
-	memcpy(&dummy->tx, key, sizeof(*key));
-	err = psp_dev_tx_key_add(psd, dummy, extack);
-	if (err)
-		goto exit_free_dummy;
-
-	memcpy(pas->drv_data, dummy->drv_data, psd->caps->assoc_drv_spc);
-	memcpy(&pas->tx, key, sizeof(*key));
-
-	WRITE_ONCE(sk->sk_validate_xmit_skb, psp_validate_xmit);
-	tcp_write_collapse_fence(sk);
-	pas->upgrade_seq = tcp_sk(sk)->rcv_nxt;
-
-	icsk = inet_csk(sk);
-	icsk->icsk_ext_hdr_len += psp_sk_overhead(sk);
-	icsk->icsk_sync_mss(sk, icsk->icsk_pmtu_cookie);
-
-exit_free_dummy:
-	kfree(dummy);
+	err = psp_sock_set_tx_key(sk, psd, pas, key, extack);
 exit_unlock:
 	release_sock(sk);
 	return err;

-- 
2.47.3

[PATCH net-next 3/9] psp: support tx rekey operation

[Daniel Zahka]

The tx rekey operation creates a new psp_assoc with the same rx state,
but with new tx state. The new psp_assoc will reference the current
assocs 'prev' assoc as its own. The assoc referenced by
'sk->psp_assoc' will have its reference dropped and be freed.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 net/psp/psp_sock.c | 51 +++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 45 insertions(+), 6 deletions(-)

diff --git a/net/psp/psp_sock.c b/net/psp/psp_sock.c
index 9b0ecce8350f..f429b8b2d8f2 100644
--- a/net/psp/psp_sock.c
+++ b/net/psp/psp_sock.c
@@ -215,6 +215,47 @@ psp_pas_set_tx_key(struct psp_dev *psd, struct psp_assoc *pas,
 	return rc;
 }
 
+static int
+psp_sock_tx_rekey(struct sock *sk, struct psp_dev *psd, struct psp_assoc *pas,
+		  struct psp_key_parsed *key, struct netlink_ext_ack *extack)
+{
+	struct psp_assoc *new;
+	size_t pas_sz;
+	int err;
+
+	pas_sz = struct_size(new, drv_data, psd->caps->assoc_drv_spc);
+	new = kzalloc(pas_sz, GFP_KERNEL_ACCOUNT);
+	if (!new)
+		return -ENOMEM;
+
+	/* don't increase refcounts until we know we won't fail */
+	new->psd         = pas->psd;
+	new->dev_id      = pas->dev_id;
+	new->generation  = pas->generation;
+	new->version     = pas->version;
+	new->peer_tx     = pas->peer_tx;
+	new->upgrade_seq = pas->upgrade_seq;
+	new->prev        = pas->prev;
+	refcount_set(&new->refcnt, 1);
+	memcpy(&new->rx, &pas->rx, sizeof(new->rx));
+
+	err = psp_pas_set_tx_key(psd, new, key, extack);
+	if (err) {
+		kfree(new);
+		return err;
+	}
+
+	psp_dev_get(new->psd);
+	if (new->prev)
+		refcount_inc(&new->prev->refcnt);
+	list_add_tail(&new->assocs_list, &psd->active_assocs);
+
+	rcu_assign_pointer(sk->psp_assoc, new);
+	psp_assoc_put(pas);
+
+	return 0;
+}
+
 static int
 psp_sock_set_tx_key(struct sock *sk, struct psp_dev *psd, struct psp_assoc *pas,
 		    struct psp_key_parsed *key, struct netlink_ext_ack *extack)
@@ -269,13 +310,11 @@ int psp_sock_assoc_set_tx(struct sock *sk, struct psp_dev *psd,
 		err = -EINVAL;
 		goto exit_unlock;
 	}
-	if (pas->tx.spi) {
-		NL_SET_ERR_MSG(extack, "Tx key already set");
-		err = -EBUSY;
-		goto exit_unlock;
-	}
+	if (pas->tx.spi)
+		err = psp_sock_tx_rekey(sk, psd, pas, key, extack);
+	else
+		err = psp_sock_set_tx_key(sk, psd, pas, key, extack);
 
-	err = psp_sock_set_tx_key(sk, psd, pas, key, extack);
 exit_unlock:
 	release_sock(sk);
 	return err;

-- 
2.47.3

[PATCH net-next 4/9] psp: refactor psp_dev_tx_key_del()

[Daniel Zahka]

This is a pure refactor. Lift the list deletion and key validation
code into callers. Subsequent patches will have different requirements
for what actions need to accompany psp dev key deletion.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 net/psp/psp_main.c |  7 +++++--
 net/psp/psp_sock.c | 11 ++++++-----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/net/psp/psp_main.c b/net/psp/psp_main.c
index a8534124f626..982c96f74ae3 100644
--- a/net/psp/psp_main.c
+++ b/net/psp/psp_main.c
@@ -132,8 +132,11 @@ void psp_dev_unregister(struct psp_dev *psd)
 
 	list_splice_init(&psd->active_assocs, &psd->prev_assocs);
 	list_splice_init(&psd->prev_assocs, &psd->stale_assocs);
-	list_for_each_entry_safe(pas, next, &psd->stale_assocs, assocs_list)
-		psp_dev_tx_key_del(psd, pas);
+	list_for_each_entry_safe(pas, next, &psd->stale_assocs, assocs_list) {
+		list_del(&pas->assocs_list);
+		if (pas->tx.spi && !pas->tx_moved)
+			psp_dev_tx_key_del(psd, pas);
+	}
 
 	rcu_assign_pointer(psd->main_netdev->psp_dev, NULL);
 
diff --git a/net/psp/psp_sock.c b/net/psp/psp_sock.c
index f429b8b2d8f2..1a97609564e7 100644
--- a/net/psp/psp_sock.c
+++ b/net/psp/psp_sock.c
@@ -85,9 +85,7 @@ static int psp_dev_tx_key_add(struct psp_dev *psd, struct psp_assoc *pas,
 
 void psp_dev_tx_key_del(struct psp_dev *psd, struct psp_assoc *pas)
 {
-	if (pas->tx.spi && !pas->tx_moved)
-		psd->ops->tx_key_del(psd, pas);
-	list_del(&pas->assocs_list);
+	psd->ops->tx_key_del(psd, pas);
 }
 
 static void psp_assoc_free(struct work_struct *work)
@@ -96,8 +94,11 @@ static void psp_assoc_free(struct work_struct *work)
 	struct psp_dev *psd = pas->psd;
 
 	mutex_lock(&psd->lock);
-	if (psd->ops)
-		psp_dev_tx_key_del(psd, pas);
+	if (psd->ops) {
+		list_del(&pas->assocs_list);
+		if (pas->tx.spi && !pas->tx_moved)
+			psp_dev_tx_key_del(psd, pas);
+	}
 	mutex_unlock(&psd->lock);
 	psp_assoc_put(pas->prev);
 	psp_dev_put(psd);

-- 
2.47.3

[PATCH net-next 5/9] psp: add driver api for deferred tx key deletion

[Daniel Zahka]

There is a race condition during tx key deletion on devices that use
key identifiers in tx metadata, as opposed to inlining keys
directly. When a key is requested to be deleted, there may be skb's in
the tx ring, whose metadata descriptors contain unaccounted references
to a tx key. Under these conditions, the skb sitting in the tx ring
can race against the request to fw to delete the tx key handle.

The solution implemented here is to provide an api for core to defer
tx key deletion operations until after the driver signals that it is
safe to do so.

Core informs the driver that a key deletion grace period is starting
via psp_dev::tx_grace_begin(). While a grace period is active, core
will queue keys to be deleted to the psp_dev::tx_del_next queue.

Core will then periodically query the driver with
psp_dev::tx_grace_end() to determine if it is safe to delete the
keys it has queued. If it is, core will then move the queued keys into
a new list and begin deleting them, allowing the psp_dev::tx_del_next
queue to begin filling for the next grace period.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 include/net/psp/types.h | 32 +++++++++++++++++
 net/psp/psp.h           |  1 +
 net/psp/psp_main.c      | 93 +++++++++++++++++++++++++++++++++++++++++++++++++
 net/psp/psp_sock.c      | 13 ++++++-
 4 files changed, 138 insertions(+), 1 deletion(-)

diff --git a/include/net/psp/types.h b/include/net/psp/types.h
index 52844ac6f870..280c7fdc713c 100644
--- a/include/net/psp/types.h
+++ b/include/net/psp/types.h
@@ -5,6 +5,7 @@
 
 #include <linux/mutex.h>
 #include <linux/refcount.h>
+#include <linux/workqueue.h>
 
 struct netlink_ext_ack;
 
@@ -58,6 +59,10 @@ struct psp_dev_config {
  * @prev_assocs:	associations which use old (but still usable)
  *			device key
  * @stale_assocs:	associations which use a rotated out key
+ * @tx_del_active:	TX keys to be deleted after the current grace period
+ * @tx_del_next:	TX keys queued for deletion during current grace period
+ * @tx_del_work:	workqueue for periodic grace period checking
+ * @tx_grace_active:	true if a grace period is in progress
  *
  * @stats:	statistics maintained by the core
  * @stats.rotations:	See stats attr key-rotations
@@ -85,6 +90,11 @@ struct psp_dev {
 	struct list_head prev_assocs;
 	struct list_head stale_assocs;
 
+	struct list_head tx_del_active;
+	struct list_head tx_del_next;
+	struct delayed_work tx_del_work;
+	bool tx_grace_active;
+
 	struct {
 		unsigned long rotations;
 		unsigned long stales;
@@ -208,6 +218,28 @@ struct psp_dev_ops {
 	 */
 	void (*tx_key_del)(struct psp_dev *psd, struct psp_assoc *pas);
 
+	/**
+	 * @tx_grace_begin: begin TX grace period tracking
+	 * Begin a new TX grace period. Core will queue tx keys for deletion
+	 * and not delete them until the grace period has elapsed.
+	 *
+	 * Return: 0 on success, error code otherwise. If the operation fails,
+	 * core will try again later and a grace period is not activated.
+	 */
+	int (*tx_grace_begin)(struct psp_dev *psd);
+
+	/**
+	 * @tx_grace_end: check if TX grace period has ended
+	 * Check if the current TX grace period has ended.
+	 *
+	 * Return:
+	 *   0: grace period has ended, core will delete queued tx keys
+	 *   -EAGAIN: grace period has not ended, core will requery later
+	 *   other error: internal error (e.g., device reconfigured), core
+	 *                should return to the tx_grace_begin() phase
+	 */
+	int (*tx_grace_end)(struct psp_dev *psd);
+
 	/**
 	 * @get_stats: get statistics from the device
 	 * Stats required by the spec must be maintained and filled in.
diff --git a/net/psp/psp.h b/net/psp/psp.h
index b86539d5137a..d09fdd956fcb 100644
--- a/net/psp/psp.h
+++ b/net/psp/psp.h
@@ -15,6 +15,7 @@ extern struct mutex psp_devs_lock;
 
 void psp_dev_free(struct psp_dev *psd);
 int psp_dev_check_access(struct psp_dev *psd, struct net *net);
+void psp_tx_key_queue_del(struct psp_dev *psd, struct psp_assoc *pas);
 
 void psp_nl_notify_dev(struct psp_dev *psd, u32 cmd);
 
diff --git a/net/psp/psp_main.c b/net/psp/psp_main.c
index 982c96f74ae3..2fb886edc384 100644
--- a/net/psp/psp_main.c
+++ b/net/psp/psp_main.c
@@ -14,6 +14,10 @@
 DEFINE_XARRAY_ALLOC1(psp_devs);
 struct mutex psp_devs_lock;
 
+#define PSP_TX_DEL_INTERVAL_MS  100
+
+static void psp_tx_del_work_fn(struct work_struct *work);
+
 /**
  * DOC: PSP locking
  *
@@ -64,6 +68,9 @@ psp_dev_create(struct net_device *netdev,
 		    !psd_ops->get_stats))
 		return ERR_PTR(-EINVAL);
 
+	if (WARN_ON(!!psd_ops->tx_grace_begin != !!psd_ops->tx_grace_end))
+		return ERR_PTR(-EINVAL);
+
 	psd = kzalloc(sizeof(*psd), GFP_KERNEL);
 	if (!psd)
 		return ERR_PTR(-ENOMEM);
@@ -77,6 +84,9 @@ psp_dev_create(struct net_device *netdev,
 	INIT_LIST_HEAD(&psd->active_assocs);
 	INIT_LIST_HEAD(&psd->prev_assocs);
 	INIT_LIST_HEAD(&psd->stale_assocs);
+	INIT_LIST_HEAD(&psd->tx_del_active);
+	INIT_LIST_HEAD(&psd->tx_del_next);
+	INIT_DELAYED_WORK(&psd->tx_del_work, psp_tx_del_work_fn);
 	refcount_set(&psd->refcnt, 1);
 
 	mutex_lock(&psp_devs_lock);
@@ -118,6 +128,8 @@ void psp_dev_unregister(struct psp_dev *psd)
 {
 	struct psp_assoc *pas, *next;
 
+	cancel_delayed_work_sync(&psd->tx_del_work);
+
 	mutex_lock(&psp_devs_lock);
 	mutex_lock(&psd->lock);
 
@@ -130,6 +142,15 @@ void psp_dev_unregister(struct psp_dev *psd)
 	xa_store(&psp_devs, psd->id, NULL, GFP_KERNEL);
 	mutex_unlock(&psp_devs_lock);
 
+	list_splice_init(&psd->tx_del_active, &psd->tx_del_next);
+	list_for_each_entry_safe(pas, next, &psd->tx_del_next, assocs_list) {
+		list_del(&pas->assocs_list);
+		psp_dev_tx_key_del(psd, pas);
+		psp_assoc_put(pas->prev);
+		psp_dev_put(psd);
+		kfree(pas);
+	}
+
 	list_splice_init(&psd->active_assocs, &psd->prev_assocs);
 	list_splice_init(&psd->prev_assocs, &psd->stale_assocs);
 	list_for_each_entry_safe(pas, next, &psd->stale_assocs, assocs_list) {
@@ -149,6 +170,78 @@ void psp_dev_unregister(struct psp_dev *psd)
 }
 EXPORT_SYMBOL(psp_dev_unregister);
 
+void psp_tx_key_queue_del(struct psp_dev *psd, struct psp_assoc *pas)
+{
+	lockdep_assert_held(&psd->lock);
+
+	list_add_tail(&pas->assocs_list, &psd->tx_del_next);
+	if (!delayed_work_pending(&psd->tx_del_work))
+		schedule_delayed_work(&psd->tx_del_work,
+				      msecs_to_jiffies(PSP_TX_DEL_INTERVAL_MS));
+}
+
+static void psp_tx_del_work_fn(struct work_struct *work)
+{
+	struct delayed_work *dwork = to_delayed_work(work);
+	struct psp_dev *psd = container_of(dwork, struct psp_dev, tx_del_work);
+	struct psp_assoc *pas, *tmp;
+	bool need_reschedule = false;
+	LIST_HEAD(to_free);
+	int err;
+
+	mutex_lock(&psd->lock);
+
+	if (!psp_dev_is_registered(psd))
+		goto out_unlock;
+
+	if (psd->tx_grace_active) {
+		err = psd->ops->tx_grace_end(psd);
+		if (err == -EAGAIN) {
+			need_reschedule = true;
+			goto out_unlock;
+		}
+		if (err) {
+			/* Driver error, restart grace period */
+			psd->tx_grace_active = false;
+			list_splice_init(&psd->tx_del_active, &psd->tx_del_next);
+			goto start_grace;
+		}
+
+		list_for_each_entry_safe(pas, tmp, &psd->tx_del_active,
+					 assocs_list) {
+			list_del(&pas->assocs_list);
+			psp_dev_tx_key_del(psd, pas);
+			list_add_tail(&pas->assocs_list, &to_free);
+		}
+
+		psd->tx_grace_active = false;
+	}
+
+start_grace:
+	if (!list_empty(&psd->tx_del_next)) {
+		err = psd->ops->tx_grace_begin(psd);
+		if (!err) {
+			list_splice_init(&psd->tx_del_next, &psd->tx_del_active);
+			psd->tx_grace_active = true;
+		}
+		need_reschedule = true;
+	}
+
+out_unlock:
+	mutex_unlock(&psd->lock);
+
+	list_for_each_entry_safe(pas, tmp, &to_free, assocs_list) {
+		list_del(&pas->assocs_list);
+		psp_assoc_put(pas->prev);
+		psp_dev_put(psd);
+		kfree(pas);
+	}
+
+	if (need_reschedule)
+		schedule_delayed_work(&psd->tx_del_work,
+				      msecs_to_jiffies(PSP_TX_DEL_INTERVAL_MS));
+}
+
 unsigned int psp_key_size(u32 version)
 {
 	switch (version) {
diff --git a/net/psp/psp_sock.c b/net/psp/psp_sock.c
index 1a97609564e7..d038da122ebb 100644
--- a/net/psp/psp_sock.c
+++ b/net/psp/psp_sock.c
@@ -88,6 +88,11 @@ void psp_dev_tx_key_del(struct psp_dev *psd, struct psp_assoc *pas)
 	psd->ops->tx_key_del(psd, pas);
 }
 
+static bool psp_dev_needs_defer(struct psp_dev *psd)
+{
+	return !!psd->ops->tx_grace_begin;
+}
+
 static void psp_assoc_free(struct work_struct *work)
 {
 	struct psp_assoc *pas = container_of(work, struct psp_assoc, work);
@@ -96,8 +101,14 @@ static void psp_assoc_free(struct work_struct *work)
 	mutex_lock(&psd->lock);
 	if (psd->ops) {
 		list_del(&pas->assocs_list);
-		if (pas->tx.spi && !pas->tx_moved)
+		if (pas->tx.spi && !pas->tx_moved) {
+			if (psp_dev_needs_defer(psd)) {
+				psp_tx_key_queue_del(psd, pas);
+				mutex_unlock(&psd->lock);
+				return;
+			}
 			psp_dev_tx_key_del(psd, pas);
+		}
 	}
 	mutex_unlock(&psd->lock);
 	psp_assoc_put(pas->prev);

-- 
2.47.3

[PATCH net-next 6/9] psp: add core tracked stats for deferred key deletion

[Daniel Zahka]

It is useful to keep track of two stats related to deferred key
deletion:

- tx_key_cnt: the number of outstanding tx keys from the hw's
  perspective
- grace_periods: the number of grace periods that core has retired for
  this dev.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 Documentation/netlink/specs/psp.yaml | 14 ++++++++++++++
 include/net/psp/types.h              |  4 ++++
 include/uapi/linux/psp.h             |  2 ++
 net/psp/psp_main.c                   |  1 +
 net/psp/psp_nl.c                     |  3 +++
 net/psp/psp_sock.c                   |  9 ++++++++-
 6 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/Documentation/netlink/specs/psp.yaml b/Documentation/netlink/specs/psp.yaml
index f3a57782d2cf..efaf3f61e37a 100644
--- a/Documentation/netlink/specs/psp.yaml
+++ b/Documentation/netlink/specs/psp.yaml
@@ -153,6 +153,18 @@ attribute-sets:
         doc: |
           Number of PSP packets for transmission with errors.
           Device statistic (from the PSP spec).
+      -
+        name: tx-key-cnt
+        type: sint
+        doc: |
+          Current number of TX keys installed on the device.
+          Kernel statistic.
+      -
+        name: grace-periods
+        type: uint
+        doc: |
+          Number of TX key deletion grace periods completed.
+          Kernel statistic.
 
 operations:
   list:
@@ -267,6 +279,8 @@ operations:
             - dev-id
             - key-rotations
             - stale-events
+            - tx-key-cnt
+            - grace-periods
         pre: psp-device-get-locked
         post: psp-device-unlock
       dump:
diff --git a/include/net/psp/types.h b/include/net/psp/types.h
index 280c7fdc713c..d790efd9d4bc 100644
--- a/include/net/psp/types.h
+++ b/include/net/psp/types.h
@@ -67,6 +67,8 @@ struct psp_dev_config {
  * @stats:	statistics maintained by the core
  * @stats.rotations:	See stats attr key-rotations
  * @stats.stales:	See stats attr stale-events
+ * @stats.tx_key_cnt:	See stats attr tx-key-cnt
+ * @stats.grace_periods:	See stats attr grace-periods
  *
  * @rcu:	RCU head for freeing the structure
  */
@@ -98,6 +100,8 @@ struct psp_dev {
 	struct {
 		unsigned long rotations;
 		unsigned long stales;
+		long tx_key_cnt;
+		unsigned long grace_periods;
 	} stats;
 
 	struct rcu_head rcu;
diff --git a/include/uapi/linux/psp.h b/include/uapi/linux/psp.h
index a3a336488dc3..208ab60e2c88 100644
--- a/include/uapi/linux/psp.h
+++ b/include/uapi/linux/psp.h
@@ -58,6 +58,8 @@ enum {
 	PSP_A_STATS_TX_PACKETS,
 	PSP_A_STATS_TX_BYTES,
 	PSP_A_STATS_TX_ERROR,
+	PSP_A_STATS_TX_KEY_CNT,
+	PSP_A_STATS_GRACE_PERIODS,
 
 	__PSP_A_STATS_MAX,
 	PSP_A_STATS_MAX = (__PSP_A_STATS_MAX - 1)
diff --git a/net/psp/psp_main.c b/net/psp/psp_main.c
index 2fb886edc384..4a7b2ffa718e 100644
--- a/net/psp/psp_main.c
+++ b/net/psp/psp_main.c
@@ -215,6 +215,7 @@ static void psp_tx_del_work_fn(struct work_struct *work)
 		}
 
 		psd->tx_grace_active = false;
+		psd->stats.grace_periods++;
 	}
 
 start_grace:
diff --git a/net/psp/psp_nl.c b/net/psp/psp_nl.c
index a540a1e77694..f0006d2e85c8 100644
--- a/net/psp/psp_nl.c
+++ b/net/psp/psp_nl.c
@@ -526,6 +526,9 @@ psp_nl_stats_fill(struct psp_dev *psd, struct sk_buff *rsp,
 	    nla_put_uint(rsp, PSP_A_STATS_KEY_ROTATIONS,
 			 psd->stats.rotations) ||
 	    nla_put_uint(rsp, PSP_A_STATS_STALE_EVENTS, psd->stats.stales) ||
+	    nla_put_sint(rsp, PSP_A_STATS_TX_KEY_CNT, psd->stats.tx_key_cnt) ||
+	    nla_put_uint(rsp, PSP_A_STATS_GRACE_PERIODS,
+			 psd->stats.grace_periods) ||
 	    nla_put_uint(rsp, PSP_A_STATS_RX_PACKETS, stats.rx_packets) ||
 	    nla_put_uint(rsp, PSP_A_STATS_RX_BYTES, stats.rx_bytes) ||
 	    nla_put_uint(rsp, PSP_A_STATS_RX_AUTH_FAIL, stats.rx_auth_fail) ||
diff --git a/net/psp/psp_sock.c b/net/psp/psp_sock.c
index d038da122ebb..bf6d089657b6 100644
--- a/net/psp/psp_sock.c
+++ b/net/psp/psp_sock.c
@@ -80,12 +80,19 @@ static struct psp_assoc *psp_assoc_dummy(struct psp_assoc *pas)
 static int psp_dev_tx_key_add(struct psp_dev *psd, struct psp_assoc *pas,
 			      struct netlink_ext_ack *extack)
 {
-	return psd->ops->tx_key_add(psd, pas, extack);
+	int rc;
+
+	rc = psd->ops->tx_key_add(psd, pas, extack);
+	if (!rc)
+		psd->stats.tx_key_cnt++;
+
+	return rc;
 }
 
 void psp_dev_tx_key_del(struct psp_dev *psd, struct psp_assoc *pas)
 {
 	psd->ops->tx_key_del(psd, pas);
+	psd->stats.tx_key_cnt--;
 }
 
 static bool psp_dev_needs_defer(struct psp_dev *psd)

-- 
2.47.3

[PATCH net-next 7/9] mlx5: psp: implement deferred tx key deletion

[Daniel Zahka]

Implement the deferred tx key deletion api. In the case of mlx5,
mlx5e_psp_tx_grace_begin() records the number of wqes retired on each
tx queue, and then mlx5e_psp_tx_grace_end() returns 0 only if
all tx queues have advanced a full ring cycle past the point where
they were snapshotted.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 .../net/ethernet/mellanox/mlx5/core/en_accel/psp.c | 101 +++++++++++++++++++--
 .../net/ethernet/mellanox/mlx5/core/en_accel/psp.h |   7 ++
 drivers/net/ethernet/mellanox/mlx5/core/en_stats.h |   1 +
 drivers/net/ethernet/mellanox/mlx5/core/en_tx.c    |   1 +
 4 files changed, 104 insertions(+), 6 deletions(-)

diff --git a/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.c b/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.c
index 9a74438ce10a..30f1dbc3fa9d 100644
--- a/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.c
+++ b/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.c
@@ -1059,13 +1059,102 @@ mlx5e_psp_get_stats(struct psp_dev *psd, struct psp_dev_stats *stats)
 	stats->tx_error = atomic_read(&priv->psp->tx_drop);
 }
 
+static int mlx5e_psp_tx_grace_begin(struct psp_dev *psd)
+{
+	struct mlx5e_priv *priv = netdev_priv(psd->main_netdev);
+	struct mlx5e_psp_tx_snapshot *snap;
+	int num_channels, num_tc, num_sqs;
+	int idx = 0, rc = 0;
+	int i, tc;
+
+	mutex_lock(&priv->state_lock);
+
+	num_channels = priv->channels.num;
+	num_tc = mlx5e_get_dcb_num_tc(&priv->channels.params);
+	num_sqs = num_channels * num_tc;
+
+	snap = kzalloc(struct_size(snap, wqes, num_sqs), GFP_KERNEL);
+	if (!snap) {
+		rc = -ENOMEM;
+		goto out_unlock;
+	}
+
+	snap->num_channels = num_channels;
+	snap->num_tc = num_tc;
+
+	for (i = 0; i < priv->channels.num; i++) {
+		struct mlx5e_channel *c = priv->channels.c[i];
+
+		for (tc = 0; tc < c->num_tc; tc++)
+			snap->wqes[idx++] = READ_ONCE(c->sq[tc].stats->wqes);
+	}
+
+	priv->psp->tx_snapshot = snap;
+
+out_unlock:
+	mutex_unlock(&priv->state_lock);
+	return rc;
+}
+
+static int mlx5e_psp_tx_grace_end(struct psp_dev *psd)
+{
+	struct mlx5e_priv *priv = netdev_priv(psd->main_netdev);
+	struct mlx5e_psp_tx_snapshot *snap;
+	int num_channels, num_tc;
+	int idx = 0, rc = 0;
+	int i, tc;
+
+	mutex_lock(&priv->state_lock);
+
+	snap = priv->psp->tx_snapshot;
+	num_channels = priv->channels.num;
+	num_tc = mlx5e_get_dcb_num_tc(&priv->channels.params);
+
+	/* If channels were reconfigured, tell core to restart grace period */
+	if (snap->num_channels != num_channels || snap->num_tc != num_tc) {
+		kfree(snap);
+		priv->psp->tx_snapshot = NULL;
+		rc = -ESTALE;
+		goto out_unlock;
+	}
+
+	for (i = 0; i < priv->channels.num; i++) {
+		struct mlx5e_channel *c = priv->channels.c[i];
+
+		for (tc = 0; tc < c->num_tc; tc++) {
+			struct mlx5e_txqsq *sq = &c->sq[tc];
+			u32 ring_size = mlx5_wq_cyc_get_size(&sq->wq);
+			u64 current_wqes = READ_ONCE(sq->stats->wqes);
+			u64 snapshot_wqes = snap->wqes[idx++];
+
+			/* If the ring has cycled, any key_id handles in tx
+			 * descriptors must have been consumed by hw and
+			 * cleaned by sw.
+			 */
+			if ((s64)(current_wqes - snapshot_wqes) < ring_size) {
+				rc = -EAGAIN;
+				goto out_unlock;
+			}
+		}
+	}
+
+	kfree(snap);
+	priv->psp->tx_snapshot = NULL;
+
+out_unlock:
+	mutex_unlock(&priv->state_lock);
+	return rc;
+}
+
 static struct psp_dev_ops mlx5_psp_ops = {
-	.set_config   = mlx5e_psp_set_config,
-	.rx_spi_alloc = mlx5e_psp_rx_spi_alloc,
-	.tx_key_add   = mlx5e_psp_assoc_add,
-	.tx_key_del   = mlx5e_psp_assoc_del,
-	.key_rotate   = mlx5e_psp_key_rotate,
-	.get_stats    = mlx5e_psp_get_stats,
+	.set_config     = mlx5e_psp_set_config,
+	.rx_spi_alloc   = mlx5e_psp_rx_spi_alloc,
+	.tx_key_add     = mlx5e_psp_assoc_add,
+	.tx_key_del     = mlx5e_psp_assoc_del,
+	.key_rotate     = mlx5e_psp_key_rotate,
+	.tx_grace_begin = mlx5e_psp_tx_grace_begin,
+	.tx_grace_end   = mlx5e_psp_tx_grace_end,
+	.get_stats      = mlx5e_psp_get_stats,
 };
 
 void mlx5e_psp_unregister(struct mlx5e_priv *priv)
diff --git a/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.h b/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.h
index 6b62fef0d9a7..29800050a331 100644
--- a/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.h
+++ b/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.h
@@ -22,10 +22,17 @@ struct mlx5e_psp_stats {
 	u64 psp_tx_bytes_drop;
 };
 
+struct mlx5e_psp_tx_snapshot {
+	int num_channels;
+	int num_tc;
+	u64 wqes[];
+};
+
 struct mlx5e_psp {
 	struct psp_dev *psp;
 	struct psp_dev_caps caps;
 	struct mlx5e_psp_fs *fs;
+	struct mlx5e_psp_tx_snapshot *tx_snapshot;
 	atomic_t tx_key_cnt;
 	atomic_t tx_drop;
 };
diff --git a/drivers/net/ethernet/mellanox/mlx5/core/en_stats.h b/drivers/net/ethernet/mellanox/mlx5/core/en_stats.h
index 09f155acb461..5173a0d3f01a 100644
--- a/drivers/net/ethernet/mellanox/mlx5/core/en_stats.h
+++ b/drivers/net/ethernet/mellanox/mlx5/core/en_stats.h
@@ -445,6 +445,7 @@ struct mlx5e_sq_stats {
 	u64 cqes ____cacheline_aligned_in_smp;
 	u64 wake;
 	u64 cqe_err;
+	u64 wqes;
 };
 
 struct mlx5e_xdpsq_stats {
diff --git a/drivers/net/ethernet/mellanox/mlx5/core/en_tx.c b/drivers/net/ethernet/mellanox/mlx5/core/en_tx.c
index a01ee656a1e7..412ebc160056 100644
--- a/drivers/net/ethernet/mellanox/mlx5/core/en_tx.c
+++ b/drivers/net/ethernet/mellanox/mlx5/core/en_tx.c
@@ -883,6 +883,7 @@ bool mlx5e_poll_tx_cq(struct mlx5e_cq *cq, int napi_budget)
 	wmb();
 
 	sq->dma_fifo_cc = dma_fifo_cc;
+	stats->wqes += (u16)(sqcc - sq->cc);
 	sq->cc = sqcc;
 
 	netdev_tx_completed_queue(sq->txq, npkts, nbytes);

-- 
2.47.3

[PATCH net-next 8/9] selftests: drv-net: psp: lift psp connection setup out of _data_basic_send() testcase

[Daniel Zahka]

This is a pure refactor, so that the key exchange sequence can be used
in other test cases that start with a key exchange, but might choose
to perform different tx/rx patterns afterwards.

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 tools/testing/selftests/drivers/net/psp.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tools/testing/selftests/drivers/net/psp.py b/tools/testing/selftests/drivers/net/psp.py
index 864d9fce1094..63dc8757ba37 100755
--- a/tools/testing/selftests/drivers/net/psp.py
+++ b/tools/testing/selftests/drivers/net/psp.py
@@ -360,10 +360,8 @@ def assoc_twice(cfg):
         s.close()
 
 
-def _data_basic_send(cfg, version, ipver):
-    """ Test basic data send """
-    _init_psp_dev(cfg)
-
+def _establish_psp_conn(cfg, version, ipver=None):
+    """Establish a PSP connection and return after key exchange"""
     # Version 0 is required by spec, don't let it skip
     if version:
         name = cfg.pspnl.consts["version"].entries_by_val[version].name
@@ -388,7 +386,14 @@ def _data_basic_send(cfg, version, ipver):
                         "version": version,
                         "tx-key": tx,
                         "sock-fd": s.fileno()})
+    return s
+
+
+def _data_basic_send(cfg, version, ipver):
+    """ Test basic data send """
+    _init_psp_dev(cfg)
 
+    s = _establish_psp_conn(cfg, version, ipver)
     data_len = _send_careful(cfg, s, 100)
     _check_data_rx(cfg, data_len)
     _close_psp_conn(cfg, s)

-- 
2.47.3

[PATCH net-next 9/9] selftests: drv-net: psp: add tests for rekeying connections

[Daniel Zahka]

Add testcases for rekeying psp connections. Some simple tx or rx only
scenarios, and also a more complicated mix of tx/rx rekeys and device
key rotations.

Note: the data echo handler does not appear here because it was added
in commit 2aeb71b2f9e8 ("selftests: drv-net: add PSP responder")

Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>
---
 tools/testing/selftests/drivers/net/psp.py         | 120 +++++++++++++++++++
 .../testing/selftests/drivers/net/psp_responder.c  | 131 +++++++++++++++++++++
 2 files changed, 251 insertions(+)

diff --git a/tools/testing/selftests/drivers/net/psp.py b/tools/testing/selftests/drivers/net/psp.py
index 63dc8757ba37..f6bce462f28a 100755
--- a/tools/testing/selftests/drivers/net/psp.py
+++ b/tools/testing/selftests/drivers/net/psp.py
@@ -19,6 +19,19 @@ from lib.py import NetDrvEpEnv, PSPFamily, NlError
 from lib.py import bkg, rand_port, wait_port_listen
 
 
+class PSPExceptShortIO(Exception):
+    pass
+
+
+def psp_ver_keylen(version):
+    """Key length for given PSP version"""
+    if version == 0 or version == 2:
+        return 16
+    elif version == 1 or version == 3:
+        return 32
+    raise Exception(f"psp_ver_keylen(): bad version: {version}")
+
+
 def _get_outq(s):
     one = b'\0' * 4
     outq = fcntl.ioctl(s.fileno(), termios.TIOCOUTQ, one)
@@ -60,6 +73,19 @@ def _close_psp_conn(cfg, s):
     _close_conn(cfg, s)
 
 
+def _provide_spi(cfg, version):
+    _send_with_ack(cfg, b'provide spi\0')
+    tx = cfg.comm_sock.recv(4 + psp_ver_keylen(version))
+    return {
+        'spi': struct.unpack('I', tx[:4])[0],
+        'key': tx[4:]
+    }
+
+
+def _use_spi(cfg, rx):
+    _send_with_ack(cfg, b'use spi\0' + struct.pack('I', rx['spi']) + rx['key'])
+
+
 def _spi_xchg(s, rx):
     s.send(struct.pack('I', rx['spi']) + rx['key'])
     tx = s.recv(4 + len(rx['key']))
@@ -110,6 +136,29 @@ def _check_data_outq(s, exp_len, force_wait=False):
     ksft_eq(outq, exp_len)
 
 
+def _recv_careful(s, target, rounds=100):
+    data = b''
+    for _ in range(rounds):
+        try:
+            data += s.recv(target - len(data), socket.MSG_DONTWAIT)
+            if len(data) == target:
+                return data
+        except BlockingIOError:
+            time.sleep(0.001)
+    raise PSPExceptShortIO(target, len(data), data)
+
+
+def _req_echo(cfg, s, expect_fail=False):
+    _send_with_ack(cfg, b'data echo\0')
+    try:
+        _recv_careful(s, 5)
+        if expect_fail:
+            raise Exception("Received unexpected echo reply")
+    except PSPExceptShortIO:
+        if not expect_fail:
+            raise
+
+
 def _get_stat(cfg, key):
     return cfg.pspnl.get_stats({'dev-id': cfg.psp_dev_id})[key]
 
@@ -399,6 +448,77 @@ def _data_basic_send(cfg, version, ipver):
     _close_psp_conn(cfg, s)
 
 
+def data_basic_rx_rekey(cfg, version=0):
+    """ Test basic rx rekey """
+    _init_psp_dev(cfg)
+
+    s = _establish_psp_conn(cfg, version, None)
+    data_len = _send_careful(cfg, s, 10)
+    _check_data_rx(cfg, data_len)
+
+    for _ in range(10):
+        rx_assoc = cfg.pspnl.rx_assoc({"version": version,
+                                       "dev-id": cfg.psp_dev_id,
+                                       "sock-fd": s.fileno()})
+        rx = rx_assoc['rx-key']
+        _use_spi(cfg, rx)
+        _req_echo(cfg, s)
+
+
+def data_basic_tx_rekey(cfg, version=0):
+    """Test basic tx rekey"""
+    _init_psp_dev(cfg)
+
+    s = _establish_psp_conn(cfg, version)
+    data_len = _send_careful(cfg, s, 10)
+    _check_data_rx(cfg, data_len)
+
+    for _ in range(10):
+        tx = _provide_spi(cfg, version)
+        cfg.pspnl.tx_assoc({"dev-id": cfg.psp_dev_id,
+                            "version": version,
+                            "tx-key": tx,
+                            "sock-fd": s.fileno()})
+        data_len += _send_careful(cfg, s, 10)
+        _check_data_rx(cfg, data_len)
+
+
+def data_rekey_and_rotate(cfg, version=0):
+    """Test a mix of key rotations and rekey operations"""
+    _init_psp_dev(cfg)
+
+    s = _establish_psp_conn(cfg, version)
+    data_len = _send_careful(cfg, s, 10)
+    _check_data_rx(cfg, data_len)
+
+    rounds = 3
+    tx_rekeys_per_round = 2
+
+    for _ in range(rounds):
+        cfg.pspnl.key_rotate({"id": cfg.psp_dev_id})
+
+        # receive data on rotated key
+        _req_echo(cfg, s)
+
+        # rekey and receive data on active key
+        rx_assoc = cfg.pspnl.rx_assoc({"version": version,
+                                       "dev-id": cfg.psp_dev_id,
+                                       "sock-fd": s.fileno()})
+        rx = rx_assoc['rx-key']
+        _use_spi(cfg, rx)
+        _req_echo(cfg, s)
+
+        # perform an arbitrary number of tx rekeys
+        for _ in range(tx_rekeys_per_round):
+            tx = _provide_spi(cfg, version)
+            cfg.pspnl.tx_assoc({"dev-id": cfg.psp_dev_id,
+                                "version": version,
+                                "tx-key": tx,
+                                "sock-fd": s.fileno()})
+            data_len += _send_careful(cfg, s, 1)
+            _check_data_rx(cfg, data_len)
+
+
 def __bad_xfer_do(cfg, s, tx, version='hdr0-aes-gcm-128'):
     # Make sure we accept the ACK for the SPI before we seal with the bad assoc
     _check_data_outq(s, 0)
diff --git a/tools/testing/selftests/drivers/net/psp_responder.c b/tools/testing/selftests/drivers/net/psp_responder.c
index a26e7628bbb1..d62982e3bc68 100644
--- a/tools/testing/selftests/drivers/net/psp_responder.c
+++ b/tools/testing/selftests/drivers/net/psp_responder.c
@@ -37,6 +37,77 @@ static struct {
 	unsigned char rx;
 } psp_vers;
 
+static int psp_ver_key_len(unsigned char version)
+{
+	switch (version) {
+	case PSP_VERSION_HDR0_AES_GCM_128:
+	case PSP_VERSION_HDR0_AES_GMAC_128:
+		return 16;
+	case PSP_VERSION_HDR0_AES_GCM_256:
+	case PSP_VERSION_HDR0_AES_GMAC_256:
+		return 32;
+	default:
+		fprintf(stderr, "ERROR: %s: bad version %d\n",
+			__func__, version);
+	}
+
+	return 0;
+}
+
+static int
+set_tx_spi(struct ynl_sock *ys, struct opts *opts, __u32 spi, char *key,
+	   int data_sock)
+{
+	struct psp_tx_assoc_rsp *tsp;
+	struct psp_tx_assoc_req *teq;
+	ssize_t sz;
+
+	teq = psp_tx_assoc_req_alloc();
+
+	psp_tx_assoc_req_set_sock_fd(teq, data_sock);
+	psp_tx_assoc_req_set_version(teq, psp_vers.tx);
+	psp_tx_assoc_req_set_tx_key_spi(teq, spi);
+	psp_tx_assoc_req_set_tx_key_key(teq, key, psp_ver_key_len(psp_vers.tx));
+
+	tsp = psp_tx_assoc(ys, teq);
+	psp_tx_assoc_req_free(teq);
+	if (!tsp) {
+		perror("ERROR: failed to Tx assoc");
+		return -1;
+	}
+	psp_tx_assoc_rsp_free(tsp);
+
+	return 0;
+}
+
+static int
+get_rx_spi(struct ynl_sock *ys, struct opts *opts, __u32 *spi, char *key,
+	   int data_sock)
+{
+	struct psp_rx_assoc_rsp *rsp;
+	struct psp_rx_assoc_req *req;
+
+	req = psp_rx_assoc_req_alloc();
+
+	psp_rx_assoc_req_set_sock_fd(req, data_sock);
+	psp_rx_assoc_req_set_version(req, psp_vers.rx);
+
+	rsp = psp_rx_assoc(ys, req);
+	psp_rx_assoc_req_free(req);
+
+	if (!rsp) {
+		perror("ERROR: failed to Rx assoc");
+		return -1;
+	}
+
+	memcpy(spi, &rsp->rx_key.spi, sizeof(*spi));
+	memcpy(key, rsp->rx_key.key, rsp->rx_key._len.key);
+
+	psp_rx_assoc_rsp_free(rsp);
+
+	return 0;
+}
+
 static int conn_setup_psp(struct ynl_sock *ys, struct opts *opts, int data_sock)
 {
 	struct psp_rx_assoc_rsp *rsp;
@@ -118,6 +189,52 @@ static void send_str(int sock, int value)
 	send(sock, buf, ret + 1, MSG_WAITALL);
 }
 
+static void
+handle_provide_spi(struct ynl_sock *ys, struct opts *opts, int data_sock,
+		   int comm_sock)
+{
+	char msg[sizeof(__u32) + psp_ver_key_len(psp_vers.tx)];
+	__u32 spi;
+
+	if (data_sock < 0) {
+		fprintf(stderr, "WARN: provide tx spi but no data sock\n");
+		send_err(comm_sock);
+		return;
+	}
+
+	if (get_rx_spi(ys, opts, &spi, &msg[sizeof(spi)], data_sock)) {
+		fprintf(stderr, "ERROR: get_rx_spi() failed\n");
+		send_err(comm_sock);
+		return;
+	}
+
+	send_ack(comm_sock);
+	memcpy(msg, &spi, sizeof(spi));
+	send(comm_sock, msg, sizeof(msg), MSG_WAITALL);
+}
+
+static void
+handle_use_spi(struct ynl_sock *ys, struct opts *opts, char *msg, int data_sock,
+	       int comm_sock)
+{
+	__u32 spi;
+
+	if (data_sock < 0) {
+		fprintf(stderr, "WARN: use tx spi but no data sock\n");
+		send_err(comm_sock);
+		return;
+	}
+
+	memcpy(&spi, msg, sizeof(spi));
+	if (set_tx_spi(ys, opts, spi, &msg[sizeof(spi)], data_sock)) {
+		fprintf(stderr, "ERROR: set_tx_spi() failed!\n");
+		send_err(comm_sock);
+		return;
+	}
+
+	send_ack(comm_sock);
+}
+
 static void
 run_session(struct ynl_sock *ys, struct opts *opts,
 	    int server_sock, int comm_sock)
@@ -224,6 +341,20 @@ run_session(struct ynl_sock *ys, struct opts *opts,
 						fprintf(stderr, "WARN: echo but no data sock\n");
 					send_ack(comm_sock);
 				}
+				if (cmd("provide spi"))
+					handle_provide_spi(ys, opts, data_sock, comm_sock);
+				if (cmd("use spi")) {
+					char msg[sizeof(__u32) + psp_ver_key_len(psp_vers.tx)];
+
+					if (off >= sizeof(msg)) {
+						memcpy(msg, buf, sizeof(msg));
+						__consume(sizeof(msg));
+						handle_use_spi(ys, opts, msg, data_sock, comm_sock);
+					} else {
+						fprintf(stderr, "WARN: short use spi command!\n");
+						send_err(comm_sock);
+					}
+				}
 				if (cmd("data close")) {
 					if (data_sock >= 0) {
 						close(data_sock);

-- 
2.47.3

Re: [PATCH net-next 0/9] psp: support rekeying psp protected tcp connections

[Willem de Bruijn]

Daniel Zahka wrote:
> The PSP architecture spec specifies the need for rekeying connections
> in the event of a device key rotation. After a device key rotation has
> occurred, a new rx derived key needs to be generated and sent out to
> the other end of the connection sometime before the next device key
> rotation occurs.
> 
> Because PSP connections involve two different keys at each endpoint,
> one for decrypting ingress traffic, and one for encrypting egress
> traffic, there are two types of rekeying events that need to be
> supported. From the perspective of one endpoint of the connection:
> 
> 1. rx rekey: we need to allocate a new spi and decryption key on the
> current device key to provide to our peer.
> 
> 2. tx rekey: our peer has provided us with a new spi + encryption key
> pair which we should use for encrypting traffic immediately.
> 
> In the case of rx rekeying, there is a period where it makes sense to
> accept packets authenticated from either the previous or current
> spi. To deal with that we allow a socket to keep a chain of a max of
> two psp assocs at any time. If authentication state does not match the
> most recent assoc, the previous one will be tried.
> 
> In the case of tx rekeying, as soon as we install the new tx key, we
> have no use for the previous one, and it can be disposed of
> immediately.

So this defines a rekey event as an instant in time. An alternative
choice is to rekey at a specific seqno.

The difference matters only for retransmits.

Not sure there is a strong reason for either. But probably good to
state the choice explicitly.

> The only catch, is in the case where hw uses a key handle
> in tx descriptor state (as opposed to inlining the key directly). If
> this is the case, psp core needs to be sure that any of these
> unaccounted for references to key state are gone by the time it tries
> to sync a deleted key to hw.
> 
> To deal with this race condition, the series includes a driver api for
> implementing deferred tx key deletion, where the driver can signal
> back to the core when it is safe to dispose of old tx keys.
> 
> Lastly, some test cases for rekeying are included that go through key
> rotations and rekeying.

Still reading the code, first pass only.

Re: [PATCH net-next 2/9] psp: move code from psp_sock_assoc_set_tx() into helper functions

[Willem de Bruijn]

Daniel Zahka wrote:
> This commit is a pure refactor. Its purpose is to lift code that needs
> to be called from both initial tx establishment and tx rekeying into
> dedicated functions.
> 
> Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>

Reviewed-by: Willem de Bruijn <willemb@google.com>

Re: [PATCH net-next 4/9] psp: refactor psp_dev_tx_key_del()

[Willem de Bruijn]

Daniel Zahka wrote:
> This is a pure refactor. Lift the list deletion and key validation
> code into callers. Subsequent patches will have different requirements
> for what actions need to accompany psp dev key deletion.
> 
> Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>

Reviewed-by: Willem de Bruijn <willemb@google.com>

Re: [PATCH net-next 8/9] selftests: drv-net: psp: lift psp connection setup out of _data_basic_send() testcase

[Willem de Bruijn]

Daniel Zahka wrote:
> This is a pure refactor, so that the key exchange sequence can be used
> in other test cases that start with a key exchange, but might choose
> to perform different tx/rx patterns afterwards.
> 
> Signed-off-by: Daniel Zahka <daniel.zahka@gmail.com>

Reviewed-by: Willem de Bruijn <willemb@google.com>

Re: [PATCH net-next 0/9] psp: support rekeying psp protected tcp connections

[Daniel Zahka]

On 2/4/26 3:45 PM, Willem de Bruijn wrote:
> Daniel Zahka wrote:
>> In the case of tx rekeying, as soon as we install the new tx key, we
>> have no use for the previous one, and it can be disposed of
>> immediately.
> So this defines a rekey event as an instant in time. An alternative
> choice is to rekey at a specific seqno.
>
> The difference matters only for retransmits.
>
> Not sure there is a strong reason for either. But probably good to
> state the choice explicitly.

I suppose if we think about a rekey as occurring at a seqno we could do 
away with the deferred key deletion, and instead just wait for data sent 
before that seqno to be ack'd before deleting the key. I would say if 
nothing else that would be a significant improvement over what I have 
right now.

Re: [PATCH net-next 0/9] psp: support rekeying psp protected tcp connections

[Jakub Kicinski]

On Wed, 4 Feb 2026 16:43:46 -0500 Daniel Zahka wrote:
> On 2/4/26 3:45 PM, Willem de Bruijn wrote:
> > Daniel Zahka wrote:  
> >> In the case of tx rekeying, as soon as we install the new tx key, we
> >> have no use for the previous one, and it can be disposed of
> >> immediately.  
> > So this defines a rekey event as an instant in time. An alternative
> > choice is to rekey at a specific seqno.
> >
> > The difference matters only for retransmits.
> >
> > Not sure there is a strong reason for either. But probably good to
> > state the choice explicitly.  
> 
> I suppose if we think about a rekey as occurring at a seqno we could do 
> away with the deferred key deletion, and instead just wait for data sent 
> before that seqno to be ack'd before deleting the key. I would say if 
> nothing else that would be a significant improvement over what I have 
> right now.

If you mean the driver xmit problem - not sure this changed much,
a (spurious) retransmission may theoretically still be queued on 
the driver ring when data is acked.

Re: [PATCH net-next 0/9] psp: support rekeying psp protected tcp connections

[Jakub Kicinski]

On Wed, 04 Feb 2026 07:20:04 -0800 Daniel Zahka wrote:
> Daniel Zahka (9):
>       psp: support rx rekey operation
>       psp: move code from psp_sock_assoc_set_tx() into helper functions
>       psp: support tx rekey operation
>       psp: refactor psp_dev_tx_key_del()
>       psp: add driver api for deferred tx key deletion
>       psp: add core tracked stats for deferred key deletion
>       mlx5: psp: implement deferred tx key deletion
>       selftests: drv-net: psp: lift psp connection setup out of _data_basic_send() testcase
>       selftests: drv-net: psp: add tests for rekeying connections

Sorry to say but my immediate reaction when seeing this on the list 
on Wednesday was "wait, no packetdrill tests?.. we never finalized 
and upstreamed packet drill tests.."

Let's get that wrapped up first, then we can move onto more features.

Re: [PATCH net-next 1/9] psp: support rx rekey operation

[Jakub Kicinski]

On Wed, 04 Feb 2026 07:20:05 -0800 Daniel Zahka wrote:
> --- a/include/net/psp/types.h
> +++ b/include/net/psp/types.h
> @@ -139,6 +139,9 @@ struct psp_assoc {
>  
>  	u32 upgrade_seq;
>  
> +	struct psp_assoc *prev;
> +
> +	bool tx_moved;
>  	struct psp_key_parsed tx;
>  	struct psp_key_parsed rx;

IIUC think tx_moved is going away in v2, but please double check 
the layout here with pahole -C psp_assoc. We the fastpath fields
in the first cacheline (or two?) The tx_moved the way it's placed
would create a hole, and hopefully for success path we don't have
to check prev at all.

Re: [PATCH net-next 7/9] mlx5: psp: implement deferred tx key deletion

[Jakub Kicinski]

On Wed, 04 Feb 2026 07:20:11 -0800 Daniel Zahka wrote:
> Implement the deferred tx key deletion api. In the case of mlx5,
> mlx5e_psp_tx_grace_begin() records the number of wqes retired on each
> tx queue, and then mlx5e_psp_tx_grace_end() returns 0 only if
> all tx queues have advanced a full ring cycle past the point where
> they were snapshotted.

Either Friday night enlightenment or I'm very tired -- but I wonder if
we can do with BQL, without touching the driver at all?

struct dql {
	/* Fields accessed in enqueue path (dql_queued) */
	unsigned int	num_queued;		/* Total ever queued */

...

	/* Fields accessed only by completion path (dql_completed) */
	unsigned int	num_completed;		/* Total ever completed */
...
};

These are free running counters of how many bytes were queued to 
the driver and how many driver finished xmit'ing.

Re: [net-next,1/9] psp: support rx rekey operation

[Jakub Kicinski]

This is an AI-generated review of your patch. The human sending this
email has considered the AI review valid, or at least plausible.

For local reproduction steps see https://netdev-ai.bots.linux.dev/ai-local.html
---
psp: support rx rekey operation

This patch adds support for rx-assoc netlink operation on an already keyed
socket, enabling rx rekey by splicing a new psp_assoc object with newly
generated rx state and reused tx state.

> diff --git a/net/psp/psp_sock.c b/net/psp/psp_sock.c
> index f785672b7df6..3a8abd023f99 100644
> --- a/net/psp/psp_sock.c
> +++ b/net/psp/psp_sock.c

[ ... ]

> @@ -129,30 +130,42 @@ void psp_sk_assoc_free(struct sock *sk)
>  	psp_assoc_put(pas);
>  }
>
> -int psp_sock_assoc_set_rx(struct sock *sk, struct psp_assoc *pas,
> -			  struct psp_key_parsed *key,
> -			  struct netlink_ext_ack *extack)
> +static void psp_sock_rx_rekey(struct psp_assoc *pas, struct psp_assoc *prev)
>  {
> -	int err;
> +	lockdep_assert_held(&pas->psd->lock);
> +
> +	pas->peer_tx = prev->peer_tx;
> +	pas->upgrade_seq = prev->upgrade_seq;
> +
> +	/* steal refcount from sk->psp_assoc */
> +	pas->prev = prev;
> +
> +	memcpy(&pas->tx, &prev->tx, sizeof(pas->tx));
> +	memcpy(pas->drv_data, prev->drv_data, pas->psd->caps->assoc_drv_spc);

Is there a check that prev->psd == pas->psd before this memcpy?

The psp_assoc_device_get_locked() pre_doit handler obtains the PSP device from
the socket's current destination route via psp_dev_get_for_sock(), but this
may differ from the device used when the existing association (prev) was
created if the route changed.

The memcpy copies pas->psd->caps->assoc_drv_spc bytes from prev->drv_data, but
if prev->psd has a smaller assoc_drv_spc, this could read beyond the bounds of
prev->drv_data.

For comparison, psp_sock_assoc_set_tx() has explicit device validation:

    if (pas->psd != psd)
        NL_SET_ERR_MSG(extack, "Rx key from different device");

Should psp_sock_rx_rekey() have a similar check?

> +	prev->tx_moved = true;

Re: [net-next,7/9] mlx5: psp: implement deferred tx key deletion

[Jakub Kicinski]

This is an AI-generated review of your patch. The human sending this
email has considered the AI review valid, or at least plausible.

For local reproduction steps see https://netdev-ai.bots.linux.dev/ai-local.html
---
mlx5: psp: implement deferred tx key deletion

This patch implements the deferred tx key deletion API for mlx5. The
mlx5e_psp_tx_grace_begin() function records WQE counts on each tx queue,
and mlx5e_psp_tx_grace_end() returns success only after all tx queues
have advanced a full ring cycle past the snapshot point.

> diff --git a/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.h b/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.h
> --- a/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.h
> +++ b/drivers/net/ethernet/mellanox/mlx5/core/en_accel/psp.h
> @@ -22,10 +22,17 @@ struct mlx5e_psp_stats {
>  	u64 psp_tx_bytes_drop;
>  };
>
> +struct mlx5e_psp_tx_snapshot {
> +	int num_channels;
> +	int num_tc;
> +	u64 wqes[];
> +};
> +
>  struct mlx5e_psp {
>  	struct psp_dev *psp;
>  	struct psp_dev_caps caps;
>  	struct mlx5e_psp_fs *fs;
> +	struct mlx5e_psp_tx_snapshot *tx_snapshot;
>  	atomic_t tx_key_cnt;
>  	atomic_t tx_drop;
>  };

Can the tx_snapshot allocation be leaked during device teardown?

If mlx5e_psp_tx_grace_begin() is called and allocates tx_snapshot, but
mlx5e_psp_cleanup() runs before mlx5e_psp_tx_grace_end() frees it:

    mlx5e_psp_cleanup()
        kfree(psp);   // psp->tx_snapshot not freed

The tx_snapshot memory would be leaked. Should mlx5e_psp_cleanup() free
psp->tx_snapshot before freeing psp?