From 690c0a719db0b08fb02dab707110dadee2756742 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thoralf=20M=C3=BCller?= <thoralf.mue@gmail.com>
Date: Wed, 30 Aug 2023 15:23:30 +0200
Subject: [PATCH] Address review comments

---
 .../stronghold_backup/stronghold_snapshot.rs  | 21 +++++++++----------
 sdk/src/wallet/migration/chrysalis.rs         |  9 ++++----
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/sdk/src/wallet/core/operations/stronghold_backup/stronghold_snapshot.rs b/sdk/src/wallet/core/operations/stronghold_backup/stronghold_snapshot.rs
index 426dc53a6b..3b4adbb6a0 100644
--- a/sdk/src/wallet/core/operations/stronghold_backup/stronghold_snapshot.rs
+++ b/sdk/src/wallet/core/operations/stronghold_backup/stronghold_snapshot.rs
@@ -6,16 +6,17 @@ use std::{collections::HashMap, path::Path, sync::atomic::Ordering};
 use crate::{
     client::{
         constants::IOTA_COIN_TYPE, secret::SecretManagerConfig, storage::StorageAdapter, stronghold::StrongholdAdapter,
+        Error as ClientError,
     },
     types::TryFromDto,
     wallet::{
         account::{AccountDetails, AccountDetailsDto},
         migration::{
-            chrysalis::{key_to_chrysalis_key, migrate_from_chrysalis_data},
+            chrysalis::{migrate_from_chrysalis_data, to_chrysalis_key},
             latest_backup_migration_version, migrate, MigrationData, MIGRATION_VERSION_KEY,
         },
         storage::constants::{CHRYSALIS_STORAGE_KEY, WALLET_INDEXATION_KEY},
-        ClientOptions, Wallet,
+        ClientOptions, Error as WalletError, Wallet,
     },
 };
 
@@ -117,18 +118,16 @@ async fn migrate_snapshot_from_chrysalis_to_stardust(
         // `iota-wallet-records` was only used in chrysalis
         Err(iota_stronghold::ClientError::ClientDataNotPresent) => return Ok(None),
         Err(e) => {
-            return Err(crate::wallet::Error::Client(Box::new(
-                crate::client::Error::Stronghold(e.into()),
-            )));
+            return Err(WalletError::Client(Box::new(ClientError::Stronghold(e.into()))));
         }
     };
 
     let stronghold_store = stronghold_client.store();
     let keys = stronghold_store
         .keys()
-        .map_err(|e| crate::wallet::Error::Client(Box::new(crate::client::Error::Stronghold(e.into()))))?;
+        .map_err(|e| WalletError::Client(Box::new(ClientError::Stronghold(e.into()))))?;
 
-    let wallet_indexation_key = key_to_chrysalis_key(b"iota-wallet-account-indexation", true);
+    let wallet_indexation_key = to_chrysalis_key(b"iota-wallet-account-indexation", true);
     // check if snapshot contains chrysalis data
     if !keys.iter().any(|k| k == &wallet_indexation_key) {
         return Ok(None);
@@ -138,10 +137,10 @@ async fn migrate_snapshot_from_chrysalis_to_stardust(
     for key in keys {
         let value = stronghold_store
             .get(&key)
-            .map_err(|e| crate::wallet::Error::Client(Box::new(crate::client::Error::Stronghold(e.into()))))?;
+            .map_err(|e| WalletError::Client(Box::new(ClientError::Stronghold(e.into()))))?;
 
         let value_utf8 =
-            String::from_utf8(value.unwrap()).map_err(|_| crate::wallet::Error::Migration("invalid utf8".into()))?;
+            String::from_utf8(value.unwrap()).map_err(|_| WalletError::Migration("invalid utf8".into()))?;
 
         chrysalis_data.insert(key, value_utf8);
     }
@@ -207,13 +206,13 @@ async fn migrate_snapshot_from_chrysalis_to_stardust(
     let stronghold = stronghold_adapter.inner().await;
     let stronghold_client = stronghold
         .get_client(b"iota-wallet-records")
-        .map_err(|e| crate::wallet::Error::Client(Box::new(crate::client::Error::Stronghold(e.into()))))?;
+        .map_err(|e| WalletError::Client(Box::new(ClientError::Stronghold(e.into()))))?;
     let stronghold_store = stronghold_client.store();
 
     for key in chrysalis_data.keys() {
         stronghold_store
             .delete(key)
-            .map_err(|_| crate::wallet::Error::Migration("couldn't delete old data".into()))?;
+            .map_err(|_| WalletError::Migration("couldn't delete old data".into()))?;
     }
 
     Ok(Some(chrysalis_data_with_string_keys))
diff --git a/sdk/src/wallet/migration/chrysalis.rs b/sdk/src/wallet/migration/chrysalis.rs
index d9bc25a809..3a03fcd9a0 100644
--- a/sdk/src/wallet/migration/chrysalis.rs
+++ b/sdk/src/wallet/migration/chrysalis.rs
@@ -61,6 +61,7 @@ pub(crate) struct AccountDetailsDto {
     native_token_foundries: HashMap<String, Value>,
 }
 
+#[cfg(feature = "rocksdb")]
 pub async fn migrate_db_chrysalis_to_stardust(
     storage_path: impl Into<PathBuf> + Send,
     password: Option<Password>,
@@ -144,11 +145,11 @@ pub(crate) fn migrate_from_chrysalis_data(
     let mut new_accounts: Vec<AccountDetailsDto> = Vec::new();
     let mut secret_manager_dto: Option<String> = None;
 
-    let account_indexation_key = key_to_chrysalis_key(b"iota-wallet-account-indexation", stronghold);
+    let account_indexation_key = to_chrysalis_key(b"iota-wallet-account-indexation", stronghold);
     if let Some(account_indexation) = chrysalis_data.get(&account_indexation_key) {
         if let Some(account_keys) = serde_json::from_str::<serde_json::Value>(account_indexation)?.as_array() {
             for account_key in account_keys {
-                let account_key = key_to_chrysalis_key(
+                let account_key = to_chrysalis_key(
                     account_key["key"].as_str().expect("key must be a string").as_bytes(),
                     stronghold,
                 );
@@ -296,10 +297,10 @@ fn decrypt_record(record_bytes: Vec<u8>, encryption_key: &[u8; 32]) -> crate::wa
     )
     .map_err(|e| Error::Migration(format!("{:?}", e)))?;
 
-    Ok(String::from_utf8_lossy(&pt).to_string())
+    String::from_utf8(pt).map_err(|e| Error::Migration(format!("{:?}", e)))
 }
 
-pub(crate) fn key_to_chrysalis_key(key: &[u8], stronghold: bool) -> Vec<u8> {
+pub(crate) fn to_chrysalis_key(key: &[u8], stronghold: bool) -> Vec<u8> {
     // key only needs to be hashed for stronghold
     if stronghold {
         let mut buf = [0; 64];