Merge pull request #4437 from nymtech/feature/freepass-expiration

expire freepass bandwidth + stagger db flushes

Cargo.lock

@@ -5756,7 +5756,9 @@ dependencies = [
  "nym-validator-client",
  "rand 0.7.3",
  "serde",
+ "si-scale",
  "thiserror",
+ "time",
  "tokio",
  "tokio-stream",
  "tokio-tungstenite",

clients/native/src/commands/import_credential.rs

@@ -8,5 +8,7 @@ use nym_client_core::cli_helpers::client_import_credential::{
 };
 
 pub(crate) async fn execute(args: CommonClientImportCredentialArgs) -> Result<(), ClientError> {
-    import_credential::<CliNativeClient, _>(args).await
+    import_credential::<CliNativeClient, _>(args).await?;
+    println!("successfully imported credential!");
+    Ok(())
 }

clients/socks5/src/commands/import_credential.rs

@@ -10,5 +10,7 @@ use nym_client_core::cli_helpers::client_import_credential::{
 pub(crate) async fn execute(
     args: CommonClientImportCredentialArgs,
 ) -> Result<(), Socks5ClientError> {
-    import_credential::<CliSocks5Client, _>(args).await
+    import_credential::<CliSocks5Client, _>(args).await?;
+    println!("successfully imported credential!");
+    Ok(())
 }

common/client-libs/gateway-client/Cargo.toml

@@ -16,6 +16,8 @@ thiserror = { workspace = true }
 url = { workspace = true }
 rand = { version = "0.7.3", features = ["wasm-bindgen"] }
 tokio = { version = "1.24.1", features = ["macros"] }
+si-scale = "0.2.2"
+time.workspace = true
 
 # internal
 nym-bandwidth-controller = { path = "../../bandwidth-controller" }

common/client-libs/gateway-client/src/socket_state.rs

@@ -10,9 +10,14 @@ use futures::stream::{SplitSink, SplitStream};
 use futures::{SinkExt, StreamExt};
 use log::*;
 use nym_gateway_requests::registration::handshake::SharedKeys;
+use nym_gateway_requests::ServerResponse;
 use nym_task::TaskClient;
+use si_scale::helpers::bibytes2;
 use std::os::raw::c_int as RawFd;
+use std::sync::atomic::{AtomicI64, Ordering};
 use std::sync::Arc;
+use std::time::Duration;
+use time::OffsetDateTime;
 use tungstenite::Message;
 
 #[cfg(unix)]
@@ -48,6 +53,21 @@ pub(crate) fn ws_fd(_conn: &WsConn) -> Option<RawFd> {
     None
 }
 
+// disgusting? absolutely, but does the trick for now
+static LAST_LOGGED_BANDWIDTH_TS: AtomicI64 = AtomicI64::new(0);
+
+fn maybe_log_bandwidth(remaining: i64) {
+    // SAFETY: this value is always populated with valid timestamps
+    let last =
+        OffsetDateTime::from_unix_timestamp(LAST_LOGGED_BANDWIDTH_TS.load(Ordering::Relaxed))
+            .unwrap();
+    let now = OffsetDateTime::now_utc();
+    if last + Duration::from_secs(10) < now {
+        log::info!("remaining bandwidth: {}", bibytes2(remaining as f64));
+        LAST_LOGGED_BANDWIDTH_TS.store(now.unix_timestamp(), Ordering::Relaxed)
+    }
+}
+
 pub(crate) struct PartiallyDelegated {
     sink_half: SplitSink<WsConn, Message>,
     delegated_stream: (SplitStreamReceiver, oneshot::Sender<()>),
@@ -55,7 +75,10 @@ pub(crate) struct PartiallyDelegated {
 }
 
 impl PartiallyDelegated {
-    fn recover_received_plaintexts(ws_msgs: Vec<Message>, shared_key: &SharedKeys) -> Vec<Vec<u8>> {
+    fn recover_received_plaintexts(
+        ws_msgs: Vec<Message>,
+        shared_key: &SharedKeys,
+    ) -> Result<Vec<Vec<u8>>, GatewayClientError> {
         let mut plaintexts = Vec::with_capacity(ws_msgs.len());
         for ws_msg in ws_msgs {
             match ws_msg {
@@ -73,15 +96,32 @@ impl PartiallyDelegated {
                 // TODO: those can return the "send confirmations" - perhaps it should be somehow worked around?
                 Message::Text(text) => {
                     trace!(
-                    "received a text message - probably a response to some previous query! - {}",
-                    text
+                    "received a text message - probably a response to some previous query! - {text}",
                 );
+                    match ServerResponse::try_from(text)
+                        .map_err(|_| GatewayClientError::MalformedResponse)?
+                    {
+                        ServerResponse::Send {
+                            remaining_bandwidth,
+                        } => maybe_log_bandwidth(remaining_bandwidth),
+                        ServerResponse::Error { message } => {
+                            error!("gateway failure: {message}");
+                            return Err(GatewayClientError::GatewayError(message));
+                        }
+                        other => {
+                            warn!(
+                                "received illegal message of type {} in an authenticated client",
+                                other.name()
+                            )
+                        }
+                    }
+
                     continue;
                 }
                 _ => continue,
             }
         }
-        plaintexts
+        Ok(plaintexts)
     }
 
     fn route_socket_messages(
@@ -89,7 +129,7 @@ impl PartiallyDelegated {
         packet_router: &PacketRouter,
         shared_key: &SharedKeys,
     ) -> Result<(), GatewayClientError> {
-        let plaintexts = Self::recover_received_plaintexts(ws_msgs, shared_key);
+        let plaintexts = Self::recover_received_plaintexts(ws_msgs, shared_key)?;
         packet_router.route_received(plaintexts)
     }
 
@@ -129,7 +169,8 @@ impl PartiallyDelegated {
                         };
 
                         if let Err(err) = Self::route_socket_messages(ws_msgs, &packet_router, shared_key.as_ref()) {
-                            log::warn!("Route socket messages failed: {err}");
+                            log::error!("Route socket messages failed: {err}");
+                            break Err(err)
                         }
                     }
                 };

common/credential-storage/migrations/20241104120000_add_unique.sql

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+ALTER TABLE coconut_credentials
+    RENAME TO old_coconut_credentials;
+
+CREATE TABLE coconut_credentials
+(
+    id                     INTEGER                                                                     NOT NULL PRIMARY KEY AUTOINCREMENT,
+
+--     introduce a way for us to introduce breaking changes in serialization
+    serialization_revision INTEGER                                                                     NOT NULL,
+
+--     the best we can do without enums
+    credential_type        TEXT CHECK ( credential_type IN ('BandwidthVoucher', 'FreeBandwidthPass') ) NOT NULL,
+    credential_data        BLOB                                                                        NOT NULL UNIQUE,
+    epoch_id               INTEGER                                                                     NOT NULL,
+
+--     this field is only really applicable to free passes
+    expired                BOOLEAN                                                                     NOT NULL
+);
+
+ALTER TABLE credential_usage
+    RENAME TO old_credential_usage;
+
+-- for bandwidth vouchers there's going to be only a single entry; for freepasses there can be as many as there are gateways
+CREATE TABLE credential_usage
+(
+    credential_id   INTEGER NOT NULL REFERENCES coconut_credentials (id),
+    gateway_id_bs58 TEXT    NOT NULL,
+
+--     no matter credential type, we can't spend the same credential with the same gateway multiple times
+    UNIQUE (credential_id, gateway_id_bs58)
+);
+
+INSERT INTO coconut_credentials
+SELECT *
+FROM old_coconut_credentials;
+
+
+INSERT INTO credential_usage
+SELECT *
+FROM old_credential_usage;
+
+DROP TABLE old_coconut_credentials;
+DROP TABLE old_credential_usage;

common/credential-storage/src/backends/sqlite.rs

@@ -44,7 +44,7 @@ impl CoconutCredentialManager {
             r#"
                 SELECT * 
                 FROM coconut_credentials
-                WHERE coconut_credentials.credential_type == "FreeBandwidthPass"
+                WHERE coconut_credentials.credential_type == "FreeBandwidthPass" AND coconut_credentials.expired = false
                       AND NOT EXISTS (SELECT 1
                                       FROM   credential_usage
                                       WHERE  credential_usage.credential_id = coconut_credentials.id

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 }

gateway/gateway-requests/src/registration/handshake/error.rs

@@ -4,7 +4,7 @@
 use nym_crypto::asymmetric::identity;
 use thiserror::Error;
 
-#[derive(Debug, Error)]
+#[derive(Debug, Clone, Error)]
 pub enum HandshakeError {
     #[error(
         "received key material of invalid length - {0}. Expected: {}",

gateway/gateway-requests/src/types.rs

@@ -190,6 +190,22 @@ impl ClientControlRequest {
         }
     }
 
+    pub fn name(&self) -> String {
+        match self {
+            ClientControlRequest::Authenticate { .. } => "Authenticate".to_string(),
+            ClientControlRequest::RegisterHandshakeInitRequest { .. } => {
+                "RegisterHandshakeInitRequest".to_string()
+            }
+            ClientControlRequest::BandwidthCredential { .. } => "BandwidthCredential".to_string(),
+            ClientControlRequest::BandwidthCredentialV2 { .. } => {
+                "BandwidthCredentialV2".to_string()
+            }
+            ClientControlRequest::ClaimFreeTestnetBandwidth => {
+                "ClaimFreeTestnetBandwidth".to_string()
+            }
+        }
+    }
+
     pub fn new_enc_coconut_bandwidth_credential_v1(
         credential: &OldV1Credential,
         shared_key: &SharedKeys,
@@ -291,6 +307,15 @@ pub enum ServerResponse {
 }
 
 impl ServerResponse {
+    pub fn name(&self) -> String {
+        match self {
+            ServerResponse::Authenticate { .. } => "Authenticate".to_string(),
+            ServerResponse::Register { .. } => "Register".to_string(),
+            ServerResponse::Bandwidth { .. } => "Bandwidth".to_string(),
+            ServerResponse::Send { .. } => "Send".to_string(),
+            ServerResponse::Error { .. } => "Error".to_string(),
+        }
+    }
     pub fn new_error<S: Into<String>>(msg: S) -> Self {
         ServerResponse::Error {
             message: msg.into(),

gateway/migrations/20240301120000_freepass_expiration.sql

@@ -0,0 +1,7 @@
+/*
+ * Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+ALTER TABLE available_bandwidth
+ADD COLUMN freepass_expiration TIMESTAMP WITHOUT TIME ZONE;

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/<id>/config`
 pub fn default_config_directory<P: AsRef<Path>>(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,
         }
     }

gateway/src/config/old_config_v1_1_31.rs

@@ -165,6 +165,7 @@ impl From<ConfigV1_1_31> 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()
             },
         }
     }

gateway/src/node/client_handling/bandwidth.rs

@@ -42,18 +42,35 @@ pub struct Bandwidth {
 }
 
 impl Bandwidth {
-    pub const fn new(value: u64) -> Bandwidth {
+    pub const fn new_unchecked(value: u64) -> Bandwidth {
         Bandwidth { value }
     }
 
-    pub fn try_from_raw_value(value: &str, typ: CredentialType) -> Result<Self, BandwidthError> {
-        let bandwidth_value =
+    pub fn new(bandwidth_value: u64) -> Result<Bandwidth, BandwidthError> {
+        if bandwidth_value > i64::MAX as u64 {
+            // note that this would have represented more than 1 exabyte,
+            // which is like 125,000 worth of hard drives, so I don't think we have
+            // to worry about it for now...
+            warn!("Somehow we received bandwidth value higher than 9223372036854775807. We don't really want to deal with this now");
+            return Err(BandwidthError::UnsupportedBandwidthValue(bandwidth_value));
+        }
+
+        Ok(Bandwidth {
+            value: bandwidth_value,
+        })
+    }
+
+    pub(crate) fn parse_raw_bandwidth(
+        value: &str,
+        typ: CredentialType,
+    ) -> Result<(u64, Option<OffsetDateTime>), BandwidthError> {
+        let (bandwidth_value, freepass_expiration) =
             match typ {
                 CredentialType::Voucher => {
                     let token_value: u64 = value
                         .parse()
                         .map_err(|source| BandwidthError::VoucherValueParsingFailure { source })?;
-                    token_value * nym_network_defaults::BYTES_PER_UTOKEN
+                    (token_value * nym_network_defaults::BYTES_PER_UTOKEN, None)
                 }
                 CredentialType::FreePass => {
                     let expiry_timestamp: i64 = value
@@ -70,21 +87,10 @@ impl Bandwidth {
                     if expiry_date < now {
                         return Err(BandwidthError::ExpiredFreePass { expiry_date });
                     }
-                    nym_network_defaults::BYTES_PER_FREEPASS
+                    (nym_network_defaults::BYTES_PER_FREEPASS, Some(expiry_date))
                 }
             };
-
-        if bandwidth_value > i64::MAX as u64 {
-            // note that this would have represented more than 1 exabyte,
-            // which is like 125,000 worth of hard drives, so I don't think we have
-            // to worry about it for now...
-            warn!("Somehow we received bandwidth value higher than 9223372036854775807. We don't really want to deal with this now");
-            return Err(BandwidthError::UnsupportedBandwidthValue(bandwidth_value));
-        }
-
-        Ok(Bandwidth {
-            value: bandwidth_value,
-        })
+        Ok((bandwidth_value, freepass_expiration))
     }
 
     pub fn value(&self) -> u64 {

gateway/src/node/client_handling/mod.rs

@@ -8,4 +8,5 @@ mod bandwidth;
 pub(crate) mod embedded_clients;
 pub(crate) mod websocket;
 
-pub(crate) const FREE_TESTNET_BANDWIDTH_VALUE: Bandwidth = Bandwidth::new(64 * 1024 * 1024 * 1024); // 64GB
+pub(crate) const FREE_TESTNET_BANDWIDTH_VALUE: Bandwidth =
+    Bandwidth::new_unchecked(64 * 1024 * 1024 * 1024); // 64GB

gateway/src/node/client_handling/websocket/common_state.rs

@@ -0,0 +1,16 @@
+// Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
+// 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 CommonHandlerState {
+    pub(crate) coconut_verifier: Arc<CoconutVerifier>,
+    pub(crate) local_identity: Arc<identity::KeyPair>,
+    pub(crate) only_coconut_credentials: bool,
+    pub(crate) bandwidth_cfg: BandwidthFlushingBehaviourConfig,
+}

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};
 
@@ -42,6 +44,11 @@ pub enum RequestHandlingError {
     #[error("Internal gateway storage error")]
     StorageError(#[from] StorageError),
 
+    #[error(
+        "the database entry for bandwidth of the registered client {client_address} is missing!"
+    )]
+    MissingClientBandwidthEntry { client_address: String },
+
     #[error("Provided bandwidth IV is malformed - {0}")]
     MalformedIV(#[from] IVConversionError),
 
@@ -51,8 +58,8 @@ pub enum RequestHandlingError {
     #[error("Provided binary request was malformed - {0}")]
     InvalidTextRequest(<ClientControlRequest as TryFrom<String>>::Error),
 
-    #[error("The received request is not valid in the current context")]
-    IllegalRequest,
+    #[error("The received request is not valid in the current context: {additional_context}")]
+    IllegalRequest { additional_context: String },
 
     #[error("Provided bandwidth credential did not verify correctly on {0}")]
     InvalidBandwidthCredential(String),
@@ -90,9 +97,18 @@ pub enum RequestHandlingError {
     #[error("the provided credential did not contain a valid type attribute")]
     InvalidTypeAttribute,
 
+    #[error("insufficient bandwidth available to process the request. required: {required}B, available: {available}B")]
+    OutOfBandwidth { required: i64, available: i64 },
+
     #[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 +137,7 @@ impl IntoWSMessage for Result<ServerResponse, RequestHandlingError> {
 pub(crate) struct AuthenticatedHandler<R, S, St> {
     inner: FreshHandler<R, S, St>,
     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 +170,32 @@ 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<R, S, St>,
         client: ClientDetails,
         mix_receiver: MixMessageReceiver,
         is_active_request_receiver: IsActiveRequestReceiver,
-    ) -> Self {
-        AuthenticatedHandler {
+    ) -> Result<Self, RequestHandlingError> {
+        // 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::MissingClientBandwidthEntry {
+                client_address: client.address.as_base58_string(),
+            })?;
+
+        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 +205,10 @@ where
             .disconnect(self.client.address)
     }
 
-    /// Checks the amount of bandwidth available for the connected client.
-    async fn get_available_bandwidth(&self) -> Result<i64, RequestHandlingError> {
-        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,11 +216,27 @@ where
     /// # Arguments
     ///
     /// * `amount`: amount to increase the available bandwidth by.
-    async fn increase_bandwidth(&self, bandwidth: Bandwidth) -> Result<(), RequestHandlingError> {
+    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
+    }
+
+    async fn set_freepass_expiration(
+        &mut self,
+        expiration: OffsetDateTime,
+    ) -> Result<(), RequestHandlingError> {
+        self.client_bandwidth.bandwidth.freepass_expiration = Some(expiration);
         self.inner
             .storage
-            .increase_bandwidth(self.client.address, bandwidth.value() as i64)
+            .set_freepass_expiration(self.client.address, expiration)
             .await?;
+        self.client_bandwidth.update_flush_data();
         Ok(())
     }
 
@@ -204,11 +245,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 +280,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
@@ -247,12 +311,6 @@ where
             credential.data.epoch_id
         );
 
-        let aggregated_verification_key = self
-            .inner
-            .coconut_verifier
-            .verification_key(credential.data.epoch_id)
-            .await?;
-
         if !credential.data.validate_type_attribute() {
             trace!("mismatch in the type attribute");
             return Err(RequestHandlingError::InvalidTypeAttribute);
@@ -264,41 +322,52 @@ where
         };
 
         // this will extract token amounts out of bandwidth vouchers and validate expiry of free passes
-        let bandwidth = Bandwidth::try_from_raw_value(bandwidth_attribute, credential.data.typ)?;
+        let (raw_bandwidth, freepass_expiration) =
+            Bandwidth::parse_raw_bandwidth(bandwidth_attribute, credential.data.typ)?;
+
+        let bandwidth = Bandwidth::new(raw_bandwidth)?;
 
         trace!("embedded bandwidth: {bandwidth:?}");
 
         // locally verify the credential
-        let params = bandwidth_credential_params();
-        if !credential.data.verify(params, &aggregated_verification_key) {
-            trace!("the credential did not verify correctly");
-            return Err(RequestHandlingError::InvalidBandwidthCredential(
-                String::from("local credential verification has failed"),
-            ));
+        {
+            let aggregated_verification_key = self
+                .inner
+                .shared_state
+                .coconut_verifier
+                .verification_key(credential.data.epoch_id)
+                .await?;
+
+            let params = bandwidth_credential_params();
+            if !credential.data.verify(params, &aggregated_verification_key) {
+                trace!("the credential did not verify correctly");
+                return Err(RequestHandlingError::InvalidBandwidthCredential(
+                    String::from("local credential verification has failed"),
+                ));
+            }
         }
 
-        let was_freepass = match credential.data.typ {
+        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?;
-                false
             }
             CredentialType::FreePass => {
                 // no need to do anything special here, we already extracted the bandwidth amount and checked expiry
                 info!("received a free pass credential");
-
-                true
             }
-        };
+        }
 
         // technically this is not atomic, i.e. checking for the spending and then marking as spent,
         // but because we have the `UNIQUE` constraint on the database table
@@ -311,12 +380,21 @@ where
         trace!("storing serial number information");
         self.inner
             .storage
-            .insert_spent_credential(serial_number, was_freepass, self.client.address)
+            .insert_spent_credential(
+                serial_number,
+                freepass_expiration.is_some(),
+                self.client.address,
+            )
             .await?;
 
         trace!("increasing client bandwidth");
         self.increase_bandwidth(bandwidth).await?;
-        let available_total = self.get_available_bandwidth().await?;
+        // set free pass expiration
+        if let Some(expiration) = freepass_expiration {
+            self.set_freepass_expiration(expiration).await?;
+        }
+
+        let available_total = self.client_bandwidth.bandwidth.bytes;
 
         Ok(ServerResponse::Bandwidth { available_total })
     }
@@ -367,17 +445,47 @@ where
     ) -> Result<ServerResponse, RequestHandlingError> {
         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<i64, RequestHandlingError> {
+        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 {
+                required: required_bandwidth,
+                available: available_bandwidth,
+            });
+        }
+
+        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 +495,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<ServerResponse, RequestHandlingError> {
-        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 +513,7 @@ where
     /// # Arguments
     ///
     /// * `bin_msg`: raw message to handle.
-    async fn handle_binary(&self, bin_msg: Vec<u8>) -> Message {
+    async fn handle_binary(&mut self, bin_msg: Vec<u8>) -> 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) {
@@ -455,7 +555,13 @@ where
                     .handle_claim_testnet_bandwidth()
                     .await
                     .into_ws_message(),
-                _ => RequestHandlingError::IllegalRequest.into_error_message(),
+                other => RequestHandlingError::IllegalRequest {
+                    additional_context: format!(
+                        "received illegal message of type {} in an authenticated client",
+                        other.name()
+                    ),
+                }
+                .into_error_message(),
             },
         }
     }

gateway/src/node/client_handling/websocket/connection_handler/fresh.rs

@@ -1,6 +1,7 @@
 // Copyright 2021-2023 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: GPL-3.0-only
 
+use crate::error::RequestHandlingError;
 use futures::{
     channel::{mpsc, oneshot},
     SinkExt, StreamExt,
@@ -20,18 +21,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::common_state::CommonHandlerState;
+use crate::node::client_handling::websocket::connection_handler::AvailableBandwidth;
 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},
         },
@@ -77,24 +79,55 @@ pub(crate) enum InitialAuthenticationError {
 
     #[error("Attempted to negotiate connection with client using incompatible protocol version. Ours is {current} and the client reports {client:?}")]
     IncompatibleProtocol { client: Option<u8>, current: u8 },
+
+    #[error("failed to send authentication error response: {source}")]
+    ErrorResponseSendFailure {
+        #[source]
+        source: WsError,
+    },
+
+    #[error("failed to send authentication response: {source}")]
+    ResponseSendFailure {
+        #[source]
+        source: WsError,
+    },
+
+    #[error("possibly received a sphinx packet without prior authentication. Request is going to be ignored")]
+    BinaryRequestWithoutAuthentication,
+
+    #[error("received a connection close message")]
+    CloseMessage,
+
+    #[error("the connection has unexpectedly closed")]
+    ClosedConnection,
+
+    #[error("failed to obtain message from websocket stream: {source}")]
+    FailedToReadMessage {
+        #[source]
+        source: WsError,
+    },
+
+    #[error("could not establish client details")]
+    EmptyClientDetails,
+
+    #[error("failed to upgrade the client handler: {source}")]
+    HandlerUpgradeFailure { source: RequestHandlingError },
 }
 
 impl InitialAuthenticationError {
     /// Converts this Error into an appropriate websocket Message.
-    fn into_error_message(self) -> Message {
+    fn to_error_message(&self) -> Message {
         ServerResponse::new_error(self.to_string()).into()
     }
 }
 
 pub(crate) struct FreshHandler<R, S, St> {
     rng: R,
-    local_identity: Arc<identity::KeyPair>,
-    pub(crate) only_coconut_credentials: bool,
+    pub(crate) shared_state: CommonHandlerState,
     pub(crate) active_clients_store: ActiveClientsStore,
     pub(crate) outbound_mix_sender: MixForwardingSender,
     pub(crate) socket_connection: SocketStream<S>,
     pub(crate) storage: St,
-    pub(crate) coconut_verifier: Arc<CoconutVerifier>,
 
     // currently unused (but populated)
     pub(crate) negotiated_protocol: Option<u8>,
@@ -113,23 +146,19 @@ where
     pub(crate) fn new(
         rng: R,
         conn: S,
-        only_coconut_credentials: bool,
         outbound_mix_sender: MixForwardingSender,
-        local_identity: Arc<identity::KeyPair>,
         storage: St,
         active_clients_store: ActiveClientsStore,
-        coconut_verifier: Arc<CoconutVerifier>,
+        shared_state: CommonHandlerState,
     ) -> 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 +200,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 +465,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 +477,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 +538,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 +547,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 +571,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.
     ///
@@ -657,7 +699,7 @@ where
     // TODO: somehow cleanup this method
     pub(crate) async fn perform_initial_authentication(
         mut self,
-    ) -> Option<AuthenticatedHandler<R, S, St>>
+    ) -> Result<AuthenticatedHandler<R, S, St>, InitialAuthenticationError>
     where
         S: AsyncRead + AsyncWrite + Unpin + Send,
     {
@@ -666,40 +708,51 @@ where
         while let Some(msg) = self.read_websocket_message().await {
             let msg = match msg {
                 Ok(msg) => msg,
-                Err(err) => {
-                    debug!("failed to obtain message from websocket stream! stopping connection handler: {err}");
-                    break;
+                Err(source) => {
+                    debug!("failed to obtain message from websocket stream! stopping connection handler: {source}");
+                    return Err(InitialAuthenticationError::FailedToReadMessage { source });
                 }
             };
 
             if msg.is_close() {
-                break;
+                return Err(InitialAuthenticationError::CloseMessage);
             }
 
             // ONLY handle 'Authenticate' or 'Register' requests, ignore everything else
             match msg {
-                Message::Close(_) => break,
+                // we have explicitly checked for close message
+                Message::Close(_) => unreachable!(),
                 Message::Text(text_msg) => {
                     let (mix_sender, mix_receiver) = mpsc::unbounded();
-                    match self.handle_initial_authentication_request(text_msg).await {
+                    return match self.handle_initial_authentication_request(text_msg).await {
                         Err(err) => {
-                            if let Err(err) =
-                                self.send_websocket_message(err.into_error_message()).await
+                            debug!("authentication failure: {err}");
+
+                            // try to send error to the client
+                            if let Err(source) =
+                                self.send_websocket_message(err.to_error_message()).await
                             {
-                                debug!("Failed to send authentication error response - {err}");
-                                return None;
+                                debug!("Failed to send authentication error response: {source}");
+                                return Err(InitialAuthenticationError::ErrorResponseSendFailure {
+                                    source,
+                                });
                             }
+                            // return the underlying error
+                            Err(err)
                         }
                         Ok(auth_result) => {
-                            if let Err(err) = self
+                            // try to send auth response back to the client
+                            if let Err(source) = self
                                 .send_websocket_message(auth_result.server_response.into())
                                 .await
                             {
-                                debug!("Failed to send authentication response - {err}");
-                                return None;
+                                debug!("Failed to send authentication response: {source}");
+                                return Err(InitialAuthenticationError::ResponseSendFailure {
+                                    source,
+                                });
                             }
 
-                            return if let Some(client_details) = auth_result.client_details {
+                            if let Some(client_details) = auth_result.client_details {
                                 // Channel for handlers to ask other handlers if they are still active.
                                 let (is_active_request_sender, is_active_request_receiver) =
                                     mpsc::unbounded();
@@ -708,23 +761,29 @@ where
                                     mix_sender,
                                     is_active_request_sender,
                                 );
-                                Some(AuthenticatedHandler::upgrade(
+                                AuthenticatedHandler::upgrade(
                                     self,
                                     client_details,
                                     mix_receiver,
                                     is_active_request_receiver,
-                                ))
+                                )
+                                .await
+                                .map_err(|source| {
+                                    InitialAuthenticationError::HandlerUpgradeFailure { source }
+                                })
                             } else {
-                                None
-                            };
+                                // honestly, it's been so long I don't remember under what conditions its possible (if at all)
+                                // to have empty client details
+                                Err(InitialAuthenticationError::EmptyClientDetails)
+                            }
                         }
-                    }
+                    };
                 }
                 Message::Binary(_) => {
                     // perhaps logging level should be reduced here, let's leave it for now and see what happens
                     // if client is working correctly, this should have never happened
                     debug!("possibly received a sphinx packet without prior authentication. Request is going to be ignored");
-                    if let Err(err) = self
+                    if let Err(source) = self
                         .send_websocket_message(
                             ServerResponse::new_error(
                                 "binary request without prior authentication",
@@ -733,15 +792,18 @@ where
                         )
                         .await
                     {
-                        debug!("Failed to send error response during authentication: {err}",)
+                        return Err(InitialAuthenticationError::ErrorResponseSendFailure {
+                            source,
+                        });
                     }
-                    return None;
+                    return Err(InitialAuthenticationError::BinaryRequestWithoutAuthentication);
                 }
 
                 _ => continue,
             };
         }
-        None
+
+        Err(InitialAuthenticationError::ClosedConnection)
     }
 
     pub(crate) async fn start_handling(self, shutdown: nym_task::TaskClient)

gateway/src/node/client_handling/websocket/connection_handler/mod.rs

@@ -1,6 +1,7 @@
 // Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
 // 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};
@@ -103,12 +106,84 @@ pub(crate) async fn handle_connection<R, S, St>(
             trace!("received shutdown signal while performing initial authentication");
             return;
         }
-        Some(None) => {
-            warn!("authentication has failed");
+        Some(Err(err)) => {
+            warn!("authentication has failed: {err}");
             return;
         }
-        Some(Some(auth_handle)) => auth_handle.listen_for_requests(shutdown).await,
+        Some(Ok(auth_handle)) => auth_handle.listen_for_requests(shutdown).await,
     }
 
     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<OffsetDateTime>,
+}
+
+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;
+    }
+}

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::common_state::CommonHandlerState;
 use crate::node::client_handling::websocket::connection_handler::FreshHandler;
 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<identity::KeyPair>,
-    only_coconut_credentials: bool,
-    pub(crate) coconut_verifier: Arc<CoconutVerifier>,
+    shared_state: CommonHandlerState,
 }
 
 impl Listener {
-    pub(crate) fn new(
-        address: SocketAddr,
-        local_identity: Arc<identity::KeyPair>,
-        only_coconut_credentials: bool,
-        coconut_verifier: Arc<CoconutVerifier>,
-    ) -> Self {
+    pub(crate) fn new(address: SocketAddr, shared_state: CommonHandlerState) -> 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 });

gateway/src/node/client_handling/websocket/mod.rs

@@ -3,6 +3,9 @@
 
 pub(crate) use listener::Listener;
 
+pub(crate) mod common_state;
 pub(crate) mod connection_handler;
 pub(crate) mod listener;
 pub(crate) mod message_receiver;
+
+pub(crate) use common_state::CommonHandlerState;

gateway/src/node/mod.rs

@@ -254,13 +254,14 @@ impl<St> Gateway<St> {
             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::CommonHandlerState {
             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,

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<Option<PersistedBandwidth>, 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,

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<Option<i64>, StorageError>;
+    ) -> Result<Option<PersistedBandwidth>, 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<Option<i64>, 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<Option<PersistedBandwidth>, 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<Option<i64>, 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<Option<PersistedBandwidth>, StorageError> {
+        todo!()
+    }
+
+    async fn set_bandwidth(
         &self,
         _client_address: DestinationAddressBytes,
         _amount: i64,

gateway/src/node/storage/models.rs

@@ -1,7 +1,9 @@
 // Copyright 2021-2024 - Nym Technologies SA <contact@nymtech.net>
 // 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<u8>,
 }
 
+#[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<OffsetDateTime>,
+}
+
+impl From<PersistedBandwidth> for AvailableBandwidth {
+    fn from(value: PersistedBandwidth) -> Self {
+        AvailableBandwidth {
+            bytes: value.available,
+            freepass_expiration: value.freepass_expiration,
+        }
+    }
+}
+
+impl From<Option<PersistedBandwidth>> for AvailableBandwidth {
+    fn from(value: Option<PersistedBandwidth>) -> Self {
+        match value {
+            None => AvailableBandwidth::default(),
+            Some(b) => b.into(),
+        }
+    }
 }
 
 #[derive(Debug, Clone, FromRow)]

nym-node/src/cli/helpers.rs

@@ -195,6 +195,14 @@ pub(crate) struct MixnetArgs {
         env = NYMNODE_NYM_APIS_ARG
     )]
     pub(crate) nym_api_urls: Option<Vec<Url>>,
+
+    /// Addresses to nyxd chain endpoint which the node will use for chain interactions.
+    #[clap(
+        long,
+        value_delimiter = ',',
+        env = NYMNODE_NYXD_URLS_ARG
+    )]
+    pub(crate) nyxd_urls: Option<Vec<Url>>,
 }
 
 impl MixnetArgs {
@@ -210,6 +218,9 @@ impl MixnetArgs {
         if let Some(nym_api_urls) = self.nym_api_urls {
             section.nym_api_urls = nym_api_urls
         }
+        if let Some(nyxd_urls) = self.nyxd_urls {
+            section.nyxd_urls = nyxd_urls
+        }
         section
     }
 }

nym-node/src/env.rs

@@ -34,6 +34,7 @@ pub mod vars {
     // mixnet:
     pub const NYMNODE_MIXNET_BIND_ADDRESS_ARG: &str = "NYMNODE_MIXNET_BIND_ADDRESS";
     pub const NYMNODE_NYM_APIS_ARG: &str = "NYMNODE_NYM_APIS";
+    pub const NYMNODE_NYXD_URLS_ARG: &str = "NYMNODE_NYXD";
 
     // wireguard:
     pub const NYMNODE_WG_ENABLED_ARG: &str = "NYMNODE_WG_ENABLED";

service-providers/ip-packet-router/src/cli/import_credential.rs

@@ -10,5 +10,7 @@ use nym_ip_packet_router::error::IpPacketRouterError;
 pub(crate) async fn execute(
     args: CommonClientImportCredentialArgs,
 ) -> Result<(), IpPacketRouterError> {
-    import_credential::<CliIpPacketRouterClient, _>(args).await
+    import_credential::<CliIpPacketRouterClient, _>(args).await?;
+    println!("successfully imported credential!");
+    Ok(())
 }

service-providers/network-requester/src/cli/import_credential.rs

@@ -10,5 +10,7 @@ use nym_client_core::cli_helpers::client_import_credential::{
 pub(crate) async fn execute(
     args: CommonClientImportCredentialArgs,
 ) -> Result<(), NetworkRequesterError> {
-    import_credential::<CliNetworkRequesterClient, _>(args).await
+    import_credential::<CliNetworkRequesterClient, _>(args).await?;
+    println!("successfully imported credential!");
+    Ok(())
 }