From 2bcbe44c79827cb4ba12f94e0a2cd8a24bfdb27d Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Tue, 9 Jun 2020 21:15:54 -0700
Subject: [PATCH] Add VoteInstruction::UpdateCommission

---
 programs/vote/src/vote_instruction.rs | 35 +++++++++++++++
 programs/vote/src/vote_state/mod.rs   | 61 +++++++++++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/programs/vote/src/vote_instruction.rs b/programs/vote/src/vote_instruction.rs
index 9cfa0ad1b0..e265a753de 100644
--- a/programs/vote/src/vote_instruction.rs
+++ b/programs/vote/src/vote_instruction.rs
@@ -93,6 +93,13 @@ pub enum VoteInstruction {
     ///   2. [SIGNER] Withdraw authority
     UpdateValidatorIdentity,
 
+    /// Update the commission for the vote account
+    ///
+    /// # Account references
+    ///   0. [WRITE] Vote account to be updated
+    ///   1. [SIGNER] Withdraw authority
+    UpdateCommission(u8),
+
     /// A Vote instruction with recent votes
     ///
     /// # Account references
@@ -190,6 +197,23 @@ pub fn update_validator_identity(
     )
 }
 
+pub fn update_commission(
+    vote_pubkey: &Pubkey,
+    authorized_withdrawer_pubkey: &Pubkey,
+    commission: u8,
+) -> Instruction {
+    let account_metas = vec![
+        AccountMeta::new(*vote_pubkey, false),
+        AccountMeta::new_readonly(*authorized_withdrawer_pubkey, true),
+    ];
+
+    Instruction::new(
+        id(),
+        &VoteInstruction::UpdateCommission(commission),
+        account_metas,
+    )
+}
+
 pub fn vote(vote_pubkey: &Pubkey, authorized_voter_pubkey: &Pubkey, vote: Vote) -> Instruction {
     let account_metas = vec![
         AccountMeta::new(*vote_pubkey, false),
@@ -271,6 +295,9 @@ pub fn process_instruction(
             next_keyed_account(keyed_accounts)?.unsigned_key(),
             &signers,
         ),
+        VoteInstruction::UpdateCommission(commission) => {
+            vote_state::update_commission(me, commission, &signers)
+        }
         VoteInstruction::Vote(vote) | VoteInstruction::VoteSwitch(vote, _) => {
             inc_new_counter_info!("vote-native", 1);
             vote_state::process_vote(
@@ -380,6 +407,14 @@ mod tests {
             )),
             Err(InstructionError::InvalidAccountData),
         );
+        assert_eq!(
+            process_instruction(&update_commission(
+                &Pubkey::default(),
+                &Pubkey::default(),
+                0,
+            )),
+            Err(InstructionError::InvalidAccountData),
+        );
 
         assert_eq!(
             process_instruction(&withdraw(
diff --git a/programs/vote/src/vote_state/mod.rs b/programs/vote/src/vote_state/mod.rs
index e9e960d2a7..26ed5e1e9d 100644
--- a/programs/vote/src/vote_state/mod.rs
+++ b/programs/vote/src/vote_state/mod.rs
@@ -597,6 +597,23 @@ pub fn update_validator_identity<S: std::hash::BuildHasher>(
     vote_account.set_state(&VoteStateVersions::Current(Box::new(vote_state)))
 }
 
+/// Update the vote account's commission
+pub fn update_commission<S: std::hash::BuildHasher>(
+    vote_account: &KeyedAccount,
+    commission: u8,
+    signers: &HashSet<Pubkey, S>,
+) -> Result<(), InstructionError> {
+    let mut vote_state: VoteState =
+        State::<VoteStateVersions>::state(vote_account)?.convert_to_current();
+
+    // current authorized withdrawer must say "yay"
+    verify_authorized_signer(&vote_state.authorized_withdrawer, signers)?;
+
+    vote_state.commission = commission;
+
+    vote_account.set_state(&VoteStateVersions::Current(Box::new(vote_state)))
+}
+
 fn verify_authorized_signer<S: std::hash::BuildHasher>(
     authorized: &Pubkey,
     signers: &HashSet<Pubkey, S>,
@@ -997,6 +1014,50 @@ mod tests {
         assert_eq!(vote_state.node_pubkey, node_pubkey);
     }
 
+    #[test]
+    fn test_vote_update_commission() {
+        let (vote_pubkey, _authorized_voter, authorized_withdrawer, vote_account) =
+            create_test_account_with_authorized();
+
+        let authorized_withdrawer_account = RefCell::new(Account::default());
+
+        let keyed_accounts = &[
+            KeyedAccount::new(&vote_pubkey, true, &vote_account),
+            KeyedAccount::new(
+                &authorized_withdrawer,
+                false,
+                &authorized_withdrawer_account,
+            ),
+        ];
+        let signers: HashSet<Pubkey> = get_signers(keyed_accounts);
+        let res = update_commission(&keyed_accounts[0], 42, &signers);
+        assert_eq!(res, Err(InstructionError::MissingRequiredSignature));
+
+        let keyed_accounts = &[
+            KeyedAccount::new(&vote_pubkey, true, &vote_account),
+            KeyedAccount::new(&authorized_withdrawer, true, &authorized_withdrawer_account),
+        ];
+        let signers: HashSet<Pubkey> = get_signers(keyed_accounts);
+        let res = update_commission(&keyed_accounts[0], 42, &signers);
+        assert_eq!(res, Ok(()));
+        let vote_state: VoteState = StateMut::<VoteStateVersions>::state(&*vote_account.borrow())
+            .unwrap()
+            .convert_to_current();
+        assert_eq!(vote_state.commission, 42);
+
+        let keyed_accounts = &[
+            KeyedAccount::new(&vote_pubkey, true, &vote_account),
+            KeyedAccount::new(&authorized_withdrawer, true, &authorized_withdrawer_account),
+        ];
+        let signers: HashSet<Pubkey> = get_signers(keyed_accounts);
+        let res = update_commission(&keyed_accounts[0], u8::MAX, &signers);
+        assert_eq!(res, Ok(()));
+        let vote_state: VoteState = StateMut::<VoteStateVersions>::state(&*vote_account.borrow())
+            .unwrap()
+            .convert_to_current();
+        assert_eq!(vote_state.commission, u8::MAX);
+    }
+
     #[test]
     fn test_vote_signature() {
         let (vote_pubkey, vote_account) = create_test_account();