From 021c2d14f24b5bc34c7b899b678096378f45137d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Tue, 15 Feb 2022 20:50:31 +0100 Subject: [PATCH] gateway: extract out PersistantStorage --- Cargo.lock | 1 + gateway/Cargo.toml | 1 + gateway/src/commands/init.rs | 60 +++- gateway/src/commands/mod.rs | 1 - gateway/src/commands/node_details.rs | 6 +- gateway/src/commands/run.rs | 4 +- gateway/src/config/mod.rs | 4 +- .../connection_handler/authenticated.rs | 13 +- .../websocket/connection_handler/fresh.rs | 13 +- .../websocket/connection_handler/mod.rs | 5 +- .../client_handling/websocket/listener.rs | 19 +- .../receiver/connection_handler.rs | 12 +- .../node/mixnet_handling/receiver/listener.rs | 11 +- gateway/src/node/mod.rs | 51 ++- gateway/src/node/storage/mod.rs | 293 +++++++++++++----- 15 files changed, 365 insertions(+), 129 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5649592080..4c13b838d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3847,6 +3847,7 @@ name = "nym-gateway" version = "0.12.1" dependencies = [ "anyhow", + "async-trait", "bandwidth-claim-contract", "bip39", "bs58", diff --git a/gateway/Cargo.toml b/gateway/Cargo.toml index 50a65ec598..147cae6444 100644 --- a/gateway/Cargo.toml +++ b/gateway/Cargo.toml @@ -13,6 +13,7 @@ rust-version = "1.56" [dependencies] anyhow = "1.0.53" +async-trait = { version = "0.1.51" } bip39 = "1.0.1" bs58 = "0.4.0" clap = { version = "3.0.10", features = ["cargo", "derive"] } diff --git a/gateway/src/commands/init.rs b/gateway/src/commands/init.rs index 6b384f1207..c923052b2a 100644 --- a/gateway/src/commands/init.rs +++ b/gateway/src/commands/init.rs @@ -4,7 +4,6 @@ use crate::{ commands::{override_config, OverrideConfig}, config::{persistence::pathfinder::GatewayPathfinder, Config}, - node::Gateway, }; use clap::Args; use config::NymConfig; @@ -68,7 +67,6 @@ pub struct Init { impl From for OverrideConfig { fn from(init_config: Init) -> Self { OverrideConfig { - id: init_config.id, host: Some(init_config.host), wallet_address: Some(init_config.wallet_address), mix_port: init_config.mix_port, @@ -93,20 +91,23 @@ impl From for OverrideConfig { } pub async fn execute(args: &Init) { - let override_config_fields = OverrideConfig::from(args.clone()); - let id = &override_config_fields.id; - println!("Initialising gateway {}...", id); + println!("Initialising gateway {}...", args.id); - let already_init = if Config::default_config_file_path(Some(id)).exists() { - println!("Gateway \"{}\" was already initialised before! Config information will be overwritten (but keys will be kept)!", id); + let already_init = if Config::default_config_file_path(Some(&args.id)).exists() { + println!( + "Gateway \"{}\" was already initialised before! Config information will be \ + overwritten (but keys will be kept)!", + args.id + ); true } else { false }; - let mut config = Config::new(id); + let override_config_fields = OverrideConfig::from(args.clone()); - config = override_config(config, override_config_fields); + // Initialising the config structure is just overriding a default constructed one + let config = override_config(Config::new(&args.id), override_config_fields); // if gateway was already initialised, don't generate new keys if !already_init { @@ -143,5 +144,44 @@ pub async fn execute(args: &Init) { println!("Saved configuration file to {:?}", config_save_location); println!("Gateway configuration completed.\n\n\n"); - Gateway::new(config).await.print_node_details(); + crate::node::create_gateway(config) + .await + .print_node_details(); +} + +#[cfg(test)] +mod tests { + use crate::node::{storage::InMemStorage, Gateway}; + + use super::*; + + #[tokio::test] + async fn create_gateway_with_in_mem_storage() { + let args = Init { + id: "foo-id".to_string(), + host: "foo-host".to_string(), + wallet_address: "nymt1z9egw0knv47nmur0p8vk4rcx59h9gg4zuxrrr9".to_string(), + mix_port: Some(42), + clients_port: Some(43), + announce_host: Some("foo-announce-host".to_string()), + datastore: Some("foo-datastore".to_string()), + validator_apis: None, + }; + + let config = Config::new(&args.id); + let config = override_config(config, OverrideConfig::from(args.clone())); + + let (identity_keys, sphinx_keys) = { + let mut rng = rand::rngs::OsRng; + ( + identity::KeyPair::new(&mut rng), + encryption::KeyPair::new(&mut rng), + ) + }; + + // The test is really if this instantiates with InMemStorage without panics + let _gateway = + Gateway::new_from_keys_and_storage(config, identity_keys, sphinx_keys, InMemStorage) + .await; + } } diff --git a/gateway/src/commands/mod.rs b/gateway/src/commands/mod.rs index a10f81b494..8ce291f83b 100644 --- a/gateway/src/commands/mod.rs +++ b/gateway/src/commands/mod.rs @@ -43,7 +43,6 @@ pub(crate) enum Commands { // Configuration that can be overridden. pub(crate) struct OverrideConfig { - id: String, host: Option, wallet_address: Option, mix_port: Option, diff --git a/gateway/src/commands/node_details.rs b/gateway/src/commands/node_details.rs index d5a2398df9..15602113c1 100644 --- a/gateway/src/commands/node_details.rs +++ b/gateway/src/commands/node_details.rs @@ -1,7 +1,7 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::{config::Config, node::Gateway}; +use crate::config::Config; use clap::Args; use config::NymConfig; use log::error; @@ -26,5 +26,7 @@ pub async fn execute(args: &NodeDetails) { } }; - Gateway::new(config).await.print_node_details(); + crate::node::create_gateway(config) + .await + .print_node_details(); } diff --git a/gateway/src/commands/run.rs b/gateway/src/commands/run.rs index c030b86076..fc7ff4c4e6 100644 --- a/gateway/src/commands/run.rs +++ b/gateway/src/commands/run.rs @@ -4,7 +4,6 @@ use crate::{ commands::{override_config, version_check, OverrideConfig}, config::Config, - node::Gateway, }; use clap::Args; use config::NymConfig; @@ -68,7 +67,6 @@ pub struct Run { impl From for OverrideConfig { fn from(run_config: Run) -> Self { OverrideConfig { - id: run_config.id, host: run_config.host, wallet_address: run_config.wallet_address, mix_port: run_config.mix_port, @@ -133,7 +131,7 @@ pub async fn execute(args: &Run) { show_binding_warning(config.get_listening_address().to_string()); } - let mut gateway = Gateway::new(config).await; + let mut gateway = crate::node::create_gateway(config).await; println!( "\nTo bond your gateway you will need to install the Nym wallet, go to https://nymtech.net/get-involved and select the Download button.\n\ Select the correct version and install it to your machine. You will need to provide the following: \n "); diff --git a/gateway/src/config/mod.rs b/gateway/src/config/mod.rs index 421703e7df..2524bfdcf0 100644 --- a/gateway/src/config/mod.rs +++ b/gateway/src/config/mod.rs @@ -416,11 +416,11 @@ impl Default for Gateway { #[derive(Debug, Default, Deserialize, PartialEq, Serialize)] #[serde(deny_unknown_fields)] -pub struct Logging {} +struct Logging {} #[derive(Debug, Deserialize, PartialEq, Serialize)] #[serde(default)] -pub struct Debug { +struct Debug { /// Initial value of an exponential backoff to reconnect to dropped TCP connection when /// forwarding sphinx packets. #[serde(with = "humantime_serde")] diff --git a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs index 5fde529264..becd4ee436 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs @@ -4,6 +4,7 @@ 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 crate::node::storage::Storage; use futures::StreamExt; use gateway_requests::iv::IVConversionError; use gateway_requests::types::{BinaryRequest, ServerResponse}; @@ -89,14 +90,14 @@ impl IntoWSMessage for Result { } } -pub(crate) struct AuthenticatedHandler { - inner: FreshHandler, +pub(crate) struct AuthenticatedHandler { + inner: FreshHandler, client: ClientDetails, mix_receiver: MixMessageReceiver, } // explicitly remove handle from the global store upon being dropped -impl Drop for AuthenticatedHandler { +impl Drop for AuthenticatedHandler { fn drop(&mut self) { self.inner .active_clients_store @@ -104,10 +105,11 @@ impl Drop for AuthenticatedHandler { } } -impl AuthenticatedHandler +impl AuthenticatedHandler where // TODO: those trait bounds here don't really make sense.... R: Rng + CryptoRng, + St: Storage, { /// 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 @@ -119,7 +121,7 @@ where /// * `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, + fresh: FreshHandler, client: ClientDetails, mix_receiver: MixMessageReceiver, ) -> Self { @@ -410,6 +412,7 @@ where pub(crate) async fn listen_for_requests(mut self) where S: AsyncRead + AsyncWrite + Unpin, + St: Storage, { trace!("Started listening for ALL incoming requests..."); diff --git a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs index 51b2729be9..745aaccb67 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -6,7 +6,7 @@ use crate::node::client_handling::websocket::connection_handler::{ AuthenticatedHandler, ClientDetails, InitialAuthResult, SocketStream, }; use crate::node::storage::error::StorageError; -use crate::node::storage::PersistentStorage; +use crate::node::storage::Storage; use crypto::asymmetric::identity; use futures::{channel::mpsc, SinkExt, StreamExt}; use gateway_requests::authentication::encrypted_address::{ @@ -68,14 +68,14 @@ impl InitialAuthenticationError { } } -pub(crate) struct FreshHandler { +pub(crate) struct FreshHandler { rng: R, local_identity: Arc, pub(crate) testnet_mode: bool, pub(crate) active_clients_store: ActiveClientsStore, pub(crate) outbound_mix_sender: MixForwardingSender, pub(crate) socket_connection: SocketStream, - pub(crate) storage: PersistentStorage, + pub(crate) storage: St, #[cfg(feature = "coconut")] pub(crate) aggregated_verification_key: VerificationKey, @@ -84,9 +84,10 @@ pub(crate) struct FreshHandler { pub(crate) erc20_bridge: Arc, } -impl FreshHandler +impl FreshHandler where R: Rng + CryptoRng, + St: Storage, { // 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 @@ -99,7 +100,7 @@ where testnet_mode: bool, outbound_mix_sender: MixForwardingSender, local_identity: Arc, - storage: PersistentStorage, + storage: St, active_clients_store: ActiveClientsStore, #[cfg(feature = "coconut")] aggregated_verification_key: VerificationKey, #[cfg(not(feature = "coconut"))] erc20_bridge: Arc, @@ -543,7 +544,7 @@ where // TODO: somehow cleanup this method pub(crate) async fn perform_initial_authentication( mut self, - ) -> Option> + ) -> Option> where S: AsyncRead + AsyncWrite + Unpin + Send, { diff --git a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs index 0dda785d38..fdd90e3aea 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs @@ -9,6 +9,8 @@ use rand::{CryptoRng, Rng}; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::WebSocketStream; +use crate::node::storage::Storage; + pub(crate) use self::authenticated::AuthenticatedHandler; pub(crate) use self::fresh::FreshHandler; @@ -64,10 +66,11 @@ impl InitialAuthResult { } } -pub(crate) async fn handle_connection(mut handle: FreshHandler) +pub(crate) async fn handle_connection(mut handle: FreshHandler) where R: Rng + CryptoRng, S: AsyncRead + AsyncWrite + Unpin + Send, + St: Storage, { if let Err(err) = handle.perform_websocket_handshake().await { warn!( diff --git a/gateway/src/node/client_handling/websocket/listener.rs b/gateway/src/node/client_handling/websocket/listener.rs index 2c227f7afb..a0c46c713c 100644 --- a/gateway/src/node/client_handling/websocket/listener.rs +++ b/gateway/src/node/client_handling/websocket/listener.rs @@ -3,7 +3,7 @@ use crate::node::client_handling::active_clients::ActiveClientsStore; use crate::node::client_handling::websocket::connection_handler::FreshHandler; -use crate::node::storage::PersistentStorage; +use crate::node::storage::Storage; use crypto::asymmetric::identity; use log::*; use mixnet_client::forwarder::MixForwardingSender; @@ -52,12 +52,14 @@ impl Listener { // TODO: change the signature to pub(crate) async fn run(&self, handler: Handler) - pub(crate) async fn run( + pub(crate) async fn run( &mut self, outbound_mix_sender: MixForwardingSender, - storage: PersistentStorage, + storage: St, active_clients_store: ActiveClientsStore, - ) { + ) where + St: Storage + Clone + 'static, + { info!("Starting websocket listener at {}", self.address); let tcp_listener = match tokio::net::TcpListener::bind(self.address).await { Ok(listener) => listener, @@ -93,12 +95,15 @@ impl Listener { } } - pub(crate) fn start( + pub(crate) fn start( mut self, outbound_mix_sender: MixForwardingSender, - storage: PersistentStorage, + storage: St, active_clients_store: ActiveClientsStore, - ) -> JoinHandle<()> { + ) -> JoinHandle<()> + where + St: Storage + Clone + 'static, + { tokio::spawn(async move { self.run(outbound_mix_sender, storage, active_clients_store) .await diff --git a/gateway/src/node/mixnet_handling/receiver/connection_handler.rs b/gateway/src/node/mixnet_handling/receiver/connection_handler.rs index df4b6cb17c..8dd85d9ded 100644 --- a/gateway/src/node/mixnet_handling/receiver/connection_handler.rs +++ b/gateway/src/node/mixnet_handling/receiver/connection_handler.rs @@ -5,7 +5,7 @@ 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::error::StorageError; -use crate::node::storage::PersistentStorage; +use crate::node::storage::Storage; use futures::StreamExt; use log::*; use mixnet_client::forwarder::MixForwardingSender; @@ -19,7 +19,7 @@ use std::net::SocketAddr; use tokio::net::TcpStream; use tokio_util::codec::Framed; -pub(crate) struct ConnectionHandler { +pub(crate) struct ConnectionHandler { packet_processor: PacketProcessor, // TODO: investigate performance trade-offs for whether this cache even makes sense @@ -28,11 +28,11 @@ pub(crate) struct ConnectionHandler { // and each `get` internally copies the channel, however, is it really that expensive? clients_store_cache: HashMap, active_clients_store: ActiveClientsStore, - storage: PersistentStorage, + storage: St, ack_sender: MixForwardingSender, } -impl Clone for ConnectionHandler { +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()); @@ -52,10 +52,10 @@ impl Clone for ConnectionHandler { } } -impl ConnectionHandler { +impl ConnectionHandler { pub(crate) fn new( packet_processor: PacketProcessor, - storage: PersistentStorage, + storage: St, ack_sender: MixForwardingSender, active_clients_store: ActiveClientsStore, ) -> Self { diff --git a/gateway/src/node/mixnet_handling/receiver/listener.rs b/gateway/src/node/mixnet_handling/receiver/listener.rs index bc78ea317c..1cb1f2850e 100644 --- a/gateway/src/node/mixnet_handling/receiver/listener.rs +++ b/gateway/src/node/mixnet_handling/receiver/listener.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use crate::node::mixnet_handling::receiver::connection_handler::ConnectionHandler; +use crate::node::storage::Storage; use log::*; use std::net::SocketAddr; use std::process; @@ -17,7 +18,10 @@ impl Listener { Listener { address } } - pub(crate) async fn run(&mut self, connection_handler: ConnectionHandler) { + pub(crate) async fn run(&mut self, connection_handler: ConnectionHandler) + where + St: Storage + Clone + 'static, + { info!("Starting mixnet listener at {}", self.address); let tcp_listener = match tokio::net::TcpListener::bind(self.address).await { Ok(listener) => listener, @@ -38,7 +42,10 @@ impl Listener { } } - pub(crate) fn start(mut self, connection_handler: ConnectionHandler) -> JoinHandle<()> { + pub(crate) fn start(mut self, connection_handler: ConnectionHandler) -> JoinHandle<()> + where + St: Storage + Clone + 'static, + { info!("Running mix listener on {:?}", self.address.to_string()); tokio::spawn(async move { self.run(connection_handler).await }) diff --git a/gateway/src/node/mod.rs b/gateway/src/node/mod.rs index e113e6aa00..72c8fbd322 100644 --- a/gateway/src/node/mod.rs +++ b/gateway/src/node/mod.rs @@ -7,7 +7,7 @@ use crate::config::Config; 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::PersistentStorage; +use crate::node::storage::Storage; use crypto::asymmetric::{encryption, identity}; use log::*; use mixnet_client::forwarder::{MixForwardingSender, PacketForwarder}; @@ -25,23 +25,44 @@ use coconut_interface::VerificationKey; #[cfg(feature = "coconut")] use credentials::obtain_aggregate_verification_key; +use self::storage::PersistentStorage; + pub(crate) mod client_handling; pub(crate) mod mixnet_handling; pub(crate) mod storage; -pub struct Gateway { +/// Wire up and create Gateway instance +pub(crate) async fn create_gateway(config: Config) -> Gateway { + let storage = initialise_storage(&config).await; + Gateway::new(config, storage).await +} + +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(crate) struct Gateway { config: Config, /// ed25519 keypair used to assert one's identity. identity_keypair: Arc, /// x25519 keypair used for Diffie-Hellman. Currently only used for sphinx key derivation. sphinx_keypair: Arc, - storage: PersistentStorage, + storage: St, } -impl Gateway { - pub async fn new(config: Config) -> Self { - let storage = Self::initialise_storage(&config).await; +impl Gateway +where + St: Storage + Clone + 'static, +{ + /// Construct from the given `Config` instance. + pub async fn new(config: Config, storage: St) -> Self { let pathfinder = GatewayPathfinder::new_from_config(&config); + // let storage = Self::initialise_storage(&config).await; Gateway { config, @@ -51,12 +72,18 @@ impl Gateway { } } - 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, + #[cfg(test)] + pub async fn new_from_keys_and_storage( + config: Config, + identity_keypair: identity::KeyPair, + sphinx_keypair: encryption::KeyPair, + storage: St, + ) -> Self { + Gateway { + config, + identity_keypair: Arc::new(identity_keypair), + sphinx_keypair: Arc::new(sphinx_keypair), + storage, } } diff --git a/gateway/src/node/storage/mod.rs b/gateway/src/node/storage/mod.rs index a74d53a389..aeed709455 100644 --- a/gateway/src/node/storage/mod.rs +++ b/gateway/src/node/storage/mod.rs @@ -6,6 +6,7 @@ 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 async_trait::async_trait; use gateway_requests::registration::handshake::SharedKeys; use log::{debug, error}; use nymsphinx::DestinationAddressBytes; @@ -18,6 +19,123 @@ mod inboxes; mod models; mod shared_keys; +#[async_trait] +pub(crate) trait Storage: Send + Sync { + /// 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. + async fn insert_shared_keys( + &self, + client_address: DestinationAddressBytes, + shared_keys: SharedKeys, + ) -> Result<(), StorageError>; + + /// Tries to retrieve shared keys stored for the particular client. + /// + /// # Arguments + /// + /// * `client_address`: address of the client + async fn get_shared_keys( + &self, + client_address: DestinationAddressBytes, + ) -> Result, StorageError>; + + /// 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)] + async fn remove_shared_keys( + &self, + client_address: DestinationAddressBytes, + ) -> Result<(), StorageError>; + + /// 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. + async fn store_message( + &self, + client_address: DestinationAddressBytes, + message: Vec, + ) -> Result<(), StorageError>; + + /// 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. + async fn retrieve_messages( + &self, + client_address: DestinationAddressBytes, + start_after: Option, + ) -> Result<(Vec, Option), StorageError>; + + /// Removes messages with the specified ids + /// + /// # Arguments + /// + /// * `ids`: ids of the messages to remove + async fn remove_messages(&self, ids: Vec) -> Result<(), StorageError>; + + /// Creates a new bandwidth entry for the particular client. + /// + /// # Arguments + /// + /// * `client_address`: address of the client + async fn create_bandwidth_entry( + &self, + client_address: DestinationAddressBytes, + ) -> Result<(), StorageError>; + + /// Tries to retrieve available bandwidth for the particular client. + /// + /// # Arguments + /// + /// * `client_address`: address of the client + async fn get_available_bandwidth( + &self, + client_address: DestinationAddressBytes, + ) -> Result, StorageError>; + + /// 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. + async fn increase_bandwidth( + &self, + client_address: DestinationAddressBytes, + amount: i64, + ) -> Result<(), StorageError>; + + /// Decreases available bandwidth of the particular client by the specified amount. + /// + /// # Arguments + /// + /// * `client_address`: address of the client + /// * `amount`: amount of available bandwidth to be removed from the client. + async fn consume_bandwidth( + &self, + client_address: DestinationAddressBytes, + amount: i64, + ) -> Result<(), StorageError>; +} + // note that clone here is fine as upon cloning the same underlying pool will be used #[derive(Clone)] pub(crate) struct PersistentStorage { @@ -33,7 +151,7 @@ impl PersistentStorage { /// /// * `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>( + pub async fn init + Send>( database_path: P, message_retrieval_limit: i64, ) -> Result { @@ -72,15 +190,11 @@ impl PersistentStorage { 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( +#[async_trait] +impl Storage for PersistentStorage { + async fn insert_shared_keys( &self, client_address: DestinationAddressBytes, shared_keys: SharedKeys, @@ -95,12 +209,7 @@ impl PersistentStorage { 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( + async fn get_shared_keys( &self, client_address: DestinationAddressBytes, ) -> Result, StorageError> { @@ -111,15 +220,8 @@ impl PersistentStorage { 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( + async fn remove_shared_keys( &self, client_address: DestinationAddressBytes, ) -> Result<(), StorageError> { @@ -129,13 +231,7 @@ impl PersistentStorage { 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( + async fn store_message( &self, client_address: DestinationAddressBytes, message: Vec, @@ -146,16 +242,7 @@ impl PersistentStorage { 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( + async fn retrieve_messages( &self, client_address: DestinationAddressBytes, start_after: Option, @@ -167,24 +254,14 @@ impl PersistentStorage { 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) -> Result<(), StorageError> { + async fn remove_messages(&self, ids: Vec) -> 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( + async fn create_bandwidth_entry( &self, client_address: DestinationAddressBytes, ) -> Result<(), StorageError> { @@ -194,12 +271,7 @@ impl PersistentStorage { Ok(()) } - /// Tries to retrieve available bandwidth for the particular client. - /// - /// # Arguments - /// - /// * `client_address`: address of the client - pub(crate) async fn get_available_bandwidth( + async fn get_available_bandwidth( &self, client_address: DestinationAddressBytes, ) -> Result, StorageError> { @@ -211,13 +283,7 @@ impl PersistentStorage { 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( + async fn increase_bandwidth( &self, client_address: DestinationAddressBytes, amount: i64, @@ -228,13 +294,7 @@ impl PersistentStorage { 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( + async fn consume_bandwidth( &self, client_address: DestinationAddressBytes, amount: i64, @@ -245,3 +305,92 @@ impl PersistentStorage { Ok(()) } } + +/// In-memory implementation of `Storage`. The intention is primarily in testing environments. +#[cfg(test)] +#[derive(Clone)] +pub(crate) struct InMemStorage; + +#[cfg(test)] +impl InMemStorage { + #[allow(unused)] + async fn init + Send>() -> Result { + todo!() + } +} + +#[cfg(test)] +#[async_trait] +impl Storage for InMemStorage { + async fn insert_shared_keys( + &self, + _client_address: DestinationAddressBytes, + _shared_keys: SharedKeys, + ) -> Result<(), StorageError> { + todo!() + } + + async fn get_shared_keys( + &self, + _client_address: DestinationAddressBytes, + ) -> Result, StorageError> { + todo!() + } + + async fn remove_shared_keys( + &self, + _client_address: DestinationAddressBytes, + ) -> Result<(), StorageError> { + todo!() + } + + async fn store_message( + &self, + _client_address: DestinationAddressBytes, + _message: Vec, + ) -> Result<(), StorageError> { + todo!() + } + + async fn retrieve_messages( + &self, + _client_address: DestinationAddressBytes, + _start_after: Option, + ) -> Result<(Vec, Option), StorageError> { + todo!() + } + + async fn remove_messages(&self, _ids: Vec) -> Result<(), StorageError> { + todo!() + } + + async fn create_bandwidth_entry( + &self, + _client_address: DestinationAddressBytes, + ) -> Result<(), StorageError> { + todo!() + } + + async fn get_available_bandwidth( + &self, + _client_address: DestinationAddressBytes, + ) -> Result, StorageError> { + todo!() + } + + async fn increase_bandwidth( + &self, + _client_address: DestinationAddressBytes, + _amount: i64, + ) -> Result<(), StorageError> { + todo!() + } + + async fn consume_bandwidth( + &self, + _client_address: DestinationAddressBytes, + _amount: i64, + ) -> Result<(), StorageError> { + todo!() + } +}