diff --git a/common/client-core/src/client/mix_traffic/transceiver.rs b/common/client-core/src/client/mix_traffic/transceiver.rs index b6197fb490..0862911d93 100644 --- a/common/client-core/src/client/mix_traffic/transceiver.rs +++ b/common/client-core/src/client/mix_traffic/transceiver.rs @@ -3,10 +3,12 @@ use async_trait::async_trait; use log::{debug, error}; +use nym_credential_storage::storage::Storage as CredentialStorage; use nym_crypto::asymmetric::identity; use nym_gateway_client::GatewayClient; pub use nym_gateway_client::{GatewayPacketRouter, PacketRouter}; use nym_sphinx::forwarding::packet::MixPacket; +use nym_validator_client::nyxd::contract_traits::DkgQueryClient; use std::fmt::Debug; use std::os::raw::c_int as RawFd; use thiserror::Error; @@ -111,8 +113,9 @@ impl RemoteGateway { impl GatewayTransceiver for RemoteGateway where - C: Send, - St: Send, + C: DkgQueryClient + Send + Sync, + St: CredentialStorage, + ::StorageError: Send + Sync + 'static, { fn gateway_identity(&self) -> identity::PublicKey { self.gateway_client.gateway_identity() @@ -126,8 +129,9 @@ where #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl GatewaySender for RemoteGateway where - C: Send, - St: Send, + C: DkgQueryClient + Send + Sync, + St: CredentialStorage, + ::StorageError: Send + Sync + 'static, { async fn send_mix_packet(&mut self, packet: MixPacket) -> Result<(), ErasedGatewayError> { self.gateway_client diff --git a/common/client-libs/gateway-client/src/client.rs b/common/client-libs/gateway-client/src/client.rs index 39172158c0..3dfc8a62c7 100644 --- a/common/client-libs/gateway-client/src/client.rs +++ b/common/client-libs/gateway-client/src/client.rs @@ -23,12 +23,13 @@ use nym_gateway_requests::{ BinaryRequest, ClientControlRequest, ServerResponse, CREDENTIAL_UPDATE_V2_PROTOCOL_VERSION, CURRENT_PROTOCOL_VERSION, }; -use nym_network_defaults::{REMAINING_BANDWIDTH_THRESHOLD, TOKENS_TO_BURN}; +use nym_network_defaults::REMAINING_BANDWIDTH_THRESHOLD; use nym_sphinx::forwarding::packet::MixPacket; use nym_task::TaskClient; use nym_validator_client::nyxd::contract_traits::DkgQueryClient; use rand::rngs::OsRng; - +use std::convert::TryFrom; +use std::sync::atomic::{AtomicI64, Ordering}; use std::sync::Arc; use std::time::Duration; use tungstenite::protocol::Message; @@ -83,7 +84,7 @@ impl GatewayConfig { pub struct GatewayClient { authenticated: bool, disabled_credentials_mode: bool, - bandwidth_remaining: i64, + bandwidth_remaining: Arc, gateway_address: String, gateway_identity: identity::PublicKey, local_identity: Arc, @@ -122,7 +123,7 @@ impl GatewayClient { GatewayClient { authenticated: false, disabled_credentials_mode: true, - bandwidth_remaining: 0, + bandwidth_remaining: Arc::new(AtomicI64::new(0)), gateway_address: config.gateway_listener, gateway_identity: config.gateway_identity, local_identity, @@ -182,7 +183,7 @@ impl GatewayClient { } pub fn remaining_bandwidth(&self) -> i64 { - self.bandwidth_remaining + self.bandwidth_remaining.load(Ordering::Acquire) } #[cfg(not(target_arch = "wasm32"))] @@ -528,7 +529,8 @@ impl GatewayClient { } => { self.check_gateway_protocol(protocol_version)?; self.authenticated = status; - self.bandwidth_remaining = bandwidth_remaining; + self.bandwidth_remaining + .store(bandwidth_remaining, Ordering::Release); self.negotiated_protocol = protocol_version; log::debug!("authenticated: {status}, bandwidth remaining: {bandwidth_remaining}"); Ok(()) @@ -564,35 +566,38 @@ impl GatewayClient { } } - async fn claim_coconut_bandwidth( + async fn claim_ecash_bandwidth( &mut self, credential: CredentialSpendingData, ) -> Result<(), GatewayClientError> { let mut rng = OsRng; let iv = IV::new_random(&mut rng); - let msg = ClientControlRequest::new_enc_coconut_bandwidth_credential_v2( + let msg = ClientControlRequest::new_enc_ecash_credential( credential, self.shared_key.as_ref().unwrap(), iv, ) .into(); - self.bandwidth_remaining = match self.send_websocket_message(msg).await? { + let bandwidth_remaining = match self.send_websocket_message(msg).await? { ServerResponse::Bandwidth { available_total } => Ok(available_total), ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)), _ => Err(GatewayClientError::UnexpectedResponse), }?; + self.bandwidth_remaining + .store(bandwidth_remaining, Ordering::Release); Ok(()) } async fn try_claim_testnet_bandwidth(&mut self) -> Result<(), GatewayClientError> { let msg = ClientControlRequest::ClaimFreeTestnetBandwidth.into(); - self.bandwidth_remaining = match self.send_websocket_message(msg).await? { + let bandwidth_remaining = match self.send_websocket_message(msg).await? { ServerResponse::Bandwidth { available_total } => Ok(available_total), ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)), _ => Err(GatewayClientError::UnexpectedResponse), }?; - + self.bandwidth_remaining + .store(bandwidth_remaining, Ordering::Release); Ok(()) } @@ -629,49 +634,45 @@ impl GatewayClient { negotiated_protocol: Some(gateway_protocol), }); } - - let gateway_id = self.gateway_identity().to_base58_string(); - let prepared_credential = self .bandwidth_controller .as_ref() .unwrap() - .prepare_bandwidth_credential(&gateway_id) + .prepare_ecash_credential(self.gateway_identity.to_bytes()) .await?; - self.claim_coconut_bandwidth(prepared_credential.data) - .await?; + self.claim_ecash_bandwidth(prepared_credential.data).await?; self.bandwidth_controller .as_ref() .unwrap() - .consume_credential(prepared_credential.credential_id, &gateway_id) + .update_ecash_wallet( + prepared_credential.updated_credential, + prepared_credential.credential_id, + ) .await?; Ok(()) } - fn estimate_required_bandwidth(&self, packets: &[MixPacket]) -> i64 { - packets - .iter() - .map(|packet| packet.packet().len()) - .sum::() as i64 - } - pub async fn batch_send_mix_packets( &mut self, packets: Vec, - ) -> Result<(), GatewayClientError> { + ) -> Result<(), GatewayClientError> + where + C: DkgQueryClient + Send + Sync, + St: CredentialStorage, + ::StorageError: Send + Sync + 'static, + { debug!("Sending {} mix packets", packets.len()); if !self.authenticated { return Err(GatewayClientError::NotAuthenticated); } - if self.estimate_required_bandwidth(&packets) > self.bandwidth_remaining { - return Err(GatewayClientError::NotEnoughBandwidth( - self.estimate_required_bandwidth(&packets), - self.bandwidth_remaining, - )); + let bandwidth_remaining = self.bandwidth_remaining.load(Ordering::Acquire); + if bandwidth_remaining < REMAINING_BANDWIDTH_THRESHOLD { + self.claim_bandwidth().await?; } + if !self.connection.is_established() { return Err(GatewayClientError::ConnectionNotEstablished); } @@ -730,19 +731,20 @@ impl GatewayClient { } // TODO: possibly make responses optional - pub async fn send_mix_packet( - &mut self, - mix_packet: MixPacket, - ) -> Result<(), GatewayClientError> { + pub async fn send_mix_packet(&mut self, mix_packet: MixPacket) -> Result<(), GatewayClientError> + where + C: DkgQueryClient + Send + Sync, + St: CredentialStorage, + ::StorageError: Send + Sync + 'static, + { if !self.authenticated { return Err(GatewayClientError::NotAuthenticated); } - if (mix_packet.packet().len() as i64) > self.bandwidth_remaining { - return Err(GatewayClientError::NotEnoughBandwidth( - mix_packet.packet().len() as i64, - self.bandwidth_remaining, - )); + let bandwidth_remaining = self.bandwidth_remaining.load(Ordering::Acquire); + if bandwidth_remaining < REMAINING_BANDWIDTH_THRESHOLD { + self.claim_bandwidth().await?; } + if !self.connection.is_established() { return Err(GatewayClientError::ConnectionNotEstablished); } @@ -796,6 +798,7 @@ impl GatewayClient { .as_ref() .expect("no shared key present even though we're authenticated!"), ), + self.bandwidth_remaining.clone(), self.shutdown.clone(), ) } @@ -836,10 +839,9 @@ impl GatewayClient { self.establish_connection().await?; } let shared_key = self.perform_initial_authentication().await?; - - if self.bandwidth_remaining < REMAINING_BANDWIDTH_THRESHOLD { - info!("Claiming more bandwidth for your tokens. This will use {} token(s) from your wallet. \ - Stop the process now if you don't want that to happen.", TOKENS_TO_BURN); + let bandwidth_remaining = self.bandwidth_remaining.load(Ordering::Acquire); + if bandwidth_remaining < REMAINING_BANDWIDTH_THRESHOLD { + info!("Claiming more bandwidth with existing credentials. Stop the process now if you don't want that to happen."); self.claim_bandwidth().await?; } @@ -876,7 +878,7 @@ impl GatewayClient { GatewayClient { authenticated: false, disabled_credentials_mode: true, - bandwidth_remaining: 0, + bandwidth_remaining: Arc::new(AtomicI64::new(0)), gateway_address: gateway_listener.to_string(), gateway_identity, local_identity, diff --git a/common/client-libs/gateway-client/src/socket_state.rs b/common/client-libs/gateway-client/src/socket_state.rs index 386ad5b147..128f6919e6 100644 --- a/common/client-libs/gateway-client/src/socket_state.rs +++ b/common/client-libs/gateway-client/src/socket_state.rs @@ -79,6 +79,7 @@ impl PartiallyDelegated { fn recover_received_plaintexts( ws_msgs: Vec, shared_key: &SharedKeys, + bandwidth_remaining: Arc, ) -> Result>, GatewayClientError> { let mut plaintexts = Vec::with_capacity(ws_msgs.len()); for ws_msg in ws_msgs { @@ -104,7 +105,11 @@ impl PartiallyDelegated { { ServerResponse::Send { remaining_bandwidth, - } => maybe_log_bandwidth(remaining_bandwidth), + } => { + maybe_log_bandwidth(remaining_bandwidth); + bandwidth_remaining + .store(remaining_bandwidth, std::sync::atomic::Ordering::Release) + } ServerResponse::Error { message } => { error!("gateway failure: {message}"); return Err(GatewayClientError::GatewayError(message)); @@ -129,8 +134,10 @@ impl PartiallyDelegated { ws_msgs: Vec, packet_router: &PacketRouter, shared_key: &SharedKeys, + bandwidth_remaining: Arc, ) -> Result<(), GatewayClientError> { - let plaintexts = Self::recover_received_plaintexts(ws_msgs, shared_key)?; + let plaintexts = + Self::recover_received_plaintexts(ws_msgs, shared_key, bandwidth_remaining)?; packet_router.route_received(plaintexts) } @@ -138,6 +145,7 @@ impl PartiallyDelegated { conn: WsConn, mut packet_router: PacketRouter, shared_key: Arc, + bandwidth_remaining: Arc, mut shutdown: TaskClient, ) -> Self { // when called for, it NEEDS TO yield back the stream so that we could merge it and @@ -169,7 +177,7 @@ impl PartiallyDelegated { Ok(msgs) => msgs }; - if let Err(err) = Self::route_socket_messages(ws_msgs, &packet_router, shared_key.as_ref()) { + if let Err(err) = Self::route_socket_messages(ws_msgs, &packet_router, shared_key.as_ref(), bandwidth_remaining.clone()) { log::error!("Route socket messages failed: {err}"); break Err(err) }