diff --git a/program/src/deposit.rs b/program/src/deposit.rs index f2cbcf0..793c190 100644 --- a/program/src/deposit.rs +++ b/program/src/deposit.rs @@ -29,12 +29,6 @@ pub fn process_deposit(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResu token_program.is_program(&spl_token::ID)?; associated_token_program.is_program(&spl_associated_token_account::ID)?; - // Whitelist - assert!( - AUTHORIZED_ACCOUNTS.contains(&signer_info.key), - "Signer not whitelisted" - ); - // Open stake account. let stake = if stake_info.data_is_empty() { create_program_account::( @@ -100,5 +94,10 @@ pub fn process_deposit(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResu .as_str(), ); + // Safety check. + let stake_tokens = + stake_tokens_info.as_associated_token_account(stake_info.key, mint_info.key)?; + assert!(stake_tokens.amount() >= stake.balance); + Ok(()) } diff --git a/program/src/lib.rs b/program/src/lib.rs index e0822a5..869b860 100644 --- a/program/src/lib.rs +++ b/program/src/lib.rs @@ -61,7 +61,7 @@ pub fn process_instruction( // Staker OreInstruction::Deposit => process_deposit(accounts, data)?, OreInstruction::Withdraw => process_withdraw(accounts, data)?, - // OreInstruction::ClaimYield => process_claim_yield(accounts, data)?, + OreInstruction::ClaimYield => process_claim_yield(accounts, data)?, // Admin OreInstruction::Bury => process_bury(accounts, data)?, diff --git a/program/src/withdraw.rs b/program/src/withdraw.rs index 8e5ea85..0537bcc 100644 --- a/program/src/withdraw.rs +++ b/program/src/withdraw.rs @@ -32,11 +32,6 @@ pub fn process_withdraw(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramRes token_program.is_program(&spl_token::ID)?; associated_token_program.is_program(&spl_associated_token_account::ID)?; - assert!( - AUTHORIZED_ACCOUNTS.contains(&signer_info.key), - "Signer not whitelisted" - ); - // Open recipient token account. if recipient_info.data_is_empty() { create_associated_token_account( @@ -72,5 +67,10 @@ pub fn process_withdraw(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramRes .as_str(), ); + // Safety check. + let stake_tokens = + stake_tokens_info.as_associated_token_account(stake_info.key, mint_info.key)?; + assert!(stake_tokens.amount() >= stake.balance); + Ok(()) }