lib.rs

extern crate jito_merkle_tree;
extern crate merkle_distributor;

use std::{path::PathBuf, str::FromStr};

use anchor_lang::{prelude::Pubkey, AccountDeserialize, InstructionData, Key, ToAccountMetas};
use anchor_spl::token;
use clap::{Parser, Subcommand};
use jito_merkle_tree::{
    airdrop_merkle_tree::AirdropMerkleTree,
    utils::{get_claim_status_pda, get_merkle_distributor_pda, get_page_account_pda},
};
use merkle_distributor::state::merkle_distributor::MerkleDistributor;
use solana_program::instruction::Instruction;
use solana_rpc_client::rpc_client::RpcClient;
use solana_sdk::{
    account::Account,
    bs58,
    commitment_config::CommitmentConfig,
    compute_budget::ComputeBudgetInstruction,
    signature::{read_keypair_file, Keypair},
    signer::Signer,
    transaction::Transaction,
};
use spl_associated_token_account::{
    get_associated_token_address, instruction::create_associated_token_account,
};

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
pub struct Args {
    #[clap(subcommand)]
    pub command: Commands,

    /// Airdrop version
    #[clap(long, default_value_t = 0)]
    pub airdrop_version: u64,

    /// SPL Mint address
    #[clap(long)]
    pub mint: Pubkey,

    /// Staking address
    #[clap(long)]
    pub staking: Pubkey,

    /// RPC url
    #[clap(long)]
    pub rpc_url: String,

    /// Program id
    #[clap(long, default_value_t = merkle_distributor::id())]
    pub program_id: Pubkey,

    /// Payer keypair
    #[clap(long)]
    pub keypair_path: PathBuf,

    /// Priority fee
    #[clap(long)]
    pub priority: Option<u64>,
}

// cargo run new-distributor --mint 4dzPmLDFSpuaCcTUoQjw71Bq8u8RWgHJQcLS63Y8ZrZp --staking 5dNR1Lav7T4WNbZUHd2EFMZ71Rfhz3sEFGyR6tgWB5uK --clawback-receiver-token-account 5pT9ijgv2Qpxn4ux4u4crCCJhgAe4w7GoeaCPJKgP4NW --start-vesting-ts 1738340679 --end-vesting-ts 1749350679 --merkle-tree-path /Users/mykyta/development/composable/mantis-staking-program/solana/merkle-tree/test_fixtures/test_csv.csv --clawback-start-ts 1749950679

// Subcommands
#[derive(Subcommand, Debug)]
pub enum Commands {
    /// Claim unlocked tokens
    Claim(ClaimArgs),
    /// Admin claim unlocked tokens
    AdminNewClaim,
    /// Create a new instance of a merkle distributor
    NewDistributor(NewDistributorArgs),
    /// Clawback tokens from merkle distributor
    #[clap(hide = true)]
    Clawback(ClawbackArgs),
    /// Create a Merkle tree, given a CSV of recipients
    CreateMerkleTree(CreateMerkleTreeArgs),
    SetAdmin(SetAdminArgs),
}

// NewClaim and Claim subcommand args
#[derive(Parser, Debug)]
pub struct ClaimArgs {
    /// Merkle distributor path
    #[clap(long)]
    pub merkle_tree_path: PathBuf,
}

// NewDistributor subcommand args
#[derive(Parser, Debug)]
pub struct NewDistributorArgs {
    /// Clawback receiver token account
    #[clap(long)]
    pub clawback_receiver_token_account: Pubkey,

    /// Lockup timestamp start
    #[clap(long)]
    pub start_vesting_ts: i64,

    /// Lockup timestamp end (unix timestamp)
    #[clap(long)]
    pub end_vesting_ts: i64,

    /// Merkle distributor path
    #[clap(long)]
    pub merkle_tree_path: PathBuf,

    /// When to make the clawback period start. Must be at least a day after the end_vesting_ts
    #[clap(long)]
    pub clawback_start_ts: i64,
}

#[derive(Parser, Debug)]
pub struct ClawbackArgs {
    #[clap(long)]
    pub clawback_keypair_path: PathBuf,
}

#[derive(Parser, Debug)]
pub struct CreateMerkleTreeArgs {
    /// CSV path
    #[clap(long)]
    pub csv_path: PathBuf,

    /// Merkle tree out path
    #[clap(long)]
    pub merkle_tree_path: PathBuf,
}

#[derive(Parser, Debug)]
pub struct SetAdminArgs {
    #[clap(long)]
    pub new_admin: Pubkey,
}

fn process_new_claim(args: &Args, claim_args: &ClaimArgs) {
    use dotenv::dotenv;
    dotenv().ok();

    let env_siner_private_key =
        std::env::var("SIGNER_PRIV_KEY").expect("SIGNER_PRIV_KEY must be set.");
    let private_key_bytes = bs58::decode(env_siner_private_key).into_vec().unwrap();

    let keypair = Keypair::from_bytes(&private_key_bytes).unwrap();
    // let keypair = read_keypair_file(&args.keypair_path).expect("Failed reading keypair file");
    let claimant = keypair.pubkey();
    println!("Claiming tokens for user {}...", claimant);

    let merkle_tree = AirdropMerkleTree::new_from_csv(&claim_args.merkle_tree_path)
        .expect("failed to load merkle tree from file");

    let (distributor, _bump) =
        get_merkle_distributor_pda(&args.program_id, &args.mint, args.airdrop_version);

    // Get user's node in claim
    let node = merkle_tree.get_node(&claimant);

    let (claim_status_pda, _bump) = get_claim_status_pda(&args.program_id, &claimant, &distributor);

    let client = RpcClient::new_with_commitment(&args.rpc_url, CommitmentConfig::confirmed());

    let claimant_ata = get_associated_token_address(&claimant, &args.mint);

    let mut ixs = vec![];

    match client.get_account(&claimant_ata) {
        Ok(_) => {}
        Err(e) => {
            // TODO: directly pattern match on error kind
            if e.to_string().contains("AccountNotFound") {
                println!("PDA does not exist. creating.");
                let ix =
                    create_associated_token_account(&claimant, &claimant, &args.mint, &token::ID);
                ixs.push(ix);
            } else {
                panic!("Error fetching PDA: {e}")
            }
        }
    }
    let x = get_associated_token_address(&distributor, &args.mint);
    println!("distributor: {distributor}");
    println!("claim_status: {claim_status_pda}");
    println!("from: {x}");
    println!("to: {claimant_ata}");
    println!("claimant: {claimant}");

    let new_claim_ix = Instruction {
        program_id: args.program_id,
        accounts: merkle_distributor::accounts::NewClaim {
            distributor,
            claim_status: claim_status_pda,
            from: get_associated_token_address(&distributor, &args.mint),
            to: claimant_ata,
            claimant,
            token_program: token::ID,
            system_program: solana_program::system_program::ID,
        }
        .to_account_metas(None),
        data: merkle_distributor::instruction::NewClaim {
            amount_unlocked: node.amount_unlocked(),
            amount_locked: node.amount_locked(),
            proof: node.proof.expect("proof not found"),
        }
        .data(),
    };

    ixs.push(new_claim_ix);

    let blockhash = client.get_latest_blockhash().unwrap();
    let tx =
        Transaction::new_signed_with_payer(&ixs, Some(&claimant.key()), &[&keypair], blockhash);

    let signature = client
        .send_and_confirm_transaction_with_spinner(&tx)
        .unwrap();
    println!("successfully created new claim with signature {signature:#?}");
}

pub fn call_process_admin_claim(
    claimant: Pubkey,
    amount_unlocked: u64,
    amount_locked: u64,
    page_index: u8,
) -> (bool, String) {
    use dotenv::dotenv;
    dotenv().ok();
    let mint = std::env::var("WHITELIST_MINT").expect("WHITELIST_MINT must be set.");
    let vesting_program =
        std::env::var("VESTING_PROGRAM_ID").expect("VESTING_PROGRAM_ID must be set.");
    let solana_rpc_url = std::env::var("SOLANA_RPC_URL").expect("SOLANA_RPC_URL must be set.");
    let args: Args = Args {
        command: Commands::AdminNewClaim,
        airdrop_version: 0,
        mint: Pubkey::from_str(&mint).unwrap(),
        staking: Pubkey::from_str(&vesting_program).unwrap(),
        rpc_url: solana_rpc_url,
        program_id: Pubkey::from_str(&vesting_program).unwrap(),
        keypair_path: PathBuf::from(""),
        priority: None,
    };
    return process_admin_claim(&args, claimant, amount_unlocked, amount_locked, page_index);
}

fn process_admin_claim(
    args: &Args,
    claimant: Pubkey,
    amount_unlocked: u64,
    amount_locked: u64,
    page_index: u8,
) -> (bool, String) {
    use dotenv::dotenv;
    dotenv().ok();

    let env_siner_private_key =
        std::env::var("SIGNER_PRIV_KEY").expect("SIGNER_PRIV_KEY must be set.");
    let admin_private_key_bytes = bs58::decode(env_siner_private_key).into_vec().unwrap();

    let keypair = Keypair::from_bytes(&admin_private_key_bytes).unwrap();
    println!("Setting airdrop/vesting schedule for user {}...", claimant);

    let (distributor, _bump) =
        get_merkle_distributor_pda(&args.program_id, &args.mint, args.airdrop_version);

    let (page_account_pda, _bump) = get_page_account_pda(&args.program_id, 1);

    let client = RpcClient::new_with_commitment(&args.rpc_url, CommitmentConfig::confirmed());

    // let claimant_ata = get_associated_token_address(&claimant, &args.mint);

    let mut ixs = vec![];

    // match client.get_account(&claimant_ata) {
    //     Ok(_) => {}
    //     Err(e) => {
    //         // TODO: directly pattern match on error kind
    //         if e.to_string().contains("AccountNotFound") {
    //             println!("PDA does not exist. creating.");
    //             let ix =
    //                 create_associated_token_account(&claimant, &claimant, &args.mint, &token::ID);
    //             ixs.push(ix);
    //         } else {
    //             panic!("Error fetching PDA: {e}")
    //         }
    //     }
    // }
    let x = get_associated_token_address(&distributor, &args.mint);
    println!("distributor: {distributor}");
    println!("claim_status: {page_account_pda}");
    println!("from: {x}");
    // println!("to: {claimant_ata}");
    println!("claimant: {claimant}");

    let new_claim_ix = Instruction {
        program_id: args.program_id,
        accounts: merkle_distributor::accounts::AdminNewClaim {
            distributor,
            admin: keypair.pubkey(),
            page_account_state: page_account_pda,
            token_program: token::ID,
            system_program: solana_program::system_program::ID,
        }
        .to_account_metas(None),
        data: merkle_distributor::instruction::AdminNewClaim {
            amount_unlocked: amount_unlocked,
            amount_locked: amount_locked,
            page_index: page_index,
            claimant: claimant,
        }
        .data(),
    };

    ixs.push(new_claim_ix);

    let blockhash = client.get_latest_blockhash().unwrap();
    let tx = Transaction::new_signed_with_payer(
        &ixs,
        Some(&keypair.pubkey().key()),
        &[&keypair],
        blockhash,
    );

    let signature = client.send_and_confirm_transaction_with_spinner(&tx);
    let msg_signature = format!("{:?}", signature);
    println!("{}", msg_signature);
    if signature.is_err() {
        return (false, msg_signature);
    }
    return (true, msg_signature);
}

fn process_claim(args: &Args, claim_args: &ClaimArgs) {
    use dotenv::dotenv;
    dotenv().ok();

    let env_siner_private_key =
        std::env::var("SIGNER_PRIV_KEY").expect("SIGNER_PRIV_KEY must be set.");
    let private_key_bytes = bs58::decode(env_siner_private_key).into_vec().unwrap();

    let keypair = Keypair::from_bytes(&private_key_bytes).unwrap();
    println!("This is pubkey {}", keypair.pubkey().to_string());
    // let keypair = read_keypair_file(&args.keypair_path).expect("Failed reading keypair file");
    let claimant = keypair.pubkey();

    let priority_fee = args.priority.unwrap_or(0);

    let (distributor, bump) =
        get_merkle_distributor_pda(&args.program_id, &args.mint, args.airdrop_version);
    println!("distributor pubkey {}", distributor);

    let (claim_status_pda, _bump) = get_claim_status_pda(&args.program_id, &claimant, &distributor);
    println!("claim pda: {claim_status_pda}, bump: {bump}");

    let client = RpcClient::new_with_commitment(&args.rpc_url, CommitmentConfig::confirmed());

    match client.get_account(&claim_status_pda) {
        Ok(_) => {}
        Err(e) => {
            // TODO: match on the error kind
            if e.to_string().contains("AccountNotFound") {
                println!("PDA does not exist. creating.");
                process_new_claim(args, claim_args);
            } else {
                panic!("error getting PDA: {e}")
            }
        }
    }

    let claimant_ata = get_associated_token_address(&claimant, &args.mint);

    let mut ixs = vec![];
    println!("distributed pubkey: {distributor}");
    println!("claim_status: {claim_status_pda}");
    let x = get_associated_token_address(&distributor, &args.mint);
    println!("from: {x}");
    println!("to: {claimant_ata}");
    println!("claimant: {claimant}");
    // return;

    let claim_ix = Instruction {
        program_id: args.program_id,
        accounts: merkle_distributor::accounts::ClaimLocked {
            distributor,
            claim_status: claim_status_pda,
            from: get_associated_token_address(&distributor, &args.mint),
            to: claimant_ata,
            claimant,
            token_program: token::ID,
        }
        .to_account_metas(None),
        data: merkle_distributor::instruction::ClaimLocked {}.data(),
    };
    ixs.push(claim_ix);

    if priority_fee > 0 {
        let instruction = ComputeBudgetInstruction::set_compute_unit_price(priority_fee);
        ixs.push(instruction);
        println!(
            "Added priority fee instruction of {} microlamports",
            priority_fee
        );
    } else {
        println!("No priority fee added. Add one with --priority <microlamports u64>");
    }

    let blockhash = client.get_latest_blockhash().unwrap();
    let tx =
        Transaction::new_signed_with_payer(&ixs, Some(&claimant.key()), &[&keypair], blockhash);

    let signature = client
        .send_and_confirm_transaction_with_spinner(&tx)
        .unwrap();
    println!("successfully claimed tokens with signature {signature:#?}",);
}

fn check_distributor_onchain_matches(
    account: &Account,
    merkle_tree: &AirdropMerkleTree,
    new_distributor_args: &NewDistributorArgs,
    pubkey: Pubkey,
) -> Result<(), &'static str> {
    if let Ok(distributor) = MerkleDistributor::try_deserialize(&mut account.data.as_slice()) {
        if distributor.root != merkle_tree.merkle_root {
            return Err("root mismatch");
        }
        if distributor.max_total_claim != merkle_tree.max_total_claim {
            return Err("max_total_claim mismatch");
        }
        if distributor.max_num_nodes != merkle_tree.max_num_nodes {
            return Err("max_num_nodes mismatch");
        }

        if distributor.start_ts != new_distributor_args.start_vesting_ts {
            return Err("start_ts mismatch");
        }
        if distributor.end_ts != new_distributor_args.end_vesting_ts {
            return Err("end_ts mismatch");
        }
        if distributor.clawback_start_ts != new_distributor_args.clawback_start_ts {
            return Err("clawback_start_ts mismatch");
        }
        if distributor.clawback_receiver != new_distributor_args.clawback_receiver_token_account {
            return Err("clawback_receiver mismatch");
        }
        if distributor.admin != pubkey {
            return Err("admin mismatch");
        }
    }
    Ok(())
}

fn process_new_distributor(args: &Args, new_distributor_args: &NewDistributorArgs) {
    use dotenv::dotenv;
    dotenv().ok();

    let env_siner_private_key =
        std::env::var("SIGNER_PRIV_KEY").expect("SIGNER_PRIV_KEY must be set.");
    let private_key_bytes = bs58::decode(env_siner_private_key).into_vec().unwrap();

    let keypair = Keypair::from_bytes(&private_key_bytes).unwrap();
    println!("This is pubkey {}", keypair.pubkey().to_string());

    let ret_back = get_associated_token_address(&keypair.pubkey(), &args.mint);

    // let keypair = read_keypair_file(&args.keypair_path).expect("Failed reading keypair file");

    let client = RpcClient::new_with_commitment(&args.rpc_url, CommitmentConfig::finalized());

    let merkle_tree = AirdropMerkleTree::new_from_csv(&new_distributor_args.merkle_tree_path)
        .expect("failed to read");
    let (distributor_pubkey, _bump) =
        get_merkle_distributor_pda(&args.program_id, &args.mint, args.airdrop_version);
    let token_vault = get_associated_token_address(&distributor_pubkey, &args.mint);

    if let Some(account) = client
        .get_account_with_commitment(&distributor_pubkey, CommitmentConfig::confirmed())
        .unwrap()
        .value
    {
        println!("merkle distributor account exists, checking parameters...");
        check_distributor_onchain_matches(
            &account,
            &merkle_tree,
            new_distributor_args,
            keypair.pubkey(),
        ).expect("merkle root on-chain does not match provided arguments! Confirm admin and clawback parameters to avoid loss of funds!");
    }

    println!("creating new distributor with args: {new_distributor_args:#?}");

    let new_distributor_ix = Instruction {
        program_id: args.program_id,
        accounts: merkle_distributor::accounts::NewDistributor {
            clawback_receiver: ret_back,
            mint: args.mint,
            token_vault,
            distributor: distributor_pubkey,
            system_program: solana_program::system_program::id(),
            associated_token_program: spl_associated_token_account::ID,
            token_program: token::ID,
            admin: keypair.pubkey(),
        }
        .to_account_metas(None),
        data: merkle_distributor::instruction::NewDistributor {
            version: args.airdrop_version,
            root: merkle_tree.merkle_root,
            max_total_claim: merkle_tree.max_total_claim,
            max_num_nodes: merkle_tree.max_num_nodes,
            start_vesting_ts: new_distributor_args.start_vesting_ts,
            end_vesting_ts: new_distributor_args.end_vesting_ts,
            clawback_start_ts: new_distributor_args.clawback_start_ts,
            staking_contract_address: None,
        }
        .data(),
    };

    let blockhash = client.get_latest_blockhash().unwrap();
    let tx = Transaction::new_signed_with_payer(
        &[new_distributor_ix],
        Some(&keypair.pubkey()),
        &[&keypair],
        blockhash,
    );

    // See comments on new_distributor instruction inside the program to ensure this transaction
    // didn't get frontrun.
    // If this fails, make sure to run it again.
    match client.send_and_confirm_transaction_with_spinner(&tx) {
        Ok(_) => {}
        Err(e) => {
            println!("Failed to create MerkleDistributor: {:?}", e);

            // double check someone didn't frontrun this transaction with a malicious merkle root
            if let Some(account) = client
                .get_account_with_commitment(&distributor_pubkey, CommitmentConfig::processed())
                .unwrap()
                .value
            {
                check_distributor_onchain_matches(
                    &account,
                    &merkle_tree,
                    new_distributor_args,
                    keypair.pubkey(),
                ).expect("merkle root on-chain does not match provided arguments! Confirm admin and clawback parameters to avoid loss of funds!");
            }
        }
    }
}

fn process_clawback(args: &Args, clawback_args: &ClawbackArgs) {
    let payer_keypair = read_keypair_file(&args.keypair_path).expect("Failed reading keypair file");
    let clawback_keypair = read_keypair_file(&clawback_args.clawback_keypair_path)
        .expect("Failed reading keypair file");

    let clawback_ata = get_associated_token_address(&clawback_keypair.pubkey(), &args.mint);

    let client = RpcClient::new_with_commitment(&args.rpc_url, CommitmentConfig::confirmed());

    let (distributor, _bump) =
        get_merkle_distributor_pda(&args.program_id, &args.mint, args.airdrop_version);

    let from = get_associated_token_address(&distributor, &args.mint);
    println!("from: {from}");

    let clawback_ix = Instruction {
        program_id: args.program_id,
        accounts: merkle_distributor::accounts::Clawback {
            distributor,
            from,
            to: clawback_ata,
            claimant: clawback_keypair.pubkey(),
            system_program: solana_program::system_program::ID,
            token_program: token::ID,
        }
        .to_account_metas(None),
        data: merkle_distributor::instruction::Clawback {}.data(),
    };

    let tx = Transaction::new_signed_with_payer(
        &[clawback_ix],
        Some(&payer_keypair.pubkey()),
        &[&payer_keypair, &clawback_keypair],
        client.get_latest_blockhash().unwrap(),
    );

    let signature = client
        .send_and_confirm_transaction_with_spinner(&tx)
        .unwrap();

    println!("Successfully clawed back funds! signature: {signature:#?}");
}

fn process_create_merkle_tree(merkle_tree_args: &CreateMerkleTreeArgs) {
    let merkle_tree = AirdropMerkleTree::new_from_csv(&merkle_tree_args.csv_path).unwrap();
    merkle_tree.write_to_file(&merkle_tree_args.merkle_tree_path);
}

fn process_set_admin(args: &Args, set_admin_args: &SetAdminArgs) {
    let keypair = read_keypair_file(&args.keypair_path).expect("Failed reading keypair file");

    let client = RpcClient::new_with_commitment(&args.rpc_url, CommitmentConfig::confirmed());

    let (distributor, _bump) =
        get_merkle_distributor_pda(&args.program_id, &args.mint, args.airdrop_version);

    let set_admin_ix = Instruction {
        program_id: args.program_id,
        accounts: merkle_distributor::accounts::SetAdmin {
            distributor,
            admin: keypair.pubkey(),
            new_admin: set_admin_args.new_admin,
        }
        .to_account_metas(None),
        data: merkle_distributor::instruction::SetAdmin {}.data(),
    };

    let tx = Transaction::new_signed_with_payer(
        &[set_admin_ix],
        Some(&keypair.pubkey()),
        &[&keypair],
        client.get_latest_blockhash().unwrap(),
    );

    let signature = client
        .send_and_confirm_transaction_with_spinner(&tx)
        .unwrap();

    println!("Successfully set admin! signature: {signature:#?}");
}