expire freepass bandwidth + stagger db flushes

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/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/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<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 +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<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::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<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,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<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);
+        }
+
+        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<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 +481,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) {

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<R, S, St> {
     rng: R,
-    local_identity: Arc<identity::KeyPair>,
-    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<S>,
     pub(crate) storage: St,
-    pub(crate) coconut_verifier: Arc<CoconutVerifier>,
 
     // currently unused (but populated)
     pub(crate) negotiated_protocol: Option<u8>,
@@ -113,23 +112,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: 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
                             };

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};
@@ -112,3 +115,75 @@ pub(crate) async fn handle_connection<R, S, St>(
 
     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::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<identity::KeyPair>,
-    only_coconut_credentials: bool,
-    pub(crate) coconut_verifier: Arc<CoconutVerifier>,
+    shared_state: SharedHandlerState,
 }
 
 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: 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 });

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;

gateway/src/node/client_handling/websocket/shared_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 SharedHandlerState {
+    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/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::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,

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)]