From 466bb97bc721c13d1401268f02129c8b0cc2dc17 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= Date: Fri, 6 Jun 2025 10:12:50 +0300 Subject: [PATCH] Use the same client bandwidth for top up (#5813) * Use the same client bandwidth for top up * Fix clippy --- .../src/bandwidth_storage_manager.rs | 4 +++ .../src/client_bandwidth.rs | 2 +- common/wireguard/src/peer_controller.rs | 20 +++++++++++ .../authenticator/src/cli/peer_handler.rs | 8 +++-- .../authenticator/src/mixnet_listener.rs | 16 ++++----- .../authenticator/src/peer_manager.rs | 33 ++++++++++++++++--- 6 files changed, 66 insertions(+), 17 deletions(-) diff --git a/common/credential-verification/src/bandwidth_storage_manager.rs b/common/credential-verification/src/bandwidth_storage_manager.rs index 19df1dba6f..c861ef123f 100644 --- a/common/credential-verification/src/bandwidth_storage_manager.rs +++ b/common/credential-verification/src/bandwidth_storage_manager.rs @@ -40,6 +40,10 @@ impl BandwidthStorageManager { } } + pub fn client_bandwidth(&self) -> ClientBandwidth { + self.client_bandwidth.clone() + } + pub async fn available_bandwidth(&self) -> i64 { self.client_bandwidth.available().await } diff --git a/common/credential-verification/src/client_bandwidth.rs b/common/credential-verification/src/client_bandwidth.rs index d98f89b511..0a0acba8c1 100644 --- a/common/credential-verification/src/client_bandwidth.rs +++ b/common/credential-verification/src/client_bandwidth.rs @@ -73,7 +73,7 @@ impl ClientBandwidth { false } - pub(crate) async fn available(&self) -> i64 { + pub async fn available(&self) -> i64 { self.inner.read().await.bandwidth.bytes } diff --git a/common/wireguard/src/peer_controller.rs b/common/wireguard/src/peer_controller.rs index ae59500429..0718d233ab 100644 --- a/common/wireguard/src/peer_controller.rs +++ b/common/wireguard/src/peer_controller.rs @@ -45,6 +45,10 @@ pub enum PeerControlRequest { key: Key, response_tx: oneshot::Sender, }, + GetClientBandwidth { + key: Key, + response_tx: oneshot::Sender, + }, } pub struct AddPeerControlResponse { @@ -65,6 +69,10 @@ pub struct QueryBandwidthControlResponse { pub bandwidth_data: Option, } +pub struct GetClientBandwidthControlResponse { + pub client_bandwidth: Option, +} + pub struct PeerController { storage: GatewayStorage, @@ -267,6 +275,14 @@ impl PeerController { })) } + async fn handle_get_client_bandwidth(&self, key: &Key) -> Option { + if let Some(Some(bandwidth_storage_manager)) = self.bw_storage_managers.get(key) { + Some(bandwidth_storage_manager.read().await.client_bandwidth()) + } else { + None + } + } + async fn update_metrics(&self, new_host: &Host) { let now = SystemTime::now(); const ACTIVITY_THRESHOLD: Duration = Duration::from_secs(180); @@ -388,6 +404,10 @@ impl PeerController { response_tx.send(QueryBandwidthControlResponse { success: false, bandwidth_data: None }).ok(); } } + Some(PeerControlRequest::GetClientBandwidth { key, response_tx }) => { + let client_bandwidth = self.handle_get_client_bandwidth(&key).await; + response_tx.send(GetClientBandwidthControlResponse { client_bandwidth }).ok(); + } None => { log::trace!("PeerController [main loop]: stopping since channel closed"); break; diff --git a/service-providers/authenticator/src/cli/peer_handler.rs b/service-providers/authenticator/src/cli/peer_handler.rs index 99a260ff7e..ebf34dd3b0 100644 --- a/service-providers/authenticator/src/cli/peer_handler.rs +++ b/service-providers/authenticator/src/cli/peer_handler.rs @@ -3,8 +3,8 @@ use nym_sdk::TaskClient; use nym_wireguard::peer_controller::{ - AddPeerControlResponse, PeerControlRequest, QueryBandwidthControlResponse, - QueryPeerControlResponse, RemovePeerControlResponse, + AddPeerControlResponse, GetClientBandwidthControlResponse, PeerControlRequest, + QueryBandwidthControlResponse, QueryPeerControlResponse, RemovePeerControlResponse, }; use tokio::sync::mpsc; @@ -43,6 +43,10 @@ impl DummyHandler { log::info!("[DUMMY] Querying bandwidth for peer {:?}", key); response_tx.send(QueryBandwidthControlResponse { success: true, bandwidth_data: None }).ok(); } + PeerControlRequest::GetClientBandwidth{key, response_tx} => { + log::info!("[DUMMY] Getting client bandwidth for peer {:?}", key); + response_tx.send(GetClientBandwidthControlResponse {client_bandwidth: None }).ok(); + } } } else { break; diff --git a/service-providers/authenticator/src/mixnet_listener.rs b/service-providers/authenticator/src/mixnet_listener.rs index 4dfdcf845b..382c0f2900 100644 --- a/service-providers/authenticator/src/mixnet_listener.rs +++ b/service-providers/authenticator/src/mixnet_listener.rs @@ -716,6 +716,7 @@ impl MixnetListener { let Some(ecash_verifier) = self.ecash_verifier.clone() else { return Err(AuthenticatorError::UnsupportedOperation); }; + let client_id = ecash_verifier .storage() .get_wireguard_peer(&msg.pub_key().to_string()) @@ -723,19 +724,16 @@ impl MixnetListener { .ok_or(AuthenticatorError::MissingClientBandwidthEntry)? .client_id .ok_or(AuthenticatorError::OldClient)?; - let bandwidth = ecash_verifier - .storage() - .get_available_bandwidth(client_id) + let client_bandwidth = self + .peer_manager + .query_client_bandwidth(msg.pub_key()) .await? - .ok_or(AuthenticatorError::InternalError( - "bandwidth entry should have just been created".to_string(), - ))?; + .ok_or(AuthenticatorError::MissingClientBandwidthEntry)?; - let available_bandwidth = if self.received_retry(&*msg) { + let available_bandwidth = if self.received_retry(msg.as_ref()) { // don't process the credential and just return the current bandwidth - bandwidth.available + client_bandwidth.available().await } else { - let client_bandwidth = ClientBandwidth::new(bandwidth.into()); let credential = msg.credential(); let mut verifier = CredentialVerifier::new( CredentialSpendingRequest::new(credential.clone()), diff --git a/service-providers/authenticator/src/peer_manager.rs b/service-providers/authenticator/src/peer_manager.rs index 00f74deb67..2cb5ac137b 100644 --- a/service-providers/authenticator/src/peer_manager.rs +++ b/service-providers/authenticator/src/peer_manager.rs @@ -8,10 +8,11 @@ use nym_authenticator_requests::{ latest::registration::{GatewayClient, RemainingBandwidthData}, traits::QueryBandwidthMessage, }; +use nym_credential_verification::ClientBandwidth; use nym_wireguard::{ peer_controller::{ - AddPeerControlResponse, PeerControlRequest, QueryBandwidthControlResponse, - QueryPeerControlResponse, RemovePeerControlResponse, + AddPeerControlResponse, GetClientBandwidthControlResponse, PeerControlRequest, + QueryBandwidthControlResponse, QueryPeerControlResponse, RemovePeerControlResponse, }, WireguardGatewayData, }; @@ -109,9 +110,9 @@ impl PeerManager { let QueryBandwidthControlResponse { success, bandwidth_data, - } = response_rx - .await - .map_err(|_| AuthenticatorError::InternalError("no response for query".to_string()))?; + } = response_rx.await.map_err(|_| { + AuthenticatorError::InternalError("no response for query bandwidth".to_string()) + })?; if !success { return Err(AuthenticatorError::InternalError( "querying bandwidth could not be performed".to_string(), @@ -119,4 +120,26 @@ impl PeerManager { } Ok(bandwidth_data) } + + pub async fn query_client_bandwidth( + &mut self, + key: PeerPublicKey, + ) -> Result> { + let key = Key::new(key.to_bytes()); + let (response_tx, response_rx) = oneshot::channel(); + let msg = PeerControlRequest::GetClientBandwidth { key, response_tx }; + self.wireguard_gateway_data + .peer_tx() + .send(msg) + .await + .map_err(|_| AuthenticatorError::PeerInteractionStopped)?; + + let GetClientBandwidthControlResponse { client_bandwidth } = + response_rx.await.map_err(|_| { + AuthenticatorError::InternalError( + "no response for query client bandwidth".to_string(), + ) + })?; + Ok(client_bandwidth) + } }