Skip to content

Commit

Permalink
fix backup/restore
Browse files Browse the repository at this point in the history
  • Loading branch information
@Alex6323
Alex6323 committed Jan 10, 2024
1 parent 6572ef5 commit 6be0c25
Show file tree
Hide file tree
Showing 7 changed files with 121 additions and 70 deletions.
2 changes: 1 addition & 1 deletion bindings/core/src/method_handler/wallet.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ pub(crate) async fn call_wallet_method_internal(wallet: &Wallet, method: WalletM
ignore_if_bech32_mismatch,
} => {
wallet
.restore_backup(
.restore_from_backup(
source,
password,
ignore_if_coin_type_mismatch,
Expand Down
21 changes: 13 additions & 8 deletions cli/src/command/wallet.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,7 @@ pub async fn mnemonic_command(output_file_name: Option<String>, output_stdout: O
}

pub async fn new_account_command(wallet: &Wallet, alias: Option<String>) -> Result<AccountIdentifier, Error> {
let alias = add_account(&wallet, alias).await?;
let alias = add_account(wallet, alias).await?;

Ok(alias)
}
Expand All @@ -220,24 +220,27 @@ pub async fn node_info_command(wallet: &Wallet) -> Result<(), Error> {

pub async fn restore_command_stronghold(
storage_path: &Path,
password: Option<Password>,
snapshot_path: &Path,
backup_path: &Path,
) -> Result<Wallet, Error> {
check_file_exists(backup_path).await?;

let mut builder = Wallet::builder();
if check_file_exists(snapshot_path).await.is_ok() {
println!(
"Detected a stronghold file at {}. Enter password to unlock:",
snapshot_path.to_str().unwrap()
);
let password = get_password("Stronghold password", false)?;
// providing a password means that a Stronghold snapshot exists (verified by the caller)
if let Some(password) = password {
println!("Detected a stronghold file at {}.", snapshot_path.to_str().unwrap());
let secret_manager = SecretManager::Stronghold(
StrongholdSecretManager::builder()
.password(password)
.build(snapshot_path)?,
);
builder = builder.with_secret_manager(secret_manager);
} else {
// If there is no db, set the placeholder so the wallet builder doesn't fail.
if check_file_exists(storage_path).await.is_err() {
builder = builder.with_secret_manager(SecretManager::Placeholder);
}
}

let wallet = builder
Expand All @@ -250,7 +253,9 @@ pub async fn restore_command_stronghold(
.await?;

let password = get_password("Stronghold backup password", false)?;
wallet.restore_backup(backup_path.into(), password, None, None).await?;
wallet
.restore_from_backup(backup_path.into(), password, None, None)
.await?;

println_log_info!(
"Wallet has been restored from the backup file \"{}\".",
Expand Down
139 changes: 92 additions & 47 deletions cli/src/wallet.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,48 +28,53 @@ pub async fn new_wallet(cli: WalletCli) -> Result<(Option<Wallet>, Option<Accoun

enum LinkedSecretManager {
Stronghold {
password: iota_sdk::client::Password,
snapshot_path: std::path::PathBuf,
snapshot_exists: Option<iota_sdk::client::Password>,
},
LedgerNano {
is_connected: bool,
},
LedgerNano,
LedgerNanoSimulator,
NotFound,
}

let wallet_and_secret_manager = if let Ok(wallet) = Wallet::builder().with_storage_path(storage_path).finish().await
{
let linked_secret_manager = match &mut *wallet.get_secret_manager().write().await {
SecretManager::Stronghold(stronghold) => {
let snapshot_path = stronghold.snapshot_path.clone();
// `set_password` will trigger writing the snapshot file, so we need to make sure it already exists,
// otherwise we could run into an inconsistent wallet.
if snapshot_path.exists() {
let password = get_password("Stronghold password", false)?;
stronghold.set_password(password.clone()).await?;
LinkedSecretManager::Stronghold {
password,
snapshot_path,
}
} else {
LinkedSecretManager::NotFound
let wallet_and_secret_manager = {
if storage_path.is_dir() {
match Wallet::builder().with_storage_path(storage_path).finish().await {
Ok(wallet) => {
let linked_secret_manager = match &mut *wallet.get_secret_manager().write().await {
SecretManager::Stronghold(stronghold) => {
let snapshot_path = stronghold.snapshot_path.clone();
// `set_password` will trigger writing the snapshot file, so we need to make sure it already
// exists, otherwise we could run into an inconsistent wallet.
if snapshot_path.exists() {
let password = get_password("Stronghold password", false)?;
stronghold.set_password(password.clone()).await?;
LinkedSecretManager::Stronghold {
snapshot_path,
snapshot_exists: Some(password),
}
} else {
LinkedSecretManager::Stronghold {
snapshot_path,
snapshot_exists: None,
}
}
}
SecretManager::LedgerNano(ledger_nano) => LinkedSecretManager::LedgerNano {
// is_simulator: ledger_nano.is_simulator,
is_connected: ledger_nano.get_ledger_nano_status().await.connected(),
},
_ => panic!("only Stronghold and LedgerNano supported at the moment."),
};
Some((wallet, linked_secret_manager))
}
}
SecretManager::LedgerNano(ledger_nano) => {
if ledger_nano.get_ledger_nano_status().await.connected() {
if ledger_nano.is_simulator {
LinkedSecretManager::LedgerNanoSimulator
} else {
LinkedSecretManager::LedgerNano
}
} else {
LinkedSecretManager::NotFound
Err(e) => {
println_log_error!("failed to load wallet db from storage: {e}");
return Ok((None, None));
}
}
_ => panic!("only Stronghold and LedgerNano supported at the moment."),
};
Some((wallet, linked_secret_manager))
} else {
None
} else {
None
}
};

let (wallet, account_id) = if let Some(command) = cli.command {
Expand Down Expand Up @@ -101,10 +106,19 @@ pub async fn new_wallet(cli: WalletCli) -> Result<(Option<Wallet>, Option<Accoun
WalletCommand::Backup { backup_path } => {
if let Some((wallet, secret_manager)) = wallet_and_secret_manager {
match secret_manager {
LinkedSecretManager::Stronghold { password, .. } => {
LinkedSecretManager::Stronghold {
snapshot_exists: Some(password),
..
} => {
backup_command_stronghold(&wallet, &password, Path::new(&backup_path)).await?;
return Ok((None, None));
}
LinkedSecretManager::Stronghold { snapshot_path, .. } => {
return Err(Error::Miscellaneous(format!(
"Stronghold snapshot does not exist at '{}'",
snapshot_path.display()
)));
}
_ => {
println_log_info!("only Stronghold backup supported");
return Ok((None, None));
Expand All @@ -120,10 +134,19 @@ pub async fn new_wallet(cli: WalletCli) -> Result<(Option<Wallet>, Option<Accoun
WalletCommand::ChangePassword => {
if let Some((wallet, secret_manager)) = wallet_and_secret_manager {
match secret_manager {
LinkedSecretManager::Stronghold { password, .. } => {
LinkedSecretManager::Stronghold {
snapshot_exists: Some(password),
..
} => {
change_password_command(&wallet, password).await?;
(Some(wallet), None)
}
LinkedSecretManager::Stronghold { snapshot_path, .. } => {
return Err(Error::Miscellaneous(format!(
"Stronghold snapshot does not exist at '{}'",
snapshot_path.display()
)));
}
_ => {
println_log_info!("only Stronghold password change supported");
return Ok((None, None));
Expand Down Expand Up @@ -195,12 +218,18 @@ pub async fn new_wallet(cli: WalletCli) -> Result<(Option<Wallet>, Option<Accoun
}
}
WalletCommand::Restore { backup_path } => {
// TODO: not sure how we do this here is correct
if let Some((_, linked_secret_manager)) = wallet_and_secret_manager {
if let Some((wallet, linked_secret_manager)) = wallet_and_secret_manager {
match linked_secret_manager {
LinkedSecretManager::Stronghold { snapshot_path, .. } => {
LinkedSecretManager::Stronghold {
snapshot_path,
snapshot_exists,
} => {
// we need to explicitly drop the current wallet here to prevent:
// "error accessing storage: IO error: lock hold by current process"
drop(wallet);
let wallet = restore_command_stronghold(
storage_path,
snapshot_exists,
snapshot_path.as_path(),
Path::new(&backup_path),
)
Expand All @@ -213,10 +242,12 @@ pub async fn new_wallet(cli: WalletCli) -> Result<(Option<Wallet>, Option<Accoun
}
}
} else {
return Err(Error::Miscellaneous(format!(
"wallet db does not exist at '{}'",
storage_path.display()
)));
// the wallet db does not exist
let init_params = InitParameters::default();
let snapshot_path = Path::new(&init_params.stronghold_snapshot_path);
let wallet =
restore_command_stronghold(storage_path, None, snapshot_path, Path::new(&backup_path)).await?;
(Some(wallet), None)
}
}
WalletCommand::MigrateStrongholdSnapshotV2ToV3 { path } => {
Expand All @@ -234,10 +265,24 @@ pub async fn new_wallet(cli: WalletCli) -> Result<(Option<Wallet>, Option<Accoun
} else {
// no wallet command provided
if let Some((wallet, linked_secret_manager)) = wallet_and_secret_manager {
if let LinkedSecretManager::NotFound = linked_secret_manager {
println_log_error!("The secret manager linked with the wallet was not found");
return Ok((None, None));
match linked_secret_manager {
LinkedSecretManager::Stronghold {
snapshot_exists: None,
snapshot_path,
} => {
println_log_error!(
"Snapshot file for Stronghold secret manager linked with the wallet not found at '{}'",
snapshot_path.display()
);
return Ok((None, None));
}
LinkedSecretManager::LedgerNano { is_connected: false } => {
println_log_error!("Ledger Nano linked with the wallet not connected");
return Ok((None, None));
}
_ => {}
}

if wallet.get_accounts().await?.is_empty() {
create_initial_account(wallet).await?
} else if let Some(alias) = cli.account {
Expand Down
9 changes: 5 additions & 4 deletions sdk/src/wallet/core/operations/stronghold_backup/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,8 @@ use crate::{
};

impl Wallet {
/// Backup the wallet data in a Stronghold file
/// stronghold_password must be the current one when Stronghold is used as SecretManager.
/// Backup the wallet data in a Stronghold file.
/// `stronghold_password` must be the current one when Stronghold is used as SecretManager.
pub async fn backup(
&self,
backup_path: PathBuf,
Expand Down Expand Up @@ -60,7 +60,7 @@ impl Wallet {
Ok(())
}

/// Restore a backup from a Stronghold file
/// Restore the wallet from a Stronghold backup file.
/// Replaces client_options, coin_type, secret_manager and accounts. Returns an error if accounts were already
/// created If Stronghold is used as secret_manager, the existing Stronghold file will be overwritten. If a
/// mnemonic was stored, it will be gone.
Expand All @@ -69,7 +69,7 @@ impl Wallet {
/// coin type doesn't match
/// if ignore_if_bech32_hrp_mismatch == Some("rms"), but addresses have something different like "smr", no accounts
/// will be restored.
pub async fn restore_backup(
pub async fn restore_from_backup(
&self,
backup_path: PathBuf,
stronghold_password: impl Into<Password> + Send,
Expand All @@ -96,6 +96,7 @@ impl Wallet {
let new_snapshot_path = if let SecretManager::Stronghold(stronghold) = &mut *secret_manager {
stronghold.snapshot_path.clone()
} else {
// TODO: move and use default constant from cli?
PathBuf::from("wallet.stronghold")
};

Expand Down
16 changes: 8 additions & 8 deletions sdk/tests/wallet/backup_restore.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ async fn backup_and_restore() -> Result<()> {

// Wrong password fails
restore_wallet
.restore_backup(
.restore_from_backup(
PathBuf::from("test-storage/backup_and_restore/backup.stronghold"),
"wrong password".to_owned(),
None,
Expand All @@ -80,7 +80,7 @@ async fn backup_and_restore() -> Result<()> {

// Correct password works, even after trying with a wrong one before
restore_wallet
.restore_backup(
.restore_from_backup(
PathBuf::from("test-storage/backup_and_restore/backup.stronghold"),
stronghold_password,
None,
Expand Down Expand Up @@ -162,7 +162,7 @@ async fn backup_and_restore_mnemonic_secret_manager() -> Result<()> {
.await?;

restore_wallet
.restore_backup(
.restore_from_backup(
PathBuf::from("test-storage/backup_and_restore_mnemonic_secret_manager/backup.stronghold"),
stronghold_password,
None,
Expand Down Expand Up @@ -247,7 +247,7 @@ async fn backup_and_restore_different_coin_type() -> Result<()> {

// restore with ignore_if_coin_type_mismatch: Some(true) to not overwrite the coin type
restore_wallet
.restore_backup(
.restore_from_backup(
PathBuf::from("test-storage/backup_and_restore_different_coin_type/backup.stronghold"),
stronghold_password,
Some(true),
Expand Down Expand Up @@ -331,7 +331,7 @@ async fn backup_and_restore_same_coin_type() -> Result<()> {

// restore with ignore_if_coin_type_mismatch: Some(true) to not overwrite the coin type
restore_wallet
.restore_backup(
.restore_from_backup(
PathBuf::from("test-storage/backup_and_restore_same_coin_type/backup.stronghold"),
stronghold_password,
Some(true),
Expand Down Expand Up @@ -413,7 +413,7 @@ async fn backup_and_restore_different_coin_type_dont_ignore() -> Result<()> {

// restore with ignore_if_coin_type_mismatch: Some(true) to not overwrite the coin type
restore_wallet
.restore_backup(
.restore_from_backup(
PathBuf::from("test-storage/backup_and_restore_different_coin_type_dont_ignore/backup.stronghold"),
stronghold_password,
Some(false),
Expand Down Expand Up @@ -498,7 +498,7 @@ async fn backup_and_restore_bech32_hrp_mismatch() -> Result<()> {
.await?;

restore_wallet
.restore_backup(
.restore_from_backup(
PathBuf::from("test-storage/backup_and_restore_bech32_hrp_mismatch/backup.stronghold"),
stronghold_password,
None,
Expand Down Expand Up @@ -550,7 +550,7 @@ async fn restore_no_secret_manager_data() -> Result<()> {
let stronghold_password = "some_hopefully_secure_password".to_owned();

restore_wallet
.restore_backup(
.restore_from_backup(
PathBuf::from("./tests/wallet/fixtures/no_secret_manager_data.stronghold"),
stronghold_password.clone(),
None,
Expand Down
2 changes: 1 addition & 1 deletion sdk/tests/wallet/chrysalis_migration.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ async fn migrate_chrysalis_stronghold() -> Result<()> {
.await?;

wallet
.restore_backup(
.restore_from_backup(
"./tests/wallet/fixtures/chrysalis-backup-work-factor-0.stronghold".into(),
Password::from("password".to_string()),
None,
Expand Down
2 changes: 1 addition & 1 deletion sdk/tests/wallet/migrate_stronghold_snapshot_v2_to_v3.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ async fn stronghold_snapshot_v2_v3_migration() {

// restore with ignore_if_coin_type_mismatch: Some(true) to not overwrite the coin type
let error = restore_manager
.restore_backup(
.restore_from_backup(
PathBuf::from("./tests/wallet/fixtures/v3.stronghold"),
"wrong_password".to_owned(),
Some(false),
Expand Down

0 comments on commit 6be0c25

Please sign in to comment.