reward-distribution-merkle-tree.rs

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;
use std::collections::HashMap;
use std::mem::size_of;

/**
 * This code demonstrates a Merkle tree rewards distribution system on Solana
 *
 * The main components of the system are:
 * 1. Updating rewards for a batch of users:
 *    - The `process_update_rewards` function handles the updating of rewards for multiple users in a single transaction.
 *    - It takes a vector of tuples containing user public keys and their newly earned reward amounts.
 *    - The function updates the Merkle tree by adding the new reward amounts to the existing reward amounts for each user.
 *    - It also maintains a `claimed_rewards` HashMap to keep track of whether a user has claimed their rewards or not.
 *
 * 2. Claiming rewards:
 *    - The `process_reward_claim` function handles the processing of reward claims by individual users.
 *    - It verifies the provided Merkle proof against the Merkle root and the user's public key.
 *    - If the proof is valid and the user hasn't claimed their rewards yet, it transfers the reward amount from the
 *      reward account to the user's account using the SPL token program.
 *    - It updates the `claimed_rewards` HashMap to mark the user's rewards as claimed, preventing them from claiming
 *      the same rewards multiple times.
 *
 * This code provides a secure and efficient way to distribute rewards to users on the Solana while preventing
 * duplicate claims and ensuring the integrity of the reward data using a Merkle tree.
 */

#[derive(BorshSerialize, BorshDeserialize, Debug)]
pub struct RewardInstruction {
    pub merkle_root: [u8; 32],
    pub proof: Vec<[u8; 32]>,
    pub amount: u64,
}

#[derive(BorshSerialize, BorshDeserialize, Debug)]
pub struct UpdateRewardsInstruction {
    pub rewards: Vec<(Pubkey, 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_reward_claim(program_id, accounts, instruction_data),
        &spl_token::ID => process_update_rewards(program_id, accounts, instruction_data),
        _ => Err(ProgramError::InvalidInstructionData.into()),
    }
}

fn process_reward_claim(
    program_id: &Pubkey,
    accounts: &[AccountInfo],
    instruction_data: &[u8],
) -> ProgramResult {
    let instruction = RewardInstruction::try_from_slice(instruction_data)?;
    let merkle_root = instruction.merkle_root;
    let proof = instruction.proof;
    let amount = instruction.amount;

    let account_info_iter = &mut accounts.iter();

    let payer = next_account_info(account_info_iter)?;
    let reward_account = next_account_info(account_info_iter)?;
    let token_program = next_account_info(account_info_iter)?;
    let merkle_tree_account = next_account_info(account_info_iter)?;
    let claimed_rewards_account = next_account_info(account_info_iter)?;

    let (reward_pda, _bump_seed) = Pubkey::find_program_address(
        &[b"reward".as_ref(), payer.key.as_ref()],
        program_id,
    );

    if reward_pda != *reward_account.key {
        msg!("Invalid reward PDA");
        return Err(ProgramError::InvalidArgument.into());
    }

    let reward_info = TokenAccount::unpack(&reward_account.data.borrow())?;
    let reward_amount = reward_info.amount;

    if !verify_merkle_proof(&proof, &merkle_root, &payer.key.to_bytes(), reward_amount) {
        msg!("Invalid Merkle proof");
        return Err(ProgramError::InvalidArgument.into());
    }

    let mut claimed_rewards = HashMap::try_from_slice(&claimed_rewards_account.data.borrow())?;

    if claimed_rewards.get(payer.key).copied().unwrap_or(false) {
        msg!("Rewards already claimed");
        return Err(ProgramError::InvalidArgument.into());
    }

    invoke(
        &spl_token::instruction::transfer(
            token_program.key,
            reward_account.key,
            payer.key,
            &[],
            amount,
        )?,
        &[reward_account.clone(), payer.clone(), token_program.clone()],
    )?;

    claimed_rewards.insert(*payer.key, true);
    claimed_rewards.serialize(&mut claimed_rewards_account.data.borrow_mut())?;

    msg!("Reward distributed successfully");

    Ok(())
}

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 merkle_tree_account = next_account_info(account_info_iter)?;
    let claimed_rewards_account = next_account_info(account_info_iter)?;

    if !admin_account.is_signer {
        return Err(ProgramError::MissingRequiredSignature.into());
    }

    let mut merkle_tree = MerkleTree::try_from_slice(&merkle_tree_account.data.borrow())?;
    let mut claimed_rewards = HashMap::try_from_slice(&claimed_rewards_account.data.borrow())?;

    for (user_pubkey, reward_amount) in rewards {
        let user_reward_leaf = UserRewardLeaf {
            user_pubkey,
            reward_amount,
        };
        merkle_tree.update_leaf(&user_reward_leaf)?;

        if !claimed_rewards.contains_key(&user_pubkey) {
            claimed_rewards.insert(user_pubkey, false);
        }
    }

    merkle_tree.serialize(&mut merkle_tree_account.data.borrow_mut())?;
    claimed_rewards.serialize(&mut claimed_rewards_account.data.borrow_mut())?;

    msg!("Rewards updated successfully");

    Ok(())
}

fn verify_merkle_proof(proof: &[[u8; 32]], root: &[u8; 32], leaf: &[u8], amount: u64) -> bool {
    let mut computed_hash = leaf.to_vec();
    computed_hash.extend_from_slice(&amount.to_le_bytes());
    computed_hash = hashv(&[&computed_hash]).to_bytes();

    for proof_element in proof.iter() {
        if computed_hash <= proof_element {
            computed_hash = hashv(&[&computed_hash, proof_element]).to_bytes();
        } else {
            computed_hash = hashv(&[proof_element, &computed_hash]).to_bytes();
        }
    }

    computed_hash == *root
}

#[derive(BorshSerialize, BorshDeserialize, Debug, PartialEq)]
pub struct UserRewardLeaf {
    pub user_pubkey: Pubkey,
    pub reward_amount: u64,
}

pub struct MerkleTree {
    pub root: [u8; 32],
    pub leaves: Vec<UserRewardLeaf>,
}

impl MerkleTree {
    pub fn update_leaf(&mut self, leaf: &UserRewardLeaf) -> ProgramResult {
        match self.leaves.iter_mut().find(|l| l.user_pubkey == leaf.user_pubkey) {
            Some(existing_leaf) => {
                existing_leaf.reward_amount += leaf.reward_amount;
            }
            None => {
                self.leaves.push(leaf.clone());
            }
        }
        self.update_root()?;
        Ok(())
    }

    fn update_root(&mut self) -> ProgramResult {
        // ... (code to update the Merkle root based on the updated leaves)
    }
}