diff --git a/common/client-libs/gateway-client/src/bandwidth.rs b/common/client-libs/gateway-client/src/bandwidth.rs index f232c004ec..1f45fdc8ea 100644 --- a/common/client-libs/gateway-client/src/bandwidth.rs +++ b/common/client-libs/gateway-client/src/bandwidth.rs @@ -2,21 +2,37 @@ // SPDX-License-Identifier: Apache-2.0 use si_scale::helpers::bibytes2; -use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; use std::sync::Arc; use std::time::Duration; use time::OffsetDateTime; -#[derive(Clone, Default)] +pub(crate) struct BandwidthClaimGuard { + inner: Arc, +} + +impl Drop for BandwidthClaimGuard { + fn drop(&mut self) { + let old = self.inner.claiming_more.swap(false, Ordering::SeqCst); + assert!( + old, + "critical failure: there were multiple BandwidthClaimGuard existing" + ) + } +} + +#[derive(Clone)] pub struct ClientBandwidth { inner: Arc, } -#[derive(Default)] struct ClientBandwidthInner { /// the actual bandwidth amount (in bytes) available available: AtomicI64, + /// flag to indicate whether this client is currently in the process of claiming additional bandwidth + claiming_more: AtomicBool, + /// defines the timestamp when the bandwidth information has been logged to the logs stream last_logged_ts: AtomicI64, @@ -29,11 +45,28 @@ impl ClientBandwidth { ClientBandwidth { inner: Arc::new(ClientBandwidthInner { available: AtomicI64::new(0), + claiming_more: AtomicBool::new(false), last_logged_ts: AtomicI64::new(0), last_updated_ts: AtomicI64::new(0), }), } } + + pub(crate) fn begin_bandwidth_claim(&self) -> Option { + if self + .inner + .claiming_more + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + { + Some(BandwidthClaimGuard { + inner: self.inner.clone(), + }) + } else { + None + } + } + pub(crate) fn remaining(&self) -> i64 { self.inner.available.load(Ordering::Acquire) } diff --git a/common/client-libs/gateway-client/src/client/mod.rs b/common/client-libs/gateway-client/src/client/mod.rs index d899ec2928..ad30cc2c8e 100644 --- a/common/client-libs/gateway-client/src/client/mod.rs +++ b/common/client-libs/gateway-client/src/client/mod.rs @@ -598,6 +598,11 @@ impl GatewayClient { return Err(GatewayClientError::NoBandwidthControllerAvailable); } + let Some(_claim_guard) = self.bandwidth.begin_bandwidth_claim() else { + debug!("there's already an existing bandwidth claim ongoing"); + return Ok(()); + }; + warn!("Not enough bandwidth. Trying to get more bandwidth, this might take a while"); if !self.cfg.bandwidth.require_tickets { info!("The client is running in disabled credentials mode - attempting to claim bandwidth without a credential"); diff --git a/common/credential-verification/src/bandwidth_storage_manager.rs b/common/credential-verification/src/bandwidth_storage_manager.rs index 3a67d43743..b2f5cefb83 100644 --- a/common/credential-verification/src/bandwidth_storage_manager.rs +++ b/common/credential-verification/src/bandwidth_storage_manager.rs @@ -1,6 +1,9 @@ // Copyright 2024 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use crate::error::*; +use crate::BandwidthFlushingBehaviourConfig; +use crate::ClientBandwidth; use nym_credentials::ecash::utils::ecash_today; use nym_credentials_interface::Bandwidth; use nym_gateway_requests::ServerResponse; @@ -9,10 +12,6 @@ use si_scale::helpers::bibytes2; use time::OffsetDateTime; use tracing::*; -use crate::error::*; -use crate::BandwidthFlushingBehaviourConfig; -use crate::ClientBandwidth; - const FREE_TESTNET_BANDWIDTH_VALUE: Bandwidth = Bandwidth::new_unchecked(64 * 1024 * 1024 * 1024); // 64GB #[derive(Clone)] @@ -43,7 +42,7 @@ impl BandwidthStorageManager { async fn sync_expiration(&mut self) -> Result<()> { self.storage - .set_expiration(self.client_id, self.client_bandwidth.bandwidth.expiration) + .set_expiration(self.client_id, self.client_bandwidth.expiration().await) .await?; Ok(()) } @@ -57,17 +56,17 @@ impl BandwidthStorageManager { self.increase_bandwidth(FREE_TESTNET_BANDWIDTH_VALUE, ecash_today()) .await?; - let available_total = self.client_bandwidth.bandwidth.bytes; + let available_total = self.client_bandwidth.available().await; Ok(ServerResponse::Bandwidth { available_total }) } #[instrument(skip_all)] pub async fn try_use_bandwidth(&mut self, required_bandwidth: i64) -> Result { - if self.client_bandwidth.bandwidth.expired() { + if self.client_bandwidth.expired().await { self.expire_bandwidth().await?; } - let available_bandwidth = self.client_bandwidth.bandwidth.bytes; + let available_bandwidth = self.client_bandwidth.available().await; if available_bandwidth < required_bandwidth { return Err(Error::OutOfBandwidth { @@ -86,8 +85,7 @@ impl BandwidthStorageManager { async fn expire_bandwidth(&mut self) -> Result<()> { self.storage.reset_bandwidth(self.client_id).await?; - self.client_bandwidth.bandwidth = Default::default(); - self.client_bandwidth.update_sync_data(); + self.client_bandwidth.expire_bandwidth().await; Ok(()) } @@ -97,31 +95,31 @@ impl BandwidthStorageManager { /// /// * `amount`: amount to decrease the available bandwidth by. async fn consume_bandwidth(&mut self, amount: i64) -> Result<()> { - self.client_bandwidth.bandwidth.bytes -= amount; - self.client_bandwidth.bytes_delta_since_sync -= amount; + self.client_bandwidth.decrease_bandwidth(amount).await; // since we're going to be operating on a fair use policy anyway, even if we crash and let extra few packets // through, that's completely fine - if self.client_bandwidth.should_sync(self.bandwidth_cfg) { - self.sync_bandwidth().await?; + if self.client_bandwidth.should_sync(self.bandwidth_cfg).await { + self.sync_storage_bandwidth().await?; } Ok(()) } #[instrument(level = "trace", skip_all)] - async fn sync_bandwidth(&mut self) -> Result<()> { + async fn sync_storage_bandwidth(&mut self) -> Result<()> { trace!("syncing client bandwidth with the underlying storage"); let updated = self .storage - .increase_bandwidth(self.client_id, self.client_bandwidth.bytes_delta_since_sync) + .increase_bandwidth( + self.client_id, + self.client_bandwidth.delta_since_sync().await, + ) .await?; - trace!(updated); - - self.client_bandwidth.bandwidth.bytes = updated; - - self.client_bandwidth.update_sync_data(); + self.client_bandwidth + .resync_bandwidth_with_storage(updated) + .await; Ok(()) } @@ -136,13 +134,14 @@ impl BandwidthStorageManager { bandwidth: Bandwidth, expiration: OffsetDateTime, ) -> Result<()> { - self.client_bandwidth.bandwidth.bytes += bandwidth.value() as i64; - self.client_bandwidth.bytes_delta_since_sync += bandwidth.value() as i64; - self.client_bandwidth.bandwidth.expiration = expiration; + self.client_bandwidth + .increase_bandwidth(bandwidth.value() as i64, expiration) + .await; // any increases to bandwidth should get flushed immediately // (we don't want to accidentally miss somebody claiming a gigabyte voucher) self.sync_expiration().await?; - self.sync_bandwidth().await + self.sync_storage_bandwidth().await?; + Ok(()) } } diff --git a/common/credential-verification/src/client_bandwidth.rs b/common/credential-verification/src/client_bandwidth.rs index 610a198e88..8b165d9353 100644 --- a/common/credential-verification/src/client_bandwidth.rs +++ b/common/credential-verification/src/client_bandwidth.rs @@ -1,10 +1,12 @@ // Copyright 2024 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use std::time::Duration; - +use nym_credentials::ecash::utils::ecash_today; use nym_credentials_interface::AvailableBandwidth; +use std::sync::Arc; +use std::time::Duration; use time::OffsetDateTime; +use tokio::sync::RwLock; #[derive(Debug, Clone, Copy)] pub struct BandwidthFlushingBehaviourConfig { @@ -15,10 +17,15 @@ pub struct BandwidthFlushingBehaviourConfig { pub client_bandwidth_max_delta_flushing_amount: i64, } -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone)] pub struct ClientBandwidth { + inner: Arc>, +} + +#[derive(Debug)] +struct ClientBandwidthInner { pub(crate) bandwidth: AvailableBandwidth, - pub(crate) last_flushed: OffsetDateTime, + pub(crate) last_synced: OffsetDateTime, /// the number of bytes the client had during the last sync. /// it is used to determine whether the current value should be synced with the storage @@ -30,28 +37,74 @@ pub struct ClientBandwidth { impl ClientBandwidth { pub fn new(bandwidth: AvailableBandwidth) -> ClientBandwidth { ClientBandwidth { - bandwidth, - last_flushed: OffsetDateTime::now_utc(), - bytes_at_last_sync: bandwidth.bytes, - bytes_delta_since_sync: 0, + inner: Arc::new(RwLock::new(ClientBandwidthInner { + bandwidth, + last_synced: OffsetDateTime::now_utc(), + bytes_at_last_sync: bandwidth.bytes, + bytes_delta_since_sync: 0, + })), } } - pub(crate) fn should_sync(&self, cfg: BandwidthFlushingBehaviourConfig) -> bool { - if self.bytes_delta_since_sync.abs() >= cfg.client_bandwidth_max_delta_flushing_amount { + pub(crate) async fn should_sync(&self, cfg: BandwidthFlushingBehaviourConfig) -> bool { + let guard = self.inner.read().await; + + if guard.bytes_delta_since_sync.abs() >= cfg.client_bandwidth_max_delta_flushing_amount { return true; } - if self.last_flushed + cfg.client_bandwidth_max_flushing_rate < OffsetDateTime::now_utc() { + if guard.last_synced + cfg.client_bandwidth_max_flushing_rate < OffsetDateTime::now_utc() { return true; } false } - pub(crate) fn update_sync_data(&mut self) { - self.last_flushed = OffsetDateTime::now_utc(); - self.bytes_at_last_sync = self.bandwidth.bytes; - self.bytes_delta_since_sync = 0; + pub(crate) async fn available(&self) -> i64 { + self.inner.read().await.bandwidth.bytes + } + + pub(crate) async fn delta_since_sync(&self) -> i64 { + self.inner.read().await.bytes_delta_since_sync + } + pub(crate) async fn expiration(&self) -> OffsetDateTime { + self.inner.read().await.bandwidth.expiration + } + + pub(crate) async fn expired(&self) -> bool { + self.expiration().await < ecash_today() + } + + pub(crate) async fn decrease_bandwidth(&self, decrease: i64) { + let mut guard = self.inner.write().await; + + guard.bandwidth.bytes -= decrease; + guard.bytes_delta_since_sync -= decrease; + } + + pub(crate) async fn increase_bandwidth(&self, increase: i64, new_expiration: OffsetDateTime) { + let mut guard = self.inner.write().await; + + guard.bandwidth.bytes += increase; + guard.bandwidth.expiration = new_expiration; + guard.bytes_delta_since_sync += increase; + } + + pub(crate) async fn expire_bandwidth(&self) { + let mut guard = self.inner.write().await; + + guard.bandwidth = AvailableBandwidth::default(); + guard.last_synced = OffsetDateTime::now_utc(); + guard.bytes_at_last_sync = 0; + guard.bytes_delta_since_sync = 0; + } + + pub(crate) async fn resync_bandwidth_with_storage(&self, stored: i64) { + let mut guard = self.inner.write().await; + + guard.bandwidth.bytes = stored; + guard.bytes_at_last_sync = stored; + guard.bytes_delta_since_sync = 0; + guard.last_synced = OffsetDateTime::now_utc(); } } diff --git a/common/credential-verification/src/lib.rs b/common/credential-verification/src/lib.rs index bb3a0cab8f..4d2e7be2da 100644 --- a/common/credential-verification/src/lib.rs +++ b/common/credential-verification/src/lib.rs @@ -150,7 +150,7 @@ impl CredentialVerifier { Ok(self .bandwidth_storage_manager .client_bandwidth - .bandwidth - .bytes) + .available() + .await) } } diff --git a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs index 12bbb16b51..6bfd95478f 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs @@ -26,7 +26,7 @@ use nym_gateway_storage::{error::StorageError, Storage}; use nym_sphinx::forwarding::packet::MixPacket; use nym_task::TaskClient; use nym_validator_client::coconut::EcashApiError; -use rand::{CryptoRng, Rng}; +use rand::{random, CryptoRng, Rng}; use std::{process, time::Duration}; use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite}; @@ -217,11 +217,7 @@ where enc_credential: Vec, iv: Vec, ) -> Result { - // TODO: change it into a span field instead once we move to tracing - debug!( - "handling e-cash bandwidth request from {}", - self.client.address - ); + debug!("handling e-cash bandwidth request"); let credential = ClientControlRequest::try_from_enc_ecash_credential( enc_credential, @@ -234,7 +230,11 @@ where self.bandwidth_storage_manager.clone(), ); - let available_total = verifier.verify().await?; + let available_total = verifier + .verify() + .await + .inspect_err(|verification_failure| debug!("{verification_failure}"))?; + trace!("available total bandwidth: {available_total}"); Ok(ServerResponse::Bandwidth { available_total }) } @@ -297,40 +297,43 @@ where /// * `raw_request`: raw message to handle. async fn handle_text(&mut self, raw_request: String) -> Message { trace!("text request"); - match ClientControlRequest::try_from(raw_request) { - Err(e) => RequestHandlingError::InvalidTextRequest(e).into_error_message(), - Ok(request) => match request { - ClientControlRequest::EcashCredential { enc_credential, iv } => self - .handle_ecash_bandwidth(enc_credential, iv) - .await - .into_ws_message(), - ClientControlRequest::BandwidthCredential { .. } => { - RequestHandlingError::IllegalRequest { - additional_context: "coconut credential are not longer supported".into(), - } - .into_error_message() - } - ClientControlRequest::BandwidthCredentialV2 { .. } => { - RequestHandlingError::IllegalRequest { - additional_context: "coconut credential are not longer supported".into(), - } - .into_error_message() - } - ClientControlRequest::ClaimFreeTestnetBandwidth => self - .bandwidth_storage_manager - .handle_claim_testnet_bandwidth() - .await - .map_err(|e| e.into()) - .into_ws_message(), - other => RequestHandlingError::IllegalRequest { - additional_context: format!( - "received illegal message of type {} in an authenticated client", - other.name() - ), - } - .into_error_message(), - }, + + let request = match ClientControlRequest::try_from(raw_request) { + Ok(req) => { + debug!("received request of type {}", req.name()); + req + } + Err(err) => { + debug!("request was malformed: {err}"); + return RequestHandlingError::InvalidTextRequest(err).into_error_message(); + } + }; + + match request { + ClientControlRequest::EcashCredential { enc_credential, iv } => { + self.handle_ecash_bandwidth(enc_credential, iv).await + } + ClientControlRequest::BandwidthCredential { .. } + | ClientControlRequest::BandwidthCredentialV2 { .. } => { + Err(RequestHandlingError::IllegalRequest { + additional_context: "coconut credential are not longer supported".into(), + }) + } + ClientControlRequest::ClaimFreeTestnetBandwidth => self + .bandwidth_storage_manager + .handle_claim_testnet_bandwidth() + .await + .map_err(|e| e.into()), + other => Err(RequestHandlingError::IllegalRequest { + additional_context: format!( + "received illegal message of type {} in an authenticated client", + other.name() + ), + }), } + .inspect(|res| debug!(response = ?res, "success")) + .inspect_err(|err| debug!(error = %err, "failure")) + .into_ws_message() } /// Handles pong message received from the client. @@ -364,12 +367,13 @@ where /// # Arguments /// /// * `raw_request`: raw received websocket message. + #[instrument(level = "debug", skip_all, + fields( + client = %self.client.address.as_base58_string() + ) + )] async fn handle_request(&mut self, raw_request: Message) -> Option { - // TODO: this should be added via tracing - debug!( - "handling request from {}", - self.client.address.as_base58_string() - ); + trace!("new request"); // apparently tungstenite auto-handles ping/pong/close messages so for now let's ignore // them and let's test that claim. If that's not the case, just copy code from @@ -390,8 +394,8 @@ where where S: AsyncRead + AsyncWrite + Unpin, { - let tag: u64 = rand::thread_rng().gen(); - debug!("Got request to ping our connection: {}", tag); + let tag: u64 = random(); + debug!("got request to ping our connection: {tag}"); self.inner .send_websocket_message(Message::Ping(tag.to_be_bytes().to_vec())) .await?; diff --git a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs index 289d7922bb..11f0ff2e34 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -419,7 +419,7 @@ where // we can't handle clients with higher protocol than ours // (perhaps we could try to negotiate downgrade on our end? sounds like a nice future improvement) if client_protocol_version <= CURRENT_PROTOCOL_VERSION { - info!("the client is using exactly the same (or older) protocol version as we are. We're good to continue!"); + debug!("the client is using exactly the same (or older) protocol version as we are. We're good to continue!"); Ok(CURRENT_PROTOCOL_VERSION) } else { let err = InitialAuthenticationError::IncompatibleProtocol { diff --git a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs index 6e1d8a9b5a..bd33b25a9a 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs @@ -102,7 +102,7 @@ where .await { Err(timeout_err) => { - warn!("websocket handshake timedout: {timeout_err}"); + warn!("websocket handshake timed out: {timeout_err}"); return; } Ok(Err(err)) => {