waffles_solana_program/message/
compiled_keys.rs

1#[cfg(not(target_os = "solana"))]
2use crate::{
3    address_lookup_table_account::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 mut 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 mut 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        pub struct KeyFlags: u8 {
191            const SIGNER   = 0b00000001;
192            const WRITABLE = 0b00000010;
193            const INVOKED  = 0b00000100;
194        }
195    }
196
197    impl From<KeyFlags> for CompiledKeyMeta {
198        fn from(flags: KeyFlags) -> Self {
199            Self {
200                is_signer: flags.contains(KeyFlags::SIGNER),
201                is_writable: flags.contains(KeyFlags::WRITABLE),
202                is_invoked: flags.contains(KeyFlags::INVOKED),
203            }
204        }
205    }
206
207    #[test]
208    fn test_compile_with_dups() {
209        let program_id0 = Pubkey::new_unique();
210        let program_id1 = Pubkey::new_unique();
211        let program_id2 = Pubkey::new_unique();
212        let program_id3 = Pubkey::new_unique();
213        let id0 = Pubkey::new_unique();
214        let id1 = Pubkey::new_unique();
215        let id2 = Pubkey::new_unique();
216        let id3 = Pubkey::new_unique();
217        let compiled_keys = CompiledKeys::compile(
218            &[
219                Instruction::new_with_bincode(
220                    program_id0,
221                    &0,
222                    vec![
223                        AccountMeta::new_readonly(id0, false),
224                        AccountMeta::new_readonly(id1, true),
225                        AccountMeta::new(id2, false),
226                        AccountMeta::new(id3, true),
227                        // duplicate the account inputs
228                        AccountMeta::new_readonly(id0, false),
229                        AccountMeta::new_readonly(id1, true),
230                        AccountMeta::new(id2, false),
231                        AccountMeta::new(id3, true),
232                        // reference program ids
233                        AccountMeta::new_readonly(program_id0, false),
234                        AccountMeta::new_readonly(program_id1, true),
235                        AccountMeta::new(program_id2, false),
236                        AccountMeta::new(program_id3, true),
237                    ],
238                ),
239                Instruction::new_with_bincode(program_id1, &0, vec![]),
240                Instruction::new_with_bincode(program_id2, &0, vec![]),
241                Instruction::new_with_bincode(program_id3, &0, vec![]),
242            ],
243            None,
244        );
245
246        assert_eq!(
247            compiled_keys,
248            CompiledKeys {
249                payer: None,
250                key_meta_map: BTreeMap::from([
251                    (id0, KeyFlags::empty().into()),
252                    (id1, KeyFlags::SIGNER.into()),
253                    (id2, KeyFlags::WRITABLE.into()),
254                    (id3, (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
255                    (program_id0, KeyFlags::INVOKED.into()),
256                    (program_id1, (KeyFlags::INVOKED | KeyFlags::SIGNER).into()),
257                    (program_id2, (KeyFlags::INVOKED | KeyFlags::WRITABLE).into()),
258                    (program_id3, KeyFlags::all().into()),
259                ]),
260            }
261        );
262    }
263
264    #[test]
265    fn test_compile_with_dup_payer() {
266        let program_id = Pubkey::new_unique();
267        let payer = Pubkey::new_unique();
268        let compiled_keys = CompiledKeys::compile(
269            &[Instruction::new_with_bincode(
270                program_id,
271                &0,
272                vec![AccountMeta::new_readonly(payer, false)],
273            )],
274            Some(payer),
275        );
276        assert_eq!(
277            compiled_keys,
278            CompiledKeys {
279                payer: Some(payer),
280                key_meta_map: BTreeMap::from([
281                    (payer, (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
282                    (program_id, KeyFlags::INVOKED.into()),
283                ]),
284            }
285        );
286    }
287
288    #[test]
289    fn test_compile_with_dup_signer_mismatch() {
290        let program_id = Pubkey::new_unique();
291        let id0 = Pubkey::new_unique();
292        let compiled_keys = CompiledKeys::compile(
293            &[Instruction::new_with_bincode(
294                program_id,
295                &0,
296                vec![AccountMeta::new(id0, false), AccountMeta::new(id0, true)],
297            )],
298            None,
299        );
300
301        // Ensure the dup writable key is a signer
302        assert_eq!(
303            compiled_keys,
304            CompiledKeys {
305                payer: None,
306                key_meta_map: BTreeMap::from([
307                    (id0, (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
308                    (program_id, KeyFlags::INVOKED.into()),
309                ]),
310            }
311        );
312    }
313
314    #[test]
315    fn test_compile_with_dup_signer_writable_mismatch() {
316        let program_id = Pubkey::new_unique();
317        let id0 = Pubkey::new_unique();
318        let compiled_keys = CompiledKeys::compile(
319            &[Instruction::new_with_bincode(
320                program_id,
321                &0,
322                vec![
323                    AccountMeta::new_readonly(id0, true),
324                    AccountMeta::new(id0, true),
325                ],
326            )],
327            None,
328        );
329
330        // Ensure the dup signer key is writable
331        assert_eq!(
332            compiled_keys,
333            CompiledKeys {
334                payer: None,
335                key_meta_map: BTreeMap::from([
336                    (id0, (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
337                    (program_id, KeyFlags::INVOKED.into()),
338                ]),
339            }
340        );
341    }
342
343    #[test]
344    fn test_compile_with_dup_nonsigner_writable_mismatch() {
345        let program_id = Pubkey::new_unique();
346        let id0 = Pubkey::new_unique();
347        let compiled_keys = CompiledKeys::compile(
348            &[
349                Instruction::new_with_bincode(
350                    program_id,
351                    &0,
352                    vec![
353                        AccountMeta::new_readonly(id0, false),
354                        AccountMeta::new(id0, false),
355                    ],
356                ),
357                Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, false)]),
358            ],
359            None,
360        );
361
362        // Ensure the dup nonsigner key is writable
363        assert_eq!(
364            compiled_keys,
365            CompiledKeys {
366                payer: None,
367                key_meta_map: BTreeMap::from([
368                    (id0, KeyFlags::WRITABLE.into()),
369                    (program_id, KeyFlags::INVOKED.into()),
370                ]),
371            }
372        );
373    }
374
375    #[test]
376    fn test_try_into_message_components() {
377        let keys = vec![
378            Pubkey::new_unique(),
379            Pubkey::new_unique(),
380            Pubkey::new_unique(),
381            Pubkey::new_unique(),
382        ];
383
384        let compiled_keys = CompiledKeys {
385            payer: None,
386            key_meta_map: BTreeMap::from([
387                (keys[0], (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
388                (keys[1], KeyFlags::SIGNER.into()),
389                (keys[2], KeyFlags::WRITABLE.into()),
390                (keys[3], KeyFlags::empty().into()),
391            ]),
392        };
393
394        let result = compiled_keys.try_into_message_components();
395        assert_eq!(result.as_ref().err(), None);
396        let (header, static_keys) = result.unwrap();
397
398        assert_eq!(static_keys, keys);
399        assert_eq!(
400            header,
401            MessageHeader {
402                num_required_signatures: 2,
403                num_readonly_signed_accounts: 1,
404                num_readonly_unsigned_accounts: 1,
405            }
406        );
407    }
408
409    #[test]
410    fn test_try_into_message_components_with_too_many_keys() {
411        const TOO_MANY_KEYS: usize = 257;
412
413        for key_flags in [
414            KeyFlags::WRITABLE | KeyFlags::SIGNER,
415            KeyFlags::SIGNER,
416            // skip writable_non_signer_keys because it isn't used for creating header values
417            KeyFlags::empty(),
418        ] {
419            let test_keys = CompiledKeys {
420                payer: None,
421                key_meta_map: BTreeMap::from_iter(
422                    (0..TOO_MANY_KEYS).map(|_| (Pubkey::new_unique(), key_flags.into())),
423                ),
424            };
425
426            assert_eq!(
427                test_keys.try_into_message_components(),
428                Err(CompileError::AccountIndexOverflow)
429            );
430        }
431    }
432
433    #[test]
434    fn test_try_extract_table_lookup() {
435        let keys = vec![
436            Pubkey::new_unique(),
437            Pubkey::new_unique(),
438            Pubkey::new_unique(),
439            Pubkey::new_unique(),
440            Pubkey::new_unique(),
441            Pubkey::new_unique(),
442        ];
443
444        let mut compiled_keys = CompiledKeys {
445            payer: None,
446            key_meta_map: BTreeMap::from([
447                (keys[0], (KeyFlags::SIGNER | KeyFlags::WRITABLE).into()),
448                (keys[1], KeyFlags::SIGNER.into()),
449                (keys[2], KeyFlags::WRITABLE.into()),
450                (keys[3], KeyFlags::empty().into()),
451                (keys[4], (KeyFlags::INVOKED | KeyFlags::WRITABLE).into()),
452                (keys[5], (KeyFlags::INVOKED).into()),
453            ]),
454        };
455
456        // add some duplicates to ensure lowest index is selected
457        let addresses = [keys.clone(), keys.clone()].concat();
458        let lookup_table_account = AddressLookupTableAccount {
459            key: Pubkey::new_unique(),
460            addresses,
461        };
462
463        assert_eq!(
464            compiled_keys.try_extract_table_lookup(&lookup_table_account),
465            Ok(Some((
466                MessageAddressTableLookup {
467                    account_key: lookup_table_account.key,
468                    writable_indexes: vec![2],
469                    readonly_indexes: vec![3],
470                },
471                LoadedAddresses {
472                    writable: vec![keys[2]],
473                    readonly: vec![keys[3]],
474                },
475            )))
476        );
477
478        assert_eq!(compiled_keys.key_meta_map.len(), 4);
479        assert!(!compiled_keys.key_meta_map.contains_key(&keys[2]));
480        assert!(!compiled_keys.key_meta_map.contains_key(&keys[3]));
481    }
482
483    #[test]
484    fn test_try_extract_table_lookup_returns_none() {
485        let mut compiled_keys = CompiledKeys {
486            payer: None,
487            key_meta_map: BTreeMap::from([
488                (Pubkey::new_unique(), KeyFlags::WRITABLE.into()),
489                (Pubkey::new_unique(), KeyFlags::empty().into()),
490            ]),
491        };
492
493        let lookup_table_account = AddressLookupTableAccount {
494            key: Pubkey::new_unique(),
495            addresses: vec![],
496        };
497
498        let expected_compiled_keys = compiled_keys.clone();
499        assert_eq!(
500            compiled_keys.try_extract_table_lookup(&lookup_table_account),
501            Ok(None)
502        );
503        assert_eq!(compiled_keys, expected_compiled_keys);
504    }
505
506    #[test]
507    fn test_try_extract_table_lookup_for_invalid_table() {
508        let writable_key = Pubkey::new_unique();
509        let mut compiled_keys = CompiledKeys {
510            payer: None,
511            key_meta_map: BTreeMap::from([
512                (writable_key, KeyFlags::WRITABLE.into()),
513                (Pubkey::new_unique(), KeyFlags::empty().into()),
514            ]),
515        };
516
517        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
518        let mut addresses = vec![Pubkey::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
519        addresses.push(writable_key);
520
521        let lookup_table_account = AddressLookupTableAccount {
522            key: Pubkey::new_unique(),
523            addresses,
524        };
525
526        let expected_compiled_keys = compiled_keys.clone();
527        assert_eq!(
528            compiled_keys.try_extract_table_lookup(&lookup_table_account),
529            Err(CompileError::AddressTableLookupIndexOverflow),
530        );
531        assert_eq!(compiled_keys, expected_compiled_keys);
532    }
533
534    #[test]
535    fn test_try_drain_keys_found_in_lookup_table() {
536        let orig_keys = vec![
537            Pubkey::new_unique(),
538            Pubkey::new_unique(),
539            Pubkey::new_unique(),
540            Pubkey::new_unique(),
541            Pubkey::new_unique(),
542        ];
543
544        let mut compiled_keys = CompiledKeys {
545            payer: None,
546            key_meta_map: BTreeMap::from([
547                (orig_keys[0], KeyFlags::empty().into()),
548                (orig_keys[1], KeyFlags::WRITABLE.into()),
549                (orig_keys[2], KeyFlags::WRITABLE.into()),
550                (orig_keys[3], KeyFlags::empty().into()),
551                (orig_keys[4], KeyFlags::empty().into()),
552            ]),
553        };
554
555        let lookup_table_addresses = vec![
556            Pubkey::new_unique(),
557            orig_keys[0],
558            Pubkey::new_unique(),
559            orig_keys[4],
560            Pubkey::new_unique(),
561            orig_keys[2],
562            Pubkey::new_unique(),
563        ];
564
565        let drain_result = compiled_keys
566            .try_drain_keys_found_in_lookup_table(&lookup_table_addresses, |meta| {
567                !meta.is_writable
568            });
569        assert_eq!(drain_result.as_ref().err(), None);
570        let (lookup_table_indexes, drained_keys) = drain_result.unwrap();
571
572        assert_eq!(
573            compiled_keys.key_meta_map.keys().collect::<Vec<&_>>(),
574            vec![&orig_keys[1], &orig_keys[2], &orig_keys[3]]
575        );
576        assert_eq!(drained_keys, vec![orig_keys[0], orig_keys[4]]);
577        assert_eq!(lookup_table_indexes, vec![1, 3]);
578    }
579
580    #[test]
581    fn test_try_drain_keys_found_in_lookup_table_with_empty_keys() {
582        let mut compiled_keys = CompiledKeys::default();
583
584        let lookup_table_addresses = vec![
585            Pubkey::new_unique(),
586            Pubkey::new_unique(),
587            Pubkey::new_unique(),
588        ];
589
590        let drain_result =
591            compiled_keys.try_drain_keys_found_in_lookup_table(&lookup_table_addresses, |_| true);
592        assert_eq!(drain_result.as_ref().err(), None);
593        let (lookup_table_indexes, drained_keys) = drain_result.unwrap();
594
595        assert!(drained_keys.is_empty());
596        assert!(lookup_table_indexes.is_empty());
597    }
598
599    #[test]
600    fn test_try_drain_keys_found_in_lookup_table_with_empty_table() {
601        let original_keys = vec![
602            Pubkey::new_unique(),
603            Pubkey::new_unique(),
604            Pubkey::new_unique(),
605        ];
606
607        let mut compiled_keys = CompiledKeys {
608            payer: None,
609            key_meta_map: BTreeMap::from_iter(
610                original_keys
611                    .iter()
612                    .map(|key| (*key, CompiledKeyMeta::default())),
613            ),
614        };
615
616        let lookup_table_addresses = vec![];
617
618        let drain_result =
619            compiled_keys.try_drain_keys_found_in_lookup_table(&lookup_table_addresses, |_| true);
620        assert_eq!(drain_result.as_ref().err(), None);
621        let (lookup_table_indexes, drained_keys) = drain_result.unwrap();
622
623        assert_eq!(compiled_keys.key_meta_map.len(), original_keys.len());
624        assert!(drained_keys.is_empty());
625        assert!(lookup_table_indexes.is_empty());
626    }
627
628    #[test]
629    fn test_try_drain_keys_found_in_lookup_table_with_too_many_addresses() {
630        let key = Pubkey::new_unique();
631        let mut compiled_keys = CompiledKeys {
632            payer: None,
633            key_meta_map: BTreeMap::from([(key, CompiledKeyMeta::default())]),
634        };
635
636        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
637        let mut lookup_table_addresses = vec![Pubkey::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
638        lookup_table_addresses.push(key);
639
640        let drain_result =
641            compiled_keys.try_drain_keys_found_in_lookup_table(&lookup_table_addresses, |_| true);
642        assert_eq!(
643            drain_result.err(),
644            Some(CompileError::AddressTableLookupIndexOverflow)
645        );
646    }
647}
waffles_solana_program/message/compiled_keys.rs

waffles_solana_program/message/
compiled_keys.rs
