From a9bf1954bc0dc4a9ce1db30258367f71cafbb62b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= Date: Mon, 8 Jun 2026 11:19:42 +0300 Subject: [PATCH] Keep peer in wg table when updating psk (#6856) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Keep peer in wg table when updating psk * Fix unit test * update handle_update_peer_psk_request --------- Co-authored-by: Jędrzej Stuczyński --- .../tests/src/v2/peer_controller.rs | 11 ++++ common/wireguard/src/peer_controller/mock.rs | 10 ++++ common/wireguard/src/peer_controller/mod.rs | 60 +++++++++++++++++++ .../wireguard/new_peer_registration/lp.rs | 17 +----- gateway/src/node/wireguard/peer_manager.rs | 38 ++++++++++++ .../src/prometheus_wrapper.rs | 8 ++- 6 files changed, 129 insertions(+), 15 deletions(-) diff --git a/common/wireguard-private-metadata/tests/src/v2/peer_controller.rs b/common/wireguard-private-metadata/tests/src/v2/peer_controller.rs index 2b29de63c6..10f0d20ace 100644 --- a/common/wireguard-private-metadata/tests/src/v2/peer_controller.rs +++ b/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( diff --git a/common/wireguard/src/peer_controller/mock.rs b/common/wireguard/src/peer_controller/mock.rs index 37f9ba1f0a..bb6451ee43 100644 --- a/common/wireguard/src/peer_controller/mock.rs +++ b/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 { 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) } diff --git a/common/wireguard/src/peer_controller/mod.rs b/common/wireguard/src/peer_controller/mod.rs index 19f936b93b..a4fbbe2c02 100644 --- a/common/wireguard/src/peer_controller/mod.rs +++ b/common/wireguard/src/peer_controller/mod.rs @@ -76,6 +76,12 @@ pub enum PeerControlRequest { peer: Peer, response_tx: oneshot::Sender, }, + /// Update PSK for an existing peer, without changing its IP allocation + UpdatePeerPsk { + peer_key: Key, + psk: Key, + response_tx: oneshot::Sender, + }, /// Attempt to allocate an IP pair from the pool PreAllocateIpPair { response_tx: oneshot::Sender, @@ -118,6 +124,7 @@ pub enum PeerControlRequest { } pub type AddPeerControlResponse = Result<()>; +pub type UpdatePeerPskControlResponse = Result<()>; pub type AllocatePeerControlResponse = Result; 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(); } diff --git a/gateway/src/node/wireguard/new_peer_registration/lp.rs b/gateway/src/node/wireguard/new_peer_registration/lp.rs index 352e517ccf..bd547bce0e 100644 --- a/gateway/src/node/wireguard/new_peer_registration/lp.rs +++ b/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 - 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(()) + self.peer_manager.update_peer_psk(peer, psk).await } fn lp_peer_to_final_response( diff --git a/gateway/src/node/wireguard/peer_manager.rs b/gateway/src/node/wireguard/peer_manager.rs index ffb25840e9..55a7445d45 100644 --- a/gateway/src/node/wireguard/peer_manager.rs +++ b/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()); diff --git a/nym-node/nym-node-metrics/src/prometheus_wrapper.rs b/nym-node/nym-node-metrics/src/prometheus_wrapper.rs index 72c4b2f4d0..b393c236af 100644 --- a/nym-node/nym-node-metrics/src/prometheus_wrapper.rs +++ b/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]