From d5b662d0665e46b9d516e075c11406f0d17910ea Mon Sep 17 00:00:00 2001
From: carla <kirkcohenc@gmail.com>
Date: Mon, 15 Jul 2019 10:27:32 -0400
Subject: [PATCH] cmd: Add retry for password capture in create wallet

---
 cmd/lncli/commands.go | 109 +++++++++++++++++++++++-------------------
 1 file changed, 60 insertions(+), 49 deletions(-)

diff --git a/cmd/lncli/commands.go b/cmd/lncli/commands.go
index 46a05d0f..45e06a51 100644
--- a/cmd/lncli/commands.go
+++ b/cmd/lncli/commands.go
@@ -1398,33 +1398,12 @@ func create(ctx *cli.Context) error {
 	client, cleanUp := getWalletUnlockerClient(ctx)
 	defer cleanUp()
 
-	// First, we'll prompt the user for their passphrase twice to ensure
-	// both attempts match up properly.
-	fmt.Printf("Input wallet password: ")
-	pw1, err := terminal.ReadPassword(int(syscall.Stdin))
+	walletPassword, err := capturePassword(
+		"Input wallet password: ", false, walletunlocker.ValidatePassword,
+	)
 	if err != nil {
 		return err
 	}
-	fmt.Println()
-
-	fmt.Printf("Confirm wallet password: ")
-	pw2, err := terminal.ReadPassword(int(syscall.Stdin))
-	if err != nil {
-		return err
-	}
-	fmt.Println()
-
-	// If the passwords don't match, then we'll return an error.
-	if !bytes.Equal(pw1, pw2) {
-		return fmt.Errorf("passwords don't match")
-	}
-
-	// If the password length is less than 8 characters, then we'll
-	// return an error.
-	pwErr := walletunlocker.ValidatePassword(pw1)
-	if pwErr != nil {
-		return pwErr
-	}
 
 	// Next, we'll see if the user has 24-word mnemonic they want to use to
 	// derive a seed within the wallet.
@@ -1539,39 +1518,23 @@ mnemonicCheck:
 		// want to use, we'll generate a fresh one with the GenSeed
 		// command.
 		fmt.Println("Your cipher seed can optionally be encrypted.")
-		fmt.Printf("Input your passphrase if you wish to encrypt it " +
+
+		instruction := "Input your passphrase if you wish to encrypt it " +
 			"(or press enter to proceed without a cipher seed " +
-			"passphrase): ")
-		aezeedPass1, err := terminal.ReadPassword(int(syscall.Stdin))
+			"passphrase): "
+		aezeedPass, err = capturePassword(
+			instruction, true, func(_ []byte) error { return nil },
+		)
 		if err != nil {
 			return err
 		}
-		fmt.Println()
-
-		if len(aezeedPass1) != 0 {
-			fmt.Printf("Confirm cipher seed passphrase: ")
-			aezeedPass2, err := terminal.ReadPassword(
-				int(syscall.Stdin),
-			)
-			if err != nil {
-				return err
-			}
-			fmt.Println()
-
-			// If the passwords don't match, then we'll return an
-			// error.
-			if !bytes.Equal(aezeedPass1, aezeedPass2) {
-				return fmt.Errorf("cipher seed pass phrases " +
-					"don't match")
-			}
-		}
 
 		fmt.Println()
 		fmt.Println("Generating fresh cipher seed...")
 		fmt.Println()
 
 		genSeedReq := &lnrpc.GenSeedRequest{
-			AezeedPassphrase: aezeedPass1,
+			AezeedPassphrase: aezeedPass,
 		}
 		seedResp, err := client.GenSeed(ctxb, genSeedReq)
 		if err != nil {
@@ -1579,7 +1542,6 @@ mnemonicCheck:
 		}
 
 		cipherSeedMnemonic = seedResp.CipherSeedMnemonic
-		aezeedPass = aezeedPass1
 	}
 
 	// Before we initialize the wallet, we'll display the cipher seed to
@@ -1635,7 +1597,7 @@ mnemonicCheck:
 	// With either the user's prior cipher seed, or a newly generated one,
 	// we'll go ahead and initialize the wallet.
 	req := &lnrpc.InitWalletRequest{
-		WalletPassword:     pw1,
+		WalletPassword:     walletPassword,
 		CipherSeedMnemonic: cipherSeedMnemonic,
 		AezeedPassphrase:   aezeedPass,
 		RecoveryWindow:     recoveryWindow,
@@ -1650,6 +1612,55 @@ mnemonicCheck:
 	return nil
 }
 
+// capturePassword returns a password value that has been entered twice by the
+// user, to ensure that the user knows what password they have entered. The user
+// will be prompted to retry until the passwords match. If the optional param is
+// true, the function may return an empty byte array if the user opts against
+// using a password.
+func capturePassword(instruction string, optional bool,
+	validate func([]byte) error) ([]byte, error) {
+
+	for {
+		fmt.Printf(instruction)
+		password, err := terminal.ReadPassword(int(syscall.Stdin))
+		if err != nil {
+			return nil, err
+		}
+		fmt.Println()
+
+		// Do not require users to repeat password if
+		// it is optional and they are not using one.
+		if len(password) == 0 && optional {
+			return nil, nil
+		}
+
+		// If the password provided is not valid, restart
+		// password capture process from the beginning.
+		if err := validate(password); err != nil {
+			fmt.Println(err.Error())
+			fmt.Println()
+			continue
+		}
+
+		fmt.Println("Confirm password:")
+		passwordConfirmed, err := terminal.ReadPassword(
+			int(syscall.Stdin),
+		)
+		if err != nil {
+			return nil, err
+		}
+		fmt.Println()
+
+		if bytes.Equal(password, passwordConfirmed) {
+			return password, nil
+		}
+
+		fmt.Println("Passwords don't match, " +
+			"please try again")
+		fmt.Println()
+	}
+}
+
 var unlockCommand = cli.Command{
 	Name:     "unlock",
 	Category: "Startup",