Skip to content

Commit

Permalink
rusk-wallet: Gracefully exit on interrupt
Browse files Browse the repository at this point in the history
Ask confimration before deleting cache
Make info message more accurate
  • Loading branch information
Daksh14 committed Oct 24, 2024
1 parent 376c9c0 commit 3351b86
Show file tree
Hide file tree
Showing 5 changed files with 76 additions and 22 deletions.
7 changes: 7 additions & 0 deletions rusk-wallet/src/bin/io/prompt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,13 @@ pub(crate) fn ask_confirm() -> anyhow::Result<bool> {
Ok(a.as_bool().expect("answer to be a bool"))
}

/// Asks the user for confirmation before deleting cache
pub(crate) fn ask_confirm_erase_cache(msg: &str) -> anyhow::Result<bool> {
let q = requestty::Question::confirm("confirm").message(msg).build();
let a = requestty::prompt_one(q)?;
Ok(a.as_bool().expect("answer to be a bool"))
}

/// Request a receiver address
pub(crate) fn request_rcvr_addr(addr_for: &str) -> anyhow::Result<Address> {
// let the user input the receiver address
Expand Down
56 changes: 46 additions & 10 deletions rusk-wallet/src/bin/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,10 @@ mod settings;
pub(crate) use command::{Command, RunResult};
pub(crate) use menu::Menu;

use std::fs::{self, File};
use std::io::Write;

use bip39::{Language, Mnemonic, MnemonicType};
use clap::Parser;
use tracing::{warn, Level};
use rocksdb::ErrorKind;
use tracing::{error, info, warn, Level};

use crate::command::TransactionHistory;
use crate::settings::{LogFormat, Settings};
Expand All @@ -34,6 +32,9 @@ use config::Config;
use io::{prompt, status};
use io::{GraphQL, WalletArgs};

use std::fs::{self, File};
use std::io::Write;

#[derive(Debug, Clone)]
pub(crate) struct WalletFile {
path: WalletPath,
Expand Down Expand Up @@ -71,7 +72,7 @@ async fn connect<F>(
mut wallet: Wallet<F>,
settings: &Settings,
status: fn(&str),
) -> Wallet<F>
) -> anyhow::Result<Wallet<F>>
where
F: SecureWalletFile + std::fmt::Debug,
{
Expand All @@ -85,12 +86,44 @@ where

// check for connection errors
match con {
Err(Error::RocksDB(e)) => panic!{"Please reset the cache! {e}"},
Err(e) => warn!("[OFFLINE MODE]: Unable to connect to Rusk, limited functionality available: {e}"),
Err(Error::RocksDB(e)) => {
wallet.close();

let msg = match e.kind() {
ErrorKind::InvalidArgument => {
format!("You seem to try access a wallet with a different seed-phrase \n\r\n\r{0: <1} delete the cache? (Alternatively specify the --profile flag to add a new wallet under the given path)", "[ALERT]")
},
ErrorKind::Corruption => {
format!("The database appears to be corrupted \n\r\n\r{0: <1} delete the cache?", "[ALERT]")
},
_ => {
format!("Unknown database error {:?} \n\r\n\r{1: <1} delete the cache?", e, "[ALERT]")
}
};

match prompt::ask_confirm_erase_cache(&msg)? {
true => {
if let Some(io_err) = wallet.delete_cache().err() {
error!("Error while deleting the cache: {io_err}");
}

info!("Restart the application to create new wallet.");
},
false => {
info!("Wallet cannot proceed will now exit");
},

}

// Exit because we cannot proceed because of db error
// wallet is already closed
std::process::exit(1);
},
Err(ref e) => warn!("[OFFLINE MODE]: Unable to connect to Rusk, limited functionality available: {e}"),
_ => {}
}
};

wallet
Ok(wallet)
}

async fn exec() -> anyhow::Result<()> {
Expand Down Expand Up @@ -284,7 +317,7 @@ async fn exec() -> anyhow::Result<()> {
false => status::interactive,
};

wallet = connect(wallet, &settings, status_cb).await;
wallet = connect(wallet, &settings, status_cb).await?;

// run command
match cmd {
Expand Down Expand Up @@ -383,5 +416,8 @@ async fn exec() -> anyhow::Result<()> {
}
}

// Gracefully close the wallet
wallet.close();

Ok(())
}
3 changes: 1 addition & 2 deletions rusk-wallet/src/cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -233,8 +233,7 @@ impl Cache {
Ok(notes)
}

pub fn close(&self) -> Result<(), Error> {
pub fn close(&self) {
self.db.cancel_all_background_work(false);
Ok(())
}
}
4 changes: 1 addition & 3 deletions rusk-wallet/src/clients.rs
Original file line number Diff line number Diff line change
Expand Up @@ -400,9 +400,7 @@ impl State {
}

pub fn close(&mut self) {
// UNWRAP: its okay to panic here because we're closing the database
// if there's an error we want an exception to happen
self.cache().close().unwrap();
self.cache().close();
let store = &mut self.store;

// if there's sync handle we abort it
Expand Down
28 changes: 21 additions & 7 deletions rusk-wallet/src/wallet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -244,13 +244,7 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
_=> {},
}

let cache_dir = {
if let Some(file) = &self.file {
file.path().cache_dir()
} else {
return Err(Error::WalletFileMissing);
}
};
let cache_dir = self.cache_path()?;

// create a state client
self.state = Some(State::new(
Expand Down Expand Up @@ -398,6 +392,19 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
derive_phoenix_vk(seed, index)
}

/// get cache database path
pub(crate) fn cache_path(&self) -> Result<PathBuf, Error> {
let cache_dir = {
if let Some(file) = &self.file {
file.path().cache_dir()
} else {
return Err(Error::WalletFileMissing);
}
};

Ok(cache_dir)
}

/// Returns the shielded key for a given index.
///
/// # Errors
Expand Down Expand Up @@ -579,6 +586,13 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
Ok(network_last_pos == db_pos)
}

/// Erase the cache directory
pub fn delete_cache(&mut self) -> Result<(), Error> {
let path = self.cache_path()?;

std::fs::remove_dir_all(path).map_err(Error::IO)
}

/// Close the wallet and zeroize the seed
pub fn close(&mut self) {
self.store.inner_mut().zeroize();
Expand Down

0 comments on commit 3351b86

Please sign in to comment.