Feature/persistent gateway storage (#784)

* Sqlx struct stub

* Initial schema

* Initial error enum

* Managed for persisted shared keys

* Initial inbox manager

* Comments

* Using new database in clients handler

* Extending gateway storage API

* tokio::main + placeholder values

* Removed old client store

* Simplified logic of async packet processing

* Renamed table + not null restriction

* BandwidthManager

* Removed sled dependency

* Using centralised storage for bandwidth

* Dead code removal

* WIP connection_handler split and simplification

Maybe it doesn't look like it right now, but once completed it will remove bunch of redundant checks for Nones etc

* Further more explicit clients handler split

* Minor cleanup

* Temporary store for active client handles

* Fixed error types

* Error trait on iv and encrypted address

* Authentication and registration moved to the handler

* Removal of clients handler

* Further logic simplification + returned explicit bandwidth values

* Further cleanup and comments

* Updated config with relevant changes

* Basic bandwidth tracking in client

* FreshHandle doc comments + fixed stagger issue

* Removed side-effects from .map

* More doc comments

* Database migration on build

* Increased default claimed bandwidth

* Renaming

* Fixed client determining available bandwidth

* Removed dead sql table that might be used in the future

* Windows workaround

* Comment

* Return error rather than cap credential

Cargo.lock

@@ -3354,7 +3354,8 @@ dependencies = [
  "pretty_env_logger",
  "rand 0.7.3",
  "serde",
- "sled",
+ "sqlx",
+ "thiserror",
  "tokio",
  "tokio-stream",
  "tokio-tungstenite",

clients/client-core/src/client/mix_traffic.rs

@@ -40,7 +40,7 @@ impl MixTrafficController {
     async fn on_messages(&mut self, mut mix_packets: Vec<MixPacket>) {
         debug_assert!(!mix_packets.is_empty());
 
-        let success = if mix_packets.len() == 1 {
+        let result = if mix_packets.len() == 1 {
             let mix_packet = mix_packets.pop().unwrap();
             self.gateway_client.send_mix_packet(mix_packet).await
         } else {
@@ -49,7 +49,7 @@ impl MixTrafficController {
                 .await
         };
 
-        match success {
+        match result {
             Err(e) => {
                 error!("Failed to send sphinx packet(s) to the gateway! - {:?}", e);
                 self.consecutive_gateway_failure_count += 1;

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

@@ -36,8 +36,7 @@ const DEFAULT_RECONNECTION_BACKOFF: Duration = Duration::from_secs(5);
 
 pub struct GatewayClient {
     authenticated: bool,
-    // TODO: This should be replaced by an actual bandwidth value, with 0 meaning no bandwidth
-    has_bandwidth: bool,
+    bandwidth_remaining: i64,
     gateway_address: String,
     gateway_identity: identity::PublicKey,
     local_identity: Arc<identity::KeyPair>,
@@ -72,7 +71,7 @@ impl GatewayClient {
     ) -> Self {
         GatewayClient {
             authenticated: false,
-            has_bandwidth: false,
+            bandwidth_remaining: 0,
             gateway_address,
             gateway_identity,
             local_identity,
@@ -117,7 +116,7 @@ impl GatewayClient {
 
         GatewayClient {
             authenticated: false,
-            has_bandwidth: false,
+            bandwidth_remaining: 0,
             gateway_address,
             gateway_identity,
             local_identity,
@@ -474,14 +473,21 @@ impl GatewayClient {
         )
         .ok_or(GatewayClientError::SerializeCredential)?
         .into();
-        self.has_bandwidth = match self.send_websocket_message(msg).await? {
-            ServerResponse::Bandwidth { status } => Ok(status),
+        self.bandwidth_remaining = match self.send_websocket_message(msg).await? {
+            ServerResponse::Bandwidth { available_total } => Ok(available_total),
             ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)),
             _ => Err(GatewayClientError::UnexpectedResponse),
         }?;
         Ok(())
     }
 
+    fn estimate_required_bandwidth(&self, packets: &[MixPacket]) -> i64 {
+        packets
+            .iter()
+            .map(|packet| packet.sphinx_packet().len())
+            .sum::<usize>() as i64
+    }
+
     pub async fn batch_send_mix_packets(
         &mut self,
         packets: Vec<MixPacket>,
@@ -489,7 +495,7 @@ impl GatewayClient {
         if !self.authenticated {
             return Err(GatewayClientError::NotAuthenticated);
         }
-        if !self.has_bandwidth {
+        if self.estimate_required_bandwidth(&packets) < self.bandwidth_remaining {
             return Err(GatewayClientError::NotEnoughBandwidth);
         }
         if !self.connection.is_established() {
@@ -550,7 +556,7 @@ impl GatewayClient {
         if !self.authenticated {
             return Err(GatewayClientError::NotAuthenticated);
         }
-        if !self.has_bandwidth {
+        if (mix_packet.sphinx_packet().len() as i64) > self.bandwidth_remaining {
             return Err(GatewayClientError::NotEnoughBandwidth);
         }
         if !self.connection.is_established() {
@@ -598,7 +604,7 @@ impl GatewayClient {
         if !self.authenticated {
             return Err(GatewayClientError::NotAuthenticated);
         }
-        if !self.has_bandwidth {
+        if self.bandwidth_remaining <= 0 {
             return Err(GatewayClientError::NotEnoughBandwidth);
         }
         if self.connection.is_partially_delegated() {

common/credentials/src/bandwidth.rs

@@ -12,7 +12,7 @@ use crate::error::Error;
 use crate::utils::{obtain_aggregate_signature, prepare_credential_for_spending};
 use coconut_interface::{hash_to_scalar, Credential, Parameters, Signature, VerificationKey};
 
-const BANDWIDTH_VALUE: u64 = 1024 * 1024; // 1 MB
+const BANDWIDTH_VALUE: u64 = 10 * 1024 * 1024 * 1024; // 10 GB
 
 pub const PUBLIC_ATTRIBUTES: u32 = 1;
 pub const PRIVATE_ATTRIBUTES: u32 = 1;

gateway/Cargo.toml

@@ -20,12 +20,13 @@ log = "0.4"
 pretty_env_logger = "0.4"
 rand = "0.7"
 serde = { version = "1.0.104", features = ["derive"] }
-sled = "0.34"
+thiserror = "1"
 tokio = { version = "1.4", features = [ "rt-multi-thread", "net", "signal", "fs" ] }
 tokio-util = { version = "0.6", features = [ "codec" ] }
 tokio-stream = { version = "0.1", features = [ "fs" ] }
 tokio-tungstenite = "0.14"
 url = { version = "2.2", features = [ "serde" ] }
+sqlx = { version = "0.5", features = ["runtime-tokio-rustls", "sqlite", "macros", "migrate"] }
 
 # internal
 coconut-interface = { path = "../common/coconut-interface" }
@@ -39,3 +40,7 @@ nymsphinx = { path = "../common/nymsphinx" }
 pemstore = { path = "../common/pemstore" }
 validator-client = { path = "../common/client-libs/validator-client" }
 version-checker = { path = "../common/version-checker" }
+
+[build-dependencies]
+tokio = { version = "1.4", features = ["rt-multi-thread", "macros"] }
+sqlx = { version = "0.5", features = ["runtime-tokio-rustls", "sqlite", "macros", "migrate"] }

gateway/build.rs

@@ -0,0 +1,25 @@
+use sqlx::{Connection, SqliteConnection};
+use std::env;
+
+#[tokio::main]
+async fn main() {
+    let out_dir = env::var("OUT_DIR").unwrap();
+    let database_path = format!("{}/gateway-example.sqlite", out_dir);
+
+    let mut conn = SqliteConnection::connect(&*format!("sqlite://{}?mode=rwc", database_path))
+        .await
+        .expect("Failed to create SQLx database connection");
+
+    sqlx::migrate!("./migrations")
+        .run(&mut conn)
+        .await
+        .expect("Failed to perform SQLx migrations");
+
+    #[cfg(target_family = "unix")]
+    println!("cargo:rustc-env=DATABASE_URL=sqlite://{}", &database_path);
+
+    #[cfg(target_family = "windows")]
+    // for some strange reason we need to add a leading `/` to the windows path even though it's
+    // not a valid windows path... but hey, it works...
+    println!("cargo:rustc-env=DATABASE_URL=sqlite:///{}", &database_path);
+}

gateway/gateway-requests/src/authentication/encrypted_address.rs

@@ -6,6 +6,7 @@ use crate::registration::handshake::shared_key::SharedKeys;
 use crypto::symmetric::stream_cipher;
 use nymsphinx::params::GatewayEncryptionAlgorithm;
 use nymsphinx::{DestinationAddressBytes, DESTINATION_ADDRESS_LENGTH};
+use thiserror::Error;
 
 pub const ENCRYPTED_ADDRESS_SIZE: usize = DESTINATION_ADDRESS_LENGTH;
 
@@ -16,9 +17,11 @@ pub const ENCRYPTED_ADDRESS_SIZE: usize = DESTINATION_ADDRESS_LENGTH;
 #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
 pub struct EncryptedAddressBytes([u8; ENCRYPTED_ADDRESS_SIZE]);
 
-#[derive(Debug)]
+#[derive(Debug, Error)]
 pub enum EncryptedAddressConversionError {
-    DecodeError(bs58::decode::Error),
+    #[error("Failed to decode the encrypted address - {0}")]
+    DecodeError(#[from] bs58::decode::Error),
+    #[error("The decoded address has invalid length")]
     StringOfInvalidLengthError,
 }
 
@@ -54,10 +57,7 @@ impl EncryptedAddressBytes {
     pub fn try_from_base58_string<S: Into<String>>(
         val: S,
     ) -> Result<Self, EncryptedAddressConversionError> {
-        let decoded = match bs58::decode(val.into()).into_vec() {
-            Ok(decoded) => decoded,
-            Err(err) => return Err(EncryptedAddressConversionError::DecodeError(err)),
-        };
+        let decoded = bs58::decode(val.into()).into_vec()?;
 
         if decoded.len() != ENCRYPTED_ADDRESS_SIZE {
             return Err(EncryptedAddressConversionError::StringOfInvalidLengthError);

gateway/gateway-requests/src/iv.rs

@@ -5,6 +5,7 @@ use crypto::generic_array::{typenum::Unsigned, GenericArray};
 use crypto::symmetric::stream_cipher::{random_iv, NewCipher, IV as CryptoIV};
 use nymsphinx::params::GatewayEncryptionAlgorithm;
 use rand::{CryptoRng, RngCore};
+use thiserror::Error;
 
 type NonceSize = <GatewayEncryptionAlgorithm as NewCipher>::NonceSize;
 
@@ -12,12 +13,17 @@ type NonceSize = <GatewayEncryptionAlgorithm as NewCipher>::NonceSize;
 #[allow(clippy::upper_case_acronyms)]
 pub struct IV(CryptoIV<GatewayEncryptionAlgorithm>);
 
-#[derive(Debug)]
+#[derive(Error, Debug)]
 // I think 'IV' looks better than 'Iv', feel free to change that.
 #[allow(clippy::upper_case_acronyms)]
 pub enum IVConversionError {
-    DecodeError(bs58::decode::Error),
+    #[error("Failed to decode the iv - {0}")]
+    DecodeError(#[from] bs58::decode::Error),
+
+    #[error("The decoded bytes iv has invalid length")]
     BytesOfInvalidLengthError,
+
+    #[error("The decoded string iv has invalid length")]
     StringOfInvalidLengthError,
 }
 
@@ -47,10 +53,7 @@ impl IV {
     }
 
     pub fn try_from_base58_string<S: Into<String>>(val: S) -> Result<Self, IVConversionError> {
-        let decoded = match bs58::decode(val.into()).into_vec() {
-            Ok(decoded) => decoded,
-            Err(err) => return Err(IVConversionError::DecodeError(err)),
-        };
+        let decoded = bs58::decode(val.into()).into_vec()?;
 
         if decoded.len() != NonceSize::to_usize() {
             return Err(IVConversionError::StringOfInvalidLengthError);

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

@@ -22,7 +22,7 @@ type EncryptionKeySize = <GatewayEncryptionAlgorithm as NewCipher>::KeySize;
 /// Shared key used when computing MAC for messages exchanged between client and its gateway.
 pub type MacKey = GenericArray<u8, MacKeySize>;
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Copy, Debug)]
 pub struct SharedKeys {
     encryption_key: CipherKey<GatewayEncryptionAlgorithm>,
     mac_key: MacKey,

gateway/gateway-requests/src/types.rs

@@ -89,6 +89,8 @@ impl fmt::Display for GatewayRequestsError {
     }
 }
 
+impl std::error::Error for GatewayRequestsError {}
+
 impl From<NymNodeRoutingAddressError> for GatewayRequestsError {
     fn from(_: NymNodeRoutingAddressError) -> Self {
         GatewayRequestsError::IncorrectlyEncodedAddress
@@ -191,9 +193,8 @@ impl TryInto<String> for ClientControlRequest {
 pub enum ServerResponse {
     Authenticate { status: bool },
     Register { status: bool },
-    // Maybe we could return the remaining bandwidth?
-    Bandwidth { status: bool },
-    Send { status: bool },
+    Bandwidth { available_total: i64 },
+    Send { remaining_bandwidth: i64 },
     Error { message: String },
 }
 

gateway/migrations/20210921120000_create_initial_tables.sql

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+CREATE TABLE shared_keys
+(
+    client_address_bs58                      TEXT NOT NULL PRIMARY KEY UNIQUE,
+    derived_aes128_ctr_blake3_hmac_keys_bs58 TEXT NOT NULL
+);
+
+CREATE TABLE message_store
+(
+    id                  INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+    client_address_bs58 TEXT    NOT NULL,
+    content             BLOB    NOT NULL
+);
+
+CREATE TABLE available_bandwidth
+(
+    client_address_bs58 TEXT    NOT NULL PRIMARY KEY UNIQUE,
+    available           INTEGER NOT NULL
+);
+
+CREATE INDEX `message_store_index` ON `message_store` (`client_address_bs58`, `content`);

gateway/src/commands/init.rs

@@ -44,15 +44,9 @@ pub fn command_args<'a, 'b>() -> clap::App<'a, 'b> {
                 .takes_value(true),
         )
         .arg(
-            Arg::with_name(INBOXES_ARG_NAME)
-                .long(INBOXES_ARG_NAME)
-                .help("Directory with inboxes where all packets for the clients are stored")
-                .takes_value(true)
-        )
-        .arg(
-            Arg::with_name(CLIENTS_LEDGER_ARG_NAME)
-                .long(CLIENTS_LEDGER_ARG_NAME)
-                .help("Ledger file containing registered clients")
+            Arg::with_name(DATASTORE_PATH)
+                .long(DATASTORE_PATH)
+                .help("Path to sqlite database containing all gateway persistent data")
                 .takes_value(true)
         )
         .arg(
@@ -115,7 +109,7 @@ fn show_bonding_info(config: &Config) {
     );
 }
 
-pub fn execute(matches: &ArgMatches) {
+pub async fn execute(matches: ArgMatches<'static>) {
     let id = matches.value_of(ID_ARG_NAME).unwrap();
     println!("Initialising gateway {}...", id);
 
@@ -128,7 +122,7 @@ pub fn execute(matches: &ArgMatches) {
 
     let mut config = Config::new(id);
 
-    config = override_config(config, matches);
+    config = override_config(config, &matches);
 
     // if gateway was already initialised, don't generate new keys
     if !already_init {

gateway/src/commands/mod.rs

@@ -15,8 +15,7 @@ pub(crate) const MIX_PORT_ARG_NAME: &str = "mix-port";
 pub(crate) const CLIENTS_PORT_ARG_NAME: &str = "clients-port";
 pub(crate) const VALIDATORS_ARG_NAME: &str = "validators";
 pub(crate) const ANNOUNCE_HOST_ARG_NAME: &str = "announce-host";
-pub(crate) const INBOXES_ARG_NAME: &str = "inboxes";
-pub(crate) const CLIENTS_LEDGER_ARG_NAME: &str = "clients-ledger";
+pub(crate) const DATASTORE_PATH: &str = "datastore";
 
 fn parse_validators(raw: &str) -> Vec<Url> {
     raw.split(',')
@@ -69,12 +68,8 @@ pub(crate) fn override_config(mut config: Config, matches: &ArgMatches) -> Confi
         config = config.with_custom_validator_apis(parse_validators(raw_validators));
     }
 
-    if let Some(inboxes_dir) = matches.value_of(INBOXES_ARG_NAME) {
-        config = config.with_custom_clients_inboxes(inboxes_dir);
-    }
-
-    if let Some(clients_ledger) = matches.value_of(CLIENTS_LEDGER_ARG_NAME) {
-        config = config.with_custom_clients_ledger(clients_ledger);
+    if let Some(datastore_path) = matches.value_of(DATASTORE_PATH) {
+        config = config.with_custom_persistent_store(datastore_path);
     }
 
     config

gateway/src/commands/run.rs

@@ -47,15 +47,9 @@ pub fn command_args<'a, 'b>() -> clap::App<'a, 'b> {
                 .takes_value(true),
         )
         .arg(
-            Arg::with_name(INBOXES_ARG_NAME)
-                .long(INBOXES_ARG_NAME)
-                .help("Directory with inboxes where all packets for the clients are stored")
-                .takes_value(true)
-        )
-        .arg(
-            Arg::with_name(CLIENTS_LEDGER_ARG_NAME)
-                .long(CLIENTS_LEDGER_ARG_NAME)
-                .help("Ledger file containing registered clients")
+            Arg::with_name(DATASTORE_PATH)
+                .long(DATASTORE_PATH)
+                .help("Path to sqlite database containing all gateway persistent data")
                 .takes_value(true)
         )
         .arg(
@@ -126,7 +120,7 @@ fn version_check(cfg: &Config) -> bool {
     }
 }
 
-pub fn execute(matches: &ArgMatches) {
+pub async fn execute(matches: ArgMatches<'static>) {
     let id = matches.value_of(ID_ARG_NAME).unwrap();
 
     println!("Starting gateway {}...", id);
@@ -139,7 +133,7 @@ pub fn execute(matches: &ArgMatches) {
         }
     };
 
-    config = override_config(config, matches);
+    config = override_config(config, &matches);
 
     if !version_check(&config) {
         error!("failed the local version check");
@@ -168,15 +162,10 @@ pub fn execute(matches: &ArgMatches) {
         config.get_announce_address()
     );
 
-    println!(
-        "Inboxes directory is: {:?}",
-        config.get_clients_inboxes_dir()
-    );
+    println!("Data store is at: {:?}", config.get_persistent_store_path());
 
-    println!(
-        "Clients ledger is stored at: {:?}",
-        config.get_clients_ledger_path()
-    );
-
-    Gateway::new(config, sphinx_keypair, identity).run();
+    Gateway::new(config, sphinx_keypair, identity)
+        .await
+        .run()
+        .await;
 }

gateway/src/commands/upgrade.rs

@@ -148,7 +148,7 @@ fn do_upgrade(mut config: Config, matches: &ArgMatches, package_version: Version
     }
 }
 
-pub fn execute(matches: &ArgMatches) {
+pub async fn execute(matches: ArgMatches<'static>) {
     let package_version = parse_package_version();
 
     let id = matches.value_of(ID_ARG_NAME).unwrap();
@@ -163,5 +163,5 @@ pub fn execute(matches: &ArgMatches) {
         process::exit(1);
     }
 
-    do_upgrade(existing_config, matches, package_version)
+    do_upgrade(existing_config, &matches, package_version)
 }

gateway/src/config/mod.rs

@@ -25,7 +25,7 @@ const DEFAULT_INITIAL_CONNECTION_TIMEOUT: Duration = Duration::from_millis(1_500
 const DEFAULT_MAXIMUM_CONNECTION_BUFFER_SIZE: usize = 128;
 
 const DEFAULT_STORED_MESSAGE_FILENAME_LENGTH: u16 = 16;
-const DEFAULT_MESSAGE_RETRIEVAL_LIMIT: u16 = 5;
+const DEFAULT_MESSAGE_RETRIEVAL_LIMIT: i64 = 100;
 
 pub fn missing_string_value() -> String {
     MISSING_VALUE.to_string()
@@ -47,8 +47,6 @@ fn default_clients_port() -> u16 {
 pub struct Config {
     gateway: Gateway,
 
-    clients_endpoint: ClientsEndpoint,
-
     #[serde(default)]
     logging: Logging,
     #[serde(default)]
@@ -115,17 +113,9 @@ impl Config {
             self.gateway.public_identity_key_file =
                 self::Gateway::default_public_identity_key_file(&id);
         }
-        if self
-            .clients_endpoint
-            .inboxes_directory
-            .as_os_str()
-            .is_empty()
-        {
-            self.clients_endpoint.inboxes_directory =
-                self::ClientsEndpoint::default_inboxes_directory(&id);
-        }
-        if self.clients_endpoint.ledger_path.as_os_str().is_empty() {
-            self.clients_endpoint.ledger_path = self::ClientsEndpoint::default_ledger_path(&id);
+
+        if self.gateway.persistent_storage.as_os_str().is_empty() {
+            self.gateway.persistent_storage = self::Gateway::default_database_path(&id)
         }
 
         self.gateway.id = id;
@@ -170,13 +160,8 @@ impl Config {
         self
     }
 
-    pub fn with_custom_clients_inboxes<S: Into<String>>(mut self, inboxes_dir: S) -> Self {
-        self.clients_endpoint.inboxes_directory = PathBuf::from(inboxes_dir.into());
-        self
-    }
-
-    pub fn with_custom_clients_ledger<S: Into<String>>(mut self, ledger_path: S) -> Self {
-        self.clients_endpoint.ledger_path = PathBuf::from(ledger_path.into());
+    pub fn with_custom_persistent_store<S: Into<String>>(mut self, store_dir: S) -> Self {
+        self.gateway.persistent_storage = PathBuf::from(store_dir.into());
         self
     }
 
@@ -226,12 +211,8 @@ impl Config {
         self.gateway.clients_port
     }
 
-    pub fn get_clients_inboxes_dir(&self) -> PathBuf {
-        self.clients_endpoint.inboxes_directory.clone()
-    }
-
-    pub fn get_clients_ledger_path(&self) -> PathBuf {
-        self.clients_endpoint.ledger_path.clone()
+    pub fn get_persistent_store_path(&self) -> PathBuf {
+        self.gateway.persistent_storage.clone()
     }
 
     pub fn get_packet_forwarding_initial_backoff(&self) -> Duration {
@@ -250,14 +231,10 @@ impl Config {
         self.debug.maximum_connection_buffer_size
     }
 
-    pub fn get_message_retrieval_limit(&self) -> u16 {
+    pub fn get_message_retrieval_limit(&self) -> i64 {
         self.debug.message_retrieval_limit
     }
 
-    pub fn get_stored_messages_filename_length(&self) -> u16 {
-        self.debug.stored_messages_filename_length
-    }
-
     pub fn get_version(&self) -> &str {
         &self.gateway.version
     }
@@ -310,6 +287,10 @@ pub struct Gateway {
     /// nym_home_directory specifies absolute path to the home nym gateways directory.
     /// It is expected to use default value and hence .toml file should not redefine this field.
     nym_root_directory: PathBuf,
+
+    /// Path to sqlite database containing all persistent data: messages for offline clients,
+    /// derived shared keys and available client bandwidths.
+    persistent_storage: PathBuf,
 }
 
 impl Gateway {
@@ -328,6 +309,10 @@ impl Gateway {
     fn default_public_identity_key_file(id: &str) -> PathBuf {
         Config::default_data_directory(Some(id)).join("public_identity.pem")
     }
+
+    fn default_database_path(id: &str) -> PathBuf {
+        Config::default_data_directory(Some(id)).join("db.sqlite")
+    }
 }
 
 impl Default for Gateway {
@@ -345,35 +330,7 @@ impl Default for Gateway {
             public_sphinx_key_file: Default::default(),
             validator_api_urls: default_api_endpoints(),
             nym_root_directory: Config::default_root_directory(),
-        }
-    }
-}
-
-#[derive(Debug, Deserialize, PartialEq, Serialize)]
-pub struct ClientsEndpoint {
-    /// Path to the directory with clients inboxes containing messages stored for them.
-    inboxes_directory: PathBuf,
-
-    /// Full path to a file containing mapping of
-    /// client addresses to their access tokens.
-    ledger_path: PathBuf,
-}
-
-impl ClientsEndpoint {
-    fn default_inboxes_directory(id: &str) -> PathBuf {
-        Config::default_data_directory(Some(id)).join("inboxes")
-    }
-
-    fn default_ledger_path(id: &str) -> PathBuf {
-        Config::default_data_directory(Some(id)).join("client_ledger.sled")
-    }
-}
-
-impl Default for ClientsEndpoint {
-    fn default() -> Self {
-        ClientsEndpoint {
-            inboxes_directory: Default::default(),
-            ledger_path: Default::default(),
+            persistent_storage: Default::default(),
         }
     }
 }
@@ -415,10 +372,8 @@ pub struct Debug {
     /// Length of filenames for new client messages.
     stored_messages_filename_length: u16,
 
-    /// Number of messages client gets on each request
-    /// if there are no real messages, dummy ones are created to always return  
-    /// `message_retrieval_limit` total messages
-    message_retrieval_limit: u16,
+    /// Number of messages from offline client that can be pulled at once from the storage.
+    message_retrieval_limit: i64,
 }
 
 impl Default for Debug {

gateway/src/config/template.rs

@@ -62,16 +62,9 @@ validator_api_urls = [
 # It is expected to use default value and hence .toml file should not redefine this field.
 nym_root_directory = '{{ gateway.nym_root_directory }}'
 
-#### Clients endpoint config options #####
-
-[clients_endpoint]
-
-# Path to the directory with clients inboxes containing messages stored for them.
-inboxes_directory = '{{ clients_endpoint.inboxes_directory }}'
-
-# Full path to a file containing mapping of client addresses to their access tokens.
-ledger_path = '{{ clients_endpoint.ledger_path }}'
-
+# Path to sqlite database containing all persistent data: messages for offline clients,
+# derived shared keys and available client bandwidths.
+persistent_storage = '{{ gateway.persistent_storage }}'
 
 ##### logging configuration options #####
 

gateway/src/main.rs

@@ -7,7 +7,8 @@ mod commands;
 mod config;
 mod node;
 
-fn main() {
+#[tokio::main]
+async fn main() {
     dotenv::dotenv().ok();
     setup_logging();
     println!("{}", banner());
@@ -21,14 +22,14 @@ fn main() {
         .subcommand(commands::upgrade::command_args())
         .get_matches();
 
-    execute(arg_matches);
+    execute(arg_matches).await;
 }
 
-fn execute(matches: ArgMatches) {
+async fn execute(matches: ArgMatches<'static>) {
     match matches.subcommand() {
-        ("init", Some(m)) => commands::init::execute(m),
-        ("run", Some(m)) => commands::run::execute(m),
-        ("upgrade", Some(m)) => commands::upgrade::execute(m),
+        ("init", Some(m)) => commands::init::execute(m.clone()).await,
+        ("run", Some(m)) => commands::run::execute(m.clone()).await,
+        ("upgrade", Some(m)) => commands::upgrade::execute(m.clone()).await,
         _ => println!("{}", usage()),
     }
 }

gateway/src/node/client_handling/active_clients.rs

@@ -0,0 +1,58 @@
+// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::node::client_handling::websocket::message_receiver::MixMessageSender;
+use dashmap::DashMap;
+use nymsphinx::DestinationAddressBytes;
+use std::sync::Arc;
+
+#[derive(Clone)]
+pub(crate) struct ActiveClientsStore(Arc<DashMap<DestinationAddressBytes, MixMessageSender>>);
+
+impl ActiveClientsStore {
+    /// Creates new instance of `ActiveClientsStore` to store in-memory handles to all currently connected clients.
+    pub(crate) fn new() -> Self {
+        ActiveClientsStore(Arc::new(DashMap::new()))
+    }
+
+    /// Tries to obtain sending channel to specified client. Note that if stale entry existed, it is
+    /// removed and a `None` is returned instead.
+    ///
+    /// # Arguments
+    ///
+    /// * `client`: address of the client for which to obtain the handle.
+    pub(crate) fn get(&self, client: DestinationAddressBytes) -> Option<MixMessageSender> {
+        let entry = self.0.get(&client)?;
+        let handle = entry.value();
+
+        // if the entry is stale, remove it from the map
+        // if handle.is_valid() {
+        if !handle.is_closed() {
+            Some(handle.clone())
+        } else {
+            // drop the reference to the map to prevent deadlocks
+            drop(entry);
+            self.0.remove(&client);
+            None
+        }
+    }
+
+    /// Indicates particular client has disconnected from the gateway and its handle should get removed.
+    ///
+    /// # Arguments
+    ///
+    /// * `client`: address of the client for which to remove the handle.
+    pub(crate) fn disconnect(&self, client: DestinationAddressBytes) {
+        self.0.remove(&client);
+    }
+
+    /// Insert new client handle into the store.
+    ///
+    /// # Arguments
+    ///
+    /// * `client`: address of the client for which to insert the handle.
+    /// * `handle`: the sender channel for all mix packets to be pushed back onto the websocket
+    pub(crate) fn insert(&self, client: DestinationAddressBytes, handle: MixMessageSender) {
+        self.0.insert(client, handle);
+    }
+}

gateway/src/node/client_handling/bandwidth.rs

@@ -1,23 +1,13 @@
 // Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use std::collections::HashMap;
 use std::convert::TryFrom;
-use std::sync::Arc;
-use tokio::sync::RwLock;
 
 use coconut_interface::Credential;
 use credentials::error::Error;
-use nymsphinx::DestinationAddressBytes;
 
 const BANDWIDTH_INDEX: usize = 0;
 
-pub type BandwidthDatabase = Arc<RwLock<HashMap<DestinationAddressBytes, u64>>>;
-
-pub fn empty_bandwidth_database() -> BandwidthDatabase {
-    Arc::new(RwLock::new(HashMap::new()))
-}
-
 pub struct Bandwidth {
     value: u64,
 }
@@ -26,45 +16,6 @@ impl Bandwidth {
     pub fn value(&self) -> u64 {
         self.value
     }
-
-    pub async fn consume_bandwidth(
-        bandwidths: &BandwidthDatabase,
-        remote_address: &DestinationAddressBytes,
-        consumed: u64,
-    ) -> Result<(), Error> {
-        if let Some(bandwidth) = bandwidths.write().await.get_mut(remote_address) {
-            if let Some(res) = bandwidth.checked_sub(consumed) {
-                *bandwidth = res;
-                Ok(())
-            } else {
-                Err(Error::BandwidthOverflow(String::from(
-                    "Allocate more bandwidth for consumption",
-                )))
-            }
-        } else {
-            Err(Error::MissingBandwidth)
-        }
-    }
-
-    pub async fn increase_bandwidth(
-        bandwidths: &BandwidthDatabase,
-        remote_address: &DestinationAddressBytes,
-        increase: u64,
-    ) -> Result<(), Error> {
-        let mut db = bandwidths.write().await;
-        if let Some(bandwidth) = db.get_mut(remote_address) {
-            if let Some(new_bandwidth) = bandwidth.checked_add(increase) {
-                *bandwidth = new_bandwidth;
-            } else {
-                return Err(Error::BandwidthOverflow(String::from(
-                    "Use some of the already allocated bandwidth",
-                )));
-            }
-        } else {
-            db.insert(*remote_address, increase);
-        }
-        Ok(())
-    }
 }
 
 impl TryFrom<Credential> for Bandwidth {

gateway/src/node/client_handling/clients_handler.rs

@@ -1,316 +0,0 @@
-// Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
-// SPDX-License-Identifier: Apache-2.0
-
-use crate::node::{
-    client_handling::websocket::message_receiver::MixMessageSender,
-    storage::{inboxes::ClientStorage, ClientLedger},
-};
-use futures::{
-    channel::{mpsc, oneshot},
-    StreamExt,
-};
-use gateway_requests::authentication::encrypted_address::EncryptedAddressBytes;
-use gateway_requests::iv::IV;
-use gateway_requests::registration::handshake::SharedKeys;
-use log::*;
-use nymsphinx::DestinationAddressBytes;
-use std::collections::HashMap;
-use tokio::task::JoinHandle;
-
-pub(crate) type ClientsHandlerRequestSender = mpsc::UnboundedSender<ClientsHandlerRequest>;
-pub(crate) type ClientsHandlerRequestReceiver = mpsc::UnboundedReceiver<ClientsHandlerRequest>;
-
-pub(crate) type ClientsHandlerResponseSender = oneshot::Sender<ClientsHandlerResponse>;
-
-// #[derive(Debug)]
-pub(crate) enum ClientsHandlerRequest {
-    // client
-    Register(
-        DestinationAddressBytes,
-        SharedKeys,
-        MixMessageSender,
-        ClientsHandlerResponseSender,
-    ),
-    Authenticate(
-        DestinationAddressBytes,
-        EncryptedAddressBytes,
-        IV,
-        MixMessageSender,
-        ClientsHandlerResponseSender,
-    ),
-    Disconnect(DestinationAddressBytes),
-
-    // mix
-    IsOnline(DestinationAddressBytes, ClientsHandlerResponseSender),
-}
-
-#[derive(Debug)]
-pub(crate) enum ClientsHandlerResponse {
-    Register(bool),
-    Authenticate(Option<SharedKeys>),
-    IsOnline(Option<MixMessageSender>),
-    Error(Box<dyn std::error::Error + Send + Sync>),
-}
-
-pub(crate) struct ClientsHandler {
-    open_connections: HashMap<DestinationAddressBytes, MixMessageSender>,
-    clients_ledger: ClientLedger,
-    clients_inbox_storage: ClientStorage,
-}
-
-impl ClientsHandler {
-    pub(crate) fn new(clients_ledger: ClientLedger, clients_inbox_storage: ClientStorage) -> Self {
-        ClientsHandler {
-            open_connections: HashMap::new(),
-            clients_ledger,
-            clients_inbox_storage,
-        }
-    }
-
-    fn make_error_response<E>(&self, err: E) -> ClientsHandlerResponse
-    where
-        E: Into<Box<dyn std::error::Error + Send + Sync>>,
-    {
-        ClientsHandlerResponse::Error(err.into())
-    }
-
-    // best effort sending error responses
-    fn send_error_response<E>(&self, err: E, res_channel: ClientsHandlerResponseSender)
-    where
-        E: Into<Box<dyn std::error::Error + Send + Sync>>,
-    {
-        if res_channel.send(self.make_error_response(err)).is_err() {
-            error!("Somehow we failed to send response back to websocket handler - there seem to be a weird bug present!");
-        }
-    }
-
-    async fn push_stored_messages_to_client_and_save_channel(
-        &mut self,
-        client_address: DestinationAddressBytes,
-        comm_channel: MixMessageSender,
-    ) {
-        // TODO: it is possible that during a small window some of client messages will be "lost",
-        // i.e. be stored on the disk rather than pushed to the client, reason for this is as follows:
-        // now we push all stored messages from client's inbox to its websocket connection
-        // however, say, at the same time there's new message to the client - it gets stored on the disk!
-        // And only after this method exists, mix receivers will become aware of the client
-        // connection going online and being able to forward traffic there.
-        //
-        // possible solution: spawn a future to empty inbox in X seconds rather than immediately
-        // JS: I will most likely do that (with including entries to config, etc.) once the
-        // basic version is up and running as not to waste time on it now
-
-        // NOTE: THIS IGNORES MESSAGE RETRIEVAL LIMIT AND TAKES EVERYTHING!
-        let all_stored_messages = match self
-            .clients_inbox_storage
-            .retrieve_all_client_messages(client_address)
-            .await
-        {
-            Ok(msgs) => msgs,
-            Err(e) => {
-                error!(
-                    "failed to retrieve client messages. {:?} inbox might be corrupted now - {:?}",
-                    client_address.as_base58_string(),
-                    e
-                );
-                return;
-            }
-        };
-
-        let (messages, paths): (Vec<_>, Vec<_>) = all_stored_messages
-            .into_iter()
-            .map(|c| c.into_tuple())
-            .unzip();
-
-        if comm_channel.unbounded_send(messages).is_err() {
-            error!("Somehow we failed to stored messages to a fresh client channel - there seem to be a weird bug present!");
-        } else {
-            // but if all went well, we can now delete it
-            if let Err(e) = self.clients_inbox_storage.delete_files(paths).await {
-                error!(
-                    "Failed to remove client ({:?}) files - {:?}",
-                    client_address.as_base58_string(),
-                    e
-                );
-            } else {
-                // finally, everything was fine - we retrieved everything, we deleted everything,
-                // we assume we can now safely delegate client message pushing
-                self.open_connections.insert(client_address, comm_channel);
-            }
-        }
-    }
-
-    async fn handle_register_request(
-        &mut self,
-        address: DestinationAddressBytes,
-        derived_shared_key: SharedKeys,
-        comm_channel: MixMessageSender,
-        res_channel: ClientsHandlerResponseSender,
-    ) {
-        debug!(
-            "Processing register new client request: {:?}",
-            address.as_base58_string()
-        );
-
-        if self.open_connections.get(&address).is_some() {
-            warn!(
-                "Tried to process register request for a client with an already opened connection!"
-            );
-            self.send_error_response("duplicate connection detected", res_channel);
-            return;
-        }
-
-        if self
-            .clients_ledger
-            .insert_shared_key(derived_shared_key, address)
-            .unwrap()
-            .is_some()
-        {
-            info!(
-                "Client {:?} was already registered before!",
-                address.as_base58_string()
-            )
-        } else if let Err(e) = self.clients_inbox_storage.create_storage_dir(address).await {
-            error!("We failed to create inbox directory for the client -{:?}\nReverting stored shared key...", e);
-            // we must revert our changes if this operation failed
-            self.clients_ledger.remove_shared_key(&address).unwrap();
-            self.send_error_response("failed to complete issuing shared key", res_channel);
-            return;
-        }
-
-        self.push_stored_messages_to_client_and_save_channel(address, comm_channel)
-            .await;
-
-        if res_channel
-            .send(ClientsHandlerResponse::Register(true))
-            .is_err()
-        {
-            error!("Somehow we failed to send response back to websocket handler - there seem to be a weird bug present!");
-        }
-    }
-
-    async fn handle_authenticate_request(
-        &mut self,
-        address: DestinationAddressBytes,
-        encrypted_address: EncryptedAddressBytes,
-        iv: IV,
-        comm_channel: MixMessageSender,
-        res_channel: ClientsHandlerResponseSender,
-    ) {
-        debug!(
-            "Processing authenticate client request: {:?}",
-            address.as_base58_string()
-        );
-
-        if self.open_connections.get(&address).is_some() {
-            warn!("Tried to process authenticate request for a client with an already opened connection!");
-            self.send_error_response("duplicate connection detected", res_channel);
-            return;
-        }
-
-        if self
-            .clients_ledger
-            .verify_shared_key(&address, &encrypted_address, &iv)
-            .unwrap()
-        {
-            // The first unwrap is due to possible db read errors, but I'm not entirely sure when could
-            // the second one happen.
-            let shared_key = self
-                .clients_ledger
-                .get_shared_key(&address)
-                .unwrap()
-                .unwrap();
-            self.push_stored_messages_to_client_and_save_channel(address, comm_channel)
-                .await;
-            if res_channel
-                .send(ClientsHandlerResponse::Authenticate(Some(shared_key)))
-                .is_err()
-            {
-                error!("Somehow we failed to send response back to websocket handler - there seem to be a weird bug present!");
-            }
-        } else if res_channel
-            .send(ClientsHandlerResponse::Authenticate(None))
-            .is_err()
-        {
-            error!("Somehow we failed to send response back to websocket handler - there seem to be a weird bug present!");
-        }
-    }
-
-    fn handle_disconnect(&mut self, address: DestinationAddressBytes) {
-        debug!(
-            "Processing disconnect client request: {:?}",
-            address.as_base58_string()
-        );
-        self.open_connections.remove(&address);
-    }
-
-    fn handle_is_online_request(
-        &self,
-        address: DestinationAddressBytes,
-        res_channel: ClientsHandlerResponseSender,
-    ) {
-        debug!(
-            "Processing is online request for: {:?}",
-            address.as_base58_string()
-        );
-
-        let response_value = self.open_connections.get(&address).cloned();
-        // if this fails, it's a critical failure, because mix handlers should ALWAYS be online
-        res_channel
-            .send(ClientsHandlerResponse::IsOnline(response_value))
-            .unwrap();
-    }
-
-    pub(crate) async fn run(
-        &mut self,
-        mut request_receiver_channel: ClientsHandlerRequestReceiver,
-    ) {
-        while let Some(request) = request_receiver_channel.next().await {
-            match request {
-                ClientsHandlerRequest::Register(
-                    address,
-                    derived_shared_key,
-                    comm_channel,
-                    res_channel,
-                ) => {
-                    self.handle_register_request(
-                        address,
-                        derived_shared_key,
-                        comm_channel,
-                        res_channel,
-                    )
-                    .await
-                }
-                ClientsHandlerRequest::Authenticate(
-                    address,
-                    encrypted_address,
-                    iv,
-                    comm_channel,
-                    res_channel,
-                ) => {
-                    self.handle_authenticate_request(
-                        address,
-                        encrypted_address,
-                        iv,
-                        comm_channel,
-                        res_channel,
-                    )
-                    .await
-                }
-                ClientsHandlerRequest::Disconnect(address) => self.handle_disconnect(address),
-                ClientsHandlerRequest::IsOnline(address, res_channel) => {
-                    self.handle_is_online_request(address, res_channel)
-                }
-            };
-        }
-        error!("Something bad has happened and we stopped listening for requests!");
-    }
-
-    pub(crate) fn start(mut self) -> (JoinHandle<()>, ClientsHandlerRequestSender) {
-        let (sender, receiver) = mpsc::unbounded();
-        (
-            tokio::spawn(async move { self.run(receiver).await }),
-            sender,
-        )
-    }
-}

gateway/src/node/client_handling/mod.rs

@@ -1,6 +1,6 @@
 // Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
+pub(crate) mod active_clients;
 mod bandwidth;
-pub(crate) mod clients_handler;
 pub(crate) mod websocket;

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

@@ -1,598 +0,0 @@
-// Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
-// SPDX-License-Identifier: Apache-2.0
-
-use crate::node::client_handling::bandwidth::{Bandwidth, BandwidthDatabase};
-use crate::node::client_handling::clients_handler::{
-    ClientsHandlerRequest, ClientsHandlerRequestSender, ClientsHandlerResponse,
-};
-use crate::node::client_handling::websocket::message_receiver::{
-    MixMessageReceiver, MixMessageSender,
-};
-use coconut_interface::VerificationKey;
-use crypto::asymmetric::identity;
-use futures::{
-    channel::{mpsc, oneshot},
-    SinkExt, StreamExt,
-};
-use gateway_requests::authentication::encrypted_address::EncryptedAddressBytes;
-use gateway_requests::iv::IV;
-use gateway_requests::registration::handshake::error::HandshakeError;
-use gateway_requests::registration::handshake::{gateway_handshake, SharedKeys};
-use gateway_requests::types::{BinaryRequest, ClientControlRequest, ServerResponse};
-use gateway_requests::BinaryResponse;
-use log::*;
-use mixnet_client::forwarder::MixForwardingSender;
-use nymsphinx::DestinationAddressBytes;
-use rand::{CryptoRng, Rng};
-use std::convert::TryFrom;
-use std::mem;
-use std::sync::Arc;
-use tokio::io::{AsyncRead, AsyncWrite};
-use tokio_tungstenite::{
-    tungstenite::{protocol::Message, Error as WsError},
-    WebSocketStream,
-};
-
-//// TODO: note for my future self to consider the following idea:
-//// split the socket connection into sink and stream
-//// stream will be for reading explicit requests
-//// and sink for pumping responses AND mix traffic
-//// but as byproduct this might (or might not) break the clean "SocketStream" enum here
-
-enum SocketStream<S> {
-    RawTcp(S),
-    UpgradedWebSocket(WebSocketStream<S>),
-    Invalid,
-}
-
-impl<S> SocketStream<S> {
-    fn is_websocket(&self) -> bool {
-        matches!(self, SocketStream::UpgradedWebSocket(_))
-    }
-}
-
-pub(crate) struct Handle<R, S> {
-    rng: R,
-    remote_address: Option<DestinationAddressBytes>,
-    shared_key: Option<SharedKeys>,
-    clients_handler_sender: ClientsHandlerRequestSender,
-    outbound_mix_sender: MixForwardingSender,
-    socket_connection: SocketStream<S>,
-    local_identity: Arc<identity::KeyPair>,
-
-    aggregated_verification_key: VerificationKey,
-    bandwidths: BandwidthDatabase,
-}
-
-impl<R, S> Handle<R, S>
-where
-    R: Rng + CryptoRng,
-{
-    // for time being we assume handle is always constructed from raw socket.
-    // if we decide we want to change it, that's not too difficult
-    pub(crate) fn new(
-        rng: R,
-        conn: S,
-        clients_handler_sender: ClientsHandlerRequestSender,
-        outbound_mix_sender: MixForwardingSender,
-        local_identity: Arc<identity::KeyPair>,
-        aggregated_verification_key: VerificationKey,
-        bandwidths: BandwidthDatabase,
-    ) -> Self {
-        Handle {
-            rng,
-            remote_address: None,
-            shared_key: None,
-            clients_handler_sender,
-            outbound_mix_sender,
-            socket_connection: SocketStream::RawTcp(conn),
-            local_identity,
-            aggregated_verification_key,
-            bandwidths,
-        }
-    }
-
-    async fn perform_websocket_handshake(&mut self) -> Result<(), WsError>
-    where
-        S: AsyncRead + AsyncWrite + Unpin,
-    {
-        self.socket_connection =
-            match std::mem::replace(&mut self.socket_connection, SocketStream::Invalid) {
-                SocketStream::RawTcp(conn) => {
-                    // TODO: perhaps in the future, rather than panic here (and uncleanly shut tcp stream)
-                    // return a result with an error?
-                    let ws_stream = tokio_tungstenite::accept_async(conn).await?;
-                    SocketStream::UpgradedWebSocket(ws_stream)
-                }
-                other => other,
-            };
-        Ok(())
-    }
-
-    async fn perform_registration_handshake(
-        &mut self,
-        init_msg: Vec<u8>,
-    ) -> Result<SharedKeys, HandshakeError>
-    where
-        S: AsyncRead + AsyncWrite + Unpin + Send,
-    {
-        debug_assert!(self.socket_connection.is_websocket());
-        match &mut self.socket_connection {
-            SocketStream::UpgradedWebSocket(ws_stream) => {
-                gateway_handshake(
-                    &mut self.rng,
-                    ws_stream,
-                    self.local_identity.as_ref(),
-                    init_msg,
-                )
-                .await
-            }
-            _ => unreachable!(),
-        }
-    }
-
-    async fn next_websocket_request(&mut self) -> Option<Result<Message, WsError>>
-    where
-        S: AsyncRead + AsyncWrite + Unpin,
-    {
-        match self.socket_connection {
-            SocketStream::UpgradedWebSocket(ref mut ws_stream) => ws_stream.next().await,
-            _ => panic!("impossible state - websocket handshake was somehow reverted"),
-        }
-    }
-
-    async fn send_websocket_response(&mut self, msg: Message) -> Result<(), WsError>
-    where
-        S: AsyncRead + AsyncWrite + Unpin,
-    {
-        match self.socket_connection {
-            // TODO: more closely investigate difference between `Sink::send` and `Sink::send_all`
-            // it got something to do with batching and flushing - it might be important if it
-            // turns out somehow we've got a bottleneck here
-            SocketStream::UpgradedWebSocket(ref mut ws_stream) => ws_stream.send(msg).await,
-            _ => panic!("impossible state - websocket handshake was somehow reverted"),
-        }
-    }
-
-    // Note that it encrypts each message and slaps a MAC on it
-    async fn send_websocket_unwrapped_sphinx_packets(
-        &mut self,
-        packets: Vec<Vec<u8>>,
-    ) -> Result<(), WsError>
-    where
-        S: AsyncRead + AsyncWrite + Unpin,
-    {
-        let shared_key = self
-            .shared_key
-            .as_ref()
-            .expect("no shared key present even though we authenticated the client!");
-
-        // note: into_ws_message encrypts the requests and adds a MAC on it. Perhaps it should
-        // be more explicit in the naming?
-        let messages: Vec<Result<Message, WsError>> = packets
-            .into_iter()
-            .map(|received_message| {
-                Ok(BinaryResponse::new_pushed_mix_message(received_message)
-                    .into_ws_message(shared_key))
-            })
-            .collect();
-        let mut send_stream = futures::stream::iter(messages);
-        match self.socket_connection {
-            SocketStream::UpgradedWebSocket(ref mut ws_stream) => {
-                ws_stream.send_all(&mut send_stream).await
-            }
-            _ => panic!("impossible state - websocket handshake was somehow reverted"),
-        }
-    }
-
-    fn disconnect(&self) {
-        // if we never established what is the address of the client, its connection was never
-        // announced hence we do not need to send 'disconnect' message
-        if let Some(addr) = self.remote_address.as_ref() {
-            self.clients_handler_sender
-                .unbounded_send(ClientsHandlerRequest::Disconnect(*addr))
-                .unwrap();
-        }
-    }
-
-    async fn handle_binary(&self, bin_msg: Vec<u8>) -> Message {
-        trace!("Handling binary message (presumably sphinx packet)");
-
-        // this function decrypts the request and checks the MAC
-        match BinaryRequest::try_from_encrypted_tagged_bytes(
-            bin_msg,
-            self.shared_key
-                .as_ref()
-                .expect("no shared key present even though we authenticated the client!"),
-        ) {
-            Err(e) => ServerResponse::new_error(e.to_string()),
-            Ok(request) => match request {
-                // currently only a single type exists
-                BinaryRequest::ForwardSphinx(mix_packet) => {
-                    let consumed_bandwidth = mem::size_of_val(&mix_packet) as u64;
-                    if let Err(e) = Bandwidth::consume_bandwidth(
-                        &self.bandwidths,
-                        &self.remote_address.unwrap(),
-                        consumed_bandwidth,
-                    )
-                    .await
-                    {
-                        ServerResponse::new_error(format!("{:?}", e))
-                    } else {
-                        self.outbound_mix_sender.unbounded_send(mix_packet).unwrap();
-                        ServerResponse::Send { status: true }
-                    }
-                }
-            },
-        }
-        .into()
-    }
-
-    async fn handle_authenticate(
-        &mut self,
-        address: String,
-        enc_address: String,
-        iv: String,
-        mix_sender: MixMessageSender,
-    ) -> ServerResponse {
-        let address = match DestinationAddressBytes::try_from_base58_string(address) {
-            Ok(address) => address,
-            Err(e) => {
-                trace!("failed to parse received DestinationAddress: {:?}", e);
-                return ServerResponse::new_error("malformed destination address");
-            }
-        };
-
-        let encrypted_address = match EncryptedAddressBytes::try_from_base58_string(enc_address) {
-            Ok(address) => address,
-            Err(e) => {
-                trace!("failed to parse received encrypted address: {:?}", e);
-                return ServerResponse::new_error("malformed encrypted address");
-            }
-        };
-
-        let iv = match IV::try_from_base58_string(iv) {
-            Ok(iv) => iv,
-            Err(e) => {
-                trace!("failed to parse received IV {:?}", e);
-                return ServerResponse::new_error("malformed iv");
-            }
-        };
-
-        let (res_sender, res_receiver) = oneshot::channel();
-        let clients_handler_request = ClientsHandlerRequest::Authenticate(
-            address,
-            encrypted_address,
-            iv,
-            mix_sender,
-            res_sender,
-        );
-        self.clients_handler_sender
-            .unbounded_send(clients_handler_request)
-            .unwrap(); // the receiver MUST BE alive
-
-        match res_receiver.await.unwrap() {
-            ClientsHandlerResponse::Authenticate(shared_key) => {
-                if shared_key.is_some() {
-                    self.remote_address = Some(address);
-                    self.shared_key = shared_key;
-                    ServerResponse::Authenticate { status: true }
-                } else {
-                    ServerResponse::Authenticate { status: false }
-                }
-            }
-            ClientsHandlerResponse::Error(e) => {
-                error!("Authentication unexpectedly failed - {}", e);
-                ServerResponse::Error {
-                    message: format!("Authentication failure - {}", e),
-                }
-            }
-            _ => panic!("received response to wrong query!"), // this should NEVER happen
-        }
-    }
-
-    fn extract_remote_identity_from_register_init(init_data: &[u8]) -> Option<identity::PublicKey> {
-        if init_data.len() < identity::PUBLIC_KEY_LENGTH {
-            None
-        } else {
-            identity::PublicKey::from_bytes(&init_data[..identity::PUBLIC_KEY_LENGTH]).ok()
-        }
-    }
-
-    async fn handle_register(
-        &mut self,
-        init_data: Vec<u8>,
-        mix_sender: MixMessageSender,
-    ) -> ServerResponse
-    where
-        S: AsyncRead + AsyncWrite + Unpin + Send,
-    {
-        // not entirely sure how to it more "nicely"...
-        // hopefully, eventually this will go away once client's identity is known beforehand
-        let remote_identity = match Self::extract_remote_identity_from_register_init(&init_data) {
-            Some(address) => address,
-            None => return ServerResponse::new_error("malformed request"),
-        };
-        let remote_address = remote_identity.derive_destination_address();
-
-        let derived_shared_key = match self.perform_registration_handshake(init_data).await {
-            Ok(shared_key) => shared_key,
-            Err(err) => {
-                return ServerResponse::new_error(format!(
-                    "failed to perform the handshake - {}",
-                    err
-                ))
-            }
-        };
-
-        let (res_sender, res_receiver) = oneshot::channel();
-        let clients_handler_request = ClientsHandlerRequest::Register(
-            remote_address,
-            derived_shared_key.clone(),
-            mix_sender,
-            res_sender,
-        );
-
-        self.clients_handler_sender
-            .unbounded_send(clients_handler_request)
-            .unwrap(); // the receiver MUST BE alive
-
-        match res_receiver.await.unwrap() {
-            // currently register can't fail (as in if all machines are working correctly and you
-            // managed to complete registration handshake)
-            ClientsHandlerResponse::Register(status) => {
-                self.remote_address = Some(remote_address);
-                if status {
-                    self.shared_key = Some(derived_shared_key);
-                }
-                ServerResponse::Register { status }
-            }
-            ClientsHandlerResponse::Error(e) => {
-                error!("Post-handshake registration unexpectedly failed - {}", e);
-                ServerResponse::Error {
-                    message: format!("Registration failure - {}", e),
-                }
-            }
-            _ => panic!("received response to wrong query!"), // this should NEVER happen
-        }
-    }
-
-    async fn handle_bandwidth(&mut self, enc_credential: Vec<u8>, iv: Vec<u8>) -> ServerResponse {
-        if self.shared_key.is_none() {
-            return ServerResponse::new_error("No shared key has been exchanged with the gateway");
-        }
-        if self.remote_address.is_none() {
-            return ServerResponse::new_error("No remote address has been set");
-        }
-        let iv = match IV::try_from_bytes(&iv) {
-            Ok(iv) => iv,
-            Err(e) => {
-                trace!("failed to parse received IV {:?}", e);
-                return ServerResponse::new_error("malformed iv");
-            }
-        };
-        let credential = match ClientControlRequest::try_from_enc_bandwidth_credential(
-            enc_credential,
-            self.shared_key.as_ref().unwrap(),
-            iv,
-        ) {
-            Ok(c) => c,
-            Err(e) => {
-                return ServerResponse::new_error(e.to_string());
-            }
-        };
-        if credential.verify(&self.aggregated_verification_key) {
-            match Bandwidth::try_from(credential) {
-                Ok(bandwidth) => {
-                    if let Err(e) = Bandwidth::increase_bandwidth(
-                        &self.bandwidths,
-                        &self.remote_address.unwrap(),
-                        bandwidth.value(),
-                    )
-                    .await
-                    {
-                        return ServerResponse::Error {
-                            message: format!("{:?}", e),
-                        };
-                    }
-                    ServerResponse::Bandwidth { status: true }
-                }
-                Err(e) => ServerResponse::Error {
-                    message: format!("{:?}", e),
-                },
-            }
-        } else {
-            ServerResponse::Bandwidth { status: false }
-        }
-    }
-
-    // currently the bandwidth credential request is the only one we can receive after
-    // authentication
-    async fn handle_text(&mut self, raw_request: String) -> Message {
-        if let Ok(request) = ClientControlRequest::try_from(raw_request) {
-            match request {
-                ClientControlRequest::BandwidthCredential { enc_credential, iv } => {
-                    self.handle_bandwidth(enc_credential, iv).await.into()
-                }
-                _ => ServerResponse::new_error("invalid request").into(),
-            }
-        } else {
-            ServerResponse::new_error("malformed request").into()
-        }
-    }
-
-    async fn handle_request(&mut self, raw_request: Message) -> Option<Message> {
-        // apparently tungstenite auto-handles ping/pong/close messages so for now let's ignore
-        // them and let's test that claim. If that's not the case, just copy code from
-        // desktop nym-client websocket as I've manually handled everything there
-        match raw_request {
-            Message::Binary(bin_msg) => Some(self.handle_binary(bin_msg).await),
-            Message::Text(text_msg) => Some(self.handle_text(text_msg).await),
-            _ => None,
-        }
-    }
-
-    /// Handles data that resembles request to either start registration handshake or perform
-    /// authentication.
-    async fn handle_initial_authentication_request(
-        &mut self,
-        mix_sender: MixMessageSender,
-        raw_request: String,
-    ) -> ServerResponse
-    where
-        S: AsyncRead + AsyncWrite + Unpin + Send,
-    {
-        if let Ok(request) = ClientControlRequest::try_from(raw_request) {
-            match request {
-                ClientControlRequest::Authenticate {
-                    address,
-                    enc_address,
-                    iv,
-                } => {
-                    self.handle_authenticate(address, enc_address, iv, mix_sender)
-                        .await
-                }
-                ClientControlRequest::RegisterHandshakeInitRequest { data } => {
-                    self.handle_register(data, mix_sender).await
-                }
-                _ => ServerResponse::new_error("invalid request"),
-            }
-        } else {
-            // TODO: is this a malformed request or rather a network error and
-            // connection should be terminated?
-            ServerResponse::new_error("malformed request")
-        }
-    }
-
-    /// Listens for only a subset of possible client requests, i.e. for those that can either
-    /// result in client getting registered or authenticated. All other requests, such as forwarding
-    /// sphinx packets are ignored.
-    async fn wait_for_initial_authentication(&mut self) -> Option<MixMessageReceiver>
-    where
-        S: AsyncRead + AsyncWrite + Unpin + Send,
-    {
-        trace!("Started waiting for authenticate/register request...");
-
-        while let Some(msg) = self.next_websocket_request().await {
-            let msg = match msg {
-                Ok(msg) => msg,
-                Err(err) => {
-                    error!("failed to obtain message from websocket stream! stopping connection handler: {}", err);
-                    break;
-                }
-            };
-
-            if msg.is_close() {
-                break;
-            }
-
-            let (mix_sender, mix_receiver) = mpsc::unbounded();
-
-            // ONLY handle 'Authenticate' or 'Register' requests, ignore everything else
-            let response = match msg {
-                Message::Close(_) => break,
-                Message::Text(text_msg) => {
-                    self.handle_initial_authentication_request(mix_sender, text_msg)
-                        .await
-                }
-                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
-                    warn!("possibly received a sphinx packet without prior authentication. Request is going to be ignored");
-                    ServerResponse::new_error("binary request without prior authentication")
-                }
-
-                _ => continue,
-            };
-
-            let is_done = response.implies_successful_authentication();
-
-            if let Err(err) = self.send_websocket_response(response.into()).await {
-                warn!(
-                    "Failed to send message over websocket: {}. Assuming the connection is dead.",
-                    err
-                );
-                break;
-            }
-
-            // it means we successfully managed to perform authentication and announce our
-            // presence to ClientsHandler
-            if is_done {
-                return Some(mix_receiver);
-            }
-        }
-        None
-    }
-
-    /// Simultaneously listens for incoming client requests, which realistically should only be
-    /// binary requests to forward sphinx packets, and for sphinx packets received from the mix
-    /// network that should be sent back to the client.
-    async fn listen_for_requests(&mut self, mut mix_receiver: MixMessageReceiver)
-    where
-        S: AsyncRead + AsyncWrite + Unpin,
-    {
-        trace!("Started listening for ALL incoming requests...");
-
-        loop {
-            tokio::select! {
-                socket_msg = self.next_websocket_request() => {
-                    if socket_msg.is_none() {
-                        break;
-                    }
-                    let socket_msg = match socket_msg.unwrap() {
-                        Ok(socket_msg) => socket_msg,
-                        Err(err) => {
-                            error!("failed to obtain message from websocket stream! stopping connection handler: {}", err);
-                            break;
-                        }
-                    };
-
-                    if socket_msg.is_close() {
-                        break;
-                    }
-
-                    if let Some(response) = self.handle_request(socket_msg).await {
-                        if let Err(err) = self.send_websocket_response(response).await {
-                            warn!(
-                                "Failed to send message over websocket: {}. Assuming the connection is dead.",
-                                err
-                            );
-                            break;
-                        }
-                    }
-                },
-                mix_messages = mix_receiver.next() => {
-                    let mix_messages = mix_messages.expect("sender was unexpectedly closed! this shouldn't have ever happened!");
-                    if let Err(e) = self.send_websocket_unwrapped_sphinx_packets(mix_messages).await {
-                        warn!("failed to send the unwrapped sphinx packets back to the client - {:?}, assuming the connection is dead", e);
-                        break;
-                    }
-                }
-            }
-        }
-
-        self.disconnect();
-        trace!("The stream was closed!");
-    }
-
-    pub(crate) async fn start_handling(&mut self)
-    where
-        S: AsyncRead + AsyncWrite + Unpin + Send,
-    {
-        if let Err(e) = self.perform_websocket_handshake().await {
-            warn!(
-                "Failed to complete WebSocket handshake - {:?}. Stopping the handler",
-                e
-            );
-            return;
-        }
-        trace!("Managed to perform websocket handshake!");
-        let mix_receiver = self.wait_for_initial_authentication().await;
-        trace!("Performed initial authentication");
-        match mix_receiver {
-            Some(receiver) => self.listen_for_requests(receiver).await,
-            None => trace!("But connection was closed during the process"),
-        }
-        trace!("The handler is done!");
-    }
-}

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

@@ -0,0 +1,332 @@
+// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::node::client_handling::bandwidth::Bandwidth;
+use crate::node::client_handling::websocket::connection_handler::{ClientDetails, FreshHandler};
+use crate::node::client_handling::websocket::message_receiver::MixMessageReceiver;
+use crate::node::storage::error::StorageError;
+use futures::StreamExt;
+use gateway_requests::iv::{IVConversionError, IV};
+use gateway_requests::types::{BinaryRequest, ServerResponse};
+use gateway_requests::{ClientControlRequest, GatewayRequestsError};
+use log::*;
+use nymsphinx::forwarding::packet::MixPacket;
+use rand::{CryptoRng, Rng};
+use std::convert::TryFrom;
+use std::process;
+use thiserror::Error;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_tungstenite::tungstenite::protocol::Message;
+
+#[derive(Debug, Error)]
+enum RequestHandlingError {
+    #[error("Internal gateway storage error")]
+    StorageError(#[from] StorageError),
+
+    #[error("Provided bandwidth IV is malformed - {0}")]
+    MalformedIV(#[from] IVConversionError),
+
+    #[error("Provided binary request was malformed - {0}")]
+    InvalidBinaryRequest(#[from] GatewayRequestsError),
+
+    #[error("Provided binary request was malformed - {0}")]
+    InvalidTextRequest(<ClientControlRequest as TryFrom<String>>::Error),
+
+    #[error("Provided bandwidth credential did not verify correctly")]
+    InvalidBandwidthCredential,
+
+    #[error("Provided bandwidth credential did not have expected structure - {0}")]
+    BandwidthCredentialError(#[from] credentials::error::Error),
+
+    #[error("Provided bandwidth credential asks for more bandwidth than it is supported to add at once (credential value: {0}, supported: {}). Try to split it before attempting again", i64::MAX)]
+    UnsupportedBandwidthValue(u64),
+
+    #[error("The received request is not valid in the current context")]
+    IllegalRequest,
+}
+
+impl RequestHandlingError {
+    fn into_error_message(self) -> Message {
+        ServerResponse::new_error(self.to_string()).into()
+    }
+}
+
+pub(crate) struct AuthenticatedHandler<R, S> {
+    inner: FreshHandler<R, S>,
+    client: ClientDetails,
+    mix_receiver: MixMessageReceiver,
+}
+
+// explicitly remove handle from the global store upon being dropped
+impl<R, S> Drop for AuthenticatedHandler<R, S> {
+    fn drop(&mut self) {
+        self.inner
+            .active_clients_store
+            .disconnect(self.client.address)
+    }
+}
+
+impl<R, S> AuthenticatedHandler<R, S>
+where
+    // TODO: those trait bounds here don't really make sense....
+    R: Rng + CryptoRng,
+{
+    /// Upgrades `FreshHandler` into the Authenticated variant implying the client is now authenticated
+    /// and thus allowed to perform more actions with the gateway, such as redeeming bandwidth or
+    /// sending sphinx packets.
+    ///
+    /// # Arguments
+    ///
+    /// * `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(
+        fresh: FreshHandler<R, S>,
+        client: ClientDetails,
+        mix_receiver: MixMessageReceiver,
+    ) -> Self {
+        AuthenticatedHandler {
+            inner: fresh,
+            client,
+            mix_receiver,
+        }
+    }
+
+    /// Explicitly removes handle from the global store.
+    fn disconnect(self) {
+        self.inner
+            .active_clients_store
+            .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)
+    }
+
+    /// Increases the amount of available bandwidth of the connected client by the specified value.
+    ///
+    /// # Arguments
+    ///
+    /// * `amount`: amount to increase the available bandwidth by.
+    async fn increase_bandwidth(&self, amount: i64) -> Result<(), RequestHandlingError> {
+        self.inner
+            .storage
+            .increase_bandwidth(self.client.address, amount)
+            .await?;
+        Ok(())
+    }
+
+    /// Decreases the amount of available bandwidth of the connected client by the specified value.
+    ///
+    /// # 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?;
+        Ok(())
+    }
+
+    /// Forwards the received mix packet from the client into the mix network.
+    ///
+    /// # Arguments
+    ///
+    /// * `mix_packet`: packet received from the client that should get forwarded into the network.
+    fn forward_packet(&self, mix_packet: MixPacket) {
+        if let Err(err) = self.inner.outbound_mix_sender.unbounded_send(mix_packet) {
+            error!("We failed to forward requested mix packet - {}. Presumably our mix forwarder has crashed. We cannot continue.", err);
+            process::exit(1);
+        }
+    }
+
+    /// Tries to handle the received bandwidth request by checking correctness of the received data
+    /// and if successful, increases client's bandwidth by an appropriate amount.
+    ///
+    /// # Arguments
+    ///
+    /// * `enc_credential`: raw encrypted bandwidth credential to verify.
+    /// * `iv`: fresh iv used for the credential.
+    async fn handle_bandwidth(
+        &mut self,
+        enc_credential: Vec<u8>,
+        iv: Vec<u8>,
+    ) -> Result<ServerResponse, RequestHandlingError> {
+        let iv = IV::try_from_bytes(&iv)?;
+        let credential = ClientControlRequest::try_from_enc_bandwidth_credential(
+            enc_credential,
+            &self.client.shared_keys,
+            iv,
+        )?;
+
+        if !credential.verify(&self.inner.aggregated_verification_key) {
+            return Err(RequestHandlingError::InvalidBandwidthCredential);
+        }
+
+        let bandwidth = Bandwidth::try_from(credential)?;
+        let bandwidth_value = bandwidth.value();
+
+        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(RequestHandlingError::UnsupportedBandwidthValue(
+                bandwidth_value,
+            ));
+        }
+
+        self.increase_bandwidth(bandwidth_value as i64).await?;
+        let available_total = self.get_available_bandwidth().await?;
+
+        Ok(ServerResponse::Bandwidth { available_total })
+    }
+
+    /// Tries to handle request to forward sphinx packet into the network. The request can only succeed
+    /// if the client has enough available bandwidth.
+    ///
+    /// Upon forwarding, client's bandwidth is decreased by the size of the forwarded packet.
+    ///
+    /// # Arguments
+    ///
+    /// * `mix_packet`: packet received from the client that should get forwarded into the network.
+    async fn handle_forward_sphinx(
+        &self,
+        mix_packet: MixPacket,
+    ) -> Result<ServerResponse, RequestHandlingError> {
+        // for now let's just use actual size of the sphinx packet. there's a tiny bit of overhead
+        // we're not including (but it's literally like 2 bytes) when the packet is framed
+        let consumed_bandwidth = mix_packet.sphinx_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?;
+        self.forward_packet(mix_packet);
+
+        Ok(ServerResponse::Send {
+            remaining_bandwidth: available_bandwidth - consumed_bandwidth,
+        })
+    }
+
+    /// Attempts to handle a binary data frame websocket message.
+    ///
+    /// # Arguments
+    ///
+    /// * `bin_msg`: raw message to handle.
+    async fn handle_binary(&self, bin_msg: Vec<u8>) -> Message {
+        // this function decrypts the request and checks the MAC
+        match BinaryRequest::try_from_encrypted_tagged_bytes(bin_msg, &self.client.shared_keys) {
+            Err(e) => RequestHandlingError::InvalidBinaryRequest(e).into_error_message(),
+            Ok(request) => match request {
+                // currently only a single type exists
+                BinaryRequest::ForwardSphinx(mix_packet) => {
+                    match self.handle_forward_sphinx(mix_packet).await {
+                        Ok(response) => response.into(),
+                        Err(err) => err.into_error_message(),
+                    }
+                }
+            },
+        }
+    }
+
+    /// Attempts to handle a text data frame websocket message.
+    ///
+    /// Currently the bandwidth credential request is the only one we can receive after authentication.
+    ///
+    /// # Arguments
+    ///
+    /// * `raw_request`: raw message to handle.
+    async fn handle_text(&mut self, raw_request: String) -> Message {
+        match ClientControlRequest::try_from(raw_request) {
+            Err(e) => RequestHandlingError::InvalidTextRequest(e).into_error_message(),
+            Ok(request) => match request {
+                ClientControlRequest::BandwidthCredential { enc_credential, iv } => {
+                    match self.handle_bandwidth(enc_credential, iv).await {
+                        Ok(response) => response.into(),
+                        Err(err) => err.into_error_message(),
+                    }
+                }
+                _ => RequestHandlingError::IllegalRequest.into_error_message(),
+            },
+        }
+    }
+
+    /// Attempts to handle websocket message received from the connected client.
+    ///
+    /// # Arguments
+    ///
+    /// * `raw_request`: raw received websocket message.
+    async fn handle_request(&mut self, raw_request: Message) -> Option<Message> {
+        // apparently tungstenite auto-handles ping/pong/close messages so for now let's ignore
+        // them and let's test that claim. If that's not the case, just copy code from
+        // desktop nym-client websocket as I've manually handled everything there
+        match raw_request {
+            Message::Binary(bin_msg) => Some(self.handle_binary(bin_msg).await),
+            Message::Text(text_msg) => Some(self.handle_text(text_msg).await),
+            _ => None,
+        }
+    }
+
+    /// Simultaneously listens for incoming client requests, which realistically should only be
+    /// binary requests to forward sphinx packets or increase bandwidth
+    /// and for sphinx packets received from the mix network that should be sent back to the client.
+    pub(crate) async fn listen_for_requests(mut self)
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        trace!("Started listening for ALL incoming requests...");
+
+        loop {
+            tokio::select! {
+                socket_msg = self.inner.read_websocket_message() => {
+                    let socket_msg = match socket_msg {
+                        None => break,
+                        Some(Ok(socket_msg)) => socket_msg,
+                        Some(Err(err)) => {
+                            error!("failed to obtain message from websocket stream! stopping connection handler: {}", err);
+                            break;
+                        }
+                    };
+
+                    if socket_msg.is_close() {
+                        break;
+                    }
+
+                    if let Some(response) = self.handle_request(socket_msg).await {
+                        if let Err(err) = self.inner.send_websocket_message(response).await {
+                            warn!(
+                                "Failed to send message over websocket: {}. Assuming the connection is dead.",
+                                err
+                            );
+                            break;
+                        }
+                    }
+                },
+                mix_messages = self.mix_receiver.next() => {
+                    let mix_messages = mix_messages.expect("sender was unexpectedly closed! this shouldn't have ever happened!");
+                    if let Err(e) = self.inner.push_packets_to_client(self.client.shared_keys, mix_messages).await {
+                        warn!("failed to send the unwrapped sphinx packets back to the client - {:?}, assuming the connection is dead", e);
+                        break;
+                    }
+                }
+            }
+        }
+
+        self.disconnect();
+        trace!("The stream was closed!");
+    }
+}

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

@@ -0,0 +1,598 @@
+// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::node::client_handling::active_clients::ActiveClientsStore;
+use crate::node::client_handling::websocket::connection_handler::{
+    AuthenticatedHandler, ClientDetails, InitialAuthResult, SocketStream,
+};
+use crate::node::storage::error::StorageError;
+use crate::node::storage::PersistentStorage;
+use coconut_interface::VerificationKey;
+use crypto::asymmetric::identity;
+use futures::{channel::mpsc, SinkExt, StreamExt};
+use gateway_requests::authentication::encrypted_address::{
+    EncryptedAddressBytes, EncryptedAddressConversionError,
+};
+use gateway_requests::iv::{IVConversionError, IV};
+use gateway_requests::registration::handshake::error::HandshakeError;
+use gateway_requests::registration::handshake::{gateway_handshake, SharedKeys};
+use gateway_requests::types::{ClientControlRequest, ServerResponse};
+use gateway_requests::BinaryResponse;
+use log::*;
+use mixnet_client::forwarder::MixForwardingSender;
+use nymsphinx::DestinationAddressBytes;
+use rand::{CryptoRng, Rng};
+use std::convert::TryFrom;
+use std::sync::Arc;
+use thiserror::Error;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError};
+
+#[derive(Debug, Error)]
+enum InitialAuthenticationError {
+    #[error("Internal gateway storage error")]
+    StorageError(#[from] StorageError),
+
+    #[error("Failed to perform registration handshake - {0}")]
+    HandshakeError(#[from] HandshakeError),
+
+    #[error("Provided client address is malformed - {0}")]
+    // sphinx error is not used here directly as it's messaging might be confusing to people
+    MalformedClientAddress(String),
+
+    #[error("Provided encrypted client address is malformed - {0}")]
+    MalformedEncryptedAddress(#[from] EncryptedAddressConversionError),
+
+    #[error("There is already an open connection to this client")]
+    DuplicateConnection,
+
+    #[error("Provided authentication IV is malformed - {0}")]
+    MalformedIV(#[from] IVConversionError),
+
+    #[error("Only 'Register' or 'Authenticate' requests are allowed")]
+    InvalidRequest,
+
+    #[error("Experienced connection error - {0}")]
+    ConnectionError(#[from] WsError),
+}
+
+impl InitialAuthenticationError {
+    /// Converts this Error into an appropriate websocket Message.
+    fn into_error_message(self) -> Message {
+        ServerResponse::new_error(self.to_string()).into()
+    }
+}
+
+pub(crate) struct FreshHandler<R, S> {
+    rng: R,
+    local_identity: Arc<identity::KeyPair>,
+    pub(crate) active_clients_store: ActiveClientsStore,
+    pub(crate) aggregated_verification_key: VerificationKey,
+    pub(crate) outbound_mix_sender: MixForwardingSender,
+    pub(crate) socket_connection: SocketStream<S>,
+    pub(crate) storage: PersistentStorage,
+}
+
+impl<R, S> FreshHandler<R, S>
+where
+    R: Rng + CryptoRng,
+{
+    // for time being we assume handle is always constructed from raw socket.
+    // if we decide we want to change it, that's not too difficult
+    pub(crate) fn new(
+        rng: R,
+        conn: S,
+        outbound_mix_sender: MixForwardingSender,
+        local_identity: Arc<identity::KeyPair>,
+        aggregated_verification_key: VerificationKey,
+        storage: PersistentStorage,
+        active_clients_store: ActiveClientsStore,
+    ) -> Self {
+        FreshHandler {
+            rng,
+            active_clients_store,
+            outbound_mix_sender,
+            socket_connection: SocketStream::RawTcp(conn),
+            local_identity,
+            aggregated_verification_key,
+            storage,
+        }
+    }
+
+    /// Attempts to perform websocket handshake with the remote and upgrades the raw TCP socket
+    /// to the framed WebSocket.
+    pub(crate) async fn perform_websocket_handshake(&mut self) -> Result<(), WsError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        self.socket_connection =
+            match std::mem::replace(&mut self.socket_connection, SocketStream::Invalid) {
+                SocketStream::RawTcp(conn) => {
+                    // TODO: perhaps in the future, rather than panic here (and uncleanly shut tcp stream)
+                    // return a result with an error?
+                    let ws_stream = tokio_tungstenite::accept_async(conn).await?;
+                    SocketStream::UpgradedWebSocket(ws_stream)
+                }
+                other => other,
+            };
+        Ok(())
+    }
+
+    /// Using received `init_msg` tries to continue the registration handshake with the connected
+    /// client to establish shared keys.
+    ///
+    /// # Arguments
+    ///
+    /// * `init_msg`: a client handshake init message which should contain its identity public key as well as an ephemeral key.
+    async fn perform_registration_handshake(
+        &mut self,
+        init_msg: Vec<u8>,
+    ) -> Result<SharedKeys, HandshakeError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin + Send,
+    {
+        debug_assert!(self.socket_connection.is_websocket());
+        match &mut self.socket_connection {
+            SocketStream::UpgradedWebSocket(ws_stream) => {
+                gateway_handshake(
+                    &mut self.rng,
+                    ws_stream,
+                    self.local_identity.as_ref(),
+                    init_msg,
+                )
+                .await
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    /// Attempts to read websocket message from the associated socket.
+    pub(crate) async fn read_websocket_message(&mut self) -> Option<Result<Message, WsError>>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        match self.socket_connection {
+            SocketStream::UpgradedWebSocket(ref mut ws_stream) => ws_stream.next().await,
+            _ => panic!("impossible state - websocket handshake was somehow reverted"),
+        }
+    }
+
+    /// Attempts to write a single websocket message on the available socket.
+    ///
+    /// # Arguments
+    ///
+    /// * `msg`: WebSocket message to write back to the client.
+    pub(crate) async fn send_websocket_message(&mut self, msg: Message) -> Result<(), WsError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        match self.socket_connection {
+            // TODO: more closely investigate difference between `Sink::send` and `Sink::send_all`
+            // it got something to do with batching and flushing - it might be important if it
+            // turns out somehow we've got a bottleneck here
+            SocketStream::UpgradedWebSocket(ref mut ws_stream) => ws_stream.send(msg).await,
+            _ => panic!("impossible state - websocket handshake was somehow reverted"),
+        }
+    }
+
+    /// Sends unwrapped sphinx packets (payloads) back to the client. Note that each message is encrypted and tagged with
+    /// the previously derived shared keys.
+    ///
+    /// # Arguments
+    ///
+    /// * `shared_keys`: keys derived between the client and gateway.
+    /// * `packets`: unwrapped packets that are to be pushed back to the client.
+    pub(crate) async fn push_packets_to_client(
+        &mut self,
+        shared_keys: SharedKeys,
+        packets: Vec<Vec<u8>>,
+    ) -> Result<(), WsError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        // note: into_ws_message encrypts the requests and adds a MAC on it. Perhaps it should
+        // be more explicit in the naming?
+        let messages: Vec<Result<Message, WsError>> = packets
+            .into_iter()
+            .map(|received_message| {
+                Ok(BinaryResponse::new_pushed_mix_message(received_message)
+                    .into_ws_message(&shared_keys))
+            })
+            .collect();
+        let mut send_stream = futures::stream::iter(messages);
+        match self.socket_connection {
+            SocketStream::UpgradedWebSocket(ref mut ws_stream) => {
+                ws_stream.send_all(&mut send_stream).await
+            }
+            _ => panic!("impossible state - websocket handshake was somehow reverted"),
+        }
+    }
+
+    /// Attempts to extract clients identity key from the received registration handshake init message.
+    ///
+    /// # Arguments
+    ///
+    /// * `init_data`: received init message that should contain, among other things, client's public key.
+    // Note: this is out of the scope of this PR, but in the future, this should be removed in favour
+    // of doing full parse of the init_data elsewhere
+    fn extract_remote_identity_from_register_init(
+        init_data: &[u8],
+    ) -> Result<identity::PublicKey, InitialAuthenticationError> {
+        if init_data.len() < identity::PUBLIC_KEY_LENGTH {
+            Err(InitialAuthenticationError::HandshakeError(
+                HandshakeError::MalformedRequest,
+            ))
+        } else {
+            identity::PublicKey::from_bytes(&init_data[..identity::PUBLIC_KEY_LENGTH]).map_err(
+                |_| InitialAuthenticationError::HandshakeError(HandshakeError::MalformedRequest),
+            )
+        }
+    }
+
+    /// Attempts to retrieve all messages currently stored in the persistent database to the client,
+    /// which was offline at the time of their receipt.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client that is going to receive the messages.
+    /// * `shared_keys`: shared keys derived between the client and the gateway used to encrypt and tag the messages.
+    async fn push_stored_messages_to_client(
+        &mut self,
+        client_address: DestinationAddressBytes,
+        shared_keys: SharedKeys,
+    ) -> Result<(), InitialAuthenticationError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        let mut start_next_after = None;
+        loop {
+            // retrieve some messages
+            let (messages, new_start_next_after) = self
+                .storage
+                .retrieve_messages(client_address, start_next_after)
+                .await?;
+
+            let (messages, ids) = messages
+                .into_iter()
+                .map(|msg| (msg.content, msg.id))
+                .unzip();
+
+            // push them to the client
+            if let Err(err) = self.push_packets_to_client(shared_keys, messages).await {
+                warn!(
+                    "We failed to send stored messages to fresh client - {}",
+                    err
+                );
+                return Err(InitialAuthenticationError::ConnectionError(err));
+            } else {
+                // if it was successful - remove them from the store
+                self.storage.remove_messages(ids).await?;
+            }
+
+            // no more messages to grab
+            if new_start_next_after.is_none() {
+                break;
+            } else {
+                start_next_after = new_start_next_after
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Checks whether the stored shared keys match the received data, i.e. whether the upon decryption
+    /// the provided encrypted address matches the expected unencrypted address.
+    ///
+    /// Returns the the retrieved shared keys if the check was successful.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client.
+    /// * `encrypted_address`: encrypted address of the client, presumably encrypted using the shared keys.
+    /// * `iv`: iv created for this particular encryption.
+    async fn verify_stored_shared_key(
+        &self,
+        client_address: DestinationAddressBytes,
+        encrypted_address: EncryptedAddressBytes,
+        iv: IV,
+    ) -> Result<Option<SharedKeys>, InitialAuthenticationError> {
+        let shared_keys = self.storage.get_shared_keys(client_address).await?;
+
+        if let Some(shared_keys) = shared_keys {
+            // the unwrap here is fine as we only ever construct persisted shared keys ourselves when inserting
+            // data to the storage. The only way it could fail is if we somehow changed implementation without
+            // performing proper migration
+            let keys = SharedKeys::try_from_base58_string(
+                shared_keys.derived_aes128_ctr_blake3_hmac_keys_bs58,
+            )
+            .unwrap();
+            // TODO: SECURITY:
+            // this is actually what we have been doing in the past, however,
+            // after looking deeper into implementation it seems that only checks the encryption
+            // key part of the shared keys. the MAC key might still be wrong
+            // (though I don't see how could this happen unless client messed with himself
+            // and I don't think it could lead to any attacks, but somebody smarter should take a look)
+            if encrypted_address.verify(&client_address, &keys, &iv) {
+                Ok(Some(keys))
+            } else {
+                Ok(None)
+            }
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Using the received challenge data, i.e. client's address as well the ciphertext of it plus
+    /// a fresh IV, attempts to authenticate the client by checking whether the ciphertext matches
+    /// the expected value if encrypted with the shared key.
+    ///
+    /// Finally, upon completion, all previously stored messages are pushed back to the client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client wishing to authenticate.
+    /// * `encrypted_address`: ciphertext of the address of the client wishing to authenticate.
+    /// * `iv`: fresh IV received with the request.
+    async fn authenticate_client(
+        &mut self,
+        client_address: DestinationAddressBytes,
+        encrypted_address: EncryptedAddressBytes,
+        iv: IV,
+    ) -> Result<Option<SharedKeys>, InitialAuthenticationError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        debug!(
+            "Processing authenticate client request for: {}",
+            client_address.as_base58_string()
+        );
+
+        let shared_keys = self
+            .verify_stored_shared_key(client_address, encrypted_address, iv)
+            .await?;
+
+        if let Some(shared_keys) = shared_keys {
+            self.push_stored_messages_to_client(client_address, shared_keys)
+                .await?;
+            Ok(Some(shared_keys))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Tries to handle the received authentication request by checking correctness of the received data.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client wishing to authenticate.
+    /// * `encrypted_address`: ciphertext of the address of the client wishing to authenticate.
+    /// * `iv`: fresh IV received with the request.
+    async fn handle_authenticate(
+        &mut self,
+        address: String,
+        enc_address: String,
+        iv: String,
+    ) -> Result<InitialAuthResult, InitialAuthenticationError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        let address = DestinationAddressBytes::try_from_base58_string(address)
+            .map_err(|err| InitialAuthenticationError::MalformedClientAddress(err.to_string()))?;
+        let encrypted_address = EncryptedAddressBytes::try_from_base58_string(enc_address)?;
+        let iv = IV::try_from_base58_string(iv)?;
+
+        if self.active_clients_store.get(address).is_some() {
+            return Err(InitialAuthenticationError::DuplicateConnection);
+        }
+
+        let shared_keys = self
+            .authenticate_client(address, encrypted_address, iv)
+            .await?;
+        let status = shared_keys.is_some();
+        let client_details =
+            shared_keys.map(|shared_keys| ClientDetails::new(address, shared_keys));
+
+        Ok(InitialAuthResult::new(
+            client_details,
+            ServerResponse::Authenticate { status },
+        ))
+    }
+
+    /// 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.
+    ///
+    /// Finally, upon completion, all previously stored messages are pushed back to the client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client`: details (i.e. address and shared keys) of the registered client
+    async fn register_client(
+        &mut self,
+        client: ClientDetails,
+    ) -> Result<bool, InitialAuthenticationError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        debug!(
+            "Processing register client request for: {}",
+            client.address.as_base58_string()
+        );
+
+        self.storage
+            .insert_shared_keys(client.address, client.shared_keys)
+            .await?;
+
+        // see if we have bandwidth entry for the client already, if not, create one with zero value
+        if self
+            .storage
+            .get_available_bandwidth(client.address)
+            .await?
+            .is_none()
+        {
+            self.storage.create_bandwidth_entry(client.address).await?;
+        }
+
+        self.push_stored_messages_to_client(client.address, client.shared_keys)
+            .await?;
+
+        Ok(true)
+    }
+
+    /// Tries to handle the received register request by checking attempting to complete registration
+    /// handshake using the received data.
+    ///
+    /// # Arguments
+    ///
+    /// * `init_data`: init payload of the registration handshake.
+    async fn handle_register(
+        &mut self,
+        init_data: Vec<u8>,
+    ) -> Result<InitialAuthResult, InitialAuthenticationError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin + Send,
+    {
+        let remote_identity = Self::extract_remote_identity_from_register_init(&init_data)?;
+        let remote_address = remote_identity.derive_destination_address();
+
+        if self.active_clients_store.get(remote_address).is_some() {
+            return Err(InitialAuthenticationError::DuplicateConnection);
+        }
+
+        let shared_keys = self.perform_registration_handshake(init_data).await?;
+        let client_details = ClientDetails::new(remote_address, shared_keys);
+
+        let status = self.register_client(client_details).await?;
+
+        Ok(InitialAuthResult::new(
+            Some(client_details),
+            ServerResponse::Register { status },
+        ))
+    }
+
+    /// Handles data that resembles request to either start registration handshake or perform
+    /// authentication.
+    ///
+    /// # Arguments
+    ///
+    /// * `raw_request`: raw text request received from the websocket.
+    async fn handle_initial_authentication_request(
+        &mut self,
+        raw_request: String,
+    ) -> Result<InitialAuthResult, InitialAuthenticationError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin + Send,
+    {
+        if let Ok(request) = ClientControlRequest::try_from(raw_request) {
+            match request {
+                ClientControlRequest::Authenticate {
+                    address,
+                    enc_address,
+                    iv,
+                } => self.handle_authenticate(address, enc_address, iv).await,
+                ClientControlRequest::RegisterHandshakeInitRequest { data } => {
+                    self.handle_register(data).await
+                }
+                // won't accept anything else (like bandwidth) without prior authentication
+                _ => Err(InitialAuthenticationError::InvalidRequest),
+            }
+        } else {
+            Err(InitialAuthenticationError::InvalidRequest)
+        }
+    }
+
+    /// Listens for only a subset of possible client requests, i.e. for those that can either
+    /// result in client getting registered or authenticated. All other requests, such as forwarding
+    /// sphinx packets considered an error and terminate the connection.
+    // TODO: somehow cleanup this method
+    pub(crate) async fn perform_initial_authentication(
+        mut self,
+    ) -> Option<AuthenticatedHandler<R, S>>
+    where
+        S: AsyncRead + AsyncWrite + Unpin + Send,
+    {
+        trace!("Started waiting for authenticate/register request...");
+
+        while let Some(msg) = self.read_websocket_message().await {
+            let msg = match msg {
+                Ok(msg) => msg,
+                Err(err) => {
+                    error!("failed to obtain message from websocket stream! stopping connection handler: {}", err);
+                    break;
+                }
+            };
+
+            if msg.is_close() {
+                break;
+            }
+
+            // ONLY handle 'Authenticate' or 'Register' requests, ignore everything else
+            match msg {
+                Message::Close(_) => break,
+                Message::Text(text_msg) => {
+                    let (mix_sender, mix_receiver) = mpsc::unbounded();
+                    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!("Failed to send authentication error response - {}", err);
+                                return None;
+                            }
+                        }
+                        Ok(auth_result) => {
+                            if let Err(err) = self
+                                .send_websocket_message(auth_result.server_response.into())
+                                .await
+                            {
+                                debug!("Failed to send authentication response - {}", err);
+                                return None;
+                            }
+
+                            return if let Some(client_details) = auth_result.client_details {
+                                self.active_clients_store
+                                    .insert(client_details.address, mix_sender);
+                                Some(AuthenticatedHandler::upgrade(
+                                    self,
+                                    client_details,
+                                    mix_receiver,
+                                ))
+                            } else {
+                                None
+                            };
+                        }
+                    }
+                }
+                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
+                    warn!("possibly received a sphinx packet without prior authentication. Request is going to be ignored");
+                    if let Err(err) = self
+                        .send_websocket_message(
+                            ServerResponse::new_error(
+                                "binary request without prior authentication",
+                            )
+                            .into(),
+                        )
+                        .await
+                    {
+                        debug!(
+                            "Failed to send error response during authentication - {}",
+                            err
+                        )
+                    }
+                    return None;
+                }
+
+                _ => continue,
+            };
+        }
+        None
+    }
+
+    pub(crate) async fn start_handling(self)
+    where
+        S: AsyncRead + AsyncWrite + Unpin + Send,
+    {
+        super::handle_connection(self).await
+    }
+}

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

@@ -0,0 +1,86 @@
+// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+use gateway_requests::registration::handshake::SharedKeys;
+use gateway_requests::ServerResponse;
+use log::{trace, warn};
+use nymsphinx::DestinationAddressBytes;
+use rand::{CryptoRng, Rng};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_tungstenite::WebSocketStream;
+
+pub(crate) use self::authenticated::AuthenticatedHandler;
+pub(crate) use self::fresh::FreshHandler;
+
+mod authenticated;
+mod fresh;
+
+//// TODO: note for my future self to consider the following idea:
+//// split the socket connection into sink and stream
+//// stream will be for reading explicit requests
+//// and sink for pumping responses AND mix traffic
+//// but as byproduct this might (or might not) break the clean "SocketStream" enum here
+
+pub(crate) enum SocketStream<S> {
+    RawTcp(S),
+    UpgradedWebSocket(WebSocketStream<S>),
+    Invalid,
+}
+
+impl<S> SocketStream<S> {
+    fn is_websocket(&self) -> bool {
+        matches!(self, SocketStream::UpgradedWebSocket(_))
+    }
+}
+
+#[derive(Copy, Clone)]
+pub(crate) struct ClientDetails {
+    pub(crate) address: DestinationAddressBytes,
+    pub(crate) shared_keys: SharedKeys,
+}
+
+impl ClientDetails {
+    pub(crate) fn new(address: DestinationAddressBytes, shared_keys: SharedKeys) -> Self {
+        ClientDetails {
+            address,
+            shared_keys,
+        }
+    }
+}
+
+pub(crate) struct InitialAuthResult {
+    pub(crate) client_details: Option<ClientDetails>,
+    pub(crate) server_response: ServerResponse,
+}
+
+impl InitialAuthResult {
+    fn new(client_details: Option<ClientDetails>, server_response: ServerResponse) -> Self {
+        InitialAuthResult {
+            client_details,
+            server_response,
+        }
+    }
+}
+
+pub(crate) async fn handle_connection<R, S>(mut handle: FreshHandler<R, S>)
+where
+    R: Rng + CryptoRng,
+    S: AsyncRead + AsyncWrite + Unpin + Send,
+{
+    if let Err(err) = handle.perform_websocket_handshake().await {
+        warn!(
+            "Failed to complete WebSocket handshake - {}. Stopping the handler",
+            err
+        );
+        return;
+    }
+
+    trace!("Managed to perform websocket handshake!");
+
+    if let Some(auth_handle) = handle.perform_initial_authentication().await {
+        auth_handle.listen_for_requests().await
+    } else {
+        warn!("Authentication has failed")
+    }
+    trace!("The handler is done!");
+}

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

@@ -1,9 +1,9 @@
 // Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::node::client_handling::bandwidth::empty_bandwidth_database;
-use crate::node::client_handling::clients_handler::ClientsHandlerRequestSender;
-use crate::node::client_handling::websocket::connection_handler::Handle;
+use crate::node::client_handling::active_clients::ActiveClientsStore;
+use crate::node::client_handling::websocket::connection_handler::FreshHandler;
+use crate::node::storage::PersistentStorage;
 use coconut_interface::VerificationKey;
 use crypto::asymmetric::identity;
 use log::*;
@@ -33,10 +33,13 @@ impl Listener {
         }
     }
 
+    // TODO: change the signature to pub(crate) async fn run(&self, handler: Handler)
+
     pub(crate) async fn run(
         &mut self,
-        clients_handler_sender: ClientsHandlerRequestSender,
         outbound_mix_sender: MixForwardingSender,
+        storage: PersistentStorage,
+        active_clients_store: ActiveClientsStore,
     ) {
         info!("Starting websocket listener at {}", self.address);
         let tcp_listener = match tokio::net::TcpListener::bind(self.address).await {
@@ -47,22 +50,20 @@ impl Listener {
             }
         };
 
-        let bandwidths = empty_bandwidth_database();
-
         loop {
             match tcp_listener.accept().await {
                 Ok((socket, remote_addr)) => {
                     trace!("received a socket connection from {}", remote_addr);
                     // TODO: I think we *REALLY* need a mechanism for having a maximum number of connected
                     // clients or spawned tokio tasks -> perhaps a worker system?
-                    let mut handle = Handle::new(
+                    let handle = FreshHandler::new(
                         OsRng,
                         socket,
-                        clients_handler_sender.clone(),
                         outbound_mix_sender.clone(),
                         Arc::clone(&self.local_identity),
                         self.aggregated_verification_key.clone(),
-                        Arc::clone(&bandwidths),
+                        storage.clone(),
+                        active_clients_store.clone(),
                     );
                     tokio::spawn(async move { handle.start_handling().await });
                 }
@@ -73,9 +74,13 @@ impl Listener {
 
     pub(crate) fn start(
         mut self,
-        clients_handler_sender: ClientsHandlerRequestSender,
         outbound_mix_sender: MixForwardingSender,
+        storage: PersistentStorage,
+        active_clients_store: ActiveClientsStore,
     ) -> JoinHandle<()> {
-        tokio::spawn(async move { self.run(clients_handler_sender, outbound_mix_sender).await })
+        tokio::spawn(async move {
+            self.run(outbound_mix_sender, storage, active_clients_store)
+                .await
+        })
     }
 }

gateway/src/node/mixnet_handling/receiver/connection_handler.rs

@@ -1,14 +1,11 @@
 // Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::node::client_handling::clients_handler::{
-    ClientsHandlerRequest, ClientsHandlerRequestSender, ClientsHandlerResponse,
-};
+use crate::node::client_handling::active_clients::ActiveClientsStore;
 use crate::node::client_handling::websocket::message_receiver::MixMessageSender;
 use crate::node::mixnet_handling::receiver::packet_processing::PacketProcessor;
-use crate::node::storage::inboxes::{ClientStorage, StoreData};
-use dashmap::DashMap;
-use futures::channel::oneshot;
+use crate::node::storage::error::StorageError;
+use crate::node::storage::PersistentStorage;
 use futures::StreamExt;
 use log::*;
 use mixnet_client::forwarder::MixForwardingSender;
@@ -17,67 +14,86 @@ use nymsphinx::forwarding::packet::MixPacket;
 use nymsphinx::framing::codec::SphinxCodec;
 use nymsphinx::framing::packet::FramedSphinxPacket;
 use nymsphinx::DestinationAddressBytes;
-use std::io;
+use std::collections::HashMap;
 use std::net::SocketAddr;
-use std::sync::Arc;
 use tokio::net::TcpStream;
 use tokio_util::codec::Framed;
 
 pub(crate) struct ConnectionHandler {
     packet_processor: PacketProcessor,
 
-    // TODO: method for cache invalidation so that we wouldn't keep all stale channel references
-    // we could use our friend DelayQueue. Alternatively we could periodically check for if the
-    // channels are closed.
-    available_socket_senders_cache: DashMap<DestinationAddressBytes, MixMessageSender>,
-    client_store: ClientStorage,
-    clients_handler_sender: ClientsHandlerRequestSender,
+    // TODO: investigate performance trade-offs for whether this cache even makes sense
+    // at this point.
+    // keep the following in mind: each action on ActiveClientsStore requires going through RwLock
+    // and each `get` internally copies the channel, however, is it really that expensive?
+    clients_store_cache: HashMap<DestinationAddressBytes, MixMessageSender>,
+    active_clients_store: ActiveClientsStore,
+    storage: PersistentStorage,
     ack_sender: MixForwardingSender,
 }
 
-impl ConnectionHandler {
-    pub(crate) fn new(
-        packet_processor: PacketProcessor,
-        clients_handler_sender: ClientsHandlerRequestSender,
-        client_store: ClientStorage,
-
-        ack_sender: MixForwardingSender,
-    ) -> Self {
-        ConnectionHandler {
-            packet_processor,
-            available_socket_senders_cache: DashMap::new(),
-            client_store,
-            clients_handler_sender,
-            ack_sender,
-        }
-    }
-
-    pub(crate) fn clone_without_cache(&self) -> Self {
-        // TODO: should this be even cloned?
-        let senders_cache = DashMap::with_capacity(self.available_socket_senders_cache.capacity());
-        for element_guard in self.available_socket_senders_cache.iter() {
-            let (k, v) = element_guard.pair();
-            // TODO: this will be made redundant once there's some cache invalidator mechanism here
+impl Clone for ConnectionHandler {
+    fn clone(&self) -> Self {
+        // remove stale entries from the cache while cloning
+        let mut clients_store_cache = HashMap::with_capacity(self.clients_store_cache.capacity());
+        for (k, v) in self.clients_store_cache.iter() {
             if !v.is_closed() {
-                senders_cache.insert(*k, v.clone());
+                clients_store_cache.insert(*k, v.clone());
             }
         }
 
         ConnectionHandler {
             packet_processor: self.packet_processor.clone(),
-            available_socket_senders_cache: senders_cache,
-            client_store: self.client_store.clone(),
-            clients_handler_sender: self.clients_handler_sender.clone(),
+            clients_store_cache,
+            active_clients_store: self.active_clients_store.clone(),
+            storage: self.storage.clone(),
             ack_sender: self.ack_sender.clone(),
         }
     }
+}
+
+impl ConnectionHandler {
+    pub(crate) fn new(
+        packet_processor: PacketProcessor,
+        storage: PersistentStorage,
+        ack_sender: MixForwardingSender,
+        active_clients_store: ActiveClientsStore,
+    ) -> Self {
+        ConnectionHandler {
+            packet_processor,
+            clients_store_cache: HashMap::new(),
+            storage,
+            active_clients_store,
+            ack_sender,
+        }
+    }
+
+    fn update_clients_store_cache_entry(&mut self, client_address: DestinationAddressBytes) {
+        if let Some(client_sender) = self.active_clients_store.get(client_address) {
+            self.clients_store_cache
+                .insert(client_address, client_sender);
+        }
+    }
+
+    fn check_cache(&mut self, client_address: DestinationAddressBytes) {
+        match self.clients_store_cache.get(&client_address) {
+            None => self.update_clients_store_cache_entry(client_address),
+            Some(entry) => {
+                if entry.is_closed() {
+                    self.update_clients_store_cache_entry(client_address)
+                }
+            }
+        }
+    }
 
     fn try_push_message_to_client(
-        &self,
-        sender_channel: Option<MixMessageSender>,
+        &mut self,
+        client_address: DestinationAddressBytes,
         message: Vec<u8>,
     ) -> Result<(), Vec<u8>> {
-        match sender_channel {
+        self.check_cache(client_address);
+
+        match self.clients_store_cache.get(&client_address) {
             None => Err(message),
             Some(sender_channel) => {
                 sender_channel
@@ -89,76 +105,17 @@ impl ConnectionHandler {
         }
     }
 
-    fn remove_stale_client_sender(&self, client_address: &DestinationAddressBytes) {
-        if self
-            .available_socket_senders_cache
-            .remove(client_address)
-            .is_none()
-        {
-            warn!(
-                "Tried to remove stale entry for non-existent client sender: {}",
-                client_address
-            )
-        }
-    }
-
-    async fn try_to_obtain_client_ws_message_sender(
-        &self,
-        client_address: DestinationAddressBytes,
-    ) -> Option<MixMessageSender> {
-        let mut should_remove_stale = false;
-        if let Some(sender_ref) = self.available_socket_senders_cache.get(&client_address) {
-            let sender = sender_ref.value();
-            if !sender.is_closed() {
-                return Some(sender.clone());
-            } else {
-                should_remove_stale = true;
-            }
-        }
-
-        // we want to do it outside the immutable borrow into the map
-        if should_remove_stale {
-            self.remove_stale_client_sender(&client_address)
-        }
-
-        // if we got here it means that either we have no sender channel for this client or it's closed
-        // so we must refresh it from the source, i.e. ClientsHandler
-        let (res_sender, res_receiver) = oneshot::channel();
-        let clients_handler_request = ClientsHandlerRequest::IsOnline(client_address, res_sender);
-        self.clients_handler_sender
-            .unbounded_send(clients_handler_request)
-            .unwrap(); // the receiver MUST BE alive
-
-        let client_sender = match res_receiver.await.unwrap() {
-            ClientsHandlerResponse::IsOnline(client_sender) => client_sender,
-            _ => panic!("received response to wrong query!"), // again, this should NEVER happen
-        }?;
-
-        // finally update the cache
-        if self
-            .available_socket_senders_cache
-            .insert(client_address, client_sender.clone())
-            .is_some()
-        {
-            // this warning is harmless, but I want to see if it's realistically for it to even occur
-            warn!("Other thread already updated cache for client sender!")
-        }
-
-        Some(client_sender)
-    }
-
     pub(crate) async fn store_processed_packet_payload(
         &self,
         client_address: DestinationAddressBytes,
         message: Vec<u8>,
-    ) -> io::Result<()> {
+    ) -> Result<(), StorageError> {
         debug!(
             "Storing received message for {} on the disk...",
             client_address
         );
 
-        let store_data = StoreData::new(client_address, message);
-        self.client_store.store_processed_data(store_data).await
+        self.storage.store_message(client_address, message).await
     }
 
     fn forward_ack(&self, forward_ack: Option<MixPacket>, client_address: DestinationAddressBytes) {
@@ -173,18 +130,14 @@ impl ConnectionHandler {
         }
     }
 
-    async fn handle_processed_packet(&self, processed_final_hop: ProcessedFinalHop) {
+    async fn handle_processed_packet(&mut self, processed_final_hop: ProcessedFinalHop) {
         let client_address = processed_final_hop.destination;
         let message = processed_final_hop.message;
         let forward_ack = processed_final_hop.forward_ack;
 
-        let client_sender = self
-            .try_to_obtain_client_ws_message_sender(client_address)
-            .await;
-
         // we failed to push message directly to the client - it's probably offline.
         // we should store it on the disk instead.
-        match self.try_push_message_to_client(client_sender, message) {
+        match self.try_push_message_to_client(client_address, message) {
             Err(unsent_plaintext) => match self
                 .store_processed_packet_payload(client_address, unsent_plaintext)
                 .await
@@ -201,7 +154,7 @@ impl ConnectionHandler {
         self.forward_ack(forward_ack, client_address);
     }
 
-    async fn handle_received_packet(self: Arc<Self>, framed_sphinx_packet: FramedSphinxPacket) {
+    async fn handle_received_packet(&mut self, framed_sphinx_packet: FramedSphinxPacket) {
         //
         // TODO: here be replay attack detection - it will require similar key cache to the one in
         // packet processor for vpn packets,
@@ -220,23 +173,19 @@ impl ConnectionHandler {
         self.handle_processed_packet(processed_final_hop).await
     }
 
-    pub(crate) async fn handle_connection(self, conn: TcpStream, remote: SocketAddr) {
+    pub(crate) async fn handle_connection(mut self, conn: TcpStream, remote: SocketAddr) {
         debug!("Starting connection handler for {:?}", remote);
-        let this = Arc::new(self);
         let mut framed_conn = Framed::new(conn, SphinxCodec);
         while let Some(framed_sphinx_packet) = framed_conn.next().await {
             match framed_sphinx_packet {
                 Ok(framed_sphinx_packet) => {
                     // TODO: benchmark spawning tokio task with full processing vs just processing it
-                    // synchronously (without delaying inside of course,
-                    // delay could be moved to a per-connection DelayQueue. The delay queue future
-                    // could automatically just forward packet that is done being delayed)
-                    // under higher load in single and multi-threaded situation.
-                    //
-                    // My gut feeling is saying that we might get some nice performance boost
-                    // if we introduced the change
-                    let this = Arc::clone(&this);
-                    tokio::spawn(this.handle_received_packet(framed_sphinx_packet));
+                    // synchronously under higher load in single and multi-threaded situation.
+
+                    // in theory we could process multiple sphinx packet from the same connection in parallel,
+                    // but we already handle multiple concurrent connections so if anything, making
+                    // that change would only slow things down
+                    self.handle_received_packet(framed_sphinx_packet).await;
                 }
                 Err(err) => {
                     error!(

gateway/src/node/mixnet_handling/receiver/listener.rs

@@ -30,7 +30,7 @@ impl Listener {
         loop {
             match tcp_listener.accept().await {
                 Ok((socket, remote_addr)) => {
-                    let handler = connection_handler.clone_without_cache();
+                    let handler = connection_handler.clone();
                     tokio::spawn(handler.handle_connection(socket, remote_addr));
                 }
                 Err(e) => warn!("failed to get client: {:?}", e),

gateway/src/node/mod.rs

@@ -2,10 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::config::Config;
-use crate::node::client_handling::clients_handler::{ClientsHandler, ClientsHandlerRequestSender};
+use crate::node::client_handling::active_clients::ActiveClientsStore;
 use crate::node::client_handling::websocket;
 use crate::node::mixnet_handling::receiver::connection_handler::ConnectionHandler;
-use crate::node::storage::{inboxes, ClientLedger};
+use crate::node::storage::PersistentStorage;
 use coconut_interface::VerificationKey;
 use credentials::obtain_aggregate_verification_key;
 use crypto::asymmetric::{encryption, identity};
@@ -16,7 +16,6 @@ use rand::thread_rng;
 use std::net::SocketAddr;
 use std::process;
 use std::sync::Arc;
-use tokio::runtime::Runtime;
 
 pub(crate) mod client_handling;
 pub(crate) mod mixnet_handling;
@@ -28,38 +27,38 @@ pub struct Gateway {
     identity: Arc<identity::KeyPair>,
     /// x25519 keypair used for Diffie-Hellman. Currently only used for sphinx key derivation.
     encryption_keys: Arc<encryption::KeyPair>,
-    registered_clients_ledger: ClientLedger,
-    client_inbox_storage: inboxes::ClientStorage,
+    storage: PersistentStorage,
 }
 
 impl Gateway {
-    pub fn new(
+    async fn initialise_storage(config: &Config) -> PersistentStorage {
+        let path = config.get_persistent_store_path();
+        let retrieval_limit = config.get_message_retrieval_limit();
+        match PersistentStorage::init(path, retrieval_limit).await {
+            Err(err) => panic!("failed to initialise gateway storage - {}", err),
+            Ok(storage) => storage,
+        }
+    }
+
+    pub async fn new(
         config: Config,
         encryption_keys: encryption::KeyPair,
         identity: identity::KeyPair,
     ) -> Self {
-        let registered_clients_ledger = match ClientLedger::load(config.get_clients_ledger_path()) {
-            Err(e) => panic!("Failed to load the ledger - {:?}", e),
-            Ok(ledger) => ledger,
-        };
-        let client_inbox_storage = inboxes::ClientStorage::new(
-            config.get_message_retrieval_limit() as usize,
-            config.get_stored_messages_filename_length(),
-            config.get_clients_inboxes_dir(),
-        );
+        let storage = Self::initialise_storage(&config).await;
+
         Gateway {
             config,
             identity: Arc::new(identity),
             encryption_keys: Arc::new(encryption_keys),
-            client_inbox_storage,
-            registered_clients_ledger,
+            storage,
         }
     }
 
     fn start_mix_socket_listener(
         &self,
-        clients_handler_sender: ClientsHandlerRequestSender,
         ack_sender: MixForwardingSender,
+        active_clients_store: ActiveClientsStore,
     ) {
         info!("Starting mix socket listener...");
 
@@ -68,9 +67,9 @@ impl Gateway {
 
         let connection_handler = ConnectionHandler::new(
             packet_processor,
-            clients_handler_sender,
-            self.client_inbox_storage.clone(),
+            self.storage.clone(),
             ack_sender,
+            active_clients_store,
         );
 
         let listening_address = SocketAddr::new(
@@ -84,8 +83,8 @@ impl Gateway {
     fn start_client_websocket_listener(
         &self,
         forwarding_channel: MixForwardingSender,
-        clients_handler_sender: ClientsHandlerRequestSender,
         verification_key: VerificationKey,
+        active_clients_store: ActiveClientsStore,
     ) {
         info!("Starting client [web]socket listener...");
 
@@ -99,7 +98,11 @@ impl Gateway {
             Arc::clone(&self.identity),
             verification_key,
         )
-        .start(clients_handler_sender, forwarding_channel);
+        .start(
+            forwarding_channel,
+            self.storage.clone(),
+            active_clients_store,
+        );
     }
 
     fn start_packet_forwarder(&self) -> MixForwardingSender {
@@ -116,16 +119,6 @@ impl Gateway {
         packet_sender
     }
 
-    fn start_clients_handler(&self) -> ClientsHandlerRequestSender {
-        info!("Starting clients handler");
-        let (_, clients_handler_sender) = ClientsHandler::new(
-            self.registered_clients_ledger.clone(),
-            self.client_inbox_storage.clone(),
-        )
-        .start();
-        clients_handler_sender
-    }
-
     async fn wait_for_interrupt(&self) {
         if let Err(e) = tokio::signal::ctrl_c().await {
             error!(
@@ -162,38 +155,41 @@ impl Gateway {
             .map(|node| node.gateway().identity_key.clone())
     }
 
-    // Rather than starting all futures with explicit `&Handle` argument, let's see how it works
-    // out if we make it implicit using `tokio::spawn` inside Runtime context.
-    // Basically more or less equivalent of using #[tokio::main] attribute.
-    pub fn run(&mut self) {
+    pub async fn run(&mut self) {
         info!("Starting nym gateway!");
 
-        let runtime = Runtime::new().unwrap();
-
-        runtime.block_on(async {
-            if let Some(duplicate_node_key) = self.check_if_same_ip_gateway_exists().await {
-                if duplicate_node_key == self.identity.public_key().to_base58_string() {
-                    warn!("We seem to have not unregistered after going offline - there's a node with identical identity and announce-host as us registered.")
-                } else {
-                    error!(
-                        "Our announce-host is identical to an existing node's announce-host! (its key is {:?})",
-                        duplicate_node_key
-                    );
-                    return;
-                }
+        if let Some(duplicate_node_key) = self.check_if_same_ip_gateway_exists().await {
+            if duplicate_node_key == self.identity.public_key().to_base58_string() {
+                warn!("We seem to have not unregistered after going offline - there's a node with identical identity and announce-host as us registered.")
+            } else {
+                error!(
+                    "Our announce-host is identical to an existing node's announce-host! (its key is {:?})",
+                    duplicate_node_key
+                );
+                return;
             }
+        }
 
-            let validators_verification_key = obtain_aggregate_verification_key(&self.config.get_validator_api_endpoints()).await.expect("failed to contact validators to obtain their verification keys");
+        let validators_verification_key =
+            obtain_aggregate_verification_key(&self.config.get_validator_api_endpoints())
+                .await
+                .expect("failed to contact validators to obtain their verification keys");
 
-            let mix_forwarding_channel = self.start_packet_forwarder();
-            let clients_handler_sender = self.start_clients_handler();
+        let mix_forwarding_channel = self.start_packet_forwarder();
 
-            self.start_mix_socket_listener(clients_handler_sender.clone(), mix_forwarding_channel.clone());
-            self.start_client_websocket_listener(mix_forwarding_channel, clients_handler_sender, validators_verification_key);
+        let active_clients_store = ActiveClientsStore::new();
+        self.start_mix_socket_listener(
+            mix_forwarding_channel.clone(),
+            active_clients_store.clone(),
+        );
+        self.start_client_websocket_listener(
+            mix_forwarding_channel,
+            validators_verification_key,
+            active_clients_store,
+        );
 
-            info!("Finished nym gateway startup procedure - it should now be able to receive mix and client traffic!");
+        info!("Finished nym gateway startup procedure - it should now be able to receive mix and client traffic!");
 
-            self.wait_for_interrupt().await
-        });
+        self.wait_for_interrupt().await
     }
 }

gateway/src/node/storage/bandwidth.rs

@@ -0,0 +1,106 @@
+// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::node::storage::models::PersistedBandwidth;
+
+#[derive(Clone)]
+pub(crate) struct BandwidthManager {
+    connection_pool: sqlx::SqlitePool,
+}
+
+impl BandwidthManager {
+    /// Creates new instance of the `BandwidthManager` with the provided sqlite connection pool.
+    ///
+    /// # Arguments
+    ///
+    /// * `connection_pool`: database connection pool to use.
+    pub(crate) fn new(connection_pool: sqlx::SqlitePool) -> Self {
+        BandwidthManager { connection_pool }
+    }
+
+    /// Creates a new bandwidth entry for the particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address_bs58`: base58-encoded address of the client.
+    pub(crate) async fn insert_new_client(
+        &self,
+        client_address_bs58: &str,
+    ) -> Result<(), sqlx::Error> {
+        sqlx::query!(
+            "INSERT INTO available_bandwidth(client_address_bs58, available) VALUES (?, 0)",
+            client_address_bs58
+        )
+        .execute(&self.connection_pool)
+        .await?;
+        Ok(())
+    }
+
+    /// Tries to retrieve available bandwidth for the particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address_bs58`: base58-encoded address of the client.
+    pub(crate) async fn get_available_bandwidth(
+        &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
+    }
+
+    /// Increases 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 added to the client.
+    pub(crate) async fn increase_available_bandwidth(
+        &self,
+        client_address_bs58: &str,
+        amount: i64,
+    ) -> Result<(), sqlx::Error> {
+        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 - ?
+                WHERE client_address_bs58 = ?
+            "#,
+            amount,
+            client_address_bs58
+        )
+        .execute(&self.connection_pool)
+        .await?;
+        Ok(())
+    }
+}

gateway/src/node/storage/error.rs

@@ -0,0 +1,13 @@
+// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub(crate) enum StorageError {
+    #[error("Database experienced an internal error - {0}")]
+    InternalDatabaseError(#[from] sqlx::Error),
+
+    #[error("Failed to perform database migration - {0}")]
+    MigrationError(#[from] sqlx::migrate::MigrateError),
+}

gateway/src/node/storage/inboxes.rs

@@ -1,191 +1,122 @@
 // Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use futures::lock::Mutex;
-use futures::StreamExt;
-use gateway_requests::DUMMY_MESSAGE_CONTENT;
-use log::*;
-use nymsphinx::DestinationAddressBytes;
-use rand::Rng;
-use std::io;
-use std::path::PathBuf;
-use std::sync::Arc;
-use tokio::fs;
-use tokio::fs::File;
-use tokio::io::AsyncWriteExt;
-use tokio_stream::wrappers::ReadDirStream;
+use crate::node::storage::models::StoredMessage;
 
-fn dummy_message() -> ClientFile {
-    ClientFile {
-        content: DUMMY_MESSAGE_CONTENT.to_vec(),
-        path: Default::default(),
-    }
+#[derive(Clone)]
+pub(crate) struct InboxManager {
+    connection_pool: sqlx::SqlitePool,
+    /// Maximum number of messages that can be obtained from the database per operation.
+    /// It is used to prevent out of memory errors in the case of client receiving a lot of data while
+    /// offline and then loading it all at once when he comes back online.
+    retrieval_limit: i64,
 }
 
-#[derive(Clone, Debug)]
-pub struct ClientFile {
-    content: Vec<u8>,
-    path: PathBuf,
-}
-
-impl ClientFile {
-    fn new(content: Vec<u8>, path: PathBuf) -> Self {
-        ClientFile { content, path }
-    }
-
-    pub(crate) fn into_tuple(self) -> (Vec<u8>, PathBuf) {
-        (self.content, self.path)
-    }
-}
-
-pub struct StoreData {
-    client_address: DestinationAddressBytes,
-    message: Vec<u8>,
-}
-
-impl StoreData {
-    pub(crate) fn new(client_address: DestinationAddressBytes, message: Vec<u8>) -> Self {
-        StoreData {
-            client_address,
-            message,
-        }
-    }
-}
-
-// TODO: replace with proper database...
-// Note: you should NEVER create more than a single instance of this using 'new()'.
-// You should always use .clone() to create additional instances
-#[derive(Clone, Debug)]
-pub struct ClientStorage {
-    inner: Arc<Mutex<ClientStorageInner>>,
-}
-
-// even though the data inside is extremely cheap to copy, we have to have a single mutex,
-// so might as well store the data behind it
-pub struct ClientStorageInner {
-    // basically part of rate limiting which does not exist anymore
-    #[allow(dead_code)]
-    message_retrieval_limit: usize,
-    filename_length: u16,
-    main_store_path_dir: PathBuf,
-}
-
-// TODO: change it to some generic implementation to inject fs (or even better - proper database)
-impl ClientStorage {
-    pub(crate) fn new(message_limit: usize, filename_len: u16, main_store_dir: PathBuf) -> Self {
-        ClientStorage {
-            inner: Arc::new(Mutex::new(ClientStorageInner {
-                message_retrieval_limit: message_limit,
-                filename_length: filename_len,
-                main_store_path_dir: main_store_dir,
-            })),
+impl InboxManager {
+    /// Creates new instance of the `InboxManager` with the provided sqlite connection pool.
+    ///
+    /// # Arguments
+    ///
+    /// * `connection_pool`: database connection pool to use.
+    pub(crate) fn new(connection_pool: sqlx::SqlitePool, retrieval_limit: i64) -> Self {
+        InboxManager {
+            connection_pool,
+            retrieval_limit,
         }
     }
 
-    pub(crate) async fn create_storage_dir(
+    /// Inserts new message to the storage for an offline client for future retrieval.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address_bs58`: base58-encoded address of the client
+    /// * `content`: raw content of the message to store.
+    pub(crate) async fn insert_message(
         &self,
-        client_address: DestinationAddressBytes,
-    ) -> io::Result<()> {
-        let inner_data = self.inner.lock().await;
-
-        let client_dir_name = client_address.as_base58_string();
-        let full_store_dir = inner_data.main_store_path_dir.join(client_dir_name);
-        fs::create_dir_all(full_store_dir).await
+        client_address_bs58: &str,
+        content: Vec<u8>,
+    ) -> Result<(), sqlx::Error> {
+        sqlx::query!(
+            "INSERT INTO message_store(client_address_bs58, content) VALUES (?, ?)",
+            client_address_bs58,
+            content,
+        )
+        .execute(&self.connection_pool)
+        .await?;
+        Ok(())
     }
 
-    pub(crate) fn generate_random_file_name(length: usize) -> String {
-        rand::thread_rng()
-            .sample_iter(&rand::distributions::Alphanumeric)
-            .take(length)
-            .collect::<String>()
-    }
-
-    pub(crate) async fn store_processed_data(&self, store_data: StoreData) -> io::Result<()> {
-        let inner_data = self.inner.lock().await;
-
-        let client_dir_name = store_data.client_address.as_base58_string();
-        let full_store_dir = inner_data.main_store_path_dir.join(client_dir_name);
-        let full_store_path = full_store_dir.join(Self::generate_random_file_name(
-            inner_data.filename_length as usize,
-        ));
-        trace!(
-            "going to store: {:?} in file: {:?}",
-            store_data.message,
-            full_store_path
-        );
-
-        let mut file = File::create(full_store_path).await?;
-        file.write_all(store_data.message.as_ref()).await
-    }
-
-    pub(crate) async fn retrieve_all_client_messages(
+    /// Retrieves messages stored for the particular client specified by the provided address.
+    ///
+    /// It also respects the specified retrieval limit. If there are more messages stored than allowed
+    /// by the limit, it returns id of the last message retrieved to indicate start of the next query.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address_bs58`: base58-encoded address of the client
+    /// * `start_after`: optional starting id of the messages to grab
+    ///
+    /// returns the retrieved messages alongside optional id of the last message retrieved if
+    /// there are more messages to retrieve.
+    pub(crate) async fn get_messages(
         &self,
-        client_address: DestinationAddressBytes,
-    ) -> io::Result<Vec<ClientFile>> {
-        let inner_data = self.inner.lock().await;
-
-        let client_dir_name = client_address.as_base58_string();
-        let full_store_dir = inner_data.main_store_path_dir.join(client_dir_name);
-
-        trace!("going to lookup: {:?}!", full_store_dir);
-        if !full_store_dir.exists() {
-            return Err(io::Error::new(
-                io::ErrorKind::NotFound,
-                "Target client does not exist",
-            ));
-        }
-
-        let mut msgs = Vec::new();
-        let mut read_dir = ReadDirStream::new(fs::read_dir(full_store_dir).await?);
-        while let Some(dir_entry) = read_dir.next().await {
-            if let Ok(dir_entry) = dir_entry {
-                if !Self::is_valid_file(&dir_entry).await {
-                    continue;
-                }
-                let client_file =
-                    ClientFile::new(fs::read(dir_entry.path()).await?, dir_entry.path());
-                msgs.push(client_file)
-            }
-        }
-        Ok(msgs)
-    }
-
-    async fn is_valid_file(entry: &fs::DirEntry) -> bool {
-        let metadata = match entry.metadata().await {
-            Ok(meta) => meta,
-            Err(e) => {
-                error!(
-                    "potentially corrupted client inbox! ({:?} - failed to read its metadata - {:?}",
-                    entry.path(),
-                    e,
-                );
-                return false;
-            }
+        client_address_bs58: &str,
+        start_after: Option<i64>,
+    ) -> Result<(Vec<StoredMessage>, Option<i64>), sqlx::Error> {
+        // get 1 additional message to check whether there will be more to grab
+        // next time
+        let limit = self.retrieval_limit + 1;
+        let mut res = if let Some(start_after) = start_after {
+            sqlx::query_as!(
+                StoredMessage,
+                r#"
+                    SELECT * FROM message_store 
+                    WHERE client_address_bs58 = ? AND id > ?
+                    ORDER BY id ASC
+                    LIMIT ?;
+                "#,
+                start_after,
+                client_address_bs58,
+                limit
+            )
+            .fetch_all(&self.connection_pool)
+            .await?
+        } else {
+            sqlx::query_as!(
+                StoredMessage,
+                r#"
+                    SELECT * FROM message_store 
+                    WHERE client_address_bs58 = ?
+                    ORDER BY id ASC
+                    LIMIT ?;
+                "#,
+                client_address_bs58,
+                limit
+            )
+            .fetch_all(&self.connection_pool)
+            .await?
         };
 
-        let is_file = metadata.is_file();
-        if !is_file {
-            error!(
-                "potentially corrupted client inbox! - found a non-file - {:?}",
-                entry.path()
-            );
+        if res.len() > self.retrieval_limit as usize {
+            res.truncate(self.retrieval_limit as usize);
+            // assuming retrieval_limit > 0, unwrap will not fail
+            let start_after = res.last().unwrap().id;
+            Ok((res, Some(start_after)))
+            //
+        } else {
+            Ok((res, None))
         }
-
-        is_file
     }
 
-    pub(crate) async fn delete_files(&self, file_paths: Vec<PathBuf>) -> io::Result<()> {
-        let dummy_message = dummy_message();
-        let _guard = self.inner.lock().await;
-
-        for file_path in file_paths {
-            if file_path == dummy_message.path {
-                continue;
-            }
-            if let Err(e) = fs::remove_file(file_path).await {
-                error!("Failed to delete client message! - {:?}", e)
-            }
-        }
+    /// Removes message with the specified id
+    ///
+    /// # Arguments
+    ///
+    /// * `id`: id of the message to remove
+    pub(crate) async fn remove_message(&self, id: i64) -> Result<(), sqlx::Error> {
+        sqlx::query!("DELETE FROM message_store WHERE id = ?", id)
+            .execute(&self.connection_pool)
+            .await?;
         Ok(())
     }
 }

gateway/src/node/storage/ledger.rs

@@ -1,141 +0,0 @@
-// Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
-// SPDX-License-Identifier: Apache-2.0
-
-use gateway_requests::authentication::encrypted_address::EncryptedAddressBytes;
-use gateway_requests::generic_array::typenum::Unsigned;
-use gateway_requests::iv::IV;
-use gateway_requests::registration::handshake::{SharedKeySize, SharedKeys};
-use log::*;
-use nymsphinx::{DestinationAddressBytes, DESTINATION_ADDRESS_LENGTH};
-use std::path::PathBuf;
-
-#[derive(Debug)]
-pub(crate) enum ClientLedgerError {
-    Read(sled::Error),
-    Write(sled::Error),
-    Open(sled::Error),
-}
-
-#[derive(Debug, Clone)]
-// Note: you should NEVER create more than a single instance of this using 'new()'.
-// You should always use .clone() to create additional instances
-pub(crate) struct ClientLedger {
-    db: sled::Db,
-}
-
-impl ClientLedger {
-    pub(crate) fn load(file: PathBuf) -> Result<Self, ClientLedgerError> {
-        let db = match sled::open(file) {
-            Err(e) => return Err(ClientLedgerError::Open(e)),
-            Ok(db) => db,
-        };
-
-        let ledger = ClientLedger { db };
-
-        ledger.db.iter().keys().for_each(|key| {
-            trace!(
-                "key: {:?}",
-                ledger
-                    .read_destination_address_bytes(key.unwrap())
-                    .as_base58_string()
-            );
-        });
-
-        Ok(ledger)
-    }
-
-    fn read_shared_key(&self, raw_key: sled::IVec) -> SharedKeys {
-        let key_bytes_ref = raw_key.as_ref();
-        // if this fails it means we have some database corruption and we
-        // absolutely can't continue
-
-        if key_bytes_ref.len() != SharedKeySize::to_usize() {
-            error!("CLIENT LEDGER DATA CORRUPTION - SHARED KEY HAS INVALID LENGTH");
-            panic!("CLIENT LEDGER DATA CORRUPTION - SHARED KEY HAS INVALID LENGTH");
-        }
-
-        // this can only fail if the bytes have invalid length but we already asserted it
-        SharedKeys::try_from_bytes(key_bytes_ref).unwrap()
-    }
-
-    fn read_destination_address_bytes(
-        &self,
-        raw_destination: sled::IVec,
-    ) -> DestinationAddressBytes {
-        let destination_ref = raw_destination.as_ref();
-        // if this fails it means we have some database corruption and we
-        // absolutely can't continue
-        if destination_ref.len() != DESTINATION_ADDRESS_LENGTH {
-            error!("CLIENT LEDGER DATA CORRUPTION - CLIENT ADDRESS HAS INVALID LENGTH");
-            panic!("CLIENT LEDGER DATA CORRUPTION - CLIENT ADDRESS HAS INVALID LENGTH");
-        }
-
-        let mut destination_bytes = [0u8; DESTINATION_ADDRESS_LENGTH];
-        destination_bytes.copy_from_slice(destination_ref);
-        DestinationAddressBytes::from_bytes(destination_bytes)
-    }
-
-    pub(crate) fn verify_shared_key(
-        &self,
-        client_address: &DestinationAddressBytes,
-        encrypted_address: &EncryptedAddressBytes,
-        iv: &IV,
-    ) -> Result<bool, ClientLedgerError> {
-        match self.db.get(client_address.as_bytes_ref()) {
-            Err(e) => Err(ClientLedgerError::Read(e)),
-            Ok(existing_key) => match existing_key {
-                Some(existing_key_ivec) => {
-                    let shared_key = &self.read_shared_key(existing_key_ivec);
-                    Ok(encrypted_address.verify(client_address, shared_key, iv))
-                }
-                None => Ok(false),
-            },
-        }
-    }
-
-    pub(crate) fn get_shared_key(
-        &self,
-        client_address: &DestinationAddressBytes,
-    ) -> Result<Option<SharedKeys>, ClientLedgerError> {
-        match self.db.get(client_address.as_bytes_ref()) {
-            Err(e) => Err(ClientLedgerError::Read(e)),
-            Ok(existing_key) => Ok(existing_key.map(|key_ivec| self.read_shared_key(key_ivec))),
-        }
-    }
-
-    pub(crate) fn insert_shared_key(
-        &mut self,
-        shared_key: SharedKeys,
-        client_address: DestinationAddressBytes,
-    ) -> Result<Option<SharedKeys>, ClientLedgerError> {
-        let insertion_result = match self
-            .db
-            .insert(client_address.as_bytes_ref(), shared_key.to_bytes())
-        {
-            Err(e) => Err(ClientLedgerError::Write(e)),
-            Ok(existing_key) => {
-                Ok(existing_key.map(|existing_key| self.read_shared_key(existing_key)))
-            }
-        };
-
-        // registration doesn't happen that often so might as well flush it to the disk to be sure
-        self.db.flush().unwrap();
-        insertion_result
-    }
-
-    pub(crate) fn remove_shared_key(
-        &mut self,
-        client_address: &DestinationAddressBytes,
-    ) -> Result<Option<SharedKeys>, ClientLedgerError> {
-        let removal_result = match self.db.remove(client_address.as_bytes_ref()) {
-            Err(e) => Err(ClientLedgerError::Write(e)),
-            Ok(existing_key) => {
-                Ok(existing_key.map(|existing_key| self.read_shared_key(existing_key)))
-            }
-        };
-
-        // removal of keys happens extremely rarely, so flush is also fine here
-        self.db.flush().unwrap();
-        removal_result
-    }
-}

gateway/src/node/storage/mod.rs

@@ -1,7 +1,247 @@
 // Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-pub(crate) mod inboxes;
-mod ledger;
+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::shared_keys::SharedKeysManager;
+use gateway_requests::registration::handshake::SharedKeys;
+use log::{debug, error};
+use nymsphinx::DestinationAddressBytes;
+use sqlx::ConnectOptions;
+use std::path::Path;
 
-pub(crate) use ledger::ClientLedger;
+mod bandwidth;
+pub(crate) mod error;
+mod inboxes;
+mod models;
+mod shared_keys;
+
+// note that clone here is fine as upon cloning the same underlying pool will be used
+#[derive(Clone)]
+pub(crate) struct PersistentStorage {
+    shared_key_manager: SharedKeysManager,
+    inbox_manager: InboxManager,
+    bandwidth_manager: BandwidthManager,
+}
+
+impl PersistentStorage {
+    /// Initialises `PersistentStorage` using the provided path.
+    ///
+    /// # Arguments
+    ///
+    /// * `database_path`: path to the database.
+    /// * `message_retrieval_limit`: maximum number of stored client messages that can be retrieved at once.
+    pub(crate) async fn init<P: AsRef<Path>>(
+        database_path: P,
+        message_retrieval_limit: i64,
+    ) -> Result<Self, StorageError> {
+        debug!(
+            "Attempting to connect to database {:?}",
+            database_path.as_ref().as_os_str()
+        );
+
+        // TODO: we can inject here more stuff based on our gateway global config
+        // struct. Maybe different pool size or timeout intervals?
+        let mut opts = sqlx::sqlite::SqliteConnectOptions::new()
+            .filename(database_path)
+            .create_if_missing(true);
+
+        // TODO: do we want auto_vacuum ?
+
+        opts.disable_statement_logging();
+
+        let connection_pool = match sqlx::SqlitePool::connect_with(opts).await {
+            Ok(db) => db,
+            Err(err) => {
+                error!("Failed to connect to SQLx database: {}", err);
+                return Err(err.into());
+            }
+        };
+
+        if let Err(err) = sqlx::migrate!("./migrations").run(&connection_pool).await {
+            error!("Failed to perform migration on the SQLx database: {}", err);
+            return Err(err.into());
+        }
+
+        // the cloning here are cheap as connection pool is stored behind an Arc
+        Ok(PersistentStorage {
+            shared_key_manager: SharedKeysManager::new(connection_pool.clone()),
+            inbox_manager: InboxManager::new(connection_pool.clone(), message_retrieval_limit),
+            bandwidth_manager: BandwidthManager::new(connection_pool),
+        })
+    }
+
+    /// Inserts provided derived shared keys into the database.
+    /// If keys previously existed for the provided client, they are overwritten with the new data.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client
+    /// * `shared_keys`: shared encryption (AES128CTR) and mac (hmac-blake3) derived shared keys to store.
+    pub(crate) async fn insert_shared_keys(
+        &self,
+        client_address: DestinationAddressBytes,
+        shared_keys: SharedKeys,
+    ) -> Result<(), StorageError> {
+        let persisted_shared_keys = PersistedSharedKeys {
+            client_address_bs58: client_address.as_base58_string(),
+            derived_aes128_ctr_blake3_hmac_keys_bs58: shared_keys.to_base58_string(),
+        };
+        self.shared_key_manager
+            .insert_shared_keys(persisted_shared_keys)
+            .await?;
+        Ok(())
+    }
+
+    /// Tries to retrieve shared keys stored for the particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client
+    pub(crate) async fn get_shared_keys(
+        &self,
+        client_address: DestinationAddressBytes,
+    ) -> Result<Option<PersistedSharedKeys>, StorageError> {
+        let keys = self
+            .shared_key_manager
+            .get_shared_keys(&client_address.as_base58_string())
+            .await?;
+        Ok(keys)
+    }
+
+    /// Removes from the database shared keys derived with the particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client
+    // currently there is no code flow that causes removal (not overwriting)
+    // of the stored keys. However, retain the function for consistency and completion sake
+    #[allow(dead_code)]
+    pub(crate) async fn remove_shared_keys(
+        &self,
+        client_address: DestinationAddressBytes,
+    ) -> Result<(), StorageError> {
+        self.shared_key_manager
+            .remove_shared_keys(&client_address.as_base58_string())
+            .await?;
+        Ok(())
+    }
+
+    /// Inserts new message to the storage for an offline client for future retrieval.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client
+    /// * `message`: raw message to store.
+    pub(crate) async fn store_message(
+        &self,
+        client_address: DestinationAddressBytes,
+        message: Vec<u8>,
+    ) -> Result<(), StorageError> {
+        self.inbox_manager
+            .insert_message(&client_address.as_base58_string(), message)
+            .await?;
+        Ok(())
+    }
+
+    /// Retrieves messages stored for the particular client specified by the provided address.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client
+    /// * `start_after`: optional starting id of the messages to grab
+    ///
+    /// returns the retrieved messages alongside optional id of the last message retrieved if
+    /// there are more messages to retrieve.
+    pub(crate) async fn retrieve_messages(
+        &self,
+        client_address: DestinationAddressBytes,
+        start_after: Option<i64>,
+    ) -> Result<(Vec<StoredMessage>, Option<i64>), StorageError> {
+        let messages = self
+            .inbox_manager
+            .get_messages(&client_address.as_base58_string(), start_after)
+            .await?;
+        Ok(messages)
+    }
+
+    /// Removes messages with the specified ids
+    ///
+    /// # Arguments
+    ///
+    /// * `ids`: ids of the messages to remove
+    pub(crate) async fn remove_messages(&self, ids: Vec<i64>) -> Result<(), StorageError> {
+        for id in ids {
+            self.inbox_manager.remove_message(id).await?;
+        }
+        Ok(())
+    }
+
+    /// Creates a new bandwidth entry for the particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client
+    pub(crate) async fn create_bandwidth_entry(
+        &self,
+        client_address: DestinationAddressBytes,
+    ) -> Result<(), StorageError> {
+        self.bandwidth_manager
+            .insert_new_client(&client_address.as_base58_string())
+            .await?;
+        Ok(())
+    }
+
+    /// Tries to retrieve available bandwidth for the particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client
+    pub(crate) async fn get_available_bandwidth(
+        &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)
+    }
+
+    /// Increases available bandwidth of the particular client by the specified amount.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address`: address of the client
+    /// * `amount`: amount of available bandwidth to be added to the client.
+    pub(crate) async fn increase_bandwidth(
+        &self,
+        client_address: DestinationAddressBytes,
+        amount: i64,
+    ) -> Result<(), StorageError> {
+        self.bandwidth_manager
+            .increase_available_bandwidth(&client_address.as_base58_string(), amount)
+            .await?;
+        Ok(())
+    }
+
+    /// 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.
+    pub(crate) async fn consume_bandwidth(
+        &self,
+        client_address: DestinationAddressBytes,
+        amount: i64,
+    ) -> Result<(), StorageError> {
+        self.bandwidth_manager
+            .decrease_available_bandwidth(&client_address.as_base58_string(), amount)
+            .await?;
+        Ok(())
+    }
+}

gateway/src/node/storage/models.rs

@@ -0,0 +1,20 @@
+// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+pub(crate) struct PersistedSharedKeys {
+    pub(crate) client_address_bs58: String,
+    pub(crate) derived_aes128_ctr_blake3_hmac_keys_bs58: String,
+}
+
+pub(crate) struct StoredMessage {
+    pub(crate) id: i64,
+    #[allow(dead_code)]
+    pub(crate) client_address_bs58: String,
+    pub(crate) content: Vec<u8>,
+}
+
+pub(crate) struct PersistedBandwidth {
+    #[allow(dead_code)]
+    pub(crate) client_address_bs58: String,
+    pub(crate) available: i64,
+}

gateway/src/node/storage/shared_keys.rs

@@ -0,0 +1,76 @@
+// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::node::storage::models::PersistedSharedKeys;
+
+#[derive(Clone)]
+pub(crate) struct SharedKeysManager {
+    connection_pool: sqlx::SqlitePool,
+}
+
+impl SharedKeysManager {
+    /// Creates new instance of the `SharedKeysManager` with the provided sqlite connection pool.
+    ///
+    /// # Arguments
+    ///
+    /// * `connection_pool`: database connection pool to use.
+    pub(crate) fn new(connection_pool: sqlx::SqlitePool) -> Self {
+        SharedKeysManager { connection_pool }
+    }
+
+    /// Inserts provided derived shared keys into the database.
+    /// If keys previously existed for the provided client, they are overwritten with the new data.
+    ///
+    /// # Arguments
+    ///
+    /// * `shared_keys`: shared encryption (AES128CTR) and mac (hmac-blake3) derived shared keys to store.
+    pub(crate) async fn insert_shared_keys(
+        &self,
+        shared_keys: PersistedSharedKeys,
+    ) -> Result<(), sqlx::Error> {
+        sqlx::query!("INSERT OR REPLACE INTO shared_keys(client_address_bs58, derived_aes128_ctr_blake3_hmac_keys_bs58) VALUES (?, ?)",
+            shared_keys.client_address_bs58,
+            shared_keys.derived_aes128_ctr_blake3_hmac_keys_bs58,
+        ).execute(&self.connection_pool).await?;
+        Ok(())
+    }
+
+    /// Tries to retrieve shared keys stored for the particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address_bs58`: base58-encoded address of the client
+    pub(crate) async fn get_shared_keys(
+        &self,
+        client_address_bs58: &str,
+    ) -> Result<Option<PersistedSharedKeys>, sqlx::Error> {
+        sqlx::query_as!(
+            PersistedSharedKeys,
+            "SELECT * FROM shared_keys WHERE client_address_bs58 = ?",
+            client_address_bs58
+        )
+        .fetch_optional(&self.connection_pool)
+        .await
+    }
+
+    /// Removes from the database shared keys derived with the particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_address_bs58`: base58-encoded address of the client
+    // currently there is no code flow that causes removal (not overwriting)
+    // of the stored keys. However, retain the function for consistency and completion sake
+    #[allow(dead_code)]
+    pub(crate) async fn remove_shared_keys(
+        &self,
+        client_address_bs58: &str,
+    ) -> Result<(), sqlx::Error> {
+        sqlx::query!(
+            "DELETE FROM shared_keys WHERE client_address_bs58 = ?",
+            client_address_bs58
+        )
+        .execute(&self.connection_pool)
+        .await?;
+        Ok(())
+    }
+}