From c8c3c4359f5bbe698398d4eea36a24ba986aeee5 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Thu, 24 Mar 2022 11:51:40 -0700
Subject: [PATCH] `vote-authorize-voter` now accepts either the vote or
 withdraw authority

---
 cli/src/stake.rs | 25 ++++++++++---------------
 cli/src/vote.rs  |  7 +++++--
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/cli/src/stake.rs b/cli/src/stake.rs
index e70f5e052..daeddfb61 100644
--- a/cli/src/stake.rs
+++ b/cli/src/stake.rs
@@ -1383,17 +1383,12 @@ pub fn process_stake_authorize(
             };
             if let Some(authorized) = authorized {
                 match authorization_type {
-                    StakeAuthorize::Staker => {
-                        // first check authorized withdrawer
-                        check_current_authority(&authorized.withdrawer, &authority.pubkey())
-                            .or_else(|_| {
-                                // ...then check authorized staker. If neither matches, error will
-                                // print the stake key as `expected`
-                                check_current_authority(&authorized.staker, &authority.pubkey())
-                            })?;
-                    }
+                    StakeAuthorize::Staker => check_current_authority(
+                        &[authorized.withdrawer, authorized.staker],
+                        &authority.pubkey(),
+                    )?,
                     StakeAuthorize::Withdrawer => {
-                        check_current_authority(&authorized.withdrawer, &authority.pubkey())?;
+                        check_current_authority(&[authorized.withdrawer], &authority.pubkey())?;
                     }
                 }
             } else {
@@ -1935,7 +1930,7 @@ pub fn process_stake_set_lockup(
         };
         if let Some(lockup) = lockup {
             if lockup.custodian != Pubkey::default() {
-                check_current_authority(&lockup.custodian, &custodian.pubkey())?;
+                check_current_authority(&[lockup.custodian], &custodian.pubkey())?;
             }
         } else {
             return Err(CliError::RpcRequestError(format!(
@@ -2119,13 +2114,13 @@ fn get_stake_account_state(
 }
 
 pub(crate) fn check_current_authority(
-    account_current_authority: &Pubkey,
+    permitted_authorities: &[Pubkey],
     provided_current_authority: &Pubkey,
 ) -> Result<(), CliError> {
-    if account_current_authority != provided_current_authority {
+    if !permitted_authorities.contains(provided_current_authority) {
         Err(CliError::RpcRequestError(format!(
-            "Invalid current authority provided: {:?}, expected {:?}",
-            provided_current_authority, account_current_authority
+            "Invalid authority provided: {:?}, expected {:?}",
+            provided_current_authority, permitted_authorities
         )))
     } else {
         Ok(())
diff --git a/cli/src/vote.rs b/cli/src/vote.rs
index 85602f3ea..b7f200ff8 100644
--- a/cli/src/vote.rs
+++ b/cli/src/vote.rs
@@ -910,7 +910,10 @@ pub fn process_vote_authorize(
                             "Invalid vote account state; no authorized voters found".to_string(),
                         )
                     })?;
-                check_current_authority(&current_authorized_voter, &authorized.pubkey())?;
+                check_current_authority(
+                    &[current_authorized_voter, vote_state.authorized_withdrawer],
+                    &authorized.pubkey(),
+                )?;
                 if let Some(signer) = new_authorized_signer {
                     if signer.is_interactive() {
                         return Err(CliError::BadParameter(format!(
@@ -927,7 +930,7 @@ pub fn process_vote_authorize(
                 (new_authorized_pubkey, "new_authorized_pubkey".to_string()),
             )?;
             if let Some(vote_state) = vote_state {
-                check_current_authority(&vote_state.authorized_withdrawer, &authorized.pubkey())?
+                check_current_authority(&[vote_state.authorized_withdrawer], &authorized.pubkey())?
             }
         }
     }