491045439

Build Type

github

Committed by GitHub

Commit Message

Merge 84bac2224 into 040b4a723

Pull Request Pull Request #18: test(lib) add unit tests

Run Details

5338 of 23249 branches covered (22.96%)

Branch coverage included in aggregate %.

461 of 477 new or added lines in 9 files covered. (96.65%)

2 existing lines in 1 file now uncovered.

2225 of 3806 relevant lines covered (58.46%)

521.99 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

29.7

/src/account_manager.rs

// Copyright 2020 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

#[allow(unused_imports)]
use crate::{
    account::{
        repost_message, Account, AccountHandle, AccountIdentifier, AccountInitialiser, RepostAction, SyncedAccount,
    },
    client::ClientOptions,
    event::{emit_balance_change, emit_confirmation_state_change, emit_transaction_event, TransactionEventType},
    message::{Message, MessageType, Transfer},
    signing::SignerType,
    storage::StorageAdapter,
};

use std::{
    collections::HashMap,
    convert::TryInto,
    fs,
    num::NonZeroU64,
    panic::AssertUnwindSafe,
    path::{Path, PathBuf},
    sync::Arc,
    thread,
    time::Duration,
};

use chrono::prelude::*;
use futures::FutureExt;
use getset::Getters;
use iota::{MessageId, Payload};
use serde::Deserialize;
use tokio::{
    sync::{
        broadcast::{channel as broadcast_channel, Receiver as BroadcastReceiver, Sender as BroadcastSender},
        Mutex, RwLock,
    },
    time::interval,
};

/// The default storage folder.
pub const DEFAULT_STORAGE_FOLDER: &str = "./storage";
/// The default stronghold storage file name.
#[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
pub const STRONGHOLD_FILENAME: &str = "wallet.stronghold";
/// The default SQLite storage file name.
#[cfg(feature = "sqlite-storage")]
pub const SQLITE_FILENAME: &str = "wallet.db";

pub(crate) type AccountStore = Arc<RwLock<HashMap<AccountIdentifier, AccountHandle>>>;

/// The storages used by default.
#[derive(Deserialize, PartialEq)]
#[serde(tag = "type", content = "data")]
pub enum DefaultStorage {
    /// Stronghold storage.
    #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
    Stronghold,
    /// Sqlite storage.
    #[cfg(feature = "sqlite-storage")]
    Sqlite,
}

#[cfg(any(feature = "stronghold", feature = "stronghold-storage", feature = "sqlite-storage"))]
fn storage_file_path(storage: &Option<DefaultStorage>, storage_path: &PathBuf) -> PathBuf {
    if storage_path.is_file() || storage_path.extension().is_some() || storage.is_none() {
        storage_path.clone()
    } else {
        #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
        if storage == &Some(DefaultStorage::Stronghold) {
            return storage_path.join(STRONGHOLD_FILENAME);
        }
        #[cfg(feature = "sqlite-storage")]
        if storage == &Some(DefaultStorage::Sqlite) {
            return storage_path.join(SQLITE_FILENAME);
        }
        storage_path.clone()
    }
}

/// Account manager builder.
pub struct AccountManagerBuilder {
    storage_path: PathBuf,
    initialised_storage: bool,
    polling_interval: Duration,
    skip_polling: bool,
    default_storage: Option<DefaultStorage>,
    storage_encryption_key: Option<[u8; 32]>,
}

impl Default for AccountManagerBuilder {
    fn default() -> Self {
        #[allow(unused_variables)]
        let default_storage: Option<DefaultStorage> = None;
        #[cfg(all(feature = "stronghold-storage", not(feature = "sqlite-storage")))]
        let default_storage = Some(DefaultStorage::Stronghold);
        #[cfg(all(feature = "sqlite-storage", not(feature = "stronghold-storage")))]
        let default_storage = Some(DefaultStorage::Sqlite);

        Self {
            storage_path: PathBuf::from(DEFAULT_STORAGE_FOLDER),
            initialised_storage: false,
            polling_interval: Duration::from_millis(30_000),
            skip_polling: false,
            default_storage,
            storage_encryption_key: None,
        }
    }
}

impl AccountManagerBuilder {
    /// Initialises a new instance of the account manager builder with the default storage adapter.
    pub fn new() -> Self {
        Default::default()
    }

    /// Use the specified storage folder when initialising the default storage adapter.
    #[cfg(any(feature = "stronghold-storage", feature = "sqlite-storage"))]
    pub fn with_storage_path(mut self, storage_path: impl AsRef<Path>) -> Self {
        self.storage_path = storage_path.as_ref().to_path_buf();
        self
    }

    /// Sets the default storage to be used.
    #[cfg(any(feature = "sqlite-storage", feature = "stronghold-storage"))]
    pub fn with_storage(mut self, storage: DefaultStorage) -> Self {
        self.default_storage = Some(storage);
        self
    }

    /// Sets a custom storage adapter to be used.
    #[cfg(not(any(feature = "sqlite-storage", feature = "stronghold-storage")))]
    pub fn with_storage<S: StorageAdapter + Sync + Send + 'static>(
        mut self,
        storage_path: impl AsRef<Path>,
        adapter: S,
    ) -> Self {
        crate::storage::set_adapter(&storage_path, adapter);
        self.storage_path = storage_path.as_ref().to_path_buf();
        self.initialised_storage = true;
        self
    }

    /// Sets the polling interval.
    pub fn with_polling_interval(mut self, polling_interval: Duration) -> Self {
        self.polling_interval = polling_interval;
        self
    }

    pub(crate) fn skip_polling(mut self) -> Self {
        self.skip_polling = true;
        self
    }

    pub(crate) fn with_storage_encryption_key(mut self, key: Option<[u8; 32]>) -> Self {
        self.storage_encryption_key = key;
        self
    }

    /// Builds the manager.
    pub async fn finish(self) -> crate::Result<AccountManager> {
        if !self.initialised_storage {
            #[cfg(any(feature = "stronghold-storage", feature = "sqlite-storage"))]
            {
                match &self.default_storage {
                    Some(storage) => match storage {
                        #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
                        DefaultStorage::Stronghold => {
                            let path = storage_file_path(&self.default_storage, &self.storage_path);
                            let storage = crate::storage::stronghold::StrongholdStorageAdapter::new(&path)?;
                            crate::storage::set_adapter(&path, storage);
                        }
                        #[cfg(feature = "sqlite-storage")]
                        DefaultStorage::Sqlite => {
                            let path = storage_file_path(&self.default_storage, &self.storage_path);
                            let storage = crate::storage::sqlite::SqliteStorageAdapter::new(&path, "accounts")?;
                            crate::storage::set_adapter(&path, storage);
                        }
                    },
                    None => {
                        return Err(crate::Error::StorageAdapterNotDefined);
                    }
                };
            }
            #[cfg(not(any(feature = "stronghold-storage", feature = "sqlite-storage")))]
            {
                return Err(crate::Error::StorageAdapterNotDefined);
            }
        }

        let storage_file_path = storage_file_path(&self.default_storage, &self.storage_path);

        // with one of the stronghold features, the accounts are loaded when the password is set
        #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
        let accounts: AccountStore = Default::default();
        #[cfg(not(any(feature = "stronghold", feature = "stronghold-storage")))]
        let accounts = AccountManager::load_accounts(&storage_file_path)
            .await
            .unwrap_or_else(|_| Default::default());

        let mut instance = AccountManager {
            storage_folder: if self.storage_path.is_file() || self.storage_path.extension().is_some() {
                match self.storage_path.parent() {
                    Some(p) => p.to_path_buf(),
                    None => self.storage_path,
                }
            } else {
                self.storage_path
            },
            storage_path: storage_file_path,
            accounts,
            stop_polling_sender: None,
            polling_handle: None,
            generated_mnemonic: None,
            storage_encryption_key: Arc::new(Mutex::new(self.storage_encryption_key)),
            encrypted_accounts: Vec::new(),
        };

        if !self.skip_polling {
            instance.start_background_sync(self.polling_interval).await;
        }

        Ok(instance)
    }
}

/// The account manager.
///
/// Used to manage multiple accounts.
#[derive(Getters)]
pub struct AccountManager {
    storage_folder: PathBuf,
    /// the path to the storage.
    #[getset(get = "pub")]
    storage_path: PathBuf,
    accounts: AccountStore,
    stop_polling_sender: Option<BroadcastSender<()>>,
    polling_handle: Option<thread::JoinHandle<()>>,
    generated_mnemonic: Option<String>,
    storage_encryption_key: Arc<Mutex<Option<[u8; 32]>>>,
    encrypted_accounts: Vec<String>,
}

impl Drop for AccountManager {
    fn drop(&mut self) {
        self.stop_background_sync();
    }
}

impl AccountManager {
    /// Initialises the account manager builder.
    pub fn builder() -> AccountManagerBuilder {
        AccountManagerBuilder::new()
    }

    async fn load_accounts(
        storage_file_path: &PathBuf,
        encryption_key: &Option<[u8; 32]>,
    ) -> crate::Result<(AccountStore, Vec<String>)> {
        let mut encrypted_accounts = Vec::new();
        let mut parsed_accounts = HashMap::new();

        let accounts = crate::storage::get(&storage_file_path)?.lock().await.get_all().await?;
        let accounts = crate::storage::parse_accounts(&storage_file_path, &accounts, encryption_key)?;
        for parsed_account in accounts {
            match parsed_account {
                crate::storage::ParsedAccount::Account(account) => {
                    parsed_accounts.insert(account.id().clone(), account.into());
                }
                crate::storage::ParsedAccount::EncryptedAccount(value) => {
                    encrypted_accounts.push(value);
                }
            }
        }

        Ok((Arc::new(RwLock::new(parsed_accounts)), encrypted_accounts))
    }

    /// Starts monitoring the accounts with the node's mqtt topics.
    async fn start_monitoring(&self) -> crate::Result<()> {
        for account in self.accounts.read().await.values() {
            crate::monitor::monitor_account_addresses_balance(account.clone()).await?;
            crate::monitor::monitor_unconfirmed_messages(account.clone()).await?;
        }
        Ok(())
    }

    /// Initialises the background polling and MQTT monitoring.
    async fn start_background_sync(&mut self, polling_interval: Duration) {
        let monitoring_disabled = self.start_monitoring().await.is_err();
        let (stop_polling_sender, stop_polling_receiver) = broadcast_channel(1);
        self.start_polling(polling_interval, monitoring_disabled, stop_polling_receiver);
        self.stop_polling_sender = Some(stop_polling_sender);
    }

    /// Stops the background polling and MQTT monitoring.
    pub fn stop_background_sync(&mut self) {
        if let Some(stop_polling_sender) = self.stop_polling_sender.take() {
            stop_polling_sender.send(()).expect("failed to stop polling process");
            self.polling_handle
                .take()
                .unwrap()
                .join()
                .expect("failed to join polling thread");
            let accounts = self.accounts.clone();
            thread::spawn(move || {
                crate::block_on(async move {
                    for account_handle in accounts.read().await.values() {
                        let _ = crate::monitor::unsubscribe(account_handle.clone());
                    }
                });
            })
            .join()
            .expect("failed to stop monitoring and polling systems");
        }
    }

    /// Sets the password for the stored accounts.
    pub async fn set_storage_password<P: AsRef<str>>(&mut self, key: P) -> crate::Result<()> {
        let mut dk = [0; 64];
        crypto::kdfs::pbkdf::PBKDF2_HMAC_SHA512(key.as_ref().as_bytes(), b"wallet.rs::storage", 100, &mut dk)?;
        let key: [u8; 32] = dk[0..32][..].try_into().unwrap();
        *self.storage_encryption_key.lock().await = Some(key);

        let mut accounts = self.accounts.write().await;
        for encrypted_account in &self.encrypted_accounts {
            let decrypted = crate::storage::decrypt_account_json(encrypted_account, &key)?;
            let account = serde_json::from_str::<Account>(&decrypted)?;
            accounts.insert(account.id().clone(), account.into());
        }
        self.encrypted_accounts.clear();

        Ok(())
    }

    /// Sets the stronghold password.
    #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
    pub async fn set_stronghold_password<P: AsRef<str>>(&mut self, password: P) -> crate::Result<()> {
        let mut dk = [0; 64];
        crypto::kdfs::pbkdf::PBKDF2_HMAC_SHA512(password.as_ref().as_bytes(), b"wallet.rs", 100, &mut dk)?;

        let stronghold_path = if self.storage_path.extension().unwrap_or_default() == "stronghold" {
            self.storage_path.clone()
        } else {
            self.storage_folder.join(STRONGHOLD_FILENAME)
        };
        crate::stronghold::load_snapshot(&stronghold_path, &dk[0..32][..].try_into().unwrap()).await?;

        // let is_empty = self.accounts.read().await.is_empty();
        if self.accounts.read().await.is_empty() {
            let (accounts, encrypted_accounts) =
                Self::load_accounts(&self.storage_path, &*self.storage_encryption_key.lock().await).await?;
            self.encrypted_accounts = encrypted_accounts;
            let mut accounts_store = self.accounts.write().await;
            for (id, account) in &*accounts.read().await {
                accounts_store.insert(id.clone(), account.clone());
            }
        }

        Ok(())
    }

    /// Starts the polling mechanism.
    fn start_polling(
        &mut self,
        polling_interval: Duration,
        is_monitoring_disabled: bool,
        mut stop: BroadcastReceiver<()>,
    ) {
        let storage_file_path = self.storage_path.clone();
        let accounts = self.accounts.clone();
        let storage_encryption_key = self.storage_encryption_key.clone();

        let handle = thread::spawn(move || {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_time()
                .build()
                .unwrap();
            runtime.block_on(async {
                let mut interval = interval(polling_interval);
                loop {
                    tokio::select! {
                        _ = async {
                            interval.tick().await;

                            let storage_file_path_ = storage_file_path.clone();

                            if let Err(error) = AssertUnwindSafe(poll(accounts.clone(), storage_file_path_, is_monitoring_disabled))
                                .catch_unwind()
                                .await {
                                    // if the error isn't a crate::Error type
                                if error.downcast_ref::<crate::Error>().is_none() {
                                    let msg = if let Some(message) = error.downcast_ref::<String>() {
                                        format!("Internal error: {}", message)
                                    } else if let Some(message) = error.downcast_ref::<&str>() {
                                        format!("Internal error: {}", message)
                                    } else {
                                        "Internal error".to_string()
                                    };
                                    log::error!("[POLLING] error: {}", msg);
                                    let _error = crate::Error::Panic(msg);
                                    // when the error is dropped, the on_error event will be triggered
                                }
                            }

                            {
                                let encryption_key = &*storage_encryption_key.lock().await;
                                for account_handle in accounts.read().await.values() {
                                    let _ = account_handle.write().await.save(&encryption_key).await;
                                }
                            }
                        } => {}
                        _ = stop.recv() => {
                            let encryption_key = &*storage_encryption_key.lock().await;
                            // before stopping the polling loop, we save the accounts
                            for account_handle in accounts.read().await.values() {
                                let _ = account_handle.write().await.save(&encryption_key).await;
                            }
                            break;
                        }
                    }
                }
                println!("DONE");
            });
        });
        self.polling_handle = Some(handle);
    }

    /// Stores a mnemonic for the given signer type.
    /// If the mnemonic is not provided, we'll generate one.
    pub async fn store_mnemonic(&mut self, signer_type: SignerType, mnemonic: Option<String>) -> crate::Result<()> {
        let mnemonic = match mnemonic {
            Some(m) => {
                self.verify_mnemonic(&m)?;
                m
            }
            None => self.generate_mnemonic()?,
        };

        let signer = crate::signing::get_signer(&signer_type).await;
        let signer = signer.lock().await;
        signer.store_mnemonic(&self.storage_path, mnemonic).await?;

        self.generated_mnemonic = None;

        Ok(())
    }

    /// Generates a new mnemonic.
    pub fn generate_mnemonic(&mut self) -> crate::Result<String> {
        let mut entropy = [0u8; 32];
        crypto::rand::fill(&mut entropy)?;
        let mnemonic = crypto::bip39::wordlist::encode(&entropy, &crypto::bip39::wordlist::ENGLISH)
            .map_err(|_| crate::Error::MnemonicEncode)?;
        self.generated_mnemonic = Some(mnemonic.clone());
        Ok(mnemonic)
    }

    /// Checks is the mnemonic is valid. If a mnemonic was generated with `generate_mnemonic()`, the mnemonic here
    /// should match the generated.
    pub fn verify_mnemonic<S: AsRef<str>>(&mut self, mnemonic: S) -> crate::Result<()> {
        // first we check if the mnemonic is valid to give meaningful errors
        crypto::bip39::wordlist::verify(mnemonic.as_ref(), &crypto::bip39::wordlist::ENGLISH)
            // TODO: crypto::bip39::wordlist::Error should impl Display
            .map_err(|e| crate::Error::InvalidMnemonic(format!("{:?}", e)))?;

        // then we check if the provided mnemonic matches the mnemonic generated with `generate_mnemonic`
        if let Some(generated_mnemonic) = &self.generated_mnemonic {
            if generated_mnemonic != mnemonic.as_ref() {
                return Err(crate::Error::InvalidMnemonic(
                    "doesn't match the generated mnemonic".to_string(),
                ));
            }
            self.generated_mnemonic = None;
        }
        Ok(())
    }

    /// Adds a new account.
    pub fn create_account(&self, client_options: ClientOptions) -> AccountInitialiser {
        AccountInitialiser::new(client_options, self.accounts.clone(), self.storage_path.clone())
    }

    /// Deletes an account.
    pub async fn remove_account(&self, account_id: &AccountIdentifier) -> crate::Result<()> {
        let mut accounts = self.accounts.write().await;

        {
            let account_handle = accounts.get(&account_id).ok_or(crate::Error::AccountNotFound)?;
            let account = account_handle.read().await;

            if !(account.messages().is_empty() && account.total_balance() == 0) {
                return Err(crate::Error::MessageNotEmpty);
            }
        }

        accounts.remove(account_id);

        crate::storage::get(&self.storage_path)?
            .lock()
            .await
            .remove(&account_id)
            .await?;

        Ok(())
    }

    /// Syncs all accounts.
    pub async fn sync_accounts(&self) -> crate::Result<Vec<SyncedAccount>> {
        sync_accounts(self.accounts.clone(), &self.storage_path, None).await
    }

    /// Transfers an amount from an account to another.
    pub async fn internal_transfer(
        &self,
        from_account_id: &AccountIdentifier,
        to_account_id: &AccountIdentifier,
        amount: NonZeroU64,
    ) -> crate::Result<Message> {
        let to_address = self
            .get_account(to_account_id)
            .await?
            .read()
            .await
            .latest_address()
            .ok_or(crate::Error::TransferDestinationEmpty)?
            .clone();

        let from_synchronized = self.get_account(from_account_id).await?.sync().await.execute().await?;
        from_synchronized
            .transfer(Transfer::builder(to_address.address().clone(), amount).finish())
            .await
    }

    /// Backups the storage to the given destination
    #[cfg(any(feature = "stronghold-storage", feature = "sqlite-storage"))]
    pub async fn backup<P: AsRef<Path>>(&self, destination: P) -> crate::Result<PathBuf> {
        let destination = destination.as_ref().to_path_buf();
        if !(destination.is_dir() && destination.exists()) {
            return Err(crate::Error::InvalidBackupDestination);
        }

        {
            let encryption_key = &*self.storage_encryption_key.lock().await;
            for account_handle in self.accounts.read().await.values() {
                account_handle.write().await.save(&encryption_key).await?;
            }
        }

        #[allow(unused_variables)]
        let (storage_path, backup_entire_directory) = (
            &self.storage_path,
            cfg!(feature = "stronghold") && cfg!(feature = "sqlite-storage"),
        );

        // if we're using SQLite for storage and stronghold for seed,
        // we'll backup only the stronghold file, copying SQLite data to its snapshot
        #[cfg(all(
            feature = "sqlite-storage",
            any(feature = "stronghold", feature = "stronghold-storage")
        ))]
        let (storage_path, backup_entire_directory) = {
            let storage_id = crate::storage::get(&&self.storage_path)?.lock().await.id();
            // if we're actually using the SQLite storage adapter
            let storage_path = if storage_id == crate::storage::sqlite::STORAGE_ID {
                let stronghold_storage = crate::storage::stronghold::StrongholdStorageAdapter::new(
                    &self.storage_folder.join(STRONGHOLD_FILENAME),
                )
                .unwrap();

                let encryption_key = &*self.storage_encryption_key.lock().await;
                for (account_id, account_handle) in &*self.accounts.read().await {
                    stronghold_storage
                        .set(
                            account_id,
                            crate::storage::get_account_string_to_save(&*account_handle.read().await, &encryption_key)?,
                        )
                        .await?;
                }
                self.storage_folder.join(STRONGHOLD_FILENAME)
            } else {
                self.storage_path.clone()
            };
            (storage_path, false)
        };

        if storage_path.exists() {
            let destination = if backup_entire_directory {
                backup_dir(&self.storage_folder, &destination)?;
                destination
            } else if let Some(filename) = storage_path.file_name() {
                let destination = destination.join(backup_filename(filename.to_str().unwrap()));
                let res = fs::copy(storage_path, &destination);

                // if we're using SQLite for storage and stronghold for seed,
                // we'll remove the accounts from stronghold after the backup
                #[cfg(all(
                    feature = "sqlite-storage",
                    any(feature = "stronghold", feature = "stronghold-storage")
                ))]
                {
                    let storage_id = crate::storage::get(&self.storage_path)?.lock().await.id();
                    // if we're actually using the SQLite storage adapter
                    if storage_id == crate::storage::sqlite::STORAGE_ID {
                        let stronghold_storage = crate::storage::stronghold::StrongholdStorageAdapter::new(
                            &self.storage_folder.join(STRONGHOLD_FILENAME),
                        )
                        .unwrap();
                        for account_id in self.accounts.read().await.keys() {
                            stronghold_storage.remove(account_id).await?;
                        }
                    }
                }

                res?;
                destination
            } else {
                return Err(crate::Error::StorageDoesntExist);
            };
            Ok(destination)
        } else {
            Err(crate::Error::StorageDoesntExist)
        }
    }

    /// Import backed up accounts.
    #[cfg(any(feature = "stronghold-storage", feature = "sqlite-storage"))]
    pub async fn import_accounts<S: AsRef<Path>>(
        &mut self,
        source: S,
        #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))] stronghold_password: impl AsRef<str>,
    ) -> crate::Result<()> {
        let source = source.as_ref();
        if source.is_dir() || !source.exists() {
            return Err(crate::Error::InvalidBackupFile);
        }

        #[allow(unused_variables)]
        #[cfg(feature = "stronghold-storage")]
        let storage_file_path = {
            let storage_file_path = self.storage_folder.join(STRONGHOLD_FILENAME);
            let storage_id = crate::storage::get(&self.storage_path)?.lock().await.id();
            if storage_id == crate::storage::stronghold::STORAGE_ID && storage_file_path.exists() {
                return Err(crate::Error::StorageExists);
            }
            storage_file_path
        };
        #[cfg(feature = "sqlite-storage")]
        let storage_file_path = {
            if !self.accounts.read().await.is_empty() {
                return Err(crate::Error::StorageExists);
            }

            self.storage_folder.join(SQLITE_FILENAME)
        };

        fs::create_dir_all(&self.storage_folder)?;

        if cfg!(feature = "sqlite-storage") && source.extension().unwrap_or_default() == "stronghold" {
            #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
            {
                let mut stronghold_manager = Self::builder()
                    .with_storage_path(&source)
                    .with_storage(DefaultStorage::Stronghold)
                    .with_storage_encryption_key(*self.storage_encryption_key.lock().await)
                    .skip_polling()
                    .finish()
                    .await?;
                stronghold_manager.set_stronghold_password(&stronghold_password).await?;
                for account_handle in stronghold_manager.accounts.read().await.values() {
                    account_handle.write().await.set_storage_path(self.storage_path.clone());
                }
                self.accounts = stronghold_manager.accounts.clone();
            }
            #[cfg(not(any(feature = "stronghold", feature = "stronghold-storage")))]
            return Err(crate::Error::InvalidBackupFile);
        } else {
            #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
            {
                // wait for stronghold to finish its tasks
                let _ = crate::stronghold::actor_runtime().lock().await;
            }
            fs::copy(source, &storage_file_path)?;
        }

        // the accounts map isn't empty when restoring SQLite from a stronghold snapshot
        #[cfg(not(any(feature = "stronghold", feature = "stronghold-storage")))]
        if self.accounts.read().await.is_empty() {
            self.accounts = Self::load_accounts(&self.storage_path).await?;
        }

        #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
        {
            // force stronghold to read the snapshot again, ignoring any previous cached value
            crate::stronghold::unload_snapshot(&self.storage_path, false).await?;
            if let Err(e) = self.set_stronghold_password(stronghold_password).await {
                fs::remove_file(&storage_file_path)?;
                return Err(e);
            }
        }

        Ok(())
    }

    /// Gets the account associated with the given identifier.
    pub async fn get_account(&self, account_id: &AccountIdentifier) -> crate::Result<AccountHandle> {
        let accounts = self.accounts.read().await;
        accounts.get(account_id).cloned().ok_or(crate::Error::AccountNotFound)
    }

    /// Gets the account associated with the given alias (case insensitive).
    pub async fn get_account_by_alias<S: AsRef<str>>(&self, alias: S) -> Option<AccountHandle> {
        let alias = alias.as_ref().to_lowercase();
        for account_handle in self.accounts.read().await.values() {
            let account = account_handle.read().await;
            if account
                .alias()
                .to_lowercase()
                .chars()
                .zip(alias.chars())
                .all(|(x, y)| x == y)
            {
                return Some(account_handle.clone());
            }
        }
        None
    }

    /// Gets all accounts from storage.
    pub async fn get_accounts(&self) -> Vec<AccountHandle> {
        let accounts = self.accounts.read().await;
        accounts.values().cloned().collect()
    }

    /// Reattaches an unconfirmed transaction.
    pub async fn reattach(&self, account_id: &AccountIdentifier, message_id: &MessageId) -> crate::Result<Message> {
        let account = self.get_account(account_id).await?;
        account.sync().await.execute().await?.reattach(message_id).await
    }

    /// Promotes an unconfirmed transaction.
    pub async fn promote(&self, account_id: &AccountIdentifier, message_id: &MessageId) -> crate::Result<Message> {
        let account = self.get_account(account_id).await?;
        account.sync().await.execute().await?.promote(message_id).await
    }

    /// Retries an unconfirmed transaction.
    pub async fn retry(&self, account_id: &AccountIdentifier, message_id: &MessageId) -> crate::Result<Message> {
        let account = self.get_account(account_id).await?;
        account.sync().await.execute().await?.retry(message_id).await
    }
}

async fn poll(accounts: AccountStore, storage_file_path: PathBuf, syncing: bool) -> crate::Result<()> {
    let retried = if syncing {
        let mut accounts_before_sync = Vec::new();
        for account_handle in accounts.read().await.values() {
            accounts_before_sync.push(account_handle.read().await.clone());
        }
        let synced_accounts = sync_accounts(accounts.clone(), &storage_file_path, Some(0)).await?;
        let accounts_after_sync = accounts.read().await;

        log::debug!("[POLLING] synced accounts");

        // compare accounts to check for balance changes and new messages
        for account_before_sync in &accounts_before_sync {
            let account_after_sync = accounts_after_sync.get(account_before_sync.id()).unwrap();
            let account_after_sync = account_after_sync.read().await;

            // balance event
            for address_before_sync in account_before_sync.addresses() {
                let address_after_sync = account_after_sync
                    .addresses()
                    .iter()
                    .find(|addr| addr == &address_before_sync)
                    .unwrap();
                if address_after_sync.balance() != address_before_sync.balance() {
                    log::debug!(
                        "[POLLING] address {} balance changed from {} to {}",
                        address_after_sync.address().to_bech32(),
                        address_before_sync.balance(),
                        address_after_sync.balance()
                    );
                    emit_balance_change(
                        account_after_sync.id(),
                        address_after_sync,
                        *address_after_sync.balance(),
                    );
                }
            }

            // new messages event
            account_after_sync
                .messages()
                .iter()
                .filter(|message| !account_before_sync.messages().contains(message))
                .for_each(|message| {
                    log::info!("[POLLING] new message: {:?}", message.id());
                    emit_transaction_event(TransactionEventType::NewTransaction, account_after_sync.id(), &message)
                });

            // confirmation state change event
            for message in account_after_sync.messages() {
                let changed = match account_before_sync.messages().iter().find(|m| m.id() == message.id()) {
                    Some(old_message) => message.confirmed() != old_message.confirmed(),
                    None => false,
                };
                if changed {
                    log::info!("[POLLING] message confirmed: {:?}", message.id());
                    emit_confirmation_state_change(account_after_sync.id(), &message, true);
                }
            }
        }
        retry_unconfirmed_transactions(synced_accounts).await?
    } else {
        log::info!("[POLLING] skipping syncing process because MQTT is running");
        let mut retried_messages = vec![];
        for account_handle in accounts.read().await.values() {
            let (account_id, unconfirmed_messages): (AccountIdentifier, Vec<(MessageId, Payload)>) = {
                let account = account_handle.read().await;
                let account_id = account.id().clone();
                let unconfirmed_messages = account
                    .list_messages(account.messages().len(), 0, Some(MessageType::Unconfirmed))
                    .iter()
                    .map(|m| (*m.id(), m.payload().clone()))
                    .collect();
                (account_id, unconfirmed_messages)
            };

            let mut promotions = vec![];
            let mut reattachments = vec![];
            for (message_id, payload) in unconfirmed_messages {
                let new_message = repost_message(account_handle.clone(), &message_id, RepostAction::Retry).await?;
                if new_message.payload() == &payload {
                    reattachments.push(new_message);
                } else {
                    log::info!("[POLLING] promoted and new message is {:?}", new_message.id());
                    promotions.push(new_message);
                }
            }

            retried_messages.push(RetriedData {
                promoted: promotions,
                reattached: reattachments,
                account_id,
            });
        }

        retried_messages
    };

    retried.iter().for_each(|retried_data| {
        retried_data.reattached.iter().for_each(|message| {
            emit_transaction_event(TransactionEventType::Reattachment, &retried_data.account_id, &message);
        });
    });
    Ok(())
}

async fn discover_accounts(
    accounts: AccountStore,
    storage_path: &PathBuf,
    client_options: &ClientOptions,
    signer_type: Option<SignerType>,
) -> crate::Result<Vec<(AccountHandle, SyncedAccount)>> {
    let mut synced_accounts = vec![];
    loop {
        let mut account_initialiser =
            AccountInitialiser::new(client_options.clone(), accounts.clone(), storage_path.clone()).skip_persistance();
        if let Some(signer_type) = &signer_type {
            account_initialiser = account_initialiser.signer_type(signer_type.clone());
        }
        let account_handle = account_initialiser.initialise().await?;
        log::debug!(
            "[SYNC] discovering account {}, signer type {:?}",
            account_handle.read().await.alias(),
            account_handle.read().await.signer_type()
        );
        let synced_account = account_handle.sync().await.execute().await?;
        let is_empty = *synced_account.is_empty();
        log::debug!("[SYNC] account is empty? {}", is_empty);
        if is_empty {
            break;
        } else {
            synced_accounts.push((account_handle, synced_account));
        }
    }
    Ok(synced_accounts)
}

async fn sync_accounts<'a>(
    accounts: AccountStore,
    storage_file_path: &PathBuf,
    address_index: Option<usize>,
) -> crate::Result<Vec<SyncedAccount>> {
    let mut synced_accounts = vec![];
    let mut last_account = None;

    {
        let accounts = accounts.read().await;
        for account_handle in accounts.values() {
            let mut sync = account_handle.sync().await;
            if let Some(index) = address_index {
                sync = sync.address_index(index);
            }
            let synced_account = sync.execute().await?;

            let account = account_handle.read().await;
            last_account = Some((
                account.messages().is_empty() || account.addresses().iter().all(|addr| *addr.balance() == 0),
                account.client_options().clone(),
                account.signer_type().clone(),
            ));
            synced_accounts.push(synced_account);
        }
    }

    let discovered_accounts_res = match last_account {
        Some((is_empty, client_options, signer_type)) => {
            if is_empty {
                log::debug!("[SYNC] running account discovery because the latest account is empty");
                discover_accounts(accounts.clone(), &storage_file_path, &client_options, Some(signer_type)).await
            } else {
                log::debug!("[SYNC] skipping account discovery because the latest account isn't empty");
                Ok(vec![])
            }
        }
        None => Ok(vec![]), /* None => discover_accounts(accounts.clone(), &storage_path, &ClientOptions::default(),
                             * None).await, */
    };

    if let Ok(discovered_accounts) = discovered_accounts_res {
        if !discovered_accounts.is_empty() {
            let mut accounts = accounts.write().await;
            for (account_handle, synced_account) in discovered_accounts {
                account_handle.write().await.set_skip_persistance(false);
                accounts.insert(account_handle.id().await, account_handle);
                synced_accounts.push(synced_account);
            }
        }
    }

    Ok(synced_accounts)
}

struct RetriedData {
    #[allow(dead_code)]
    promoted: Vec<Message>,
    reattached: Vec<Message>,
    account_id: AccountIdentifier,
}

async fn retry_unconfirmed_transactions(synced_accounts: Vec<SyncedAccount>) -> crate::Result<Vec<RetriedData>> {
    let mut retried_messages = vec![];
    for synced in synced_accounts {
        let account = synced.account_handle().read().await;

        let unconfirmed_messages = account.list_messages(account.messages().len(), 0, Some(MessageType::Unconfirmed));
        let mut reattachments = vec![];
        let mut promotions = vec![];
        for message in unconfirmed_messages {
            log::debug!("[POLLING] retrying {:?}", message);
            let new_message = synced.retry(message.id()).await?;
            // if the payload is the same, it was reattached; otherwise it was promoted
            if new_message.payload() == message.payload() {
                log::debug!("[POLLING] rettached and new message is {:?}", new_message);
                reattachments.push(new_message);
            } else {
                log::debug!("[POLLING] promoted and new message is {:?}", new_message);
                promotions.push(new_message);
            }
        }
        retried_messages.push(RetriedData {
            promoted: promotions,
            reattached: reattachments,
            account_id: account.id().clone(),
        });
    }
    Ok(retried_messages)
}

fn backup_dir<U: AsRef<Path>, V: AsRef<Path>>(from: U, to: V) -> Result<(), std::io::Error> {
    let mut stack = Vec::new();
    stack.push(PathBuf::from(from.as_ref()));

    let output_root = PathBuf::from(to.as_ref());
    let input_root = PathBuf::from(from.as_ref()).components().count();

    while let Some(working_path) = stack.pop() {
        let src: PathBuf = working_path.components().skip(input_root).collect();

        let dest = if src.components().count() == 0 {
            output_root.clone()
        } else {
            output_root.join(&src)
        };
        if fs::metadata(&dest).is_err() {
            fs::create_dir_all(&dest)?;
        }

        for entry in fs::read_dir(working_path)? {
            let entry = entry?;
            let path = entry.path();
            if path.is_dir() {
                stack.push(path);
            } else if let Some(filename) = path.file_name() {
                let dest_path = dest.join(backup_filename(filename.to_str().unwrap()));
                fs::copy(&path, &dest_path)?;
            }
        }
    }

    Ok(())
}

fn backup_filename(original: &str) -> String {
    let date = Local::now();
    format!(
        "{}-iota-wallet-backup{}",
        date.format("%FT%H-%M-%S").to_string(),
        if original.is_empty() {
            "".to_string()
        } else {
            format!("-{}", original)
        }
    )
}

#[cfg(test)]
mod tests {
    use crate::{
        address::{AddressBuilder, IotaAddress},
        client::ClientOptionsBuilder,
        message::Message,
    };
    use iota::{Ed25519Address, IndexationPayload, MessageBuilder, MessageId, Payload};

    #[tokio::test]
    async fn store_accounts() {
        let mut manager = crate::test_utils::get_account_manager().await;
        let account_handle = crate::test_utils::AccountCreator::new(&manager).create().await;
        manager.stop_background_sync();

        manager
            .remove_account(account_handle.read().await.id())
            .await
            .expect("failed to remove account");
    }

    #[tokio::test]
    async fn remove_account_with_message_history() {
        let manager = crate::test_utils::get_account_manager().await;

        let client_options = ClientOptionsBuilder::node("https://nodes.devnet.iota.org:443")
            .expect("invalid node URL")
            .build();

        let account_handle = manager
            .create_account(client_options)
            .messages(vec![Message::from_iota_message(
                MessageId::new([0; 32]),
                &[],
                &MessageBuilder::new()
                    .with_nonce_provider(crate::test_utils::NoopNonceProvider {}, 4000f64)
                    .with_parent1(MessageId::new([0; 32]))
                    .with_parent2(MessageId::new([0; 32]))
                    .with_payload(Payload::Indexation(Box::new(
                        IndexationPayload::new("index".to_string(), &[0; 16]).unwrap(),
                    )))
                    .with_network_id(0)
                    .finish()
                    .unwrap(),
                Some(true),
            )
            .unwrap()])
            .initialise()
            .await
            .unwrap();

        let remove_response = manager.remove_account(account_handle.read().await.id()).await;
        assert!(remove_response.is_err());
    }

    #[tokio::test]
    async fn remove_account_with_balance() {
        let manager = crate::test_utils::get_account_manager().await;

        let client_options = ClientOptionsBuilder::node("https://nodes.devnet.iota.org:443")
            .expect("invalid node URL")
            .build();

        let account_handle = manager
            .create_account(client_options)
            .addresses(vec![AddressBuilder::new()
                .balance(5)
                .key_index(0)
                .address(IotaAddress::Ed25519(Ed25519Address::new([0; 32])))
                .outputs(vec![])
                .build()
                .unwrap()])
            .initialise()
            .await
            .unwrap();

        let remove_response = manager.remove_account(account_handle.read().await.id()).await;
        assert!(remove_response.is_err());
    }

    #[tokio::test]
    async fn create_account_with_latest_without_history() {
        let manager = crate::test_utils::get_account_manager().await;

        let client_options = ClientOptionsBuilder::node("https://nodes.devnet.iota.org:443")
            .expect("invalid node URL")
            .build();

        manager
            .create_account(client_options.clone())
            .alias("alias")
            .initialise()
            .await
            .expect("failed to add account");

        let create_response = manager.create_account(client_options).initialise().await;
        assert!(create_response.is_err());
    }

    #[tokio::test]
    async fn backup_and_restore_happy_path() {
        let backup_path = "./backup/happy-path";
        let _ = std::fs::remove_dir_all(backup_path);
        std::fs::create_dir_all(backup_path).unwrap();
        let manager = crate::test_utils::get_account_manager().await;

        let account_handle = crate::test_utils::AccountCreator::new(&manager).create().await;

        // backup the stored accounts to ./backup/happy-path/${backup_name}
        let backup_path = manager.backup(backup_path).await.unwrap();
        let backup_file_path = if backup_path.is_dir() {
            std::fs::read_dir(backup_path)
                .unwrap()
                .next()
                .unwrap()
                .unwrap()
                .path()
                .to_path_buf()
        } else {
            backup_path
        };

        // get another manager instance so we can import the accounts to a different storage
        #[allow(unused_mut)]
        let mut manager = crate::test_utils::get_account_manager().await;

        #[cfg(all(
            not(feature = "sqlite-storage"),
            any(feature = "stronghold", feature = "stronghold-storage")
        ))]
        {
            // wait for stronghold to finish pending operations and delete the storage file
            crate::stronghold::unload_snapshot(manager.storage_path(), false)
                .await
                .unwrap();
            let _ = crate::stronghold::actor_runtime().lock().await;
            std::fs::remove_file(manager.storage_path()).unwrap();
        }

        manager.set_storage_password("password").await.unwrap();

        // import the accounts from the backup and assert that it's the same

        #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
        manager.import_accounts(backup_file_path, "password").await.unwrap();

        #[cfg(not(any(feature = "stronghold", feature = "stronghold-storage")))]
        manager.import_accounts(backup_file_path).await.unwrap();

        let imported_account = manager.get_account(account_handle.read().await.id()).await.unwrap();
        // set the account storage path field so the assert works
        account_handle
            .write()
            .await
            .set_storage_path(manager.storage_path().clone());
        assert_eq!(&*account_handle.read().await, &*imported_account.read().await);
    }

    #[tokio::test]
    async fn backup_and_restore_storage_already_exists() {
        let backup_path = "./backup/account-exists";
        let _ = std::fs::remove_dir_all(backup_path);
        std::fs::create_dir_all(backup_path).unwrap();
        let mut manager = crate::test_utils::get_account_manager().await;

        // first we'll create an example account
        let address = crate::test_utils::generate_random_iota_address();
        let address = AddressBuilder::new()
            .address(address.clone())
            .key_index(0)
            .balance(0)
            .outputs(vec![])
            .build()
            .unwrap();
        crate::test_utils::AccountCreator::new(&manager)
            .addresses(vec![address])
            .create()
            .await;

        let backup_path = manager.backup(backup_path).await.unwrap();
        let backup_file_path = if backup_path.is_dir() {
            std::fs::read_dir(backup_path).unwrap().next().unwrap().unwrap().path()
        } else {
            backup_path
        };

        #[cfg(any(feature = "stronghold", feature = "stronghold-storage"))]
        let response = manager.import_accounts(backup_file_path, "password").await;

        #[cfg(not(any(feature = "stronghold", feature = "stronghold-storage")))]
        let response = manager.import_accounts(backup_file_path).await;

        assert!(response.is_err());
        let err = response.unwrap_err();
        assert_eq!(
            err.to_string(),
            "failed to restore backup: storage file already exists".to_string()
        );
    }
}

iotaledger / wallet.rs / 491045439

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous