test: Allow replay of live transaction in program test (#911)

This commit is contained in:
Christian Kamm 2024-03-11 15:08:58 +01:00 committed by GitHub
parent 2cd4376466
commit 4fcaf09c09
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 246 additions and 6 deletions

3
Cargo.lock generated
View File

@ -3385,11 +3385,13 @@ version = "0.24.0"
dependencies = [
"anchor-lang",
"anchor-spl",
"anyhow",
"arrayref",
"async-trait",
"base64 0.13.1",
"bincode",
"borsh 0.10.3",
"bs58 0.5.0",
"bytemuck",
"default-env",
"derivative",
@ -3405,6 +3407,7 @@ dependencies = [
"rand 0.8.5",
"regex",
"serde",
"serde_json",
"serum_dex 0.5.10 (git+https://github.com/openbook-dex/program.git)",
"solana-address-lookup-table-program",
"solana-logger",

View File

@ -67,3 +67,6 @@ rand = "0.8.4"
lazy_static = "1.4.0"
num = "0.4.0"
regex = "1"
serde_json = "1"
bs58 = "0.5"
anyhow = "1"

View File

@ -13,6 +13,10 @@ pub trait AccountReader {
fn data(&self) -> &[u8];
}
pub trait AccountDataWriter {
fn data_as_mut_slice(&mut self) -> &mut [u8];
}
/// Like AccountReader, but can also get the account pubkey
pub trait KeyedAccountReader: AccountReader {
fn key(&self) -> &Pubkey;
@ -99,6 +103,12 @@ impl<'info, 'a> KeyedAccountReader for AccountInfoRefMut<'info, 'a> {
}
}
impl<'info, 'a> AccountDataWriter for AccountInfoRefMut<'info, 'a> {
fn data_as_mut_slice(&mut self) -> &mut [u8] {
&mut self.data
}
}
#[cfg(feature = "solana-sdk")]
impl<T: solana_sdk::account::ReadableAccount> AccountReader for T {
fn owner(&self) -> &Pubkey {
@ -110,6 +120,13 @@ impl<T: solana_sdk::account::ReadableAccount> AccountReader for T {
}
}
#[cfg(feature = "solana-sdk")]
impl<T: solana_sdk::account::WritableAccount> AccountDataWriter for T {
fn data_as_mut_slice(&mut self) -> &mut [u8] {
self.data_as_mut_slice()
}
}
#[cfg(feature = "solana-sdk")]
#[derive(Clone)]
pub struct KeyedAccount {
@ -232,28 +249,29 @@ impl<A: AccountReader> LoadZeroCopy for A {
}
}
impl<'info, 'a> LoadMutZeroCopy for AccountInfoRefMut<'info, 'a> {
impl<A: AccountReader + AccountDataWriter> LoadMutZeroCopy for A {
fn load_mut<T: ZeroCopy + Owner>(&mut self) -> Result<&mut T> {
if self.owner != &T::owner() {
if self.owner() != &T::owner() {
return Err(ErrorCode::AccountOwnedByWrongProgram.into());
}
if self.data.len() < 8 {
let data = self.data_as_mut_slice();
if data.len() < 8 {
return Err(ErrorCode::AccountDiscriminatorNotFound.into());
}
let disc_bytes = array_ref![self.data, 0, 8];
let disc_bytes = array_ref![data, 0, 8];
if disc_bytes != &T::discriminator() {
return Err(ErrorCode::AccountDiscriminatorMismatch.into());
}
Ok(bytemuck::from_bytes_mut(
&mut self.data[8..mem::size_of::<T>() + 8],
&mut data[8..mem::size_of::<T>() + 8],
))
}
fn load_mut_fully_unchecked<T: ZeroCopy + Owner>(&mut self) -> Result<&mut T> {
Ok(bytemuck::from_bytes_mut(
&mut self.data[8..mem::size_of::<T>() + 8],
&mut self.data_as_mut_slice()[8..mem::size_of::<T>() + 8],
))
}
}

View File

@ -36,6 +36,7 @@ mod test_perp_settle;
mod test_perp_settle_fees;
mod test_position_lifetime;
mod test_reduce_only;
mod test_replay;
mod test_serum;
mod test_stale_oracles;
mod test_token_conditional_swap;

View File

@ -0,0 +1,215 @@
use super::*;
use solana_address_lookup_table_program::state::AddressLookupTable;
use solana_program::program_pack::Pack;
use solana_sdk::account::{Account, ReadableAccount};
use solana_sdk::instruction::AccountMeta;
use solana_sdk::instruction::Instruction;
use solana_sdk::message::v0::LoadedAddresses;
use solana_sdk::message::SanitizedMessage;
use solana_sdk::message::SanitizedVersionedMessage;
use solana_sdk::message::SimpleAddressLoader;
use solana_sdk::transaction::VersionedTransaction;
use anyhow::Context;
use mango_v4::accounts_zerocopy::LoadMutZeroCopy;
use std::str::FromStr;
fn read_json_file<P: AsRef<std::path::Path>>(path: P) -> anyhow::Result<serde_json::Value> {
let file_contents = std::fs::read_to_string(path)?;
let json: serde_json::Value = serde_json::from_str(&file_contents)?;
Ok(json)
}
fn account_from_snapshot(snapshot_path: &str, pk: Pubkey) -> anyhow::Result<Account> {
let file_path = format!("{}/{}.json", snapshot_path, pk);
let json = read_json_file(&file_path).with_context(|| format!("reading {file_path}"))?;
let account = json.get("account").unwrap();
let data = base64::decode(
account.get("data").unwrap().as_array().unwrap()[0]
.as_str()
.unwrap(),
)
.unwrap();
let owner = Pubkey::from_str(account.get("owner").unwrap().as_str().unwrap()).unwrap();
let mut account = Account::new(u64::MAX, data.len(), &owner);
account.data = data;
Ok(account)
}
fn find_tx(block_file: &str, txsig: &str) -> Option<(u64, i64, Vec<u8>)> {
let txsig = bs58::decode(txsig).into_vec().unwrap();
let block = read_json_file(block_file).unwrap();
let slot = block.get("parentSlot").unwrap().as_u64().unwrap();
let time = block.get("blockTime").unwrap().as_i64().unwrap();
let txs = block.get("transactions").unwrap().as_array().unwrap();
for tx_obj in txs {
let tx_bytes = base64::decode(
tx_obj.get("transaction").unwrap().as_array().unwrap()[0]
.as_str()
.unwrap(),
)
.unwrap();
let sig = &tx_bytes[1..65];
if sig == txsig {
return Some((slot, time, tx_bytes.to_vec()));
}
}
None
}
#[tokio::test]
async fn test_replay() -> anyhow::Result<()> {
// Path to a directory generated with cli save-snapshot, containing <pubkey>.json files
let snapshot_path = &"path/to/directory";
// Path to the block data, retrieved with `solana block 252979760 --output json`
let block_file = &"path/to/block";
// TX signature in the block that should be looked at
let txsig = &"";
// 0-based index of instuction in the tx to try replaying
let ix_index = 3;
if txsig.is_empty() {
return Ok(());
}
let mut test_builder = TestContextBuilder::new();
test_builder.test().set_compute_max_units(400_000);
let context = test_builder.start_default().await;
let solana = &context.solana.clone();
let signer = context.users[0].key;
let known_accounts = [
"ComputeBudget111111111111111111111111111111",
"ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL",
"4MangoMjqJ2firMokCjjGgoK8d4MXcrgL7XJaL3w6fVg",
"11111111111111111111111111111111",
"TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA",
]
.iter()
.map(|s| Pubkey::from_str(s).unwrap())
.collect_vec();
// Load block, find tx
let (slot, time, tx_bytes) = find_tx(block_file, txsig).unwrap();
let tx: VersionedTransaction = bincode::deserialize(&tx_bytes).unwrap();
// Lookup ALTs so we can decompile
let loaded_addresses: LoadedAddresses = tx
.message
.address_table_lookups()
.unwrap_or_default()
.iter()
.map(|alt_lookup| {
let alt_account = account_from_snapshot(snapshot_path, alt_lookup.account_key).unwrap();
let alt = AddressLookupTable::deserialize(&alt_account.data()).unwrap();
LoadedAddresses {
readonly: alt_lookup
.readonly_indexes
.iter()
.map(|i| alt.addresses[*i as usize])
.collect_vec(),
writable: alt_lookup
.writable_indexes
.iter()
.map(|i| alt.addresses[*i as usize])
.collect_vec(),
}
})
.collect();
let alt_loader = SimpleAddressLoader::Enabled(loaded_addresses);
// decompile instructions, looking up alts at the same time
let sv_message = SanitizedVersionedMessage::try_from(tx.message).unwrap();
let s_message = SanitizedMessage::try_new(sv_message, alt_loader).unwrap();
let bix = &s_message.decompile_instructions()[ix_index];
let ix = Instruction {
program_id: *bix.program_id,
accounts: bix
.accounts
.iter()
.map(|m| AccountMeta {
pubkey: *m.pubkey,
is_writable: m.is_writable,
is_signer: m.is_signer,
})
.collect(),
data: bix.data.to_vec(),
};
// since we can't retain the original signer/blockhash, replace it
let mut replaced_signers = vec![];
let mut replaced_ix = ix.clone();
for meta in &mut replaced_ix.accounts {
if meta.is_signer {
replaced_signers.push(meta.pubkey);
meta.pubkey = signer.pubkey();
}
}
// Load all accounts, reporting missing ones, add found to context
let mut missing_accounts = vec![];
for pubkey in replaced_ix.accounts.iter().map(|m| m.pubkey) {
if known_accounts.contains(&pubkey) || pubkey == signer.pubkey() {
continue;
}
let mut account = match account_from_snapshot(snapshot_path, pubkey) {
Ok(a) => a,
Err(e) => {
println!("error reading account from snapshot: {pubkey}, error {e:?}");
missing_accounts.push(pubkey);
continue;
}
};
// Override where the previous signer was an owner
if replaced_signers.contains(&account.owner) {
account.owner = signer.pubkey();
}
// Override mango account owners or delegates
if let Ok(mut ma) = account.load_mut::<MangoAccountFixed>() {
if replaced_signers.contains(&ma.owner) {
ma.owner = signer.pubkey();
}
if replaced_signers.contains(&ma.delegate) {
ma.delegate = signer.pubkey();
}
}
// Override token account owners
if account.owner == spl_token::id() {
if let Ok(mut ta) = spl_token::state::Account::unpack(&account.data) {
if replaced_signers.contains(&ta.owner) {
ta.owner = signer.pubkey();
}
spl_token::state::Account::pack(ta, &mut account.data).unwrap();
}
}
let mut program_test_context = solana.context.borrow_mut();
program_test_context.set_account(&pubkey, &account.into());
}
if !missing_accounts.is_empty() {
println!("There were account reading errors, maybe fetch them:");
for a in &missing_accounts {
println!("solana account {a} --output json -o {snapshot_path}/{a}.json");
}
anyhow::bail!("accounts were missing");
}
// update slot/time to roughly match
let mut clock = solana.clock().await;
clock.slot = slot;
clock.unix_timestamp = time;
solana.set_clock(&clock);
// Send transaction
solana
.process_transaction(&[replaced_ix], Some(&[signer]))
.await
.unwrap();
Ok(())
}