fix: wallet unlock rendering (#145)

* fix: wallet unlock rendering

* combine the functions into one

* function desc

src/ui/identities/mod.rs

@@ -1,3 +1,23 @@
+use std::sync::{Arc, RwLock};
+
+use dash_sdk::{
+    dpp::{
+        data_contract::accessors::v0::DataContractV0Getters,
+        identity::identity_public_key::accessors::v0::IdentityPublicKeyGettersV0,
+    },
+    platform::IdentityPublicKey,
+};
+
+use crate::{
+    context::AppContext,
+    model::{
+        qualified_identity::{
+            encrypted_key_storage::PrivateKeyData, PrivateKeyTarget, QualifiedIdentity,
+        },
+        wallet::Wallet,
+    },
+};
+
 pub mod add_existing_identity_screen;
 pub mod add_new_identity_screen;
 mod funding_common;
@@ -7,3 +27,99 @@ pub mod register_dpns_name_screen;
 pub mod top_up_identity_screen;
 pub mod transfers;
 pub mod withdraw_from_identity_screen;
+
+/// Retrieves the appropriate wallet (if any) associated with the given identity.
+///
+/// # Description
+///
+/// This function tries to determine which wallet should be used, either via:
+///
+/// - The DPNS-based approach (if [`AppContext`] is provided), which looks up
+///   the `preorder` document type in the DPNS contract and retrieves the
+///   document-signing key from the given [`QualifiedIdentity`].
+/// - The fallback approach (if `app_context` is `None`), which relies on a
+///   directly provided key (`selected_key`).
+///
+/// # Parameters
+///
+/// - `qualified_identity`: A reference to the [`QualifiedIdentity`], which holds
+///   the identity, keys, and associated wallets.
+/// - `app_context`: Optional reference to the [`AppContext`] which contains the
+///   DPNS contract. When present, DPNS logic is used to find the public key.
+/// - `selected_key`: An optional reference to a chosen [`IdentityPublicKey`].
+///   When `app_context` is not provided, this is required to get the wallet.
+/// - `error_message`: A mutable optional string where any error message will
+///   be written if the function fails to retrieve a wallet.
+///
+/// # Returns
+///
+/// Returns `Some(Arc<RwLock<Wallet>>)` if a matching wallet is found, or `None`
+/// otherwise. If an error is encountered, an explanatory message is placed in
+/// `error_message`.
+///
+/// # Errors
+///
+/// - If the DPNS document type can't be found or the identity is missing the
+///   required DPNS signing key (when `app_context` is provided).
+/// - If no `selected_key` is provided (when `app_context` is `None`).
+/// - If the derived wallet derivation path is missing from the
+///   [`QualifiedIdentity`].
+pub fn get_selected_wallet(
+    qualified_identity: &QualifiedIdentity,
+    app_context: Option<&AppContext>, // Used for DPNS-based logic (the first scenario).
+    selected_key: Option<&IdentityPublicKey>, // Used for direct-key logic (the fallback scenario).
+    error_message: &mut Option<String>,
+) -> Option<Arc<RwLock<Wallet>>> {
+    // If `app_context` is provided, use the DPNS-based approach.
+    let public_key = if let Some(context) = app_context {
+        let dpns_contract = &context.dpns_contract;
+
+        // Attempt to fetch the `preorder` document type from the DPNS contract.
+        let preorder_document_type = match dpns_contract.document_type_for_name("preorder") {
+            Ok(doc_type) => doc_type,
+            Err(e) => {
+                *error_message = Some(format!("DPNS preorder document type not found: {}", e));
+                return None;
+            }
+        };
+
+        // Attempt to retrieve the public key from the identity.
+        match qualified_identity.document_signing_key(&preorder_document_type) {
+            Some(key) => key,
+            None => {
+                *error_message = Some(
+                    "Identity doesn't have an authentication key for signing document transitions"
+                        .to_string(),
+                );
+                return None;
+            }
+        }
+    } else {
+        // Fallback: directly use the provided selected key.
+        match selected_key {
+            Some(key) => key,
+            None => {
+                *error_message = Some("No key provided when getting selected wallet".to_string());
+                return None;
+            }
+        }
+    };
+
+    // Once we have the public key (either from DPNS or directly), look up
+    // the matching private key data in `qualified_identity`.
+    let key_lookup = (PrivateKeyTarget::PrivateKeyOnMainIdentity, public_key.id());
+    if let Some((_, PrivateKeyData::AtWalletDerivationPath(wallet_derivation_path))) =
+        qualified_identity
+            .private_keys
+            .private_keys
+            .get(&key_lookup)
+    {
+        // If found, return the associated wallet (cloned to preserve Arc).
+        qualified_identity
+            .associated_wallets
+            .get(&wallet_derivation_path.wallet_seed_hash)
+            .cloned()
+    } else {
+        None
+    }
+}

src/ui/identities/register_dpns_name_screen.rs

@@ -2,8 +2,7 @@ use crate::app::AppAction;
 use crate::backend_task::identity::{IdentityTask, RegisterDpnsNameInput};
 use crate::backend_task::BackendTask;
 use crate::context::AppContext;
-use crate::model::qualified_identity::encrypted_key_storage::PrivateKeyData;
-use crate::model::qualified_identity::{PrivateKeyTarget, QualifiedIdentity};
+use crate::model::qualified_identity::QualifiedIdentity;
 use crate::model::wallet::Wallet;
 use crate::ui::components::top_panel::add_top_panel;
 use crate::ui::components::wallet_unlock::ScreenWithWalletUnlock;
@@ -21,6 +20,8 @@ use std::sync::Arc;
 use std::sync::RwLock;
 use std::time::{SystemTime, UNIX_EPOCH};
 
+use super::get_selected_wallet;
+
 #[derive(PartialEq)]
 pub enum RegisterDpnsNameStatus {
     NotStarted,
@@ -78,7 +79,7 @@ impl RegisterDpnsNameScreen {
 
         let mut error_message: Option<String> = None;
         let selected_wallet = if let Some(ref identity) = selected_qualified_identity {
-            get_selected_wallet(&identity.0, app_context, &mut error_message)
+            get_selected_wallet(&identity.0, Some(app_context), None, &mut error_message)
         } else {
             None
         };
@@ -108,8 +109,12 @@ impl RegisterDpnsNameScreen {
             // Set the selected_qualified_identity to the found identity
             self.selected_qualified_identity = Some(qi.clone());
             // Update the selected wallet
-            self.selected_wallet =
-                get_selected_wallet(&qi.0, &self.app_context, &mut self.error_message);
+            self.selected_wallet = get_selected_wallet(
+                &qi.0,
+                Some(&self.app_context),
+                None,
+                &mut self.error_message,
+            );
         } else {
             // If not found, you might want to handle this case
             // For now, we'll set selected_qualified_identity to None
@@ -179,7 +184,8 @@ impl RegisterDpnsNameScreen {
                                 self.selected_qualified_identity = Some(qualified_identity.clone());
                                 self.selected_wallet = get_selected_wallet(
                                     &qualified_identity.0,
-                                    &self.app_context,
+                                    Some(&self.app_context),
+                                    None,
                                     &mut self.error_message,
                                 );
                             }
@@ -466,43 +472,3 @@ pub fn is_contested_name(name: &str) -> bool {
     }
     true
 }
-
-pub fn get_selected_wallet(
-    qualified_identity: &QualifiedIdentity,
-    app_context: &AppContext,
-    error_message: &mut Option<String>,
-) -> Option<Arc<RwLock<Wallet>>> {
-    let dpns_contract = &app_context.dpns_contract;
-
-    let preorder_document_type = match dpns_contract.document_type_for_name("preorder") {
-        Ok(doc_type) => doc_type,
-        Err(e) => {
-            *error_message = Some(format!("DPNS preorder document type not found: {}", e));
-            return None;
-        }
-    };
-
-    let public_key = match qualified_identity.document_signing_key(&preorder_document_type) {
-        Some(key) => key,
-        None => {
-            *error_message = Some(
-                "Identity doesn't have an authentication key for signing document transitions"
-                    .to_string(),
-            );
-            return None;
-        }
-    };
-
-    let key = (PrivateKeyTarget::PrivateKeyOnMainIdentity, public_key.id());
-
-    if let Some((_, PrivateKeyData::AtWalletDerivationPath(wallet_derivation_path))) =
-        qualified_identity.private_keys.private_keys.get(&key)
-    {
-        qualified_identity
-            .associated_wallets
-            .get(&wallet_derivation_path.wallet_seed_hash)
-            .cloned()
-    } else {
-        None
-    }
-}

src/ui/identities/transfers/mod.rs

@@ -20,7 +20,8 @@ use std::sync::{Arc, RwLock};
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use crate::ui::components::wallet_unlock::ScreenWithWalletUnlock;
-use crate::ui::identities::register_dpns_name_screen::get_selected_wallet;
+
+use super::get_selected_wallet;
 
 pub enum TransferCreditsStatus {
     NotStarted,
@@ -47,20 +48,19 @@ pub struct TransferScreen {
 impl TransferScreen {
     pub fn new(identity: QualifiedIdentity, app_context: &Arc<AppContext>) -> Self {
         let max_amount = identity.identity.balance();
-        let selected_key = identity
-            .identity
-            .get_first_public_key_matching(
-                Purpose::TRANSFER,
-                SecurityLevel::full_range().into(),
-                KeyType::all_key_types().into(),
-                false,
-            )
-            .cloned();
+        let identity_clone = identity.identity.clone();
+        let selected_key = identity_clone.get_first_public_key_matching(
+            Purpose::TRANSFER,
+            SecurityLevel::full_range().into(),
+            KeyType::all_key_types().into(),
+            false,
+        );
         let mut error_message = None;
-        let selected_wallet = get_selected_wallet(&identity, app_context, &mut error_message);
+        let selected_wallet =
+            get_selected_wallet(&identity, None, selected_key, &mut error_message);
         Self {
             identity,
-            selected_key,
+            selected_key: selected_key.cloned(),
             receiver_identity_id: String::new(),
             amount: String::new(),
             transfer_credits_status: TransferCreditsStatus::NotStarted,

src/ui/identities/withdraw_from_identity_screen.rs

@@ -21,6 +21,7 @@ use std::str::FromStr;
 use std::sync::{Arc, RwLock};
 use std::time::{SystemTime, UNIX_EPOCH};
 
+use super::get_selected_wallet;
 use super::keys::key_info_screen::KeyInfoScreen;
 
 pub enum WithdrawFromIdentityStatus {
@@ -48,19 +49,29 @@ pub struct WithdrawalScreen {
 impl WithdrawalScreen {
     pub fn new(identity: QualifiedIdentity, app_context: &Arc<AppContext>) -> Self {
         let max_amount = identity.identity.balance();
+        let identity_clone = identity.identity.clone();
+        let selected_key = identity_clone.get_first_public_key_matching(
+            Purpose::TRANSFER,
+            SecurityLevel::full_range().into(),
+            KeyType::all_key_types().into(),
+            false,
+        );
+        let mut error_message = None;
+        let selected_wallet =
+            get_selected_wallet(&identity, None, selected_key, &mut error_message);
         Self {
             identity,
-            selected_key: None,
+            selected_key: selected_key.cloned(),
             withdrawal_address: String::new(),
             withdrawal_amount: String::new(),
             max_amount,
             app_context: app_context.clone(),
             confirmation_popup: false,
             withdraw_from_identity_status: WithdrawFromIdentityStatus::NotStarted,
-            selected_wallet: None,
+            selected_wallet,
             wallet_password: String::new(),
             show_password: false,
-            error_message: None,
+            error_message,
         }
     }