From 22ccdb6c23965bdb231ec93fa7c709ae555ff42e Mon Sep 17 00:00:00 2001
From: Pawan Dhananjay <pawandhananjay@gmail.com>
Date: Mon, 12 Aug 2024 17:16:21 -0700
Subject: [PATCH] Reuse password option prompts again on a wrong password
 (#4380)

* Prompt for password if incorrect in import

* lint and fmt

* Use if instead of match

* Fix issue raised by @chong-he

* Merge branch 'unstable' into reuse-pw-break

* Remove unused function
---
 account_manager/src/validator/import.rs | 50 ++++++++++++++++++-------
 1 file changed, 36 insertions(+), 14 deletions(-)

diff --git a/account_manager/src/validator/import.rs b/account_manager/src/validator/import.rs
index a7c72679f74..8f04e9059a3 100644
--- a/account_manager/src/validator/import.rs
+++ b/account_manager/src/validator/import.rs
@@ -178,7 +178,13 @@ pub fn cli_run(matches: &ArgMatches, validator_dir: PathBuf) -> Result<(), Strin
         let password_opt = loop {
             if let Some(password) = previous_password.clone() {
                 eprintln!("Reuse previous password.");
-                break Some(password);
+                if check_password_on_keystore(&keystore, &password)? {
+                    break Some(password);
+                } else {
+                    eprintln!("Reused password incorrect. Retry!");
+                    previous_password = None;
+                    continue;
+                }
             }
             eprintln!();
             eprintln!("{}", PASSWORD_PROMPT);
@@ -201,20 +207,12 @@ pub fn cli_run(matches: &ArgMatches, validator_dir: PathBuf) -> Result<(), Strin
                 }
             };
 
-            match keystore.decrypt_keypair(password.as_ref()) {
-                Ok(_) => {
-                    eprintln!("Password is correct.");
-                    eprintln!();
-                    sleep(Duration::from_secs(1)); // Provides nicer UX.
-                    if reuse_password {
-                        previous_password = Some(password.clone());
-                    }
-                    break Some(password);
-                }
-                Err(eth2_keystore::Error::InvalidPassword) => {
-                    eprintln!("Invalid password");
+            // Check if the password unlocks the keystore
+            if check_password_on_keystore(&keystore, &password)? {
+                if reuse_password {
+                    previous_password = Some(password.clone());
                 }
-                Err(e) => return Err(format!("Error whilst decrypting keypair: {:?}", e)),
+                break Some(password);
             }
         };
 
@@ -317,3 +315,27 @@ pub fn cli_run(matches: &ArgMatches, validator_dir: PathBuf) -> Result<(), Strin
 
     Ok(())
 }
+
+/// Checks if the given password unlocks the keystore.
+///
+/// Returns `Ok(true)` if password unlocks the keystore successfully.
+/// Returns `Ok(false` if password is incorrect.
+/// Otherwise, returns the keystore error.
+fn check_password_on_keystore(
+    keystore: &Keystore,
+    password: &ZeroizeString,
+) -> Result<bool, String> {
+    match keystore.decrypt_keypair(password.as_ref()) {
+        Ok(_) => {
+            eprintln!("Password is correct.");
+            eprintln!();
+            sleep(Duration::from_secs(1)); // Provides nicer UX.
+            Ok(true)
+        }
+        Err(eth2_keystore::Error::InvalidPassword) => {
+            eprintln!("Invalid password");
+            Ok(false)
+        }
+        Err(e) => Err(format!("Error whilst decrypting keypair: {:?}", e)),
+    }
+}