solana_program/message/
compiled_keys.rs

1#[cfg(not(target_os = "solana"))]
2use crate::{
3    address_lookup_table::AddressLookupTableAccount,
4    message::v0::{LoadedAddresses, MessageAddressTableLookup},
5};
6use {
7    crate::{instruction::Instruction, message::MessageHeader, pubkey::Pubkey},
8    std::collections::BTreeMap,
9    thiserror::Error,
10};
11
12/// A helper struct to collect pubkeys compiled for a set of instructions
13#[derive(Default, Debug, Clone, PartialEq, Eq)]
14pub(crate) struct CompiledKeys {
15    payer: Option<Pubkey>,
16    key_meta_map: BTreeMap<Pubkey, CompiledKeyMeta>,
17}
18
19#[cfg_attr(target_os = "solana", allow(dead_code))]
20#[derive(PartialEq, Debug, Error, Eq, Clone)]
21pub enum CompileError {
22    #[error("account index overflowed during compilation")]
23    AccountIndexOverflow,
24    #[error("address lookup table index overflowed during compilation")]
25    AddressTableLookupIndexOverflow,
26    #[error("encountered unknown account key `{0}` during instruction compilation")]
27    UnknownInstructionKey(Pubkey),
28}
29
30#[derive(Default, Debug, Clone, PartialEq, Eq)]
31struct CompiledKeyMeta {
32    is_signer: bool,
33    is_writable: bool,
34    is_invoked: bool,
35}
36
37impl CompiledKeys {
38    /// Compiles the pubkeys referenced by a list of instructions and organizes by
39    /// signer/non-signer and writable/readonly.
40    pub(crate) fn compile(instructions: &[Instruction], payer: Option<Pubkey>) -> Self {
41        let mut key_meta_map = BTreeMap::<Pubkey, CompiledKeyMeta>::new();
42        for ix in instructions {
43            let meta = key_meta_map.entry(ix.program_id).or_default();
44            meta.is_invoked = true;
45            for account_meta in &ix.accounts {
46                let meta = key_meta_map.entry(account_meta.pubkey).or_default();
47                meta.is_signer |= account_meta.is_signer;
48                meta.is_writable |= account_meta.is_writable;
49            }
50        }
51        if let Some(payer) = &payer {
52            let meta = key_meta_map.entry(*payer).or_default();
53            meta.is_signer = true;
54            meta.is_writable = true;
55        }
56        Self {
57            payer,
58            key_meta_map,
59        }
60    }
61
62    pub(crate) fn try_into_message_components(
63        self,
64    ) -> Result<(MessageHeader, Vec<Pubkey>), CompileError> {
65        let try_into_u8 = |num: usize| -> Result<u8, CompileError> {
66            u8::try_from(num).map_err(|_| CompileError::AccountIndexOverflow)
67        };
68
69        let Self {
70            payer,
71            mut key_meta_map,
72        } = self;
73
74        if let Some(payer) = &payer {
75            key_meta_map.remove_entry(payer);
76        }
77
78        let writable_signer_keys: Vec<Pubkey> = payer
79            .into_iter()
80            .chain(
81                key_meta_map
82                    .iter()
83                    .filter_map(|(key, meta)| (meta.is_signer && meta.is_writable).then_some(*key)),
84            )
85            .collect();
86        let readonly_signer_keys: Vec<Pubkey> = key_meta_map
87            .iter()
88            .filter_map(|(key, meta)| (meta.is_signer && !meta.is_writable).then_some(*key))
89            .collect();
90        let writable_non_signer_keys: Vec<Pubkey> = key_meta_map
91            .iter()
92            .filter_map(|(key, meta)| (!meta.is_signer && meta.is_writable).then_some(*key))
93            .collect();
94        let readonly_non_signer_keys: Vec<Pubkey> = key_meta_map
95            .iter()
96            .filter_map(|(key, meta)| (!meta.is_signer && !meta.is_writable).then_some(*key))
97            .collect();
98
99        let signers_len = writable_signer_keys
100            .len()
101            .saturating_add(readonly_signer_keys.len());
102
103        let header = MessageHeader {
104            num_required_signatures: try_into_u8(signers_len)?,
105            num_readonly_signed_accounts: try_into_u8(readonly_signer_keys.len())?,
106            num_readonly_unsigned_accounts: try_into_u8(readonly_non_signer_keys.len())?,
107        };
108
109        let static_account_keys = std::iter::empty()
110            .chain(writable_signer_keys)
111            .chain(readonly_signer_keys)
112            .chain(writable_non_signer_keys)
113            .chain(readonly_non_signer_keys)
114            .collect();
115
116        Ok((header, static_account_keys))
117    }
118
119    #[cfg(not(target_os = "solana"))]
120    pub(crate) fn try_extract_table_lookup(
121        &mut self,
122        lookup_table_account: &AddressLookupTableAccount,
123    ) -> Result<Option<(MessageAddressTableLookup, LoadedAddresses)>, CompileError> {
124        let (writable_indexes, drained_writable_keys) = self
125            .try_drain_keys_found_in_lookup_table(&lookup_table_account.addresses, |meta| {
126                !meta.is_signer && !meta.is_invoked && meta.is_writable
127            })?;
128        let (readonly_indexes, drained_readonly_keys) = self
129            .try_drain_keys_found_in_lookup_table(&lookup_table_account.addresses, |meta| {
130                !meta.is_signer && !meta.is_invoked && !meta.is_writable
131            })?;
132
133        // Don't extract lookup if no keys were found
134        if writable_indexes.is_empty() && readonly_indexes.is_empty() {
135            return Ok(None);
136        }
137
138        Ok(Some((
139            MessageAddressTableLookup {
140                account_key: lookup_table_account.key,
141                writable_indexes,
142                readonly_indexes,
143            },
144            LoadedAddresses {
145                writable: drained_writable_keys,
146                readonly: drained_readonly_keys,
147            },
148        )))
149    }
150
151    #[cfg(not(target_os = "solana"))]
152    fn try_drain_keys_found_in_lookup_table(
153        &mut self,
154        lookup_table_addresses: &[Pubkey],
155        key_meta_filter: impl Fn(&CompiledKeyMeta) -> bool,
156    ) -> Result<(Vec<u8>, Vec<Pubkey>), CompileError> {
157        let mut lookup_table_indexes = Vec::new();
158        let mut drained_keys = Vec::new();
159
160        for search_key in self
161            .key_meta_map
162            .iter()
163            .filter_map(|(key, meta)| key_meta_filter(meta).then_some(key))
164        {
165            for (key_index, key) in lookup_table_addresses.iter().enumerate() {
166                if key == search_key {
167                    let lookup_table_index = u8::try_from(key_index)
168                        .map_err(|_| CompileError::AddressTableLookupIndexOverflow)?;
169
170                    lookup_table_indexes.push(lookup_table_index);
171                    drained_keys.push(*search_key);
172                    break;
173                }
174            }
175        }
176
177        for key in &drained_keys {
178            self.key_meta_map.remove_entry(key);
179        }
180
181        Ok((lookup_table_indexes, drained_keys))
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use {super::*, crate::instruction::AccountMeta, bitflags::bitflags};
188
189    bitflags! {
190        #[derive(Clone, Copy)]
191        pub struct KeyFlags: u8 {
192            const SIGNER   = 0b00000001;
193            const WRITABLE = 0b00000010;
194            const INVOKED  = 0b00000100;
195        }
196    }
197
198    impl From<KeyFlags> for CompiledKeyMeta {
199        fn from(flags: KeyFlags) -> Self {
200            Self {
201                is_signer: flags.contains(KeyFlags::SIGNER),
202                is_writable: flags.contains(KeyFlags::WRITABLE),
203                is_invoked: flags.contains(KeyFlags::INVOKED),
204            }
205        }
206    }
207
208    #[test]
209    fn test_compile_with_dups() {
210        let program_id0 = Pubkey::new_unique();
211        let program_id1 = Pubkey::new_unique();
212        let program_id2 = Pubkey::new_unique();
213        let program_id3 = Pubkey::new_unique();
214        let id0 = Pubkey::new_unique();
215        let id1 = Pubkey::new_unique();
216        let id2 = Pubkey::new_unique();
217        let id3 = Pubkey::new_unique();
218        let compiled_keys = CompiledKeys::compile(
219            &[
220                Instruction::new_with_bincode(
221                    program_id0,
222                    &0,
223                    vec![
224                        AccountMeta::new_readonly(id0, false),
225                        AccountMeta::new_readonly(id1, true),
226                        AccountMeta::new(id2, false),
227                        AccountMeta::new(id3, true),
228                        // duplicate the account inputs
229                        AccountMeta::new_readonly(id0, false),
230                        AccountMeta::new_readonly(id1, true),
231                        AccountMeta::new(id2, false),
232                        AccountMeta::new(id3, true),
233                        // reference program ids
234                        AccountMeta::new_readonly(program_id0, false),
235                        AccountMeta::new_readonly(program_id1, true),
236                        AccountMeta::new(program_id2, false),
237                        AccountMeta::new(program_id3, true),
238                    ],
239                ),
240                Instruction::new_with_bincode(program_id1, &0, vec![]),
241                Instruction::new_with_bincode(program_id2, &0, vec![]),
242                Instruction::new_with_bincode(program_id3, &0, vec![]),
243            ],
244            None,
245        );
246
247        assert_eq!(
248            compiled_keys,
249            CompiledKeys {
250                payer: None,
251                key_meta_map: BTreeMap::from([
252                    (id0, KeyFlags::empty().into()),
253                    (id1, KeyFlags::SIGNER.into()),
254                    (id2, KeyFlags::WRITABLE.into()),
255                    (id3, (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
256                    (program_id0, KeyFlags::INVOKED.into()),
257                    (program_id1, (KeyFlags::INVOKED | KeyFlags::SIGNER).into()),
258                    (program_id2, (KeyFlags::INVOKED | KeyFlags::WRITABLE).into()),
259                    (program_id3, KeyFlags::all().into()),
260                ]),
261            }
262        );
263    }
264
265    #[test]
266    fn test_compile_with_dup_payer() {
267        let program_id = Pubkey::new_unique();
268        let payer = Pubkey::new_unique();
269        let compiled_keys = CompiledKeys::compile(
270            &[Instruction::new_with_bincode(
271                program_id,
272                &0,
273                vec![AccountMeta::new_readonly(payer, false)],
274            )],
275            Some(payer),
276        );
277        assert_eq!(
278            compiled_keys,
279            CompiledKeys {
280                payer: Some(payer),
281                key_meta_map: BTreeMap::from([
282                    (payer, (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
283                    (program_id, KeyFlags::INVOKED.into()),
284                ]),
285            }
286        );
287    }
288
289    #[test]
290    fn test_compile_with_dup_signer_mismatch() {
291        let program_id = Pubkey::new_unique();
292        let id0 = Pubkey::new_unique();
293        let compiled_keys = CompiledKeys::compile(
294            &[Instruction::new_with_bincode(
295                program_id,
296                &0,
297                vec![AccountMeta::new(id0, false), AccountMeta::new(id0, true)],
298            )],
299            None,
300        );
301
302        // Ensure the dup writable key is a signer
303        assert_eq!(
304            compiled_keys,
305            CompiledKeys {
306                payer: None,
307                key_meta_map: BTreeMap::from([
308                    (id0, (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
309                    (program_id, KeyFlags::INVOKED.into()),
310                ]),
311            }
312        );
313    }
314
315    #[test]
316    fn test_compile_with_dup_signer_writable_mismatch() {
317        let program_id = Pubkey::new_unique();
318        let id0 = Pubkey::new_unique();
319        let compiled_keys = CompiledKeys::compile(
320            &[Instruction::new_with_bincode(
321                program_id,
322                &0,
323                vec![
324                    AccountMeta::new_readonly(id0, true),
325                    AccountMeta::new(id0, true),
326                ],
327            )],
328            None,
329        );
330
331        // Ensure the dup signer key is writable
332        assert_eq!(
333            compiled_keys,
334            CompiledKeys {
335                payer: None,
336                key_meta_map: BTreeMap::from([
337                    (id0, (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
338                    (program_id, KeyFlags::INVOKED.into()),
339                ]),
340            }
341        );
342    }
343
344    #[test]
345    fn test_compile_with_dup_nonsigner_writable_mismatch() {
346        let program_id = Pubkey::new_unique();
347        let id0 = Pubkey::new_unique();
348        let compiled_keys = CompiledKeys::compile(
349            &[
350                Instruction::new_with_bincode(
351                    program_id,
352                    &0,
353                    vec![
354                        AccountMeta::new_readonly(id0, false),
355                        AccountMeta::new(id0, false),
356                    ],
357                ),
358                Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, false)]),
359            ],
360            None,
361        );
362
363        // Ensure the dup nonsigner key is writable
364        assert_eq!(
365            compiled_keys,
366            CompiledKeys {
367                payer: None,
368                key_meta_map: BTreeMap::from([
369                    (id0, KeyFlags::WRITABLE.into()),
370                    (program_id, KeyFlags::INVOKED.into()),
371                ]),
372            }
373        );
374    }
375
376    #[test]
377    fn test_try_into_message_components() {
378        let keys = vec![
379            Pubkey::new_unique(),
380            Pubkey::new_unique(),
381            Pubkey::new_unique(),
382            Pubkey::new_unique(),
383        ];
384
385        let compiled_keys = CompiledKeys {
386            payer: None,
387            key_meta_map: BTreeMap::from([
388                (keys[0], (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
389                (keys[1], KeyFlags::SIGNER.into()),
390                (keys[2], KeyFlags::WRITABLE.into()),
391                (keys[3], KeyFlags::empty().into()),
392            ]),
393        };
394
395        let result = compiled_keys.try_into_message_components();
396        assert_eq!(result.as_ref().err(), None);
397        let (header, static_keys) = result.unwrap();
398
399        assert_eq!(static_keys, keys);
400        assert_eq!(
401            header,
402            MessageHeader {
403                num_required_signatures: 2,
404                num_readonly_signed_accounts: 1,
405                num_readonly_unsigned_accounts: 1,
406            }
407        );
408    }
409
410    #[test]
411    fn test_try_into_message_components_with_too_many_keys() {
412        const TOO_MANY_KEYS: usize = 257;
413
414        for key_flags in [
415            KeyFlags::WRITABLE | KeyFlags::SIGNER,
416            KeyFlags::SIGNER,
417            // skip writable_non_signer_keys because it isn't used for creating header values
418            KeyFlags::empty(),
419        ] {
420            let test_keys = CompiledKeys {
421                payer: None,
422                key_meta_map: BTreeMap::from_iter(
423                    (0..TOO_MANY_KEYS).map(|_| (Pubkey::new_unique(), key_flags.into())),
424                ),
425            };
426
427            assert_eq!(
428                test_keys.try_into_message_components(),
429                Err(CompileError::AccountIndexOverflow)
430            );
431        }
432    }
433
434    #[test]
435    fn test_try_extract_table_lookup() {
436        let keys = vec![
437            Pubkey::new_unique(),
438            Pubkey::new_unique(),
439            Pubkey::new_unique(),
440            Pubkey::new_unique(),
441            Pubkey::new_unique(),
442            Pubkey::new_unique(),
443        ];
444
445        let mut compiled_keys = CompiledKeys {
446            payer: None,
447            key_meta_map: BTreeMap::from([
448                (keys[0], (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
449                (keys[1], KeyFlags::SIGNER.into()),
450                (keys[2], KeyFlags::WRITABLE.into()),
451                (keys[3], KeyFlags::empty().into()),
452                (keys[4], (KeyFlags::INVOKED | KeyFlags::WRITABLE).into()),
453                (keys[5], (KeyFlags::INVOKED).into()),
454            ]),
455        };
456
457        // add some duplicates to ensure lowest index is selected
458        let addresses = [keys.clone(), keys.clone()].concat();
459        let lookup_table_account = AddressLookupTableAccount {
460            key: Pubkey::new_unique(),
461            addresses,
462        };
463
464        assert_eq!(
465            compiled_keys.try_extract_table_lookup(&lookup_table_account),
466            Ok(Some((
467                MessageAddressTableLookup {
468                    account_key: lookup_table_account.key,
469                    writable_indexes: vec![2],
470                    readonly_indexes: vec![3],
471                },
472                LoadedAddresses {
473                    writable: vec![keys[2]],
474                    readonly: vec![keys[3]],
475                },
476            )))
477        );
478
479        assert_eq!(compiled_keys.key_meta_map.len(), 4);
480        assert!(!compiled_keys.key_meta_map.contains_key(&keys[2]));
481        assert!(!compiled_keys.key_meta_map.contains_key(&keys[3]));
482    }
483
484    #[test]
485    fn test_try_extract_table_lookup_returns_none() {
486        let mut compiled_keys = CompiledKeys {
487            payer: None,
488            key_meta_map: BTreeMap::from([
489                (Pubkey::new_unique(), KeyFlags::WRITABLE.into()),
490                (Pubkey::new_unique(), KeyFlags::empty().into()),
491            ]),
492        };
493
494        let lookup_table_account = AddressLookupTableAccount {
495            key: Pubkey::new_unique(),
496            addresses: vec![],
497        };
498
499        let expected_compiled_keys = compiled_keys.clone();
500        assert_eq!(
501            compiled_keys.try_extract_table_lookup(&lookup_table_account),
502            Ok(None)
503        );
504        assert_eq!(compiled_keys, expected_compiled_keys);
505    }
506
507    #[test]
508    fn test_try_extract_table_lookup_for_invalid_table() {
509        let writable_key = Pubkey::new_unique();
510        let mut compiled_keys = CompiledKeys {
511            payer: None,
512            key_meta_map: BTreeMap::from([
513                (writable_key, KeyFlags::WRITABLE.into()),
514                (Pubkey::new_unique(), KeyFlags::empty().into()),
515            ]),
516        };
517
518        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
519        let mut addresses = vec![Pubkey::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
520        addresses.push(writable_key);
521
522        let lookup_table_account = AddressLookupTableAccount {
523            key: Pubkey::new_unique(),
524            addresses,
525        };
526
527        let expected_compiled_keys = compiled_keys.clone();
528        assert_eq!(
529            compiled_keys.try_extract_table_lookup(&lookup_table_account),
530            Err(CompileError::AddressTableLookupIndexOverflow),
531        );
532        assert_eq!(compiled_keys, expected_compiled_keys);
533    }
534
535    #[test]
536    fn test_try_drain_keys_found_in_lookup_table() {
537        let orig_keys = [
538            Pubkey::new_unique(),
539            Pubkey::new_unique(),
540            Pubkey::new_unique(),
541            Pubkey::new_unique(),
542            Pubkey::new_unique(),
543        ];
544
545        let mut compiled_keys = CompiledKeys {
546            payer: None,
547            key_meta_map: BTreeMap::from([
548                (orig_keys[0], KeyFlags::empty().into()),
549                (orig_keys[1], KeyFlags::WRITABLE.into()),
550                (orig_keys[2], KeyFlags::WRITABLE.into()),
551                (orig_keys[3], KeyFlags::empty().into()),
552                (orig_keys[4], KeyFlags::empty().into()),
553            ]),
554        };
555
556        let lookup_table_addresses = vec![
557            Pubkey::new_unique(),
558            orig_keys[0],
559            Pubkey::new_unique(),
560            orig_keys[4],
561            Pubkey::new_unique(),
562            orig_keys[2],
563            Pubkey::new_unique(),
564        ];
565
566        let drain_result = compiled_keys
567            .try_drain_keys_found_in_lookup_table(&lookup_table_addresses, |meta| {
568                !meta.is_writable
569            });
570        assert_eq!(drain_result.as_ref().err(), None);
571        let (lookup_table_indexes, drained_keys) = drain_result.unwrap();
572
573        assert_eq!(
574            compiled_keys.key_meta_map.keys().collect::<Vec<&_>>(),
575            vec![&orig_keys[1], &orig_keys[2], &orig_keys[3]]
576        );
577        assert_eq!(drained_keys, vec![orig_keys[0], orig_keys[4]]);
578        assert_eq!(lookup_table_indexes, vec![1, 3]);
579    }
580
581    #[test]
582    fn test_try_drain_keys_found_in_lookup_table_with_empty_keys() {
583        let mut compiled_keys = CompiledKeys::default();
584
585        let lookup_table_addresses = vec![
586            Pubkey::new_unique(),
587            Pubkey::new_unique(),
588            Pubkey::new_unique(),
589        ];
590
591        let drain_result =
592            compiled_keys.try_drain_keys_found_in_lookup_table(&lookup_table_addresses, |_| true);
593        assert_eq!(drain_result.as_ref().err(), None);
594        let (lookup_table_indexes, drained_keys) = drain_result.unwrap();
595
596        assert!(drained_keys.is_empty());
597        assert!(lookup_table_indexes.is_empty());
598    }
599
600    #[test]
601    fn test_try_drain_keys_found_in_lookup_table_with_empty_table() {
602        let original_keys = [
603            Pubkey::new_unique(),
604            Pubkey::new_unique(),
605            Pubkey::new_unique(),
606        ];
607
608        let mut compiled_keys = CompiledKeys {
609            payer: None,
610            key_meta_map: BTreeMap::from_iter(
611                original_keys
612                    .iter()
613                    .map(|key| (*key, CompiledKeyMeta::default())),
614            ),
615        };
616
617        let lookup_table_addresses = vec![];
618
619        let drain_result =
620            compiled_keys.try_drain_keys_found_in_lookup_table(&lookup_table_addresses, |_| true);
621        assert_eq!(drain_result.as_ref().err(), None);
622        let (lookup_table_indexes, drained_keys) = drain_result.unwrap();
623
624        assert_eq!(compiled_keys.key_meta_map.len(), original_keys.len());
625        assert!(drained_keys.is_empty());
626        assert!(lookup_table_indexes.is_empty());
627    }
628
629    #[test]
630    fn test_try_drain_keys_found_in_lookup_table_with_too_many_addresses() {
631        let key = Pubkey::new_unique();
632        let mut compiled_keys = CompiledKeys {
633            payer: None,
634            key_meta_map: BTreeMap::from([(key, CompiledKeyMeta::default())]),
635        };
636
637        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
638        let mut lookup_table_addresses = vec![Pubkey::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
639        lookup_table_addresses.push(key);
640
641        let drain_result =
642            compiled_keys.try_drain_keys_found_in_lookup_table(&lookup_table_addresses, |_| true);
643        assert_eq!(
644            drain_result.err(),
645            Some(CompileError::AddressTableLookupIndexOverflow)
646        );
647    }
648}
solana_program/message/compiled_keys.rs

solana_program/message/
compiled_keys.rs
