From da8c95362d46429d314e4b2bf65d12d020bb223e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= Date: Fri, 5 Jun 2026 12:15:55 +0300 Subject: [PATCH] Keep peer in wg table when updating psk --- .../tests/src/v2/peer_controller.rs | 11 +++++ common/wireguard/src/peer_controller/mock.rs | 10 +++++ common/wireguard/src/peer_controller/mod.rs | 44 +++++++++++++++++++ .../wireguard/new_peer_registration/lp.rs | 12 +++-- gateway/src/node/wireguard/peer_manager.rs | 38 ++++++++++++++++ .../src/prometheus_wrapper.rs | 6 +++ 6 files changed, 114 insertions(+), 7 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..4aa1daf5ab 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,34 @@ 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(()); + }; + peer.preshared_key = Some(psk); + + // 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()); + }; + + // 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; + } + + nym_metrics::inc!("wg_peer_addition_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 +548,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 9e7a45d057..9d8f059953 100644 --- a/gateway/src/node/wireguard/new_peer_registration/lp.rs +++ b/gateway/src/node/wireguard/new_peer_registration/lp.rs @@ -20,14 +20,12 @@ impl PeerRegistrator { 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(); + + // 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)) 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..a782266758 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)) }