[solana] Fix vesting logic (#216)

* Remove the ability to vest multiple different mint types

* Add vestingConfig and vestingBalance consistency checking in Delegate instruction

* Fix overwriting vest.amount when transferring vest

* Refactoring transfer_vesting: remove update_balance function

* Add test for TransferVesting instruction if another vester with an existing vest of the same kind

* Fix new_vesting_balance.vester in TransferVesting instruction

* Refactoring transfer_vesting: update set_inner

* Rename whTokenMint to votingTokenMint

* Fix vesting tests

* Apple yarn format

* Fix Underflow/Overflow error name

Author: zpoken

solana/app/StakeConnection.ts

@@ -335,7 +335,10 @@ export class StakeConnection {
     ).amount;
 
     const totalSupply = (
-      await getMint(this.program.provider.connection, this.config.whTokenMint)
+      await getMint(
+        this.program.provider.connection,
+        this.config.votingTokenMint,
+      )
     ).supply;
 
     return new StakeAccount(
@@ -366,7 +369,7 @@ export class StakeConnection {
       await this.program.methods
         .createStakeAccount()
         .accounts({
-          mint: this.config.whTokenMint,
+          mint: this.config.votingTokenMint,
         })
         .instruction(),
     );
@@ -389,7 +392,7 @@ export class StakeConnection {
       await this.program.methods
         .createStakeAccount()
         .accounts({
-          mint: this.config.whTokenMint,
+          mint: this.config.votingTokenMint,
         })
         .instruction(),
     );
@@ -402,7 +405,7 @@ export class StakeConnection {
     amount: BN,
   ): Promise<TransactionInstruction> {
     const from_account = await getAssociatedTokenAddress(
-      this.config.whTokenMint,
+      this.config.votingTokenMint,
       this.provider.wallet.publicKey,
       true,
     );
@@ -515,7 +518,7 @@ export class StakeConnection {
             delegateeStakeAccountCheckpointsAddress,
           vestingConfig: vestingConfigAccount,
           vestingBalance: vestingBalanceAccount,
-          mint: this.config.whTokenMint,
+          mint: this.config.votingTokenMint,
         })
         .instruction(),
     );
@@ -602,7 +605,7 @@ export class StakeConnection {
             delegateeStakeAccountCheckpointsAddress,
           vestingConfig: null,
           vestingBalance: null,
-          mint: this.config.whTokenMint,
+          mint: this.config.votingTokenMint,
         })
         .instruction(),
     );
@@ -739,7 +742,7 @@ export class StakeConnection {
     }
 
     const toAccount = await getAssociatedTokenAddress(
-      this.config.whTokenMint,
+      this.config.votingTokenMint,
       this.provider.wallet.publicKey,
       true,
     );
@@ -751,7 +754,7 @@ export class StakeConnection {
           this.provider.wallet.publicKey,
           toAccount,
           this.provider.wallet.publicKey,
-          this.config.whTokenMint,
+          this.config.votingTokenMint,
         ),
       );
     }

solana/app/deploy/1_init_staking.ts

@@ -19,7 +19,7 @@ async function main() {
   const globalConfig = {
     bump: 255,
     governanceAuthority: DEPLOYER_AUTHORITY_KEYPAIR.publicKey,
-    whTokenMint: WORMHOLE_TOKEN,
+    votingTokenMint: WORMHOLE_TOKEN,
     vestingAdmin: DEPLOYER_AUTHORITY_KEYPAIR.publicKey,
     maxCheckpointsAccountLimit: CHECKPOINTS_ACCOUNT_LIMIT,
   };

solana/programs/staking/src/context.rs

@@ -131,7 +131,7 @@ pub struct Delegate<'info> {
     #[account(seeds = [CONFIG_SEED.as_bytes()], bump = config.bump)]
     pub config: Box<Account<'info, global_config::GlobalConfig>>,
     // Wormhole token mint:
-    #[account(address = config.wh_token_mint)]
+    #[account(address = config.voting_token_mint)]
     pub mint: Account<'info, Mint>,
     pub system_program: Program<'info, System>,
 }
@@ -446,7 +446,7 @@ pub struct CreateStakeAccount<'info> {
     #[account(seeds = [CONFIG_SEED.as_bytes()], bump = config.bump)]
     pub config: Box<Account<'info, global_config::GlobalConfig>>,
     // Wormhole token mint:
-    #[account(address = config.wh_token_mint)]
+    #[account(address = config.voting_token_mint)]
     pub mint: Account<'info, Mint>,
     #[account(
         init,

solana/programs/staking/src/contexts/claim_vesting.rs

@@ -95,16 +95,6 @@ impl<'info> ClaimVesting<'info> {
                 );
                 drop(loaded_checkpoints);
 
-                require!(
-                    self.config.mint == self.global_config.wh_token_mint,
-                    // This error can never happen here, because for the condition above
-                    // (self.vesting_balance.stake_account_metadata != Pubkey::default())
-                    // to be met, the delegate instruction must be executed.
-                    // However, delegate cannot be executed when self.config.mint !=
-                    // self.global_config.wh_token_mint.
-                    VestingError::InvalidVestingMint
-                );
-
                 // Additional checks to ensure the owner matches
                 require!(
                     stake_account_metadata.owner == self.vesting_balance.vester,

solana/programs/staking/src/contexts/initialize.rs

@@ -11,6 +11,7 @@ use anchor_spl::token_interface::{Mint, TokenAccount, TokenInterface};
 pub struct Initialize<'info> {
     #[account(mut)]
     admin: Signer<'info>,
+    #[account(address = global_config.voting_token_mint)]
     mint: InterfaceAccount<'info, Mint>,
     // Initialize a vault for us to store our money in escrow for vesting
     #[account(

solana/programs/staking/src/contexts/transfer_vesting.rs

@@ -87,16 +87,11 @@ pub struct TransferVesting<'info> {
 }
 
 impl<'info> crate::contexts::TransferVesting<'info> {
-    pub fn transfer_vesting(&mut self, bump: u8) -> Result<()> {
-        fn update_balance(balance: &mut u64, amount: u64, is_subtract: bool) -> Result<()> {
-            if is_subtract {
-                *balance = balance.checked_sub(amount).ok_or(VestingError::Underflow)?;
-            } else {
-                *balance = balance.checked_add(amount).ok_or(VestingError::Underflow)?;
-            }
-            Ok(())
-        }
-
+    pub fn transfer_vesting(
+        &mut self,
+        new_vest_bump: u8,
+        new_vesting_balance_bump: u8,
+    ) -> Result<()> {
         if self.new_vester_ta.owner.key() == self.vester_ta.owner.key() {
             return err!(VestingError::TransferVestToMyself);
         }
@@ -149,11 +144,10 @@ impl<'info> crate::contexts::TransferVesting<'info> {
                     VestingError::InvalidStakeAccountMetadataPDA
                 );
 
-                update_balance(
-                    &mut stake_account_metadata.recorded_vesting_balance,
-                    self.vest.amount,
-                    true,
-                )?;
+                stake_account_metadata.recorded_vesting_balance = stake_account_metadata
+                    .recorded_vesting_balance
+                    .checked_sub(self.vest.amount)
+                    .ok_or(VestingError::Underflow)?;
 
                 // Update checkpoints
                 let current_delegate_checkpoints_account_info =
@@ -232,11 +226,10 @@ impl<'info> crate::contexts::TransferVesting<'info> {
                     VestingError::InvalidStakeAccountMetadataPDA
                 );
 
-                update_balance(
-                    &mut new_stake_account_metadata.recorded_vesting_balance,
-                    self.vest.amount,
-                    false,
-                )?;
+                new_stake_account_metadata.recorded_vesting_balance = new_stake_account_metadata
+                    .recorded_vesting_balance
+                    .checked_add(self.vest.amount)
+                    .ok_or(VestingError::Overflow)?;
 
                 let current_delegate_checkpoints_account_info =
                     new_stake_account_checkpoints.to_account_info();
@@ -267,21 +260,31 @@ impl<'info> crate::contexts::TransferVesting<'info> {
         self.new_vest.set_inner(Vesting {
             vester_ta: self.new_vester_ta.key(),
             config: self.vest.config,
-            amount: self.vest.amount,
+            amount: self
+                .new_vest
+                .amount
+                .checked_add(self.vest.amount)
+                .ok_or(VestingError::Overflow)?,
             maturation: self.vest.maturation,
-            bump,
+            bump: new_vest_bump,
+        });
+
+        self.new_vesting_balance.set_inner(VestingBalance {
+            vester: self.new_vester_ta.owner.key(),
+            stake_account_metadata: self.new_vesting_balance.stake_account_metadata,
+            total_vesting_balance: self
+                .new_vesting_balance
+                .total_vesting_balance
+                .checked_add(self.vest.amount)
+                .ok_or(VestingError::Overflow)?,
+            bump: new_vesting_balance_bump,
         });
 
-        update_balance(
-            &mut self.vesting_balance.total_vesting_balance,
-            self.vest.amount,
-            true,
-        )?;
-        update_balance(
-            &mut self.new_vesting_balance.total_vesting_balance,
-            self.vest.amount,
-            false,
-        )?;
+        self.vesting_balance.total_vesting_balance = self
+            .vesting_balance
+            .total_vesting_balance
+            .checked_sub(self.vest.amount)
+            .ok_or(VestingError::Underflow)?;
 
         Ok(())
     }

solana/programs/staking/src/error.rs

@@ -7,6 +7,8 @@ pub enum ErrorCode {
     TooManyCheckpoints,
     #[msg("An arithmetic operation unexpectedly overflowed")]
     GenericOverflow,
+    #[msg("An arithmetic operation unexpectedly underflowed")]
+    GenericUnderflow,
     #[msg("Error deserializing checkpoint")]
     CheckpointSerDe,
     #[msg("Checkpoint out of bounds")]
@@ -69,8 +71,6 @@ pub enum VestingError {
     InvalidStakeAccountCheckpoints,
     #[msg("Error parsing stake_account_metadata and stake_account_checkpoints")]
     ErrorOfStakeAccountParsing,
-    #[msg("Invalid vesting mint")]
-    InvalidVestingMint,
     #[msg("Invalid stake account owner")]
     InvalidStakeAccountOwner,
     #[msg("Invalid vesting admin")]
@@ -81,6 +81,8 @@ pub enum VestingError {
     InvalidStakeAccountMetadataPDA,
     #[msg("Invalid stake account checkpoints PDA")]
     InvalidStakeAccountCheckpointsPDA,
+    #[msg("Invalid vesting balance PDA")]
+    InvalidVestingBalancePDA,
     #[msg("Transfer vest to myself")]
     TransferVestToMyself,
 }

solana/programs/staking/src/lib.rs

@@ -78,7 +78,7 @@ pub mod staking {
         let config_account = &mut ctx.accounts.config_account;
         config_account.bump = ctx.bumps.config_account;
         config_account.governance_authority = global_config.governance_authority;
-        config_account.wh_token_mint = global_config.wh_token_mint;
+        config_account.voting_token_mint = global_config.voting_token_mint;
         config_account.vesting_admin = global_config.vesting_admin;
         config_account.max_checkpoints_account_limit = global_config.max_checkpoints_account_limit;
 
@@ -188,14 +188,19 @@ pub mod staking {
         if let Some(vesting_config) = &mut ctx.accounts.vesting_config {
             if vesting_config.finalized {
                 if let Some(vesting_balance) = &mut ctx.accounts.vesting_balance {
-                    require!(
-                        vesting_balance.vester == stake_account_metadata.owner,
-                        VestingError::InvalidStakeAccountOwner
+                    let (expected_vesting_balance_pda, _) = Pubkey::find_program_address(
+                        &[
+                            VESTING_BALANCE_SEED.as_bytes(),
+                            vesting_config.key().as_ref(),
+                            stake_account_metadata.owner.as_ref(),
+                        ],
+                        &crate::ID,
                     );
                     require!(
-                        vesting_config.mint == config.wh_token_mint,
-                        VestingError::InvalidVestingMint
+                        expected_vesting_balance_pda == vesting_balance.key(),
+                        VestingError::InvalidVestingBalancePDA
                     );
+
                     vesting_balance.stake_account_metadata = stake_account_metadata.key();
                     stake_account_metadata.recorded_vesting_balance =
                         vesting_balance.total_vesting_balance;
@@ -664,7 +669,8 @@ pub mod staking {
 
     // Transfer Vesting from and send to new Vester
     pub fn transfer_vesting(ctx: Context<TransferVesting>) -> Result<()> {
-        ctx.accounts.transfer_vesting(ctx.bumps.new_vest)
+        ctx.accounts
+            .transfer_vesting(ctx.bumps.new_vest, ctx.bumps.new_vesting_balance)
     }
 
     // Cancel and close a Vesting account for a non-finalized Config

solana/programs/staking/src/state/checkpoints.rs

@@ -280,7 +280,7 @@ fn calc_new_checkpoint(
             .ok_or_else(|| error!(ErrorCode::GenericOverflow))?,
         Operation::Subtract => current_value
             .checked_sub(amount_delta)
-            .ok_or_else(|| error!(ErrorCode::GenericOverflow))?,
+            .ok_or_else(|| error!(ErrorCode::GenericUnderflow))?,
     };
 
     let new_checkpoint = Checkpoint {

solana/programs/staking/src/state/global_config.rs

@@ -9,7 +9,7 @@ pub struct GlobalConfig {
     // Maximum number of checkpoints that can be stored in a single account
     pub max_checkpoints_account_limit: u32,
     pub governance_authority: Pubkey,
-    pub wh_token_mint: Pubkey,
+    pub voting_token_mint: Pubkey,
     pub vesting_admin: Pubkey,
 }
 

solana/programs/staking/src/state/vesting.rs

@@ -2,6 +2,7 @@ use anchor_lang::prelude::*;
 use std::mem::size_of;
 
 #[account]
+#[derive(Default)]
 pub struct Vesting {
     pub vester_ta: Pubkey,
     pub config: Pubkey,

solana/programs/staking/src/state/vesting_balance.rs

@@ -4,6 +4,7 @@ use std::mem::size_of;
 /// Used to store the total vesting balance of a single vester
 /// It is also used to delegate vesting
 #[account]
+#[derive(Default)]
 pub struct VestingBalance {
     pub vester: Pubkey,
     pub total_vesting_balance: u64,

solana/tests/api_test.ts

@@ -632,7 +632,7 @@ describe("api", async () => {
               user3StakeAccountCheckpointsAddress,
             vestingConfig: null,
             vestingBalance: null,
-            mint: stakeConnection.config.whTokenMint,
+            mint: stakeConnection.config.votingTokenMint,
           })
           .rpc();
 
@@ -677,7 +677,7 @@ describe("api", async () => {
               user3StakeAccountCheckpointsAddress,
             vestingConfig: null,
             vestingBalance: null,
-            mint: stakeConnection.config.whTokenMint,
+            mint: stakeConnection.config.votingTokenMint,
           })
           .rpc();
 
@@ -726,7 +726,7 @@ describe("api", async () => {
               ownerStakeAccountCheckpointsAddress, // Invalid delegatee checkpoints account
             vestingConfig: null,
             vestingBalance: null,
-            mint: stakeConnection.config.whTokenMint,
+            mint: stakeConnection.config.votingTokenMint,
           })
           .rpc();
 
@@ -801,7 +801,7 @@ describe("api", async () => {
         );
 
       const toAccount = await getAssociatedTokenAddress(
-        stakeConnection.config.whTokenMint,
+        stakeConnection.config.votingTokenMint,
         owner,
         true,
       );
@@ -842,7 +842,7 @@ describe("api", async () => {
       assert.equal(currentDelegate.toBase58(), user2.toBase58());
 
       const toAccount = await getAssociatedTokenAddress(
-        stakeConnection.config.whTokenMint,
+        stakeConnection.config.votingTokenMint,
         owner,
         true,
       );