Keep peer in wg table when updating psk (#6856)

* Keep peer in wg table when updating psk

* Fix unit test

* update handle_update_peer_psk_request

---------

Co-authored-by: Jędrzej Stuczyński <jedrzej.stuczynski@gmail.com>

common/wireguard-private-metadata/tests/src/v2/peer_controller.rs

@@ -14,6 +14,7 @@ use tokio::sync::mpsc::Receiver;
 #[derive(Hash, PartialOrd, PartialEq, Clone, Debug, Eq, Copy)]
 pub enum PeerControlRequestTypeV2 {
     AddPeer,
+    UpdatePeerPsk,
     RemovePeer,
     QueryPeer,
     GetClientBandwidthByKey,
@@ -26,6 +27,7 @@ impl From<&PeerControlRequest> for PeerControlRequestTypeV2 {
     fn from(req: &PeerControlRequest) -> Self {
         match req {
             PeerControlRequest::AddPeer { .. } => PeerControlRequestTypeV2::AddPeer,
+            PeerControlRequest::UpdatePeerPsk { .. } => PeerControlRequestTypeV2::UpdatePeerPsk,
             PeerControlRequest::PreAllocateIpPair { .. } => PeerControlRequestTypeV2::AddPeer,
             PeerControlRequest::RemovePeer { .. } => PeerControlRequestTypeV2::RemovePeer,
             PeerControlRequest::QueryPeer { .. } => PeerControlRequestTypeV2::QueryPeer,
@@ -115,6 +117,15 @@ impl MockPeerControllerV2 {
                     )
                     .unwrap();
             }
+            PeerControlRequest::UpdatePeerPsk { response_tx, .. } => {
+                response_tx
+                    .send(
+                        *response
+                            .downcast()
+                            .expect("registered response has mismatched type"),
+                    )
+                    .unwrap();
+            }
             PeerControlRequest::PreAllocateIpPair { response_tx, .. } => {
                 response_tx
                     .send(

common/wireguard/src/peer_controller/mock.rs

@@ -71,6 +71,7 @@ impl From<&Key> for KeyWrapper {
 #[derive(Hash, PartialOrd, PartialEq, Clone, Debug, Eq)]
 pub enum PeerControlRequestType {
     AddPeer { public_key: KeyWrapper },
+    UpdatePeerPsk { peer_key: KeyWrapper },
     AllocatePeerIpPair {},
     ReleaseIpPair { ip_pair: IpPair },
     RemovePeer { key: KeyWrapper },
@@ -86,6 +87,7 @@ impl PeerControlRequestType {
     pub fn peer_key(&self) -> Option<KeyWrapper> {
         match self {
             PeerControlRequestType::AddPeer { public_key } => Some(public_key.clone()),
+            PeerControlRequestType::UpdatePeerPsk { peer_key } => Some(peer_key.clone()),
             PeerControlRequestType::AllocatePeerIpPair {} => None,
             PeerControlRequestType::ReleaseIpPair { .. } => None,
             PeerControlRequestType::RemovePeer { key } => Some(key.clone()),
@@ -109,6 +111,11 @@ impl From<&PeerControlRequest> for PeerControlRequestType {
             PeerControlRequest::AddPeer { peer, .. } => PeerControlRequestType::AddPeer {
                 public_key: (&peer.public_key).into(),
             },
+            PeerControlRequest::UpdatePeerPsk { peer_key, .. } => {
+                PeerControlRequestType::UpdatePeerPsk {
+                    peer_key: peer_key.into(),
+                }
+            }
             PeerControlRequest::PreAllocateIpPair { .. } => {
                 PeerControlRequestType::AllocatePeerIpPair {}
             }
@@ -271,6 +278,9 @@ impl MockPeerController {
                 }
                 response_tx.send_downcasted(response.content)
             }
+            PeerControlRequest::UpdatePeerPsk { response_tx, .. } => {
+                response_tx.send_downcasted(response.content)
+            }
             PeerControlRequest::PreAllocateIpPair { response_tx, .. } => {
                 response_tx.send_downcasted(response.content)
             }

common/wireguard/src/peer_controller/mod.rs

@@ -76,6 +76,12 @@ pub enum PeerControlRequest {
         peer: Peer,
         response_tx: oneshot::Sender<AddPeerControlResponse>,
     },
+    /// Update PSK for an existing peer, without changing its IP allocation
+    UpdatePeerPsk {
+        peer_key: Key,
+        psk: Key,
+        response_tx: oneshot::Sender<UpdatePeerPskControlResponse>,
+    },
     /// Attempt to allocate an IP pair from the pool
     PreAllocateIpPair {
         response_tx: oneshot::Sender<AllocatePeerControlResponse>,
@@ -118,6 +124,7 @@ pub enum PeerControlRequest {
 }
 
 pub type AddPeerControlResponse = Result<()>;
+pub type UpdatePeerPskControlResponse = Result<()>;
 pub type AllocatePeerControlResponse = Result<IpPair>;
 pub type ReleaseIpPairControlResponse = Result<()>;
 pub type RemovePeerControlResponse = Result<()>;
@@ -317,6 +324,50 @@ impl PeerController {
         Ok(())
     }
 
+    async fn handle_update_peer_psk_request(&mut self, peer_key: &Key, psk: Key) -> Result<()> {
+        // observation will get automatically added once dropped
+        let _metric_timer =
+            PROMETHEUS_METRICS.start_timer(PrometheusMetric::WireguardDefguardPeerPskUpdate);
+
+        nym_metrics::inc!("wg_peer_update_psk_attempts");
+
+        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");
+            nym_metrics::inc!("wg_config_errors_total");
+            return Err(e.into());
+        };
+
+        // 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_update_psk_success");
+        Ok(())
+    }
+
     /// Allocate IP pair from pool for a new peer registration
     ///
     /// This only allocates IPs - the caller must handle database storage and
@@ -513,6 +564,15 @@ impl PeerController {
             PeerControlRequest::AddPeer { peer, response_tx } => {
                 response_tx.send(self.handle_add_request(&peer).await).ok();
             }
+            PeerControlRequest::UpdatePeerPsk {
+                peer_key,
+                psk,
+                response_tx,
+            } => {
+                response_tx
+                    .send(self.handle_update_peer_psk_request(&peer_key, psk).await)
+                    .ok();
+            }
             PeerControlRequest::PreAllocateIpPair { response_tx } => {
                 response_tx.send(self.handle_ip_allocation_request()).ok();
             }

gateway/src/node/wireguard/new_peer_registration/lp.rs

@@ -15,25 +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> {
-        // 1. check if the peer is currently being handled
+        self.peer_manager.update_peer_psk(peer, psk).await
-        if self.peer_manager.check_active_peer(peer).await? {
-            // 2. if so, force disconnect it (as we're handling new request from the same peer)
-            self.peer_manager.remove_peer(peer).await?;
-        }
-
-        // 3. update the on-disk PSK
-        let encoded_psk = psk.to_lower_hex();
-        self.ecash_verifier
-            .storage()
-            .update_peer_psk(&peer.to_string(), Some(&encoded_psk))
-            .await?;
-
-        Ok(())
     }
 
     fn lp_peer_to_final_response(

gateway/src/node/wireguard/peer_manager.rs

@@ -125,6 +125,44 @@ impl PeerManager {
         res
     }
 
+    pub async fn update_peer_psk(
+        &self,
+        pub_key: PeerPublicKey,
+        psk: Key,
+    ) -> Result<(), GatewayWireguardError> {
+        let controller_start = Instant::now();
+        let peer_key = Key::new(pub_key.to_bytes());
+        let (response_tx, response_rx) = oneshot::channel();
+        let msg = PeerControlRequest::UpdatePeerPsk {
+            peer_key,
+            psk,
+            response_tx,
+        };
+        self.wireguard_gateway_data
+            .peer_tx()
+            .send(msg)
+            .await
+            .map_err(|_| GatewayWireguardError::PeerInteractionStopped)?;
+
+        let res = response_rx
+            .await
+            .map_err(|_| GatewayWireguardError::internal("no response for update peer psk"))?
+            .map_err(|err| {
+                GatewayWireguardError::InternalError(format!(
+                    "updating peer psk could not be performed: {err:?}"
+                ))
+            });
+
+        let latency = controller_start.elapsed().as_secs_f64();
+        add_histogram_obs!(
+            "wg_peer_controller_channel_latency_seconds",
+            latency,
+            WG_CONTROLLER_LATENCY_BUCKETS
+        );
+
+        res
+    }
+
     pub async fn remove_peer(&self, pub_key: PeerPublicKey) -> Result<(), GatewayWireguardError> {
         let controller_start = Instant::now();
         let key = Key::new(pub_key.to_bytes());

nym-node/nym-node-metrics/src/prometheus_wrapper.rs

@@ -157,6 +157,9 @@ pub enum PrometheusMetric {
     #[strum(props(help = "The distribution of defguard peer creation time"))]
     WireguardDefguardPeerCreation,
 
+    #[strum(props(help = "The distribution of defguard peer psk update time"))]
+    WireguardDefguardPeerPskUpdate,
+
     #[strum(props(
         help = "The distribution of time it takes to verify a credential during peer registration"
     ))]
@@ -320,6 +323,9 @@ impl PrometheusMetric {
             PrometheusMetric::WireguardDefguardPeerCreation => {
                 Metric::new_histogram(&name, help, Some(REG_LATENCY_BUCKETS))
             }
+            PrometheusMetric::WireguardDefguardPeerPskUpdate => {
+                Metric::new_histogram(&name, help, Some(REG_LATENCY_BUCKETS))
+            }
             PrometheusMetric::DvpnAuthenticatorClientRegistrationMsg1 => {
                 Metric::new_histogram(&name, help, Some(REG_LATENCY_BUCKETS))
             }
@@ -452,7 +458,7 @@ mod tests {
         // a sanity check for anyone adding new metrics. if this test fails,
         // make sure any methods on `PrometheusMetric` enum don't need updating
         // or require custom Display impl
-        assert_eq!(46, PrometheusMetric::COUNT)
+        assert_eq!(47, PrometheusMetric::COUNT)
     }
 
     #[test]