stake-pool: Ensure zero pool token supply on init (#1572)

stake-pool/program/src/error.rs

@@ -82,6 +82,9 @@ pub enum StakePoolError {
     /// Wrong pool staker account.
     #[error("WrongStaker")]
     WrongStaker,
+    /// Pool token supply is not zero on initialization
+    #[error("NonZeroPoolTokenSupply")]
+    NonZeroPoolTokenSupply,
 }
 impl From<StakePoolError> for ProgramError {
     fn from(e: StakePoolError) -> Self {

stake-pool/program/src/instruction.rs

@@ -33,7 +33,7 @@ pub enum StakePoolInstruction {
     ///   1. `[s]` Manager
     ///   2. `[]` Staker
     ///   3. `[w]` Uninitialized validator stake list storage account
-    ///   4. `[]` Pool token mint. Must be non zero, owned by withdraw authority.
+    ///   4. `[]` Pool token mint. Must have zero supply, owned by withdraw authority.
     ///   5. `[]` Pool account to deposit the generated fee for manager.
     ///   6. `[]` Clock sysvar
     ///   7. `[]` Rent sysvar

stake-pool/program/src/processor.rs

@@ -281,7 +281,7 @@ impl Processor {
         }
 
         if manager_fee_info.owner != token_program_info.key {
-            return Err(StakePoolError::InvalidFeeAccount.into());
+            return Err(ProgramError::IncorrectProgramId);
         }
 
         if pool_mint_info.owner != token_program_info.key {
@@ -301,6 +301,10 @@ impl Processor {
 
         let pool_mint = Mint::unpack_from_slice(&pool_mint_info.data.borrow())?;
 
+        if pool_mint.supply != 0 {
+            return Err(StakePoolError::NonZeroPoolTokenSupply.into());
+        }
+
         if !pool_mint.mint_authority.contains(&withdraw_authority_key) {
             return Err(StakePoolError::WrongMintingAuthority.into());
         }
@@ -1155,6 +1159,7 @@ impl PrintProgramError for StakePoolError {
             StakePoolError::WrongMintingAuthority => msg!("Error: Wrong minting authority set for mint pool account"),
             StakePoolError::UnexpectedValidatorListAccountSize=> msg!("Error: The size of the given validator stake list does match the expected amount"),
             StakePoolError::WrongStaker=> msg!("Error: Wrong pool staker account"),
+            StakePoolError::NonZeroPoolTokenSupply => msg!("Error: Pool token supply is not zero on initialization"),
         }
     }
 }

stake-pool/program/tests/helpers/mod.rs

@@ -126,6 +126,33 @@ pub async fn create_token_account(
     Ok(())
 }
 
+pub async fn mint_tokens(
+    banks_client: &mut BanksClient,
+    payer: &Keypair,
+    recent_blockhash: &Hash,
+    mint: &Pubkey,
+    account: &Pubkey,
+    mint_authority: &Keypair,
+    amount: u64,
+) -> Result<(), TransportError> {
+    let transaction = Transaction::new_signed_with_payer(
+        &[spl_token::instruction::mint_to(
+            &spl_token::id(),
+            mint,
+            account,
+            &mint_authority.pubkey(),
+            &[],
+            amount,
+        )
+        .unwrap()],
+        Some(&payer.pubkey()),
+        &[payer, mint_authority],
+        *recent_blockhash,
+    );
+    banks_client.process_transaction(transaction).await?;
+    Ok(())
+}
+
 pub async fn get_token_balance(banks_client: &mut BanksClient, token: &Pubkey) -> u64 {
     let token_account = banks_client
         .get_account(token.clone())

stake-pool/program/tests/initialize.rs

@@ -52,7 +52,7 @@ async fn create_mint_and_token_account(
 }
 
 #[tokio::test]
-async fn test_stake_pool_initialize() {
+async fn success_initialize() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
     stake_pool_accounts
@@ -77,7 +77,7 @@ async fn test_stake_pool_initialize() {
 }
 
 #[tokio::test]
-async fn test_initialize_already_initialized_stake_pool() {
+async fn fail_double_initialize() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
     stake_pool_accounts
@@ -108,7 +108,7 @@ async fn test_initialize_already_initialized_stake_pool() {
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_already_initialized_stake_list_storage() {
+async fn fail_initialize_with_already_initialized_validator_list() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
     stake_pool_accounts
@@ -139,7 +139,7 @@ async fn test_initialize_stake_pool_with_already_initialized_stake_list_storage(
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_high_fee() {
+async fn fail_initialize_with_high_fee() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let mut stake_pool_accounts = StakePoolAccounts::new();
     stake_pool_accounts.fee = instruction::Fee {
@@ -165,7 +165,7 @@ async fn test_initialize_stake_pool_with_high_fee() {
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_wrong_max_validators() {
+async fn fail_initialize_with_wrong_max_validators() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
 
@@ -245,7 +245,7 @@ async fn test_initialize_stake_pool_with_wrong_max_validators() {
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_wrong_mint_authority() {
+async fn fail_initialize_with_wrong_mint_authority() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
     let wrong_mint = Keypair::new();
@@ -299,7 +299,7 @@ async fn test_initialize_stake_pool_with_wrong_mint_authority() {
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_wrong_token_program_id() {
+async fn fail_initialize_with_wrong_token_program_id() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
 
@@ -399,7 +399,7 @@ async fn test_initialize_stake_pool_with_wrong_token_program_id() {
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_wrong_fee_accounts_manager() {
+async fn fail_initialize_with_wrong_fee_account() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
 
@@ -446,24 +446,17 @@ async fn test_initialize_stake_pool_with_wrong_fee_accounts_manager() {
     )
     .await
     .err()
+    .unwrap()
     .unwrap();
 
-    match transaction_error {
-        TransportError::TransactionError(TransactionError::InstructionError(
-            _,
-            InstructionError::Custom(error_index),
-        )) => {
-            let program_error = error::StakePoolError::InvalidFeeAccount as u32;
-            assert_eq!(error_index, program_error);
-        }
-        _ => panic!(
-            "Wrong error occurs while try to initialize stake pool with wrong fee account's manager"
-        ),
-    }
+    assert_eq!(
+        transaction_error,
+        TransactionError::InstructionError(2, InstructionError::IncorrectProgramId)
+    );
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_wrong_withdraw_authority() {
+async fn fail_initialize_with_wrong_withdraw_authority() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let mut stake_pool_accounts = StakePoolAccounts::new();
 
@@ -490,7 +483,7 @@ async fn test_initialize_stake_pool_with_wrong_withdraw_authority() {
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_not_rent_exempt_pool() {
+async fn fail_initialize_with_not_rent_exempt_pool() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
 
@@ -563,7 +556,7 @@ async fn test_initialize_stake_pool_with_not_rent_exempt_pool() {
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_with_not_rent_exempt_validator_list() {
+async fn fail_initialize_with_not_rent_exempt_validator_list() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
 
@@ -638,7 +631,7 @@ async fn test_initialize_stake_pool_with_not_rent_exempt_validator_list() {
 }
 
 #[tokio::test]
-async fn test_initialize_stake_pool_without_manager_signature() {
+async fn fail_initialize_without_manager_signature() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
 
@@ -727,3 +720,71 @@ async fn test_initialize_stake_pool_without_manager_signature() {
         ),
     }
 }
+
+#[tokio::test]
+async fn fail_initialize_with_pre_minted_pool_tokens() {
+    let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
+    let stake_pool_accounts = StakePoolAccounts::new();
+    let mint_authority = Keypair::new();
+
+    create_mint(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts.pool_mint,
+        &mint_authority.pubkey(),
+    )
+    .await
+    .unwrap();
+
+    create_token_account(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts.pool_fee_account,
+        &stake_pool_accounts.pool_mint.pubkey(),
+        &stake_pool_accounts.manager.pubkey(),
+    )
+    .await
+    .unwrap();
+
+    mint_tokens(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts.pool_mint.pubkey(),
+        &stake_pool_accounts.pool_fee_account.pubkey(),
+        &mint_authority,
+        1,
+    )
+    .await
+    .unwrap();
+
+    let transaction_error = create_stake_pool(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts.stake_pool,
+        &stake_pool_accounts.validator_list,
+        &stake_pool_accounts.pool_mint.pubkey(),
+        &stake_pool_accounts.pool_fee_account.pubkey(),
+        &stake_pool_accounts.manager,
+        &stake_pool_accounts.staker.pubkey(),
+        &stake_pool_accounts.fee,
+        stake_pool_accounts.max_validators,
+    )
+    .await
+    .err()
+    .unwrap();
+
+    match transaction_error {
+        TransportError::TransactionError(TransactionError::InstructionError(
+            _,
+            InstructionError::Custom(error_index),
+        )) => {
+            let program_error = error::StakePoolError::NonZeroPoolTokenSupply as u32;
+            assert_eq!(error_index, program_error);
+        }
+        _ => panic!("Wrong error occurs while try to initialize stake pool with wrong mint authority of pool fee account"),
+    }
+}