From f6f0f94e17d8c5d7a9c1d51463e3e0ad56c6cd08 Mon Sep 17 00:00:00 2001
From: Tyera Eulberg <teulberg@gmail.com>
Date: Wed, 26 Feb 2020 15:24:44 -0700
Subject: [PATCH] Add flag to confirm key on device (#8478)

---
 clap-utils/src/keypair.rs           |  1 +
 cli/src/cli.rs                      | 13 +++++++++++--
 remote-wallet/src/ledger.rs         | 10 +++++++---
 remote-wallet/src/remote_keypair.rs |  5 ++++-
 remote-wallet/src/remote_wallet.rs  |  6 +++++-
 5 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/clap-utils/src/keypair.rs b/clap-utils/src/keypair.rs
index 501538962..2b625ca9a 100644
--- a/clap-utils/src/keypair.rs
+++ b/clap-utils/src/keypair.rs
@@ -86,6 +86,7 @@ pub fn signer_from_path(
                     path,
                     derivation_of(matches, "derivation_path"),
                     wallet_manager,
+                    matches.is_present("confirm_key"),
                 )?))
             } else {
                 Err(RemoteWalletError::NoDeviceFound.into())
diff --git a/cli/src/cli.rs b/cli/src/cli.rs
index a36bedebb..d3ae64b31 100644
--- a/cli/src/cli.rs
+++ b/cli/src/cli.rs
@@ -620,7 +620,7 @@ pub fn parse_command(
         ),
         ("vote-account", Some(matches)) => parse_vote_get_account_command(matches),
         // Wallet Commands
-        ("address", Some(_matches)) => Ok(CliCommandInfo {
+        ("address", Some(matches)) => Ok(CliCommandInfo {
             command: CliCommand::Address,
             signers: vec![signer_from_path(
                 matches,
@@ -2064,7 +2064,16 @@ pub fn app<'ab, 'v>(name: &str, about: &'ab str, version: &'v str) -> App<'ab, '
         .about(about)
         .version(version)
         .setting(AppSettings::SubcommandRequiredElseHelp)
-        .subcommand(SubCommand::with_name("address").about("Get your public key"))
+        .subcommand(
+            SubCommand::with_name("address")
+                .about("Get your public key")
+                .arg(
+                    Arg::with_name("confirm_key")
+                        .long("confirm-key")
+                        .takes_value(false)
+                        .help("Confirm key on device; only relevant if using remote wallet"),
+                ),
+        )
         .cluster_query_subcommands()
         .nonce_subcommands()
         .stake_subcommands()
diff --git a/remote-wallet/src/ledger.rs b/remote-wallet/src/ledger.rs
index fffc4201f..671a6daf2 100644
--- a/remote-wallet/src/ledger.rs
+++ b/remote-wallet/src/ledger.rs
@@ -261,7 +261,7 @@ impl RemoteWallet for LedgerWallet {
             .serial_number
             .clone()
             .unwrap_or_else(|| "Unknown".to_owned());
-        self.get_pubkey(&DerivationPath::default())
+        self.get_pubkey(&DerivationPath::default(), false)
             .map(|pubkey| RemoteWalletInfo {
                 model,
                 manufacturer,
@@ -270,12 +270,16 @@ impl RemoteWallet for LedgerWallet {
             })
     }
 
-    fn get_pubkey(&self, derivation_path: &DerivationPath) -> Result<Pubkey, RemoteWalletError> {
+    fn get_pubkey(
+        &self,
+        derivation_path: &DerivationPath,
+        confirm_key: bool,
+    ) -> Result<Pubkey, RemoteWalletError> {
         let derivation_path = extend_and_serialize(derivation_path);
 
         let key = self.send_apdu(
             commands::GET_PUBKEY,
-            0, // In the naive implementation, default request is for no device confirmation
+            if confirm_key { 1 } else { 0 },
             0,
             &derivation_path,
         )?;
diff --git a/remote-wallet/src/remote_keypair.rs b/remote-wallet/src/remote_keypair.rs
index d6c740f20..7cd656f3b 100644
--- a/remote-wallet/src/remote_keypair.rs
+++ b/remote-wallet/src/remote_keypair.rs
@@ -20,9 +20,10 @@ impl RemoteKeypair {
     pub fn new(
         wallet_type: RemoteWalletType,
         derivation_path: DerivationPath,
+        confirm_key: bool,
     ) -> Result<Self, RemoteWalletError> {
         let pubkey = match &wallet_type {
-            RemoteWalletType::Ledger(wallet) => wallet.get_pubkey(&derivation_path)?,
+            RemoteWalletType::Ledger(wallet) => wallet.get_pubkey(&derivation_path, confirm_key)?,
         };
 
         Ok(Self {
@@ -51,6 +52,7 @@ pub fn generate_remote_keypair(
     path: String,
     explicit_derivation_path: Option<DerivationPath>,
     wallet_manager: &RemoteWalletManager,
+    confirm_key: bool,
 ) -> Result<RemoteKeypair, RemoteWalletError> {
     let (remote_wallet_info, mut derivation_path) = RemoteWalletInfo::parse_path(path)?;
     if let Some(derivation) = explicit_derivation_path {
@@ -61,6 +63,7 @@ pub fn generate_remote_keypair(
         Ok(RemoteKeypair::new(
             RemoteWalletType::Ledger(ledger),
             derivation_path,
+            confirm_key,
         )?)
     } else {
         Err(RemoteWalletError::DeviceTypeMismatch)
diff --git a/remote-wallet/src/remote_wallet.rs b/remote-wallet/src/remote_wallet.rs
index e9ccdea24..991555b4f 100644
--- a/remote-wallet/src/remote_wallet.rs
+++ b/remote-wallet/src/remote_wallet.rs
@@ -174,7 +174,11 @@ pub trait RemoteWallet {
     ) -> Result<RemoteWalletInfo, RemoteWalletError>;
 
     /// Get solana pubkey from a RemoteWallet
-    fn get_pubkey(&self, derivation_path: &DerivationPath) -> Result<Pubkey, RemoteWalletError>;
+    fn get_pubkey(
+        &self,
+        derivation_path: &DerivationPath,
+        confirm_key: bool,
+    ) -> Result<Pubkey, RemoteWalletError>;
 
     /// Sign transaction data with wallet managing pubkey at derivation path m/44'/501'/<account>'/<change>'.
     fn sign_message(