diff --git a/common/wireguard/src/error.rs b/common/wireguard/src/error.rs index c261ac123a..ebf4928f15 100644 --- a/common/wireguard/src/error.rs +++ b/common/wireguard/src/error.rs @@ -3,9 +3,6 @@ #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("peers in wireguard don't match with in-memory ")] - PeerMismatch, - #[error("traffic byte data needs to be increasing")] InconsistentConsumedBytes, diff --git a/common/wireguard/src/peer_controller.rs b/common/wireguard/src/peer_controller.rs index c6316a720e..fea175066e 100644 --- a/common/wireguard/src/peer_controller.rs +++ b/common/wireguard/src/peer_controller.rs @@ -160,13 +160,10 @@ impl PeerController { .ok_or(Error::MissingClientBandwidthEntry)? .client_id { - let bandwidth = storage - .get_available_bandwidth(client_id) - .await? - .ok_or(Error::MissingClientBandwidthEntry)?; + storage.create_bandwidth_entry(client_id).await?; Ok(Some(BandwidthStorageManager::new( storage, - ClientBandwidth::new(bandwidth.into()), + ClientBandwidth::new(Default::default()), client_id, BandwidthFlushingBehaviourConfig::default(), true, @@ -228,14 +225,10 @@ impl PeerController { .available_bandwidth() .await } else { - let peer = self - .host_information - .read() - .await - .peers - .get(key) - .ok_or(Error::PeerMismatch)? - .clone(); + let Some(peer) = self.host_information.read().await.peers.get(key).cloned() else { + // host information not updated yet + return Ok(None); + }; BANDWIDTH_CAP_PER_DAY.saturating_sub((peer.rx_bytes + peer.tx_bytes) as i64) }; diff --git a/common/wireguard/src/peer_handle.rs b/common/wireguard/src/peer_handle.rs index abf367b3c9..85f2a612b2 100644 --- a/common/wireguard/src/peer_handle.rs +++ b/common/wireguard/src/peer_handle.rs @@ -3,6 +3,7 @@ use crate::error::Error; use crate::peer_controller::PeerControlRequest; +use defguard_wireguard_rs::host::Peer; use defguard_wireguard_rs::{host::Host, key::Key}; use futures::channel::oneshot; use nym_authenticator_requests::v2::registration::BANDWIDTH_CAP_PER_DAY; @@ -71,15 +72,11 @@ impl PeerHandle { Ok(success) } - async fn active_peer(&mut self, storage_peer: WireguardPeer) -> Result { - let kernel_peer = self - .host_information - .read() - .await - .peers - .get(&self.public_key) - .ok_or(Error::PeerMismatch)? - .clone(); + async fn active_peer( + &mut self, + storage_peer: WireguardPeer, + kernel_peer: Peer, + ) -> Result { if let Some(bandwidth_manager) = &self.bandwidth_storage_manager { let spent_bandwidth = (kernel_peer.rx_bytes + kernel_peer.tx_bytes) .checked_sub(storage_peer.rx_bytes as u64 + storage_peer.tx_bytes as u64) @@ -111,11 +108,21 @@ impl PeerHandle { while !self.task_client.is_shutdown() { tokio::select! { _ = self.timeout_check_interval.next() => { - let Some(peer) = self.storage.get_wireguard_peer(&self.public_key.to_string()).await? else { + let Some(kernel_peer) = self + .host_information + .read() + .await + .peers + .get(&self.public_key) + .cloned() else { + // the host information hasn't beed updated yet + continue; + }; + let Some(storage_peer) = self.storage.get_wireguard_peer(&self.public_key.to_string()).await? else { log::debug!("Peer {:?} not in storage anymore, shutting down handle", self.public_key); return Ok(()); }; - if !self.active_peer(peer).await? { + if !self.active_peer(storage_peer, kernel_peer).await? { log::debug!("Peer {:?} doesn't have bandwidth anymore, shutting down handle", self.public_key); return Ok(()); } diff --git a/service-providers/authenticator/src/mixnet_listener.rs b/service-providers/authenticator/src/mixnet_listener.rs index a68b86ff7e..71864047e2 100644 --- a/service-providers/authenticator/src/mixnet_listener.rs +++ b/service-providers/authenticator/src/mixnet_listener.rs @@ -288,10 +288,6 @@ impl MixnetListener { credential: CredentialSpendingData, client_id: i64, ) -> Result { - ecash_verifier - .storage() - .create_bandwidth_entry(client_id) - .await?; let bandwidth = ecash_verifier .storage() .get_available_bandwidth(client_id)