stake

program/src/deposit.rs

@@ -3,6 +3,8 @@ use solana_program::log::sol_log;
 use spl_token::amount_to_ui_amount;
 use steel::*;
 
+use crate::AUTHORIZED_ACCOUNTS;
+
 /// Deposits ORE into the staking contract.
 pub fn process_deposit(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult {
     // Parse data.
@@ -11,22 +13,27 @@ pub fn process_deposit(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResu
 
     // Load accounts.
     let clock = Clock::get()?;
-    let [signer_info, sender_info, stake_info, treasury_info, treasury_tokens_info, system_program, token_program] =
+    let [signer_info, mint_info, sender_info, stake_info, stake_tokens_info, treasury_info, system_program, token_program, associated_token_program] =
         accounts
     else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
     signer_info.is_signer()?;
+    mint_info.has_address(&MINT_ADDRESS)?.as_mint()?;
     let sender = sender_info
         .is_writable()?
         .as_associated_token_account(&signer_info.key, &MINT_ADDRESS)?;
     stake_info.is_writable()?;
     let treasury = treasury_info.as_account_mut::<Treasury>(&ore_api::ID)?;
-    treasury_tokens_info
-        .is_writable()?
-        .as_associated_token_account(&treasury_info.key, &MINT_ADDRESS)?;
     system_program.is_program(&system_program::ID)?;
     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() {
@@ -54,6 +61,21 @@ pub fn process_deposit(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResu
             .assert_mut(|s| s.authority == *signer_info.key)?
     };
 
+    // Create stake tokens account.
+    if stake_tokens_info.data_is_empty() {
+        create_associated_token_account(
+            signer_info,
+            stake_info,
+            stake_tokens_info,
+            mint_info,
+            system_program,
+            token_program,
+            associated_token_program,
+        )?;
+    } else {
+        stake_tokens_info.as_associated_token_account(stake_info.key, mint_info.key)?;
+    }
+
     // Only allow deposits from seekers.
     // assert!(stake.is_seeker == 1, "Only seekers can deposit stake");
 
@@ -64,7 +86,7 @@ pub fn process_deposit(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResu
     transfer(
         signer_info,
         sender_info,
-        treasury_tokens_info,
+        stake_tokens_info,
         token_program,
         amount,
     )?;

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/migrate_staker.rs

@@ -0,0 +1,59 @@
+use ore_api::prelude::*;
+use steel::*;
+
+/// Sets the admin.
+pub fn process_migrate_staker(accounts: &[AccountInfo<'_>], _data: &[u8]) -> ProgramResult {
+    // Load accounts.
+    let [signer_info, config_info, mint_info, stake_info, stake_tokens_info, treasury_info, treasury_tokens_info, system_program, token_program, associated_token_program] =
+        accounts
+    else {
+        return Err(ProgramError::NotEnoughAccountKeys);
+    };
+    signer_info.is_signer()?;
+    config_info.as_account::<Config>(&ore_api::ID)?.assert_err(
+        |c| c.admin == *signer_info.key,
+        OreError::NotAuthorized.into(),
+    )?;
+    mint_info.has_address(&MINT_ADDRESS)?.as_mint()?;
+    let stake = stake_info.as_account_mut::<Stake>(&ore_api::ID)?;
+    stake_tokens_info.is_empty()?.is_writable()?;
+    treasury_info.as_account_mut::<Treasury>(&ore_api::ID)?;
+    treasury_tokens_info.as_associated_token_account(&treasury_info.key, &MINT_ADDRESS)?;
+    system_program.is_program(&system_program::ID)?;
+    token_program.is_program(&spl_token::ID)?;
+    associated_token_program.is_program(&spl_associated_token_account::ID)?;
+
+    // Create stake tokens account.
+    if stake_tokens_info.data_is_empty() {
+        create_associated_token_account(
+            signer_info,
+            stake_info,
+            stake_tokens_info,
+            mint_info,
+            system_program,
+            token_program,
+            associated_token_program,
+        )?;
+    }
+    let stake_tokens =
+        stake_tokens_info.as_associated_token_account(stake_info.key, mint_info.key)?;
+
+    // Move tokens from treasury to stake tokens.
+    if stake_tokens.amount() == 0 {
+        transfer_signed(
+            treasury_info,
+            treasury_tokens_info,
+            stake_tokens_info,
+            token_program,
+            stake.balance,
+            &[TREASURY],
+        )?;
+    }
+
+    // Safety check.
+    let stake_tokens =
+        stake_tokens_info.as_associated_token_account(stake_info.key, mint_info.key)?;
+    assert_eq!(stake_tokens.amount(), stake.balance);
+
+    Ok(())
+}

program/src/withdraw.rs

@@ -3,6 +3,8 @@ use solana_program::log::sol_log;
 use spl_token::amount_to_ui_amount;
 use steel::*;
 
+use crate::AUTHORIZED_ACCOUNTS;
+
 /// Withdraws ORE from the staking contract.
 pub fn process_withdraw(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult {
     // Parse data.
@@ -11,7 +13,7 @@ pub fn process_withdraw(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramRes
 
     // Load accounts.
     let clock = Clock::get()?;
-    let [signer_info, mint_info, recipient_info, stake_info, treasury_info, treasury_tokens_info, system_program, token_program, associated_token_program] =
+    let [signer_info, mint_info, recipient_info, stake_info, stake_tokens_info, treasury_info, system_program, token_program, associated_token_program] =
         accounts
     else {
         return Err(ProgramError::NotEnoughAccountKeys);
@@ -24,14 +26,17 @@ pub fn process_withdraw(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramRes
     let stake = stake_info
         .as_account_mut::<Stake>(&ore_api::ID)?
         .assert_mut(|s| s.authority == *signer_info.key)?;
+    stake_tokens_info.as_associated_token_account(stake_info.key, mint_info.key)?;
     let treasury = treasury_info.as_account_mut::<Treasury>(&ore_api::ID)?;
-    treasury_tokens_info
-        .is_writable()?
-        .as_associated_token_account(&treasury_info.key, &mint_info.key)?;
     system_program.is_program(&system_program::ID)?;
     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(
@@ -50,12 +55,12 @@ pub fn process_withdraw(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramRes
 
     // Transfer ORE to recipient.
     transfer_signed(
-        treasury_info,
-        treasury_tokens_info,
+        stake_info,
+        stake_tokens_info,
         recipient_info,
         token_program,
         amount,
-        &[TREASURY],
+        &[STAKE, &stake.authority.to_bytes()],
     )?;
 
     // Log withdraw.