Skip to content

Instantly share code, notes, and snippets.

@lanvidr
Last active December 30, 2024 01:51
Show Gist options
  • Save lanvidr/4595f7b02236ffb2a3fb3ce9347ca044 to your computer and use it in GitHub Desktop.
Save lanvidr/4595f7b02236ffb2a3fb3ce9347ca044 to your computer and use it in GitHub Desktop.
use borsh::{BorshDeserialize, BorshSerialize};
use solana_program::{
account_info::{next_account_info, AccountInfo},
entrypoint,
entrypoint::ProgramResult,
msg,
program::invoke,
program_error::ProgramError,
pubkey::Pubkey,
rent::Rent,
system_instruction,
sysvar::Sysvar,
};
use spl_token::state::{Account as TokenAccount, Mint};
use std::collections::HashMap;
use std::mem::size_of;
use zk_token_sdk::instruction::{compress_token_account, transfer_compressed};
use zk_token_sdk::state::{CompressedTokenAccount, CompressedTokenAccountInfo};
/**
* In this example, we use ZK compressed token accounts to store and distribute rewards. The program defines two instructions: UpdateRewardsInstruction and ClaimRewardInstruction.
* The process_update_rewards function is responsible for updating the rewards for a batch of users. It takes the UpdateRewardsInstruction data, which contains a vector of tuples representing user
* public keys and their corresponding reward amounts. The function performs the following steps:
*
* Retrieves the necessary accounts (admin account, mint account, compressed token account, and token program).
* Verifies that the admin account is a signer.
* Unpacks the mint information and the compressed token account information.
* Iterates over the rewards and updates the compressed token account balances for each user. If a user's compressed account doesn't exist, it creates a new one.
* Packs the updated compressed token account information back into the account data.
*
* The process_claim_reward function allows a user to claim their rewards. It takes the ClaimRewardInstruction data, which contains the amount to be claimed. The function performs the following steps:
*
* 1. Retrieves the necessary accounts (user account, mint account, compressed token account, user token account, and token program).
* 2. Unpacks the mint information and the compressed token account information.
* 3. Finds the user's compressed token account address based on their public key and the mint authority.
* 4. Checks if the user has sufficient rewards to claim the requested amount.
* 5. Invokes the transfer_compressed instruction to transfer the reward tokens from the compressed account to the user's token account.
* 6. Updates the user's compressed account balance by subtracting the claimed amount.
* 7. Packs the updated compressed token account information back into the account data.
*
* To prevent users from claiming rewards multiple times, the program maintains a separate compressed token account for each user. When a user claims their rewards, the claimed amount is deducted from
* their compressed account balance. If a user attempts to claim more rewards than they have available, the program returns an error.
**/
#[derive(BorshSerialize, BorshDeserialize, Debug)]
pub struct UpdateRewardsInstruction {
pub rewards: Vec<(Pubkey, u64)>,
}
#[derive(BorshSerialize, BorshDeserialize, Debug)]
pub struct ClaimRewardInstruction {
pub amount: u64,
}
entrypoint!(process_instruction);
pub fn process_instruction(
program_id: &Pubkey,
accounts: &[AccountInfo],
instruction_data: &[u8],
) -> ProgramResult {
let account_info_iter = &mut accounts.iter();
let instruction_account = next_account_info(account_info_iter)?;
match instruction_account.key {
&system_program::ID => process_update_rewards(program_id, accounts, instruction_data),
&spl_token::ID => process_claim_reward(program_id, accounts, instruction_data),
_ => Err(ProgramError::InvalidInstructionData.into()),
}
}
fn process_update_rewards(
program_id: &Pubkey,
accounts: &[AccountInfo],
instruction_data: &[u8],
) -> ProgramResult {
let instruction = UpdateRewardsInstruction::try_from_slice(instruction_data)?;
let rewards = instruction.rewards;
let account_info_iter = &mut accounts.iter();
let admin_account = next_account_info(account_info_iter)?;
let mint_account = next_account_info(account_info_iter)?;
let compressed_token_account = next_account_info(account_info_iter)?;
let token_program = next_account_info(account_info_iter)?;
if !admin_account.is_signer {
return Err(ProgramError::MissingRequiredSignature.into());
}
let mint_info = Mint::unpack(&mint_account.data.borrow())?;
let mut compressed_account_info =
CompressedTokenAccountInfo::unpack(&compressed_token_account.data.borrow())?;
for (user_pubkey, reward_amount) in rewards {
let (user_compressed_account, _) = Pubkey::find_program_address(
&[
b"compressed_account",
user_pubkey.as_ref(),
mint_info.authority.as_ref(),
],
program_id,
);
if compressed_account_info
.compressed_token_accounts
.contains_key(&user_compressed_account)
{
let mut user_compressed_account_data = compressed_account_info
.compressed_token_accounts
.get_mut(&user_compressed_account)
.unwrap();
user_compressed_account_data.amount += reward_amount;
} else {
compressed_account_info
.compressed_token_accounts
.insert(user_compressed_account, CompressedTokenAccount {
mint: *mint_account.key,
owner: user_pubkey,
amount: reward_amount,
});
}
}
CompressedTokenAccountInfo::pack(
compressed_account_info,
&mut compressed_token_account.data.borrow_mut(),
)?;
msg!("Rewards updated successfully");
Ok(())
}
fn process_claim_reward(
program_id: &Pubkey,
accounts: &[AccountInfo],
instruction_data: &[u8],
) -> ProgramResult {
let instruction = ClaimRewardInstruction::try_from_slice(instruction_data)?;
let amount = instruction.amount;
let account_info_iter = &mut accounts.iter();
let user_account = next_account_info(account_info_iter)?;
let mint_account = next_account_info(account_info_iter)?;
let compressed_token_account = next_account_info(account_info_iter)?;
let user_token_account = next_account_info(account_info_iter)?;
let token_program = next_account_info(account_info_iter)?;
let mint_info = Mint::unpack(&mint_account.data.borrow())?;
let mut compressed_account_info =
CompressedTokenAccountInfo::unpack(&compressed_token_account.data.borrow())?;
let (user_compressed_account, _) = Pubkey::find_program_address(
&[
b"compressed_account",
user_account.key.as_ref(),
mint_info.authority.as_ref(),
],
program_id,
);
if let Some(user_compressed_account_data) = compressed_account_info
.compressed_token_accounts
.get_mut(&user_compressed_account)
{
if user_compressed_account_data.amount < amount {
return Err(ProgramError::InsufficientFunds.into());
}
invoke(
&transfer_compressed(
token_program.key,
compressed_token_account.key,
user_compressed_account,
user_token_account.key,
user_account.key,
&[],
amount,
)?,
&[
user_account.clone(),
compressed_token_account.clone(),
user_token_account.clone(),
token_program.clone(),
],
)?;
user_compressed_account_data.amount -= amount;
CompressedTokenAccountInfo::pack(
compressed_account_info,
&mut compressed_token_account.data.borrow_mut(),
)?;
msg!("Reward claimed successfully");
Ok(())
} else {
msg!("User does not have any rewards to claim");
Err(ProgramError::InvalidArgument.into())
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment