hacker-challenge/src/solana.rs

#![allow(dead_code)]
use crate::{grpc::challenge::Keys, persistence::KeysFile};
use solana_client::nonblocking::rpc_client::RpcClient;
use solana_program::program_pack::Pack;
use solana_sdk::{
    pubkey::Pubkey, signature::keypair::Keypair, signer::Signer, system_instruction,
    transaction::Transaction,
};
use spl_associated_token_account::{
    get_associated_token_address, instruction::create_associated_token_account,
};
use spl_token::{
    instruction::{initialize_mint, mint_to},
    state::Mint,
};
use tokio::time::{sleep, Duration};

#[cfg(feature = "test")]
const RPC_URL: &str = "https://api.devnet.solana.com";
#[cfg(not(feature = "test"))]
const RPC_URL: &str = "https://api.mainnet-beta.solana.com";

pub struct SolClient {
    client: RpcClient,
    keypair: Keypair,
    token: Pubkey,
}

impl SolClient {
    pub async fn create_new_token() -> Self {
        let client = RpcClient::new(RPC_URL.to_string());
        let keypair = Keypair::new();
        let token = create_token(&client, &keypair).await;
        Self { client, keypair, token }
    }

    pub fn get_keys(&self) -> Keys {
        Keys { keypair: self.get_keypair_bytes(), token_address: self.get_token_address() }
    }

    pub fn get_keys_file(&self) -> KeysFile {
        self.get_keys().into()
    }

    pub async fn mint(&self, recipient: &str) -> Result<String, Box<dyn std::error::Error>> {
        use std::str::FromStr;
        let recipient = solana_sdk::pubkey::Pubkey::from_str(recipient)?;
        let associated_token_address = self.create_token_account(&recipient).await?;
        let mint_to_instruction = mint_to(
            &spl_token::id(),
            &self.token,
            &associated_token_address,
            &self.keypair.pubkey(),
            &[],
            1_000_000_000,
        )?;

        let transaction = Transaction::new_signed_with_payer(
            &[mint_to_instruction],
            Some(&self.keypair.pubkey()),
            &[&self.keypair],
            self.client.get_latest_blockhash().await?,
        );
        let signature = self.client.send_and_confirm_transaction(&transaction).await?;
        Ok(signature.to_string())
    }

    async fn create_token_account(
        &self,
        recipient: &Pubkey,
    ) -> Result<Pubkey, Box<dyn std::error::Error>> {
        let address = get_associated_token_address(recipient, &self.token);
        if self.client.get_account(&address).await.is_err() {
            let create_token_account_instruction = create_associated_token_account(
                &self.keypair.pubkey(),
                recipient,
                &self.token,
                &spl_token::id(),
            );
            let recent_blockhash = self.client.get_latest_blockhash().await?;
            let tx = Transaction::new_signed_with_payer(
                &[create_token_account_instruction],
                Some(&self.keypair.pubkey()),
                &[&self.keypair],
                recent_blockhash,
            );
            self.client.send_and_confirm_transaction(&tx).await?;
        }
        Ok(address)
    }

    pub fn get_wallet_pubkey(&self) -> String {
        // Return the base58 string representation of the public key
        self.keypair.pubkey().to_string()
    }

    pub fn get_token_address(&self) -> String {
        self.token.to_string()
    }

    pub fn get_keypair_bytes(&self) -> Vec<u8> {
        self.keypair.to_bytes().to_vec()
    }
}

impl TryFrom<Keys> for SolClient {
    type Error = String;

    fn try_from(keys: Keys) -> Result<Self, Self::Error> {
        use std::str::FromStr;
        let keypair = match Keypair::from_bytes(&keys.keypair) {
            Ok(k) => k,
            Err(_) => return Err("Could not parse keypair.".into()),
        };
        let token = Pubkey::from_str(&keys.token_address)
            .map_err(|_| "Could not parse wallet address.".to_string())?;
        Ok(Self { client: RpcClient::new(RPC_URL.to_string()), keypair, token })
    }
}

impl TryFrom<KeysFile> for SolClient {
    type Error = String;

    fn try_from(keys_file: KeysFile) -> Result<Self, Self::Error> {
        let keys: Keys = keys_file.into();
        Self::try_from(keys)
    }
}

async fn wait_for_sol(client: &RpcClient, pubkey: &Pubkey) {
    println!("Waiting to receive 0.01 SOL in address {pubkey}");
    loop {
        sleep(Duration::from_secs(30)).await;
        match client.get_balance(pubkey).await {
            Ok(balance) => {
                println!("Got {balance} lamports.");
                if balance > 10_000_000 {
                    return;
                }
            }
            Err(e) => println!("Could not get balance: {e:?}"),
        }
    }
}

async fn create_token(client: &RpcClient, keypair: &Keypair) -> Pubkey {
    wait_for_sol(client, &keypair.pubkey()).await;

    let mint_keypair = Keypair::new();
    let payer = Keypair::from_base58_string(&keypair.to_base58_string());
    let mint_rent = client.get_minimum_balance_for_rent_exemption(Mint::LEN).await.unwrap();

    let create_mint_account_ix = system_instruction::create_account(
        &payer.pubkey(),
        &mint_keypair.pubkey(),
        mint_rent,
        Mint::LEN as u64,
        &spl_token::id(),
    );

    let init_mint_ix =
        initialize_mint(&spl_token::id(), &mint_keypair.pubkey(), &payer.pubkey(), None, 9)
            .unwrap();

    let recent_blockhash = client.get_latest_blockhash().await.unwrap();
    let tx = Transaction::new_signed_with_payer(
        &[create_mint_account_ix, init_mint_ix],
        Some(&payer.pubkey()),
        &[&payer, &mint_keypair],
        recent_blockhash,
    );

    let signature = client.send_and_confirm_transaction(&tx).await.unwrap();

    println!("Mint created: {}", mint_keypair.pubkey());
    println!("Transaction signature: {}", signature);

    mint_keypair.pubkey()
}