solana_message/
account_keys.rs

1use {
2    crate::{compiled_instruction::CompiledInstruction, v0::LoadedAddresses, CompileError},
3    solana_address::Address,
4    solana_instruction::Instruction,
5    std::{collections::BTreeMap, iter::zip, ops::Index},
6};
7
8/// Collection of static and dynamically loaded keys used to load accounts
9/// during transaction processing.
10#[derive(Clone, Default, Debug, Eq)]
11pub struct AccountKeys<'a> {
12    static_keys: &'a [Address],
13    dynamic_keys: Option<&'a LoadedAddresses>,
14}
15
16impl Index<usize> for AccountKeys<'_> {
17    type Output = Address;
18    #[inline]
19    fn index(&self, index: usize) -> &Self::Output {
20        self.get(index).expect("index is invalid")
21    }
22}
23
24impl<'a> AccountKeys<'a> {
25    pub fn new(static_keys: &'a [Address], dynamic_keys: Option<&'a LoadedAddresses>) -> Self {
26        Self {
27            static_keys,
28            dynamic_keys,
29        }
30    }
31
32    /// Returns an iterator of account key segments. The ordering of segments
33    /// affects how account indexes from compiled instructions are resolved and
34    /// so should not be changed.
35    #[inline]
36    fn key_segment_iter(&self) -> impl Iterator<Item = &'a [Address]> + Clone {
37        if let Some(dynamic_keys) = self.dynamic_keys {
38            [
39                self.static_keys,
40                &dynamic_keys.writable,
41                &dynamic_keys.readonly,
42            ]
43            .into_iter()
44        } else {
45            // empty segments added for branch type compatibility
46            [self.static_keys, &[], &[]].into_iter()
47        }
48    }
49
50    /// Returns the address of the account at the specified index of the list of
51    /// message account keys constructed from static keys, followed by dynamically
52    /// loaded writable addresses, and lastly the list of dynamically loaded
53    /// readonly addresses.
54    #[inline]
55    pub fn get(&self, mut index: usize) -> Option<&'a Address> {
56        for key_segment in self.key_segment_iter() {
57            if index < key_segment.len() {
58                return Some(&key_segment[index]);
59            }
60            index = index.saturating_sub(key_segment.len());
61        }
62
63        None
64    }
65
66    /// Returns the total length of loaded accounts for a message
67    #[inline]
68    pub fn len(&self) -> usize {
69        let mut len = 0usize;
70        for key_segment in self.key_segment_iter() {
71            len = len.saturating_add(key_segment.len());
72        }
73        len
74    }
75
76    /// Returns true if this collection of account keys is empty
77    pub fn is_empty(&self) -> bool {
78        self.len() == 0
79    }
80
81    /// Iterator for the addresses of the loaded accounts for a message
82    #[inline]
83    pub fn iter(&self) -> impl Iterator<Item = &'a Address> + Clone {
84        self.key_segment_iter().flatten()
85    }
86
87    /// Compile instructions using the order of account keys to determine
88    /// compiled instruction account indexes.
89    ///
90    /// # Panics
91    ///
92    /// Panics when compiling fails. See [`AccountKeys::try_compile_instructions`]
93    /// for a full description of failure scenarios.
94    pub fn compile_instructions(&self, instructions: &[Instruction]) -> Vec<CompiledInstruction> {
95        self.try_compile_instructions(instructions)
96            .expect("compilation failure")
97    }
98
99    /// Compile instructions using the order of account keys to determine
100    /// compiled instruction account indexes.
101    ///
102    /// # Errors
103    ///
104    /// Compilation will fail if any `instructions` use account keys which are not
105    /// present in this account key collection.
106    ///
107    /// Compilation will fail if any `instructions` use account keys which are located
108    /// at an index which cannot be cast to a `u8` without overflow.
109    pub fn try_compile_instructions(
110        &self,
111        instructions: &[Instruction],
112    ) -> Result<Vec<CompiledInstruction>, CompileError> {
113        let mut account_index_map = BTreeMap::<&Address, u8>::new();
114        for (index, key) in self.iter().enumerate() {
115            let index = u8::try_from(index).map_err(|_| CompileError::AccountIndexOverflow)?;
116            account_index_map.insert(key, index);
117        }
118
119        let get_account_index = |key: &Address| -> Result<u8, CompileError> {
120            account_index_map
121                .get(key)
122                .cloned()
123                .ok_or(CompileError::UnknownInstructionKey(*key))
124        };
125
126        instructions
127            .iter()
128            .map(|ix| {
129                let accounts: Vec<u8> = ix
130                    .accounts
131                    .iter()
132                    .map(|account_meta| get_account_index(&account_meta.pubkey))
133                    .collect::<Result<Vec<u8>, CompileError>>()?;
134
135                Ok(CompiledInstruction {
136                    program_id_index: get_account_index(&ix.program_id)?,
137                    data: ix.data.clone(),
138                    accounts,
139                })
140            })
141            .collect()
142    }
143}
144
145impl PartialEq for AccountKeys<'_> {
146    fn eq(&self, other: &Self) -> bool {
147        zip(self.iter(), other.iter()).all(|(a, b)| a == b)
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use {super::*, solana_instruction::AccountMeta};
154
155    fn test_account_keys() -> [Address; 6] {
156        let key0 = Address::new_unique();
157        let key1 = Address::new_unique();
158        let key2 = Address::new_unique();
159        let key3 = Address::new_unique();
160        let key4 = Address::new_unique();
161        let key5 = Address::new_unique();
162
163        [key0, key1, key2, key3, key4, key5]
164    }
165
166    #[test]
167    fn test_key_segment_iter() {
168        let keys = test_account_keys();
169
170        let static_keys = vec![keys[0], keys[1], keys[2]];
171        let dynamic_keys = LoadedAddresses {
172            writable: vec![keys[3], keys[4]],
173            readonly: vec![keys[5]],
174        };
175        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
176
177        let expected_segments = [
178            vec![keys[0], keys[1], keys[2]],
179            vec![keys[3], keys[4]],
180            vec![keys[5]],
181        ];
182
183        assert!(account_keys.key_segment_iter().eq(expected_segments.iter()));
184    }
185
186    #[test]
187    fn test_len() {
188        let keys = test_account_keys();
189
190        let static_keys = vec![keys[0], keys[1], keys[2], keys[3], keys[4], keys[5]];
191        let account_keys = AccountKeys::new(&static_keys, None);
192
193        assert_eq!(account_keys.len(), keys.len());
194    }
195
196    #[test]
197    fn test_len_with_dynamic_keys() {
198        let keys = test_account_keys();
199
200        let static_keys = vec![keys[0], keys[1], keys[2]];
201        let dynamic_keys = LoadedAddresses {
202            writable: vec![keys[3], keys[4]],
203            readonly: vec![keys[5]],
204        };
205        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
206
207        assert_eq!(account_keys.len(), keys.len());
208    }
209
210    #[test]
211    fn test_iter() {
212        let keys = test_account_keys();
213
214        let static_keys = vec![keys[0], keys[1], keys[2], keys[3], keys[4], keys[5]];
215        let account_keys = AccountKeys::new(&static_keys, None);
216
217        assert!(account_keys.iter().eq(keys.iter()));
218    }
219
220    #[test]
221    fn test_iter_with_dynamic_keys() {
222        let keys = test_account_keys();
223
224        let static_keys = vec![keys[0], keys[1], keys[2]];
225        let dynamic_keys = LoadedAddresses {
226            writable: vec![keys[3], keys[4]],
227            readonly: vec![keys[5]],
228        };
229        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
230
231        assert!(account_keys.iter().eq(keys.iter()));
232    }
233
234    #[test]
235    fn test_get() {
236        let keys = test_account_keys();
237
238        let static_keys = vec![keys[0], keys[1], keys[2], keys[3]];
239        let account_keys = AccountKeys::new(&static_keys, None);
240
241        assert_eq!(account_keys.get(0), Some(&keys[0]));
242        assert_eq!(account_keys.get(1), Some(&keys[1]));
243        assert_eq!(account_keys.get(2), Some(&keys[2]));
244        assert_eq!(account_keys.get(3), Some(&keys[3]));
245        assert_eq!(account_keys.get(4), None);
246        assert_eq!(account_keys.get(5), None);
247    }
248
249    #[test]
250    fn test_get_with_dynamic_keys() {
251        let keys = test_account_keys();
252
253        let static_keys = vec![keys[0], keys[1], keys[2]];
254        let dynamic_keys = LoadedAddresses {
255            writable: vec![keys[3], keys[4]],
256            readonly: vec![keys[5]],
257        };
258        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
259
260        assert_eq!(account_keys.get(0), Some(&keys[0]));
261        assert_eq!(account_keys.get(1), Some(&keys[1]));
262        assert_eq!(account_keys.get(2), Some(&keys[2]));
263        assert_eq!(account_keys.get(3), Some(&keys[3]));
264        assert_eq!(account_keys.get(4), Some(&keys[4]));
265        assert_eq!(account_keys.get(5), Some(&keys[5]));
266    }
267
268    #[test]
269    fn test_try_compile_instructions() {
270        let keys = test_account_keys();
271
272        let static_keys = vec![keys[0]];
273        let dynamic_keys = LoadedAddresses {
274            writable: vec![keys[1]],
275            readonly: vec![keys[2]],
276        };
277        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
278
279        let instruction = Instruction {
280            program_id: keys[0],
281            accounts: vec![
282                AccountMeta::new(keys[1], true),
283                AccountMeta::new(keys[2], true),
284            ],
285            data: vec![0],
286        };
287
288        assert_eq!(
289            account_keys.try_compile_instructions(&[instruction]),
290            Ok(vec![CompiledInstruction {
291                program_id_index: 0,
292                accounts: vec![1, 2],
293                data: vec![0],
294            }]),
295        );
296    }
297
298    #[test]
299    fn test_try_compile_instructions_with_unknown_key() {
300        let static_keys = test_account_keys();
301        let account_keys = AccountKeys::new(&static_keys, None);
302
303        let unknown_key = Address::new_unique();
304        let test_instructions = [
305            Instruction {
306                program_id: unknown_key,
307                accounts: vec![],
308                data: vec![],
309            },
310            Instruction {
311                program_id: static_keys[0],
312                accounts: vec![
313                    AccountMeta::new(static_keys[1], false),
314                    AccountMeta::new(unknown_key, false),
315                ],
316                data: vec![],
317            },
318        ];
319
320        for ix in test_instructions {
321            assert_eq!(
322                account_keys.try_compile_instructions(&[ix]),
323                Err(CompileError::UnknownInstructionKey(unknown_key))
324            );
325        }
326    }
327
328    #[test]
329    fn test_try_compile_instructions_with_too_many_account_keys() {
330        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
331        let static_keys = vec![Address::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
332        let dynamic_keys = LoadedAddresses {
333            writable: vec![Address::default()],
334            readonly: vec![],
335        };
336        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
337        assert_eq!(
338            account_keys.try_compile_instructions(&[]),
339            Err(CompileError::AccountIndexOverflow)
340        );
341    }
342}
solana_message/account_keys.rs

solana_message/
account_keys.rs
