Use the same client bandwidth for top up (#5813)

* Use the same client bandwidth for top up

* Fix clippy

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
     }

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
     }
 

common/wireguard/src/peer_controller.rs

@@ -45,6 +45,10 @@ pub enum PeerControlRequest {
         key: Key,
         response_tx: oneshot::Sender<QueryBandwidthControlResponse>,
     },
+    GetClientBandwidth {
+        key: Key,
+        response_tx: oneshot::Sender<GetClientBandwidthControlResponse>,
+    },
 }
 
 pub struct AddPeerControlResponse {
@@ -65,6 +69,10 @@ pub struct QueryBandwidthControlResponse {
     pub bandwidth_data: Option<RemainingBandwidthData>,
 }
 
+pub struct GetClientBandwidthControlResponse {
+    pub client_bandwidth: Option<ClientBandwidth>,
+}
+
 pub struct PeerController {
     storage: GatewayStorage,
 
@@ -267,6 +275,14 @@ impl PeerController {
         }))
     }
 
+    async fn handle_get_client_bandwidth(&self, key: &Key) -> Option<ClientBandwidth> {
+        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;

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;

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()),

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<Option<ClientBandwidth>> {
+        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)
+    }
 }