From d78ad2ef15486806ee6fcf97de5843d0b081b822 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?=
 <jedrzej.stuczynski@gmail.com>
Date: Fri, 5 Jun 2026 18:54:29 +0100
Subject: [PATCH] update handle_update_peer_psk_request

---
 common/wireguard/src/peer_controller/mod.rs   | 24 +++++++++++++++----
 .../wireguard/new_peer_registration/lp.rs     | 15 +++---------
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/common/wireguard/src/peer_controller/mod.rs b/common/wireguard/src/peer_controller/mod.rs
index 4aa1daf5ab..a4fbbe2c02 100644
--- a/common/wireguard/src/peer_controller/mod.rs
+++ b/common/wireguard/src/peer_controller/mod.rs
@@ -334,8 +334,17 @@ impl PeerController {
         let Ok(Some(mut peer)) = self.handle_query_peer_by_key(peer_key).await else {
             return Ok(());
         };
+        let encoded_psk = psk.to_lower_hex();
         peer.preshared_key = Some(psk);
 
+        // Account for bandwidth used so far *before* reconfiguring: `configure_peer`
+        // isn't guaranteed to preserve the kernel rx/tx counters, so fold the
+        // accrued bytes into the metrics first to avoid losing them on a reset.
+        if let Ok(host) = self.wg_api.read_interface_data() {
+            self.update_metrics(&host).await;
+            *self.host_information.write().await = host;
+        }
+
         // Try to update WireGuard peer
         if let Err(e) = self.wg_api.configure_peer(&peer) {
             nym_metrics::inc!("wg_peer_update_psk_failed");
@@ -343,12 +352,19 @@ impl PeerController {
             return Err(e.into());
         };
 
-        // try to immediately update the host information, to eliminate races
-        if let Ok(host_information) = self.wg_api.read_interface_data() {
-            *self.host_information.write().await = host_information;
+        // Persist the new PSK to disk so it survives a restart. Kernel-first: a
+        // failure here leaves the live session working, only risking drift on restart.
+        self.ecash_verifier
+            .storage()
+            .update_peer_psk(&peer_key.to_string(), Some(&encoded_psk))
+            .await?;
+
+        // Refresh again so the cached host information reflects the post-update state
+        if let Ok(host) = self.wg_api.read_interface_data() {
+            *self.host_information.write().await = host;
         }
 
-        nym_metrics::inc!("wg_peer_addition_success");
+        nym_metrics::inc!("wg_peer_update_psk_success");
         Ok(())
     }
 
diff --git a/gateway/src/node/wireguard/new_peer_registration/lp.rs b/gateway/src/node/wireguard/new_peer_registration/lp.rs
index 9d8f059953..fa19aec01d 100644
--- a/gateway/src/node/wireguard/new_peer_registration/lp.rs
+++ b/gateway/src/node/wireguard/new_peer_registration/lp.rs
@@ -15,23 +15,14 @@ use std::time::Instant;
 
 impl PeerRegistrator {
     /// In the case of an already registered WG peer, update its PSK.
+    ///
+    /// The peer controller keeps the active config and the on-disk PSK in sync.
     pub(super) async fn update_peer_psk(
         &self,
         peer: PeerPublicKey,
         psk: Key,
     ) -> Result<(), GatewayWireguardError> {
-        let encoded_psk = psk.to_lower_hex();
-
-        // 1. update the PSK in the active configuration
-        self.peer_manager.update_peer_psk(peer, psk).await?;
-
-        // 2. update the on-disk PSK
-        self.ecash_verifier
-            .storage()
-            .update_peer_psk(&peer.to_string(), Some(&encoded_psk))
-            .await?;
-
-        Ok(())
+        self.peer_manager.update_peer_psk(peer, psk).await
     }
 
     fn lp_peer_to_final_response(