Split accounts 3 (#247)

* Draft * Add authority * Clippy * Checkpoint * Do it * Rename error * Update stuff * Use epoch_of_snapshot * Checkpoint * Scaffolding * Checkpoint * Cleanup * Cleanup test * Another round * Cleanup * Restore all tests * add todos * Cleanup idls * Throw error since it's not implemented * Add some comments * Box everything * Add soruce * Add more comments * Add another comment * First implementation * Delete current request * add bumps * Fix bug * Tests works * Update idls * Clippy * Clippy * Add actual tests * Clippy * Cleanup * Fix tests * Split accounts test (#258) * vesting tests * more tests * better invariant tests * minor cleanups * cleanup implementation * idl * ok fix the ts tests * refactor * refactor * fix * cleanup * pr comments * minor * Cleanup * Comment * Add another test with full amount * Sorry --------- Co-authored-by: Jayant Krishnamurthy <[email protected]>
pyth-network · Nov 10, 2023 · 2ac18a1 · 2ac18a1
1 parent deb1121
commit 2ac18a1
Show file tree

Hide file tree

Showing 14 changed files with 802 additions and 94 deletions.
diff --git a/staking/app/StakeConnection.ts b/staking/app/StakeConnection.ts
@@ -122,6 +122,11 @@ export class StakeConnection {
     );
   }
 
+  /** The public key of the user of the staking program. This connection sends transactions as this user. */
+  public userPublicKey(): PublicKey {
+    return this.provider.wallet.publicKey;
+  }
+
   public async getAllStakeAccountAddresses(): Promise<PublicKey[]> {
     // Use the raw web3.js connection so that anchor doesn't try to borsh deserialize the zero-copy serialized account
     const allAccts = await this.provider.connection.getProgramAccounts(
@@ -539,9 +544,17 @@ export class StakeConnection {
     ) {
       throw Error(`Unexpected account state ${vestingAccountState}`);
     }
-    const owner: PublicKey = stakeAccount.stakeAccountMetadata.owner;
+
     const balanceSummary = stakeAccount.getBalanceSummary(await this.getTime());
-    const amountBN = balanceSummary.unvested.unlocked.toBN();
+    await this.lockTokens(stakeAccount, balanceSummary.unvested.unlocked);
+  }
+
+  /**
+   * Locks the specified amount of tokens in governance.
+   */
+  public async lockTokens(stakeAccount: StakeAccount, amount: PythBalance) {
+    const owner: PublicKey = stakeAccount.stakeAccountMetadata.owner;
+    const amountBN = amount.toBN();
 
     const transaction: Transaction = new Transaction();
 
@@ -860,7 +873,11 @@ export class StakeConnection {
       .rpc();
   }
 
-  public async acceptSplit(stakeAccount: StakeAccount) {
+  public async acceptSplit(
+    stakeAccount: StakeAccount,
+    amount: PythBalance,
+    recipient: PublicKey
+  ) {
     const newStakeAccountKeypair = new Keypair();
 
     const instructions = [];
@@ -872,7 +889,7 @@ export class StakeConnection {
     );
 
     await this.program.methods
-      .acceptSplit()
+      .acceptSplit(amount.toBN(), recipient)
       .accounts({
         sourceStakeAccountPositions: stakeAccount.address,
         newStakeAccountPositions: newStakeAccountKeypair.publicKey,

diff --git a/staking/programs/staking/src/context.rs b/staking/programs/staking/src/context.rs
@@ -311,6 +311,7 @@ pub struct RequestSplit<'info> {
 }
 
 #[derive(Accounts)]
+#[instruction(amount: u64, recipient: Pubkey)]
 pub struct AcceptSplit<'info> {
     // Native payer:
     #[account(mut, address = config.pda_authority)]

diff --git a/staking/programs/staking/src/error.rs b/staking/programs/staking/src/error.rs
@@ -65,6 +65,15 @@ pub enum ErrorCode {
     NotLlcMember,
     #[msg("Invalid LLC agreement")] // 6030
     InvalidLlcAgreement,
-    #[msg("Other")] //6031
+    #[msg("Can't split 0 tokens from an account")] // 6031
+    SplitZeroTokens,
+    #[msg("Can't split more tokens than are in the account")] // 6032
+    SplitTooManyTokens,
+    #[msg("Can't split a token account with staking positions. Unstake your tokens first.")]
+    // 6033
+    SplitWithStake,
+    #[msg("The approval arguments do not match the split request.")] // 6034
+    InvalidApproval,
+    #[msg("Other")] //6035
     Other,
 }
diff --git a/staking/programs/staking/src/lib.rs b/staking/programs/staking/src/lib.rs
@@ -2,6 +2,7 @@
 #![allow(dead_code)]
 #![allow(clippy::upper_case_acronyms)]
 #![allow(clippy::result_large_err)]
+#![allow(clippy::too_many_arguments)]
 // Objects of type Result must be used, otherwise we might
 // call a function that returns a Result and not handle the error
 
@@ -53,6 +54,7 @@ pub mod staking {
 
     /// Creates a global config for the program
     use super::*;
+
     pub fn init_config(ctx: Context<InitConfig>, global_config: GlobalConfig) -> Result<()> {
         let config_account = &mut ctx.accounts.config_account;
         config_account.bump = *ctx.bumps.get("config_account").unwrap();
@@ -124,25 +126,20 @@ pub mod staking {
         config.check_frozen()?;
 
         let stake_account_metadata = &mut ctx.accounts.stake_account_metadata;
-        stake_account_metadata.metadata_bump = *ctx.bumps.get("stake_account_metadata").unwrap();
-        stake_account_metadata.custody_bump = *ctx.bumps.get("stake_account_custody").unwrap();
-        stake_account_metadata.authority_bump = *ctx.bumps.get("custody_authority").unwrap();
-        stake_account_metadata.voter_bump = *ctx.bumps.get("voter_record").unwrap();
-        stake_account_metadata.owner = owner;
-        stake_account_metadata.next_index = 0;
-
-        stake_account_metadata.lock = lock;
-        stake_account_metadata.transfer_epoch = None;
-        stake_account_metadata.signed_agreement_hash = None;
+        stake_account_metadata.initialize(
+            *ctx.bumps.get("stake_account_metadata").unwrap(),
+            *ctx.bumps.get("stake_account_custody").unwrap(),
+            *ctx.bumps.get("custody_authority").unwrap(),
+            *ctx.bumps.get("voter_record").unwrap(),
+            &owner,
+        );
+        stake_account_metadata.set_lock(lock);
 
         let stake_account_positions = &mut ctx.accounts.stake_account_positions.load_init()?;
-        stake_account_positions.owner = owner;
+        stake_account_positions.initialize(&owner);
 
         let voter_record = &mut ctx.accounts.voter_record;
-
-        voter_record.realm = config.pyth_governance_realm;
-        voter_record.governing_token_mint = config.pyth_token_mint;
-        voter_record.governing_token_owner = owner;
+        voter_record.initialize(config, &owner);
 
         Ok(())
     }
@@ -573,32 +570,131 @@ pub mod staking {
      * the config account. If accepted, `amount` tokens are transferred to a new stake account
      * owned by the `recipient` and the split request is reset (by setting `amount` to 0).
      * The recipient of a transfer can't vote during the epoch of the transfer.
+     *
+     * The `pda_authority` must explicitly approve both the amount of tokens and recipient, and
+     * these parameters must match the request (in the `split_request` account).
      */
-    pub fn accept_split(ctx: Context<AcceptSplit>) -> Result<()> {
-        // TODO : Split vesting schedule between both accounts
+    pub fn accept_split(ctx: Context<AcceptSplit>, amount: u64, recipient: Pubkey) -> Result<()> {
+        let config = &ctx.accounts.config;
+        config.check_frozen()?;
 
-        // TODO : Transfer stake positions to the new account if need
+        let current_epoch = get_current_epoch(config)?;
 
-        // TODO Check both accounts are valid after the transfer
+        let split_request = &ctx.accounts.source_stake_account_split_request;
+        require!(
+            split_request.amount == amount && split_request.recipient == recipient,
+            ErrorCode::InvalidApproval
+        );
+
+        // Initialize new accounts
+        ctx.accounts.new_stake_account_metadata.initialize(
+            *ctx.bumps.get("new_stake_account_metadata").unwrap(),
+            *ctx.bumps.get("new_stake_account_custody").unwrap(),
+            *ctx.bumps.get("new_custody_authority").unwrap(),
+            *ctx.bumps.get("new_voter_record").unwrap(),
+            &split_request.recipient,
+        );
+
+        let new_stake_account_positions =
+            &mut ctx.accounts.new_stake_account_positions.load_init()?;
+        new_stake_account_positions.initialize(&split_request.recipient);
+
+        let new_voter_record = &mut ctx.accounts.new_voter_record;
+        new_voter_record.initialize(config, &split_request.recipient);
+
+        // Pre-check invariants
+        // Note that the accept operation requires the positions account to be empty, which should trivially
+        // pass this invariant check. However, we explicitly check invariants everywhere else, so may
+        // as well check in this operation also.
+        let source_stake_account_positions =
+            &mut ctx.accounts.source_stake_account_positions.load_mut()?;
+        utils::risk::validate(
+            source_stake_account_positions,
+            ctx.accounts.source_stake_account_custody.amount,
+            ctx.accounts
+                .source_stake_account_metadata
+                .lock
+                .get_unvested_balance(
+                    utils::clock::get_current_time(config),
+                    config.pyth_token_list_time,
+                )?,
+            current_epoch,
+            config.unlocking_duration,
+        )?;
 
-        // Transfer tokens
-        {
-            let split_request = &ctx.accounts.source_stake_account_split_request;
-            transfer(
-                CpiContext::from(&*ctx.accounts).with_signer(&[&[
-                    AUTHORITY_SEED.as_bytes(),
-                    ctx.accounts.source_stake_account_positions.key().as_ref(),
-                    &[ctx.accounts.source_stake_account_metadata.authority_bump],
-                ]]),
+        // Check that there aren't any positions (i.e., staked tokens) in the source account.
+        // This check allows us to create an empty positions account on behalf of the recipient and
+        // not worry about moving positions from the source account to the new account.
+        require!(
+            ctx.accounts.source_stake_account_metadata.next_index == 0,
+            ErrorCode::SplitWithStake
+        );
+
+        require!(split_request.amount > 0, ErrorCode::SplitZeroTokens);
+
+        // Split vesting account
+        let (source_vesting_schedule, new_vesting_schedule) = ctx
+            .accounts
+            .source_stake_account_metadata
+            .lock
+            .split_vesting_schedule(
                 split_request.amount,
+                ctx.accounts.source_stake_account_custody.amount,
             )?;
-        }
+        ctx.accounts
+            .source_stake_account_metadata
+            .set_lock(source_vesting_schedule);
+        ctx.accounts
+            .new_stake_account_metadata
+            .set_lock(new_vesting_schedule);
+
+
+        transfer(
+            CpiContext::from(&*ctx.accounts).with_signer(&[&[
+                AUTHORITY_SEED.as_bytes(),
+                ctx.accounts.source_stake_account_positions.key().as_ref(),
+                &[ctx.accounts.source_stake_account_metadata.authority_bump],
+            ]]),
+            split_request.amount,
+        )?;
+
+        ctx.accounts.source_stake_account_custody.reload()?;
+        ctx.accounts.new_stake_account_custody.reload()?;
+
+
+        // Post-check
+        utils::risk::validate(
+            source_stake_account_positions,
+            ctx.accounts.source_stake_account_custody.amount,
+            ctx.accounts
+                .source_stake_account_metadata
+                .lock
+                .get_unvested_balance(
+                    utils::clock::get_current_time(config),
+                    config.pyth_token_list_time,
+                )?,
+            current_epoch,
+            config.unlocking_duration,
+        )?;
+
+        utils::risk::validate(
+            new_stake_account_positions,
+            ctx.accounts.new_stake_account_custody.amount,
+            ctx.accounts
+                .new_stake_account_metadata
+                .lock
+                .get_unvested_balance(
+                    utils::clock::get_current_time(config),
+                    config.pyth_token_list_time,
+                )?,
+            current_epoch,
+            config.unlocking_duration,
+        )?;
 
         // Delete current request
-        {
-            ctx.accounts.source_stake_account_split_request.amount = 0;
-        }
-        err!(ErrorCode::NotImplemented)
+        ctx.accounts.source_stake_account_split_request.amount = 0;
+
+        Ok(())
     }
 
     /**

diff --git a/staking/programs/staking/src/state/positions.rs b/staking/programs/staking/src/state/positions.rs
@@ -47,6 +47,10 @@ impl Default for PositionData {
     }
 }
 impl PositionData {
+    pub fn initialize(&mut self, owner: &Pubkey) {
+        self.owner = *owner;
+    }
+
     /// Finds first index available for a new position, increments the internal counter
     pub fn reserve_new_index(&mut self, next_index: &mut u8) -> Result<usize> {
         let res = *next_index as usize;
@@ -403,7 +407,6 @@ pub mod tests {
         }
     }
 
-
     #[quickcheck]
     fn prop(input: Vec<DataOperation>) -> bool {
         let mut position_data = PositionData::default();

diff --git a/staking/programs/staking/src/state/stake_account.rs b/staking/programs/staking/src/state/stake_account.rs
@@ -41,6 +41,30 @@ impl StakeAccountMetadataV2 {
     }
 }
 
+impl StakeAccountMetadataV2 {
+    pub fn initialize(
+        &mut self,
+        metadata_bump: u8,
+        custody_bump: u8,
+        authority_bump: u8,
+        voter_bump: u8,
+        owner: &Pubkey,
+    ) {
+        self.metadata_bump = metadata_bump;
+        self.custody_bump = custody_bump;
+        self.authority_bump = authority_bump;
+        self.voter_bump = voter_bump;
+        self.owner = *owner;
+        self.next_index = 0;
+        self.transfer_epoch = None;
+        self.signed_agreement_hash = None;
+    }
+
+    pub fn set_lock(&mut self, lock: VestingSchedule) {
+        self.lock = lock;
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
     use {