GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

client ask for more bandwidth if it runs out

Browse Source
This commit is contained in:
Simon Wicky
2024-05-07 13:48:14 +02:00
parent 0d133187e0
commit 0d69bfbef7
3 changed files with 66 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/client-core/src/client/mix_traffic/transceiver.rs
+8 -4
View File
@@ -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<C, St> RemoteGateway<C, St> {
impl<C, St> GatewayTransceiver for RemoteGateway<C, St>
where
C: Send,
St: Send,
C: DkgQueryClient + Send + Sync,
St: CredentialStorage,
<St as 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<C, St> GatewaySender for RemoteGateway<C, St>
where
C: Send,
St: Send,
C: DkgQueryClient + Send + Sync,
St: CredentialStorage,
<St as CredentialStorage>::StorageError: Send + Sync + 'static,
{
async fn send_mix_packet(&mut self, packet: MixPacket) -> Result<(), ErasedGatewayError> {
self.gateway_client
common/client-libs/gateway-client/src/client.rs
+47 -45
View File
@@ -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<C, St = EphemeralCredentialStorage> {
authenticated: bool,
disabled_credentials_mode: bool,
bandwidth_remaining: i64,
bandwidth_remaining: Arc<AtomicI64>,
gateway_address: String,
gateway_identity: identity::PublicKey,
local_identity: Arc<identity::KeyPair>,
@@ -122,7 +123,7 @@ impl<C, St> GatewayClient<C, St> {
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<C, St> GatewayClient<C, St> {
}
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<C, St> GatewayClient<C, St> {
} => {
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<C, St> GatewayClient<C, St> {
}
}
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<C, St> GatewayClient<C, St> {
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::<usize>() as i64
}
pub async fn batch_send_mix_packets(
&mut self,
packets: Vec<MixPacket>,
) -> Result<(), GatewayClientError> {
) -> Result<(), GatewayClientError>
where
C: DkgQueryClient + Send + Sync,
St: CredentialStorage,
<St as 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<C, St> GatewayClient<C, St> {
}
// 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,
<St as 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<C, St> GatewayClient<C, St> {
.as_ref()
.expect("no shared key present even though we're authenticated!"),
),
self.bandwidth_remaining.clone(),
self.shutdown.clone(),
)
}
@@ -836,10 +839,9 @@ impl<C, St> GatewayClient<C, St> {
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<InitOnly, EphemeralCredentialStorage> {
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,
common/client-libs/gateway-client/src/socket_state.rs
+11 -3
View File
@@ -79,6 +79,7 @@ impl PartiallyDelegated {
fn recover_received_plaintexts(
ws_msgs: Vec<Message>,
shared_key: &SharedKeys,
bandwidth_remaining: Arc<AtomicI64>,
) -> Result<Vec<Vec<u8>>, 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<Message>,
packet_router: &PacketRouter,
shared_key: &SharedKeys,
bandwidth_remaining: Arc<AtomicI64>,
) -> 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<SharedKeys>,
bandwidth_remaining: Arc<AtomicI64>,
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)
}