fix

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::<Stake>(
@@ -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(())
 }

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)?,

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(())
 }