diff --git a/gateway/Cargo.toml b/gateway/Cargo.toml index 4ec9d74844..21f622c0ba 100644 --- a/gateway/Cargo.toml +++ b/gateway/Cargo.toml @@ -38,6 +38,7 @@ sqlx = { workspace = true, features = [ "sqlite", "macros", "migrate", + "time" ] } subtle-encoding = { version = "0.5", features = ["bech32-preview"] } thiserror = { workspace = true } diff --git a/gateway/migrations/20240301120000_freepass_expiration.sql b/gateway/migrations/20240301120000_freepass_expiration.sql new file mode 100644 index 0000000000..6538723ac6 --- /dev/null +++ b/gateway/migrations/20240301120000_freepass_expiration.sql @@ -0,0 +1,7 @@ +/* + * Copyright 2024 - Nym Technologies SA + * SPDX-License-Identifier: Apache-2.0 + */ + +ALTER TABLE available_bandwidth +ADD COLUMN freepass_expiration TIMESTAMP WITHOUT TIME ZONE; \ No newline at end of file diff --git a/gateway/src/config/mod.rs b/gateway/src/config/mod.rs index 5a345279d1..7932d2456f 100644 --- a/gateway/src/config/mod.rs +++ b/gateway/src/config/mod.rs @@ -42,6 +42,9 @@ const DEFAULT_MAXIMUM_CONNECTION_BUFFER_SIZE: usize = 2000; const DEFAULT_STORED_MESSAGE_FILENAME_LENGTH: u16 = 16; const DEFAULT_MESSAGE_RETRIEVAL_LIMIT: i64 = 100; +const DEFAULT_CLIENT_BANDWIDTH_MAX_FLUSHING_RATE: Duration = Duration::from_millis(5); +const DEFAULT_CLIENT_BANDWIDTH_MAX_DELTA_FLUSHING_AMOUNT: i64 = 512 * 1024; // 512kB + /// Derive default path to gateway's config directory. /// It should get resolved to `$HOME/.nym/gateways//config` pub fn default_config_directory>(id: P) -> PathBuf { @@ -516,6 +519,13 @@ pub struct Debug { /// Number of messages from offline client that can be pulled at once from the storage. pub message_retrieval_limit: i64, + /// Defines maximum delay between client bandwidth information being flushed to the persistent storage. + #[serde(with = "humantime_serde")] + pub client_bandwidth_max_flushing_rate: Duration, + + /// Defines a maximum change in client bandwidth before it gets flushed to the persistent storage. + pub client_bandwidth_max_delta_flushing_amount: i64, + /// Specifies whether the mixnode should be using the legacy framing for the sphinx packets. // it's set to true by default. The reason for that decision is to preserve compatibility with the // existing nodes whilst everyone else is upgrading and getting the code for handling the new field. @@ -533,6 +543,9 @@ impl Default for Debug { maximum_connection_buffer_size: DEFAULT_MAXIMUM_CONNECTION_BUFFER_SIZE, stored_messages_filename_length: DEFAULT_STORED_MESSAGE_FILENAME_LENGTH, message_retrieval_limit: DEFAULT_MESSAGE_RETRIEVAL_LIMIT, + client_bandwidth_max_flushing_rate: DEFAULT_CLIENT_BANDWIDTH_MAX_FLUSHING_RATE, + client_bandwidth_max_delta_flushing_amount: + DEFAULT_CLIENT_BANDWIDTH_MAX_DELTA_FLUSHING_AMOUNT, use_legacy_framed_packet_version: false, } } diff --git a/gateway/src/config/old_config_v1_1_31.rs b/gateway/src/config/old_config_v1_1_31.rs index b0797a9eaf..0d57cfe777 100644 --- a/gateway/src/config/old_config_v1_1_31.rs +++ b/gateway/src/config/old_config_v1_1_31.rs @@ -165,6 +165,7 @@ impl From for Config { stored_messages_filename_length: value.debug.stored_messages_filename_length, message_retrieval_limit: value.debug.message_retrieval_limit, use_legacy_framed_packet_version: value.debug.use_legacy_framed_packet_version, + ..Default::default() }, } } 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 9ff8882a3a..7c785547c1 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: GPL-3.0-only use crate::node::client_handling::bandwidth::BandwidthError; +use crate::node::client_handling::websocket::connection_handler::ClientBandwidth; use crate::node::{ client_handling::{ bandwidth::Bandwidth, @@ -34,6 +35,7 @@ use nym_validator_client::coconut::CoconutApiError; use rand::{CryptoRng, Rng}; use std::{process, time::Duration}; use thiserror::Error; +use time::OffsetDateTime; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError}; @@ -90,9 +92,18 @@ pub enum RequestHandlingError { #[error("the provided credential did not contain a valid type attribute")] InvalidTypeAttribute, + #[error("insufficient bandwidth available to process the request")] + OutOfBandwidth, + #[error("the provided credential did not have a bandwidth attribute")] MissingBandwidthAttribute, + #[error("attempted to claim a bandwidth voucher for an account using a free pass (it expires on {expiration})")] + BandwidthVoucherForFreePassAccount { expiration: OffsetDateTime }, + + #[error("attempted to claim another free pass for the account while another free pass is still active (it expires on {expiration})")] + PreexistingFreePass { expiration: OffsetDateTime }, + #[error("the DKG contract is unavailable")] UnavailableDkgContract, } @@ -121,6 +132,7 @@ impl IntoWSMessage for Result { pub(crate) struct AuthenticatedHandler { inner: FreshHandler, client: ClientDetails, + client_bandwidth: ClientBandwidth, mix_receiver: MixMessageReceiver, // Occasionally the handler is requested to ping the connected client for confirm that it's // active, such as when a duplicate connection is detected. This hashmap stores the oneshot @@ -153,19 +165,30 @@ where /// * `fresh`: fresh, unauthenticated, connection handler. /// * `client`: details (i.e. address and shared keys) of the registered client /// * `mix_receiver`: channel used for receiving messages from the mixnet destined for this client. - pub(crate) fn upgrade( + pub(crate) async fn upgrade( fresh: FreshHandler, client: ClientDetails, mix_receiver: MixMessageReceiver, is_active_request_receiver: IsActiveRequestReceiver, - ) -> Self { - AuthenticatedHandler { + ) -> Result { + // note: the `upgrade` function can only be called after registering or authenticating the client, + // meaning the appropriate database rows must have been created + // so in theory we could just unwrap the value here, but since we're returning a Result anyway, + // we might as well return a failure response instead + let bandwidth = fresh + .storage + .get_available_bandwidth(client.address) + .await? + .ok_or(RequestHandlingError::IllegalRequest)?; + + Ok(AuthenticatedHandler { inner: fresh, client, + client_bandwidth: ClientBandwidth::new(bandwidth.into()), mix_receiver, is_active_request_receiver, is_active_ping_pending_reply: None, - } + }) } /// Explicitly removes handle from the global store. @@ -175,15 +198,10 @@ where .disconnect(self.client.address) } - /// Checks the amount of bandwidth available for the connected client. - async fn get_available_bandwidth(&self) -> Result { - let bandwidth = self - .inner - .storage - .get_available_bandwidth(self.client.address) - .await? - .unwrap_or_default(); - Ok(bandwidth) + async fn expire_freepass(&mut self) -> Result<(), RequestHandlingError> { + self.client_bandwidth.bandwidth = Default::default(); + self.client_bandwidth.update_flush_data(); + Ok(self.inner.expire_freepass(self.client.address).await?) } /// Increases the amount of available bandwidth of the connected client by the specified value. @@ -191,12 +209,15 @@ where /// # Arguments /// /// * `amount`: amount to increase the available bandwidth by. - async fn increase_bandwidth(&self, bandwidth: Bandwidth) -> Result<(), RequestHandlingError> { - self.inner - .storage - .increase_bandwidth(self.client.address, bandwidth.value() as i64) - .await?; - Ok(()) + async fn increase_bandwidth( + &mut self, + bandwidth: Bandwidth, + ) -> Result<(), RequestHandlingError> { + self.client_bandwidth.bandwidth.bytes += bandwidth.value() as i64; + + // any increases to bandwidth should get flushed immediately + // (we don't want to accidentally miss somebody claiming a gigabyte voucher) + self.flush_bandwidth().await } /// Decreases the amount of available bandwidth of the connected client by the specified value. @@ -204,11 +225,18 @@ where /// # Arguments /// /// * `amount`: amount to decrease the available bandwidth by. - async fn consume_bandwidth(&self, amount: i64) -> Result<(), RequestHandlingError> { - self.inner - .storage - .consume_bandwidth(self.client.address, amount) - .await?; + async fn consume_bandwidth(&mut self, amount: i64) -> Result<(), RequestHandlingError> { + self.client_bandwidth.bandwidth.bytes -= amount; + + // 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_flush(self.inner.shared_state.bandwidth_cfg) + { + self.flush_bandwidth().await?; + } + Ok(()) } @@ -232,6 +260,22 @@ where let serial_number = credential.data.blinded_serial_number(); trace!("processing credential {}", serial_number.to_bs58()); + // if we already have had received a free pass (that's not expired, don't accept any additional bandwidth) + if self.client_bandwidth.bandwidth.freepass_expired() { + // the free pass we used before has expired -> reset our state and handle the request as normal + self.expire_freepass().await?; + } else if let Some(expiration) = self.client_bandwidth.bandwidth.freepass_expiration { + // the free pass is still valid -> return error + return match credential.data.typ { + CredentialType::Voucher => { + Err(RequestHandlingError::BandwidthVoucherForFreePassAccount { expiration }) + } + CredentialType::FreePass => { + Err(RequestHandlingError::PreexistingFreePass { expiration }) + } + }; + } + let already_spent = self .inner .storage @@ -249,6 +293,7 @@ where let aggregated_verification_key = self .inner + .shared_state .coconut_verifier .verification_key(credential.data.epoch_id) .await?; @@ -277,16 +322,20 @@ where )); } + drop(aggregated_verification_key); + let was_freepass = match credential.data.typ { CredentialType::Voucher => { trace!("the credential is a bandwidth voucher. attempting to release the funds"); let api_clients = self .inner + .shared_state .coconut_verifier .api_clients(credential.data.epoch_id) .await?; self.inner + .shared_state .coconut_verifier .release_bandwidth_voucher_funds(&api_clients, credential) .await?; @@ -316,7 +365,7 @@ where trace!("increasing client bandwidth"); self.increase_bandwidth(bandwidth).await?; - let available_total = self.get_available_bandwidth().await?; + let available_total = self.client_bandwidth.bandwidth.bytes; Ok(ServerResponse::Bandwidth { available_total }) } @@ -367,17 +416,44 @@ where ) -> Result { debug!("handling testnet bandwidth request"); - if self.inner.only_coconut_credentials { + if self.inner.shared_state.only_coconut_credentials { return Err(RequestHandlingError::OnlyCoconutCredentials); } self.increase_bandwidth(FREE_TESTNET_BANDWIDTH_VALUE) .await?; - let available_total = self.get_available_bandwidth().await?; + let available_total = self.client_bandwidth.bandwidth.bytes; Ok(ServerResponse::Bandwidth { available_total }) } + async fn flush_bandwidth(&mut self) -> Result<(), RequestHandlingError> { + trace!("flushing client bandwidth to the underlying storage"); + self.inner + .storage + .set_bandwidth(self.client.address, self.client_bandwidth.bandwidth.bytes) + .await?; + self.client_bandwidth.update_flush_data(); + Ok(()) + } + + async fn try_use_bandwidth( + &mut self, + required_bandwidth: i64, + ) -> Result { + if self.client_bandwidth.bandwidth.freepass_expired() { + self.expire_freepass().await?; + } + let available_bandwidth = self.client_bandwidth.bandwidth.bytes; + + if available_bandwidth < required_bandwidth { + return Err(RequestHandlingError::OutOfBandwidth); + } + + self.consume_bandwidth(required_bandwidth).await?; + Ok(self.client_bandwidth.bandwidth.bytes) + } + /// Tries to handle request to forward sphinx packet into the network. The request can only succeed /// if the client has enough available bandwidth. /// @@ -387,24 +463,16 @@ where /// /// * `mix_packet`: packet received from the client that should get forwarded into the network. async fn handle_forward_sphinx( - &self, + &mut self, mix_packet: MixPacket, ) -> Result { - let consumed_bandwidth = mix_packet.packet().len() as i64; + let required_bandwidth = mix_packet.packet().len() as i64; - let available_bandwidth = self.get_available_bandwidth().await?; - - if available_bandwidth < consumed_bandwidth { - return Ok(ServerResponse::new_error( - "Insufficient bandwidth available", - )); - } - - self.consume_bandwidth(consumed_bandwidth).await?; + let remaining_bandwidth = self.try_use_bandwidth(required_bandwidth).await?; self.forward_packet(mix_packet); Ok(ServerResponse::Send { - remaining_bandwidth: available_bandwidth - consumed_bandwidth, + remaining_bandwidth, }) } @@ -413,7 +481,7 @@ where /// # Arguments /// /// * `bin_msg`: raw message to handle. - async fn handle_binary(&self, bin_msg: Vec) -> Message { + async fn handle_binary(&mut self, bin_msg: Vec) -> Message { trace!("binary request"); // this function decrypts the request and checks the MAC match BinaryRequest::try_from_encrypted_tagged_bytes(bin_msg, &self.client.shared_keys) { 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 539026ad5a..aaf8e39518 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -20,18 +20,19 @@ use nym_gateway_requests::{ use nym_mixnet_client::forwarder::MixForwardingSender; use nym_sphinx::DestinationAddressBytes; use rand::{CryptoRng, Rng}; -use std::{sync::Arc, time::Duration}; +use std::time::Duration; use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError}; +use crate::node::client_handling::websocket::connection_handler::AvailableBandwidth; +use crate::node::client_handling::websocket::shared_state::SharedHandlerState; use crate::node::{ client_handling::{ active_clients::ActiveClientsStore, websocket::{ connection_handler::{ - coconut::CoconutVerifier, AuthenticatedHandler, ClientDetails, InitialAuthResult, - SocketStream, + AuthenticatedHandler, ClientDetails, InitialAuthResult, SocketStream, }, message_receiver::{IsActive, IsActiveRequestSender}, }, @@ -88,13 +89,11 @@ impl InitialAuthenticationError { pub(crate) struct FreshHandler { rng: R, - local_identity: Arc, - pub(crate) only_coconut_credentials: bool, + pub(crate) shared_state: SharedHandlerState, pub(crate) active_clients_store: ActiveClientsStore, pub(crate) outbound_mix_sender: MixForwardingSender, pub(crate) socket_connection: SocketStream, pub(crate) storage: St, - pub(crate) coconut_verifier: Arc, // currently unused (but populated) pub(crate) negotiated_protocol: Option, @@ -113,23 +112,19 @@ where pub(crate) fn new( rng: R, conn: S, - only_coconut_credentials: bool, outbound_mix_sender: MixForwardingSender, - local_identity: Arc, storage: St, active_clients_store: ActiveClientsStore, - coconut_verifier: Arc, + shared_state: SharedHandlerState, ) -> Self { FreshHandler { rng, active_clients_store, - only_coconut_credentials, outbound_mix_sender, socket_connection: SocketStream::RawTcp(conn), - local_identity, storage, - coconut_verifier, negotiated_protocol: None, + shared_state, } } @@ -171,7 +166,7 @@ where gateway_handshake( &mut self.rng, ws_stream, - self.local_identity.as_ref(), + self.shared_state.local_identity.as_ref(), init_msg, ) .await @@ -436,7 +431,7 @@ where // Ask the other connection to ping if they are still active. // Use a oneshot channel to return the result to us let (ping_result_sender, ping_result_receiver) = oneshot::channel(); - log::debug!("Asking other connection handler to ping the connected client to see if it is still active"); + debug!("Asking other connection handler to ping the connected client to see if it is still active"); if let Err(err) = is_active_request_tx.send(ping_result_sender).await { warn!("Failed to send ping request to other handler: {err}"); } @@ -448,31 +443,31 @@ where IsActive::NotActive => { // The other handler reported that the client is not active, so we can // disconnect the other client and continue with this connection. - log::debug!("Other handler reports it is not active"); + debug!("Other handler reports it is not active"); self.active_clients_store.disconnect(address); } IsActive::Active => { // The other handled reported a positive reply, so we have to assume it's // still active and disconnect this connection. - log::info!("Other handler reports it is active"); + info!("Other handler reports it is active"); return Err(InitialAuthenticationError::DuplicateConnection); } IsActive::BusyPinging => { // The other handler is already busy pinging the client, so we have to // assume it's still active and disconnect this connection. - log::debug!("Other handler reports it is already busy pinging the client"); + debug!("Other handler reports it is already busy pinging the client"); return Err(InitialAuthenticationError::DuplicateConnection); } } } Ok(Err(_)) => { // Other channel failed to reply (the channel sender probably dropped) - log::info!("Other connection failed to reply, disconnecting it in favour of this new connection"); + info!("Other connection failed to reply, disconnecting it in favour of this new connection"); self.active_clients_store.disconnect(address); } Err(_) => { // Timeout waiting for reply - log::warn!( + warn!( "Other connection timed out, disconnecting it in favour of this new connection" ); self.active_clients_store.disconnect(address); @@ -509,7 +504,7 @@ where // Check for duplicate clients if let Some(client_tx) = self.active_clients_store.get_remote_client(address) { - log::warn!("Detected duplicate connection for client: {address}"); + warn!("Detected duplicate connection for client: {address}"); self.handle_duplicate_client(address, client_tx.is_active_request_sender) .await?; } @@ -518,11 +513,17 @@ where .authenticate_client(address, encrypted_address, iv) .await?; let status = shared_keys.is_some(); - let bandwidth_remaining = self - .storage - .get_available_bandwidth(address) - .await? - .unwrap_or(0); + + let available_bandwidth: AvailableBandwidth = + self.storage.get_available_bandwidth(address).await?.into(); + + let bandwidth_remaining = if available_bandwidth.freepass_expired() { + self.expire_freepass(address).await?; + 0 + } else { + available_bandwidth.bytes + }; + let client_details = shared_keys.map(|shared_keys| ClientDetails::new(address, shared_keys)); @@ -536,6 +537,13 @@ where )) } + pub(crate) async fn expire_freepass( + &self, + client: DestinationAddressBytes, + ) -> Result<(), StorageError> { + self.storage.reset_freepass_bandwidth(client).await + } + /// Attempts to finalize registration of the client by storing the derived shared keys in the /// persistent store as well as creating entry for its bandwidth allocation. /// @@ -708,12 +716,14 @@ where mix_sender, is_active_request_sender, ); - Some(AuthenticatedHandler::upgrade( + AuthenticatedHandler::upgrade( self, client_details, mix_receiver, is_active_request_receiver, - )) + ) + .await + .ok() } else { None }; 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 67ec3e4316..84f61d6652 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs @@ -1,6 +1,7 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: GPL-3.0-only +use crate::config::Config; use crate::node::storage::Storage; use log::{trace, warn}; use nym_gateway_requests::registration::handshake::SharedKeys; @@ -8,6 +9,8 @@ use nym_gateway_requests::ServerResponse; use nym_sphinx::DestinationAddressBytes; use nym_task::TaskClient; use rand::{CryptoRng, Rng}; +use std::time::Duration; +use time::OffsetDateTime; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::WebSocketStream; use zeroize::{Zeroize, ZeroizeOnDrop}; @@ -112,3 +115,75 @@ pub(crate) async fn handle_connection( trace!("The handler is done!"); } + +#[derive(Debug, Clone, Copy)] +pub(crate) struct BandwidthFlushingBehaviourConfig { + /// Defines maximum delay between client bandwidth information being flushed to the persistent storage. + pub client_bandwidth_max_flushing_rate: Duration, + + /// Defines a maximum change in client bandwidth before it gets flushed to the persistent storage. + pub client_bandwidth_max_delta_flushing_amount: i64, +} + +impl<'a> From<&'a Config> for BandwidthFlushingBehaviourConfig { + fn from(value: &'a Config) -> Self { + BandwidthFlushingBehaviourConfig { + client_bandwidth_max_flushing_rate: value.debug.client_bandwidth_max_flushing_rate, + client_bandwidth_max_delta_flushing_amount: value + .debug + .client_bandwidth_max_delta_flushing_amount, + } + } +} + +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct AvailableBandwidth { + pub(crate) bytes: i64, + pub(crate) freepass_expiration: Option, +} + +impl AvailableBandwidth { + pub(crate) fn freepass_expired(&self) -> bool { + if let Some(expiration) = self.freepass_expiration { + if expiration < OffsetDateTime::now_utc() { + return true; + } + } + false + } +} + +pub(crate) struct ClientBandwidth { + pub(crate) bandwidth: AvailableBandwidth, + pub(crate) last_flushed: OffsetDateTime, + pub(crate) bytes_at_last_flush: i64, +} + +impl ClientBandwidth { + pub(crate) fn new(bandwidth: AvailableBandwidth) -> ClientBandwidth { + ClientBandwidth { + bandwidth, + last_flushed: OffsetDateTime::now_utc(), + bytes_at_last_flush: bandwidth.bytes, + } + } + + pub(crate) fn should_flush(&self, cfg: BandwidthFlushingBehaviourConfig) -> bool { + if (self.bytes_at_last_flush - self.bandwidth.bytes).abs() + >= cfg.client_bandwidth_max_delta_flushing_amount + { + return true; + } + + if self.last_flushed + cfg.client_bandwidth_max_flushing_rate < OffsetDateTime::now_utc() { + return true; + } + + false + } + + pub(crate) fn update_flush_data(&mut self) { + self.last_flushed = OffsetDateTime::now_utc(); + self.bytes_at_last_flush = self.bandwidth.bytes; + } +} diff --git a/gateway/src/node/client_handling/websocket/listener.rs b/gateway/src/node/client_handling/websocket/listener.rs index 27a03aacfc..9b073dbc5c 100644 --- a/gateway/src/node/client_handling/websocket/listener.rs +++ b/gateway/src/node/client_handling/websocket/listener.rs @@ -2,37 +2,26 @@ // SPDX-License-Identifier: GPL-3.0-only use crate::node::client_handling::active_clients::ActiveClientsStore; -use crate::node::client_handling::websocket::connection_handler::coconut::CoconutVerifier; use crate::node::client_handling::websocket::connection_handler::FreshHandler; +use crate::node::client_handling::websocket::shared_state::SharedHandlerState; use crate::node::storage::Storage; use log::*; -use nym_crypto::asymmetric::identity; use nym_mixnet_client::forwarder::MixForwardingSender; use rand::rngs::OsRng; use std::net::SocketAddr; use std::process; -use std::sync::Arc; use tokio::task::JoinHandle; pub(crate) struct Listener { address: SocketAddr, - local_identity: Arc, - only_coconut_credentials: bool, - pub(crate) coconut_verifier: Arc, + shared_state: SharedHandlerState, } impl Listener { - pub(crate) fn new( - address: SocketAddr, - local_identity: Arc, - only_coconut_credentials: bool, - coconut_verifier: Arc, - ) -> Self { + pub(crate) fn new(address: SocketAddr, shared_state: SharedHandlerState) -> Self { Listener { address, - local_identity, - only_coconut_credentials, - coconut_verifier, + shared_state, } } @@ -71,12 +60,10 @@ impl Listener { let handle = FreshHandler::new( OsRng, socket, - self.only_coconut_credentials, outbound_mix_sender.clone(), - Arc::clone(&self.local_identity), storage.clone(), active_clients_store.clone(), - Arc::clone(&self.coconut_verifier), + self.shared_state.clone(), ); let shutdown = shutdown.clone().named(format!("ClientConnectionHandler_{remote_addr}")); tokio::spawn(async move { handle.start_handling(shutdown).await }); diff --git a/gateway/src/node/client_handling/websocket/mod.rs b/gateway/src/node/client_handling/websocket/mod.rs index 673f14b8a8..6d0a990956 100644 --- a/gateway/src/node/client_handling/websocket/mod.rs +++ b/gateway/src/node/client_handling/websocket/mod.rs @@ -6,3 +6,6 @@ pub(crate) use listener::Listener; pub(crate) mod connection_handler; pub(crate) mod listener; pub(crate) mod message_receiver; +pub(crate) mod shared_state; + +pub(crate) use shared_state::SharedHandlerState; diff --git a/gateway/src/node/client_handling/websocket/shared_state.rs b/gateway/src/node/client_handling/websocket/shared_state.rs new file mode 100644 index 0000000000..1f19476e97 --- /dev/null +++ b/gateway/src/node/client_handling/websocket/shared_state.rs @@ -0,0 +1,16 @@ +// Copyright 2024 - Nym Technologies SA +// SPDX-License-Identifier: GPL-3.0-only + +use crate::node::client_handling::websocket::connection_handler::coconut::CoconutVerifier; +use crate::node::client_handling::websocket::connection_handler::BandwidthFlushingBehaviourConfig; +use nym_crypto::asymmetric::identity; +use std::sync::Arc; + +// I can see this being possible expanded with say storage or client store +#[derive(Clone)] +pub(crate) struct SharedHandlerState { + pub(crate) coconut_verifier: Arc, + pub(crate) local_identity: Arc, + pub(crate) only_coconut_credentials: bool, + pub(crate) bandwidth_cfg: BandwidthFlushingBehaviourConfig, +} diff --git a/gateway/src/node/mod.rs b/gateway/src/node/mod.rs index 53607a3dfb..5f5f57e164 100644 --- a/gateway/src/node/mod.rs +++ b/gateway/src/node/mod.rs @@ -254,13 +254,14 @@ impl Gateway { self.config.gateway.clients_port, ); - websocket::Listener::new( - listening_address, - Arc::clone(&self.identity_keypair), - self.config.gateway.only_coconut_credentials, + let shared_state = websocket::SharedHandlerState { coconut_verifier, - ) - .start( + local_identity: Arc::clone(&self.identity_keypair), + only_coconut_credentials: self.config.gateway.only_coconut_credentials, + bandwidth_cfg: (&self.config).into(), + }; + + websocket::Listener::new(listening_address, shared_state).start( forwarding_channel, self.storage.clone(), active_clients_store, diff --git a/gateway/src/node/storage/bandwidth.rs b/gateway/src/node/storage/bandwidth.rs index 3eb4dab057..1ebcccea85 100644 --- a/gateway/src/node/storage/bandwidth.rs +++ b/gateway/src/node/storage/bandwidth.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: GPL-3.0-only use crate::node::storage::models::{PersistedBandwidth, SpentCredential}; +use time::OffsetDateTime; #[derive(Clone)] pub(crate) struct BandwidthManager { @@ -36,6 +37,53 @@ impl BandwidthManager { Ok(()) } + /// Set the freepass expiration date of the particular client to the provided date. + /// + /// # Arguments + /// + /// * `client_address_bs58`: base58-encoded address of the client. + /// * `freepass_expiration`: the expiration date of the associated free pass. + pub(crate) async fn set_freepass_expiration( + &self, + client_address_bs58: &str, + freepass_expiration: OffsetDateTime, + ) -> Result<(), sqlx::Error> { + sqlx::query!( + r#" + UPDATE available_bandwidth + SET freepass_expiration = ? + WHERE client_address_bs58 = ? + "#, + freepass_expiration, + client_address_bs58 + ) + .execute(&self.connection_pool) + .await?; + Ok(()) + } + + /// Reset all the bandwidth associated with the freepass and reset its expiration date + /// + /// # Arguments + /// + /// * `client_address_bs58`: base58-encoded address of the client. + pub(crate) async fn reset_freepass_bandwidth( + &self, + client_address_bs58: &str, + ) -> Result<(), sqlx::Error> { + sqlx::query!( + r#" + UPDATE available_bandwidth + SET available = 0, freepass_expiration = NULL + WHERE client_address_bs58 = ? + "#, + client_address_bs58 + ) + .execute(&self.connection_pool) + .await?; + Ok(()) + } + /// Tries to retrieve available bandwidth for the particular client. /// /// # Arguments @@ -45,22 +93,19 @@ impl BandwidthManager { &self, client_address_bs58: &str, ) -> Result, sqlx::Error> { - sqlx::query_as!( - PersistedBandwidth, - "SELECT * FROM available_bandwidth WHERE client_address_bs58 = ?", - client_address_bs58 - ) - .fetch_optional(&self.connection_pool) - .await + sqlx::query_as("SELECT * FROM available_bandwidth WHERE client_address_bs58 = ?") + .bind(client_address_bs58) + .fetch_optional(&self.connection_pool) + .await } - /// Increases available bandwidth of the particular client by the specified amount. + /// Sets available bandwidth of the particular client to the provided amount; /// /// # Arguments /// - /// * `client_address_bs58`: base58-encoded address of the client. - /// * `amount`: amount of available bandwidth to be added to the client. - pub(crate) async fn increase_available_bandwidth( + /// * `client_address`: address of the client + /// * `amount`: the updated client bandwidth amount. + pub(crate) async fn set_available_bandwidth( &self, client_address_bs58: &str, amount: i64, @@ -68,32 +113,7 @@ impl BandwidthManager { sqlx::query!( r#" UPDATE available_bandwidth - SET available = available + ? - WHERE client_address_bs58 = ? - "#, - amount, - client_address_bs58 - ) - .execute(&self.connection_pool) - .await?; - Ok(()) - } - - /// Decreases available bandwidth of the particular client by the specified amount. - /// - /// # Arguments - /// - /// * `client_address_bs58`: base58-encoded address of the client. - /// * `amount`: amount of available bandwidth to be removed from the client. - pub(crate) async fn decrease_available_bandwidth( - &self, - client_address_bs58: &str, - amount: i64, - ) -> Result<(), sqlx::Error> { - sqlx::query!( - r#" - UPDATE available_bandwidth - SET available = available - ? + SET available = ? WHERE client_address_bs58 = ? "#, amount, diff --git a/gateway/src/node/storage/mod.rs b/gateway/src/node/storage/mod.rs index bb7330c2f6..e7123a44a6 100644 --- a/gateway/src/node/storage/mod.rs +++ b/gateway/src/node/storage/mod.rs @@ -4,7 +4,7 @@ use crate::node::storage::bandwidth::BandwidthManager; use crate::node::storage::error::StorageError; use crate::node::storage::inboxes::InboxManager; -use crate::node::storage::models::{PersistedSharedKeys, StoredMessage}; +use crate::node::storage::models::{PersistedBandwidth, PersistedSharedKeys, StoredMessage}; use crate::node::storage::shared_keys::SharedKeysManager; use async_trait::async_trait; use log::{debug, error}; @@ -13,6 +13,7 @@ use nym_gateway_requests::registration::handshake::SharedKeys; use nym_sphinx::DestinationAddressBytes; use sqlx::ConnectOptions; use std::path::Path; +use time::OffsetDateTime; mod bandwidth; pub(crate) mod error; @@ -102,6 +103,28 @@ pub trait Storage: Send + Sync { client_address: DestinationAddressBytes, ) -> Result<(), StorageError>; + /// Set the freepass expiration date of the particular client to the provided date. + /// + /// # Arguments + /// + /// * `client_address`: address of the client + /// * `freepass_expiration`: the expiration date of the associated free pass. + async fn set_freepass_expiration( + &self, + client_address: DestinationAddressBytes, + freepass_expiration: OffsetDateTime, + ) -> Result<(), StorageError>; + + /// Reset all the bandwidth associated with the freepass and reset its expiration date + /// + /// # Arguments + /// + /// * `client_address`: address of the client + async fn reset_freepass_bandwidth( + &self, + client_address: DestinationAddressBytes, + ) -> Result<(), StorageError>; + /// Tries to retrieve available bandwidth for the particular client. /// /// # Arguments @@ -110,27 +133,15 @@ pub trait Storage: Send + Sync { async fn get_available_bandwidth( &self, client_address: DestinationAddressBytes, - ) -> Result, StorageError>; + ) -> Result, StorageError>; - /// Increases available bandwidth of the particular client by the specified amount. + /// Sets available bandwidth of the particular client to the provided amount; /// /// # Arguments /// /// * `client_address`: address of the client - /// * `amount`: amount of available bandwidth to be added to the client. - async fn increase_bandwidth( - &self, - client_address: DestinationAddressBytes, - amount: i64, - ) -> Result<(), StorageError>; - - /// Decreases available bandwidth of the particular client by the specified amount. - /// - /// # Arguments - /// - /// * `client_address`: address of the client - /// * `amount`: amount of available bandwidth to be removed from the client. - async fn consume_bandwidth( + /// * `amount`: the updated client bandwidth amount. + async fn set_bandwidth( &self, client_address: DestinationAddressBytes, amount: i64, @@ -295,36 +306,44 @@ impl Storage for PersistentStorage { Ok(()) } - async fn get_available_bandwidth( + async fn set_freepass_expiration( &self, client_address: DestinationAddressBytes, - ) -> Result, StorageError> { - let res = self - .bandwidth_manager - .get_available_bandwidth(&client_address.as_base58_string()) - .await - .map(|bandwidth_option| bandwidth_option.map(|bandwidth| bandwidth.available))?; - Ok(res) - } - - async fn increase_bandwidth( - &self, - client_address: DestinationAddressBytes, - amount: i64, + freepass_expiration: OffsetDateTime, ) -> Result<(), StorageError> { self.bandwidth_manager - .increase_available_bandwidth(&client_address.as_base58_string(), amount) + .set_freepass_expiration(&client_address.as_base58_string(), freepass_expiration) .await?; Ok(()) } - async fn consume_bandwidth( + async fn reset_freepass_bandwidth( + &self, + client_address: DestinationAddressBytes, + ) -> Result<(), StorageError> { + self.bandwidth_manager + .reset_freepass_bandwidth(&client_address.as_base58_string()) + .await?; + Ok(()) + } + + async fn get_available_bandwidth( + &self, + client_address: DestinationAddressBytes, + ) -> Result, StorageError> { + Ok(self + .bandwidth_manager + .get_available_bandwidth(&client_address.as_base58_string()) + .await?) + } + + async fn set_bandwidth( &self, client_address: DestinationAddressBytes, amount: i64, ) -> Result<(), StorageError> { self.bandwidth_manager - .decrease_available_bandwidth(&client_address.as_base58_string(), amount) + .set_available_bandwidth(&client_address.as_base58_string(), amount) .await?; Ok(()) } @@ -422,22 +441,29 @@ impl Storage for InMemStorage { todo!() } - async fn get_available_bandwidth( + async fn set_freepass_expiration( &self, _client_address: DestinationAddressBytes, - ) -> Result, StorageError> { - todo!() - } - - async fn increase_bandwidth( - &self, - _client_address: DestinationAddressBytes, - _amount: i64, + _freepass_expiration: OffsetDateTime, ) -> Result<(), StorageError> { todo!() } - async fn consume_bandwidth( + async fn reset_freepass_bandwidth( + &self, + _client_address: DestinationAddressBytes, + ) -> Result<(), StorageError> { + todo!() + } + + async fn get_available_bandwidth( + &self, + _client_address: DestinationAddressBytes, + ) -> Result, StorageError> { + todo!() + } + + async fn set_bandwidth( &self, _client_address: DestinationAddressBytes, _amount: i64, diff --git a/gateway/src/node/storage/models.rs b/gateway/src/node/storage/models.rs index b3fde7d37a..999c9f827f 100644 --- a/gateway/src/node/storage/models.rs +++ b/gateway/src/node/storage/models.rs @@ -1,7 +1,9 @@ // Copyright 2021-2024 - Nym Technologies SA // SPDX-License-Identifier: GPL-3.0-only +use crate::node::client_handling::websocket::connection_handler::AvailableBandwidth; use sqlx::FromRow; +use time::OffsetDateTime; pub struct PersistedSharedKeys { pub(crate) client_address_bs58: String, @@ -15,10 +17,30 @@ pub struct StoredMessage { pub(crate) content: Vec, } +#[derive(Debug, Clone, FromRow)] pub struct PersistedBandwidth { #[allow(dead_code)] pub(crate) client_address_bs58: String, pub(crate) available: i64, + pub(crate) freepass_expiration: Option, +} + +impl From for AvailableBandwidth { + fn from(value: PersistedBandwidth) -> Self { + AvailableBandwidth { + bytes: value.available, + freepass_expiration: value.freepass_expiration, + } + } +} + +impl From> for AvailableBandwidth { + fn from(value: Option) -> Self { + match value { + None => AvailableBandwidth::default(), + Some(b) => b.into(), + } + } } #[derive(Debug, Clone, FromRow)]