Skip to content

Commit dfad1af

Browse files
authored
fix: transaction account dedupe and roles (#29)
1 parent f8e3414 commit dfad1af

File tree

4 files changed

+214
-71
lines changed

4 files changed

+214
-71
lines changed

harness/src/keys.rs

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
//! Instruction <-> Transaction key deduplication and privilege handling.
2+
//!
3+
//! Solana instructions and transactions are designed to be intentionally
4+
//! verbosely declarative, to provide the runtime with granular directives
5+
//! for manipulating chain state.
6+
//!
7+
//! As a result, when a transaction is _compiled_, many steps occur:
8+
//! * Ensuring there is a fee payer.
9+
//! * Ensuring there is a signature.
10+
//! * Deduplicating account keys.
11+
//! * Configuring the highest role awarded to each account key.
12+
//! * ...
13+
//!
14+
//! Since Mollusk does not use transactions or fee payers, the deduplication
15+
//! of account keys and handling of roles are the only two steps necessary
16+
//! to perform under the hood within the harness.
17+
//!
18+
//! This implementation closely follows the implementation in the Anza SDK
19+
//! for `Message::new_with_blockhash`. For more information, see:
20+
//! <https://github.com/anza-xyz/agave/blob/c6e8239843af8e6301cd198e39d0a44add427bef/sdk/program/src/message/legacy.rs#L357>.
21+
22+
use {
23+
solana_sdk::{
24+
account::{AccountSharedData, WritableAccount},
25+
instruction::Instruction,
26+
pubkey::Pubkey,
27+
transaction_context::{IndexOfAccount, InstructionAccount, TransactionAccount},
28+
},
29+
std::collections::HashMap,
30+
};
31+
32+
struct KeyMap(HashMap<Pubkey, (bool, bool)>);
33+
34+
impl KeyMap {
35+
fn compile(instruction: &Instruction) -> Self {
36+
let mut map: HashMap<Pubkey, (bool, bool)> = HashMap::new();
37+
map.entry(instruction.program_id).or_default();
38+
for meta in instruction.accounts.iter() {
39+
let entry = map.entry(meta.pubkey).or_default();
40+
entry.0 |= meta.is_signer;
41+
entry.1 |= meta.is_writable;
42+
}
43+
Self(map)
44+
}
45+
46+
fn is_signer(&self, key: &Pubkey) -> bool {
47+
self.0.get(key).map(|(s, _)| *s).unwrap_or(false)
48+
}
49+
50+
fn is_writable(&self, key: &Pubkey) -> bool {
51+
self.0.get(key).map(|(_, w)| *w).unwrap_or(false)
52+
}
53+
}
54+
55+
struct Keys<'a> {
56+
keys: Vec<&'a Pubkey>,
57+
key_map: &'a KeyMap,
58+
}
59+
60+
impl<'a> Keys<'a> {
61+
fn new(key_map: &'a KeyMap) -> Self {
62+
Self {
63+
keys: key_map.0.keys().collect(),
64+
key_map,
65+
}
66+
}
67+
68+
fn position(&self, key: &Pubkey) -> u8 {
69+
self.keys.iter().position(|k| *k == key).unwrap() as u8
70+
}
71+
72+
fn is_signer(&self, index: usize) -> bool {
73+
self.key_map.is_signer(self.keys[index])
74+
}
75+
76+
fn is_writable(&self, index: usize) -> bool {
77+
self.key_map.is_writable(self.keys[index])
78+
}
79+
}
80+
81+
// Helper struct so Mollusk doesn't have to clone instruction data.
82+
struct CompiledInstructionWithoutData {
83+
program_id_index: u8,
84+
accounts: Vec<u8>,
85+
}
86+
87+
pub struct CompiledAccounts {
88+
pub program_id_index: u16,
89+
pub instruction_accounts: Vec<InstructionAccount>,
90+
pub transaction_accounts: Vec<TransactionAccount>,
91+
}
92+
93+
pub fn compile_accounts(
94+
instruction: &Instruction,
95+
accounts: &[(Pubkey, AccountSharedData)],
96+
loader_key: Pubkey,
97+
) -> CompiledAccounts {
98+
let key_map = KeyMap::compile(instruction);
99+
let keys = Keys::new(&key_map);
100+
101+
let compiled_instruction = CompiledInstructionWithoutData {
102+
program_id_index: keys.position(&instruction.program_id),
103+
accounts: instruction
104+
.accounts
105+
.iter()
106+
.map(|account_meta| keys.position(&account_meta.pubkey))
107+
.collect(),
108+
};
109+
110+
let instruction_accounts: Vec<InstructionAccount> = compiled_instruction
111+
.accounts
112+
.iter()
113+
.enumerate()
114+
.map(|(ix_account_index, &index_in_transaction)| {
115+
let index_in_callee = compiled_instruction
116+
.accounts
117+
.get(0..ix_account_index)
118+
.unwrap()
119+
.iter()
120+
.position(|&account_index| account_index == index_in_transaction)
121+
.unwrap_or(ix_account_index) as IndexOfAccount;
122+
let index_in_transaction = index_in_transaction as usize;
123+
InstructionAccount {
124+
index_in_transaction: index_in_transaction as IndexOfAccount,
125+
index_in_caller: index_in_transaction as IndexOfAccount,
126+
index_in_callee,
127+
is_signer: keys.is_signer(index_in_transaction),
128+
is_writable: keys.is_writable(index_in_transaction),
129+
}
130+
})
131+
.collect();
132+
133+
let transaction_accounts: Vec<TransactionAccount> = keys
134+
.keys
135+
.iter()
136+
.map(|key| {
137+
if *key == &instruction.program_id {
138+
(**key, stub_out_program_account(loader_key))
139+
} else {
140+
let account = accounts
141+
.iter()
142+
.find(|(k, _)| k == *key)
143+
.map(|(_, account)| account.clone())
144+
.unwrap_or_else(|| {
145+
panic!(
146+
" [mollusk]: An account required by the instruction was not \
147+
provided: {:?}",
148+
key,
149+
)
150+
});
151+
(**key, account)
152+
}
153+
})
154+
.collect();
155+
156+
CompiledAccounts {
157+
program_id_index: compiled_instruction.program_id_index as u16,
158+
instruction_accounts,
159+
transaction_accounts,
160+
}
161+
}
162+
163+
fn stub_out_program_account(loader_key: Pubkey) -> AccountSharedData {
164+
let mut program_account = AccountSharedData::default();
165+
program_account.set_owner(loader_key);
166+
program_account.set_executable(true);
167+
program_account
168+
}

harness/src/lib.rs

Lines changed: 36 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
//! series of checks on the result, panicking if any checks fail.
2929
3030
pub mod file;
31+
mod keys;
3132
pub mod program;
3233
pub mod result;
3334
pub mod sysvar;
@@ -38,30 +39,23 @@ use {
3839
result::{Check, InstructionResult},
3940
sysvar::Sysvars,
4041
},
42+
keys::CompiledAccounts,
4143
solana_compute_budget::compute_budget::ComputeBudget,
4244
solana_program_runtime::{
4345
invoke_context::{EnvironmentConfig, InvokeContext},
4446
sysvar_cache::SysvarCache,
4547
timings::ExecuteTimings,
4648
},
4749
solana_sdk::{
48-
account::AccountSharedData,
49-
bpf_loader_upgradeable,
50-
feature_set::FeatureSet,
51-
fee::FeeStructure,
52-
hash::Hash,
53-
instruction::Instruction,
54-
pubkey::Pubkey,
55-
transaction_context::{InstructionAccount, TransactionContext},
50+
account::AccountSharedData, bpf_loader_upgradeable, feature_set::FeatureSet,
51+
fee::FeeStructure, hash::Hash, instruction::Instruction, pubkey::Pubkey,
52+
transaction_context::TransactionContext,
5653
},
57-
std::{collections::HashMap, sync::Arc},
54+
std::sync::Arc,
5855
};
5956

6057
const DEFAULT_LOADER_KEY: Pubkey = bpf_loader_upgradeable::id();
6158

62-
const PROGRAM_ACCOUNTS_LEN: usize = 1;
63-
const PROGRAM_INDICES: &[u16] = &[0];
64-
6559
/// The Mollusk API, providing a simple interface for testing Solana programs.
6660
///
6761
/// All fields can be manipulated through a handful of helper methods, but
@@ -203,40 +197,22 @@ impl Mollusk {
203197
let mut compute_units_consumed = 0;
204198
let mut timings = ExecuteTimings::default();
205199

206-
let instruction_accounts = {
207-
// For ensuring each account has the proper privilege level (dedupe).
208-
// <pubkey, (is_signer, is_writable)>
209-
let mut privileges: HashMap<Pubkey, (bool, bool)> = HashMap::new();
210-
211-
for meta in instruction.accounts.iter() {
212-
let entry = privileges.entry(meta.pubkey).or_default();
213-
entry.0 |= meta.is_signer;
214-
entry.1 |= meta.is_writable;
215-
}
216-
217-
instruction
218-
.accounts
219-
.iter()
220-
.enumerate()
221-
.map(|(i, meta)| {
222-
// Guaranteed by the last iteration.
223-
let (is_signer, is_writable) = privileges.get(&meta.pubkey).unwrap();
224-
InstructionAccount {
225-
index_in_callee: i as u16,
226-
index_in_caller: i as u16,
227-
index_in_transaction: (i + PROGRAM_ACCOUNTS_LEN) as u16,
228-
is_signer: *is_signer,
229-
is_writable: *is_writable,
230-
}
231-
})
232-
.collect::<Vec<_>>()
233-
};
200+
let loader_key = self
201+
.program_cache
202+
.load_program(&instruction.program_id)
203+
.unwrap_or_else(|| {
204+
panic!(
205+
" [mollusk]: Program targeted by instruction is missing from cache: {:?}",
206+
instruction.program_id,
207+
)
208+
})
209+
.account_owner();
234210

235-
let transaction_accounts = [(self.program_id, self.program_account.clone())]
236-
.iter()
237-
.chain(accounts)
238-
.cloned()
239-
.collect::<Vec<_>>();
211+
let CompiledAccounts {
212+
program_id_index,
213+
instruction_accounts,
214+
transaction_accounts,
215+
} = crate::keys::compile_accounts(instruction, accounts, loader_key);
240216

241217
let mut transaction_context = TransactionContext::new(
242218
transaction_accounts,
@@ -264,20 +240,26 @@ impl Mollusk {
264240
.process_instruction(
265241
&instruction.data,
266242
&instruction_accounts,
267-
PROGRAM_INDICES,
243+
&[program_id_index],
268244
&mut compute_units_consumed,
269245
&mut timings,
270246
)
271247
};
272248

273-
let resulting_accounts = transaction_context
274-
.deconstruct_without_keys()
275-
.unwrap()
276-
.into_iter()
277-
.skip(PROGRAM_ACCOUNTS_LEN)
278-
.zip(instruction.accounts.iter())
279-
.map(|(account, meta)| (meta.pubkey, account))
280-
.collect::<Vec<_>>();
249+
let resulting_accounts: Vec<(Pubkey, AccountSharedData)> = (0..transaction_context
250+
.get_number_of_accounts())
251+
.filter_map(|index| {
252+
let key = transaction_context
253+
.get_key_of_account_at_index(index)
254+
.unwrap();
255+
let account = transaction_context.get_account_at_index(index).unwrap();
256+
if *key != instruction.program_id {
257+
Some((*key, account.take()))
258+
} else {
259+
None
260+
}
261+
})
262+
.collect();
281263

282264
InstructionResult {
283265
compute_units_consumed,

harness/src/program.rs

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,13 @@ impl ProgramCache {
4242
&self.cache
4343
}
4444

45+
/// Add a builtin program to the cache.
46+
pub fn add_builtin(&mut self, builtin: Builtin) {
47+
let program_id = builtin.program_id;
48+
let entry = builtin.program_cache_entry();
49+
self.cache.write().unwrap().replenish(program_id, entry);
50+
}
51+
4552
/// Add a program to the cache.
4653
pub fn add_program(
4754
&mut self,
@@ -72,11 +79,9 @@ impl ProgramCache {
7279
);
7380
}
7481

75-
/// Add a builtin program to the cache.
76-
pub fn add_builtin(&mut self, builtin: Builtin) {
77-
let program_id = builtin.program_id;
78-
let entry = builtin.program_cache_entry();
79-
self.cache.write().unwrap().replenish(program_id, entry);
82+
/// Load a program from the cache.
83+
pub fn load_program(&self, program_id: &Pubkey) -> Option<Arc<ProgramCacheEntry>> {
84+
self.cache.read().unwrap().find(program_id)
8085
}
8186
}
8287

harness/tests/bpf_program.rs

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -254,18 +254,6 @@ fn test_cpi() {
254254
)
255255
};
256256

257-
// Fail CPI target program account not provided.
258-
{
259-
mollusk.process_and_validate_instruction(
260-
&instruction,
261-
&[(key, account.clone())],
262-
&[
263-
Check::err(ProgramError::NotEnoughAccountKeys),
264-
Check::compute_units(0), // No compute units used.
265-
],
266-
);
267-
}
268-
269257
// Fail CPI target program not added to test environment.
270258
{
271259
mollusk.process_and_validate_instruction(

0 commit comments

Comments
 (0)