squads_multisig/vault_transaction/
vault_transaction_message.rs

1use crate::pda::get_ephemeral_signer_pda;
2use squads_multisig_program::{
3    CompiledInstruction, MessageAddressTableLookup, TransactionMessage, VaultTransactionMessage,
4};
5use std::collections::HashMap;
6
7use super::compiled_keys::CompiledKeys;
8use crate::solana_program::address_lookup_table_account::AddressLookupTableAccount;
9use crate::solana_program::instruction::{AccountMeta, Instruction};
10use crate::solana_program::message::{AccountKeys, CompileError};
11use crate::solana_program::pubkey::Pubkey;
12
13#[derive(thiserror::Error, Debug)]
14pub enum Error {
15    #[error("Invalid AddressLookupTableAccount")]
16    InvalidAddressLookupTableAccount,
17    #[error("Invalid TransactionMessage")]
18    InvalidTransactionMessage,
19}
20
21pub trait VaultTransactionMessageExt {
22    fn as_transaction_message(&self) -> &TransactionMessage;
23
24    /// This implementation is mostly a copy-paste from `solana_program::message::v0::Message::try_compile()`,
25    /// but it constructs a `TransactionMessage` meant to be passed to `vault_transaction_create`.
26    fn try_compile(
27        vault_key: &Pubkey,
28        instructions: &[Instruction],
29        address_lookup_table_accounts: &[AddressLookupTableAccount],
30    ) -> Result<TransactionMessage, CompileError> {
31        let mut compiled_keys = CompiledKeys::compile(instructions, Some(*vault_key));
32
33        let mut address_table_lookups = Vec::with_capacity(address_lookup_table_accounts.len());
34        let mut loaded_addresses_list = Vec::with_capacity(address_lookup_table_accounts.len());
35        for lookup_table_account in address_lookup_table_accounts {
36            if let Some((lookup, loaded_addresses)) =
37                compiled_keys.try_extract_table_lookup(lookup_table_account)?
38            {
39                address_table_lookups.push(lookup);
40                loaded_addresses_list.push(loaded_addresses);
41            }
42        }
43
44        let (header, static_keys) = compiled_keys.try_into_message_components()?;
45        let dynamic_keys = loaded_addresses_list.into_iter().collect();
46        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
47        let instructions = account_keys.try_compile_instructions(instructions)?;
48
49        let num_static_keys: u8 = static_keys
50            .len()
51            .try_into()
52            .map_err(|_| CompileError::AccountIndexOverflow)?;
53
54        Ok(TransactionMessage {
55            num_signers: header.num_required_signatures,
56            num_writable_signers: header.num_required_signatures
57                - header.num_readonly_signed_accounts,
58            num_writable_non_signers: num_static_keys
59                - header.num_required_signatures
60                - header.num_readonly_unsigned_accounts,
61            account_keys: static_keys.into(),
62            instructions: instructions
63                .into_iter()
64                .map(|ix| CompiledInstruction {
65                    program_id_index: ix.program_id_index,
66                    account_indexes: ix.accounts.into(),
67                    data: ix.data.into(),
68                })
69                .collect::<Vec<CompiledInstruction>>()
70                .into(),
71            address_table_lookups: address_table_lookups
72                .into_iter()
73                .map(|lookup| MessageAddressTableLookup {
74                    account_key: lookup.account_key,
75                    writable_indexes: lookup.writable_indexes.into(),
76                    readonly_indexes: lookup.readonly_indexes.into(),
77                })
78                .collect::<Vec<MessageAddressTableLookup>>()
79                .into(),
80        })
81    }
82
83    fn get_accounts_for_execute(
84        &self,
85        vault_pda: &Pubkey,
86        transaction_pda: &Pubkey,
87        address_lookup_table_accounts: &[AddressLookupTableAccount],
88        num_ephemeral_signers: u8,
89        program_id: &Pubkey,
90    ) -> Result<Vec<AccountMeta>, Error> {
91        let message = VaultTransactionMessage::try_from(self.as_transaction_message().to_owned())
92            .map_err(|_| Error::InvalidTransactionMessage)?;
93
94        let ephemeral_signer_pdas: Vec<Pubkey> = (0..num_ephemeral_signers)
95            .into_iter()
96            .map(|ephemeral_signer_index| {
97                get_ephemeral_signer_pda(transaction_pda, ephemeral_signer_index, Some(program_id))
98                    .0
99            })
100            .collect();
101
102        // region: -- address_lookup_tables map --
103
104        let address_lookup_tables = address_lookup_table_accounts
105            .into_iter()
106            .map(|alt| (alt.key, alt))
107            .collect::<HashMap<_, _>>();
108
109        // endregion: -- address_lookup_tables map --
110
111        // region: -- Account Metas --
112
113        // First go the lookup table accounts used by the transaction. They are needed for on-chain validation.
114        let lookup_table_account_metas = address_lookup_table_accounts
115            .iter()
116            .map(|alt| AccountMeta {
117                pubkey: alt.key,
118                is_writable: false,
119                is_signer: false,
120            })
121            .collect::<Vec<_>>();
122
123        // Then come static account keys included into the message.
124        let static_account_metas = message
125            .account_keys
126            .iter()
127            .enumerate()
128            .map(|(index, &pubkey)| {
129                AccountMeta {
130                    pubkey,
131                    is_writable: message.is_static_writable_index(index),
132                    // NOTE: vaultPda and ephemeralSignerPdas cannot be marked as signers,
133                    // because they are PDAs and hence won't have their signatures on the transaction.
134                    is_signer: message.is_signer_index(index)
135                        && &pubkey != vault_pda
136                        && !ephemeral_signer_pdas.contains(&pubkey),
137                }
138            })
139            .collect::<Vec<_>>();
140
141        // And the last go the accounts that will be loaded with address lookup tables.
142        let loaded_account_metas = message
143            .address_table_lookups
144            .iter()
145            .map(|lookup| {
146                let lookup_table_account = address_lookup_tables
147                    .get(&lookup.account_key)
148                    .ok_or(Error::InvalidAddressLookupTableAccount)?;
149
150                // For each lookup, fist list writable, then readonly account metas.
151                lookup
152                    .writable_indexes
153                    .iter()
154                    .map(|&index| {
155                        let pubkey = lookup_table_account
156                            .addresses
157                            .get(index as usize)
158                            .ok_or(Error::InvalidAddressLookupTableAccount)?
159                            .to_owned();
160
161                        Ok(AccountMeta {
162                            pubkey,
163                            is_writable: true,
164                            is_signer: false,
165                        })
166                    })
167                    .chain(lookup.readonly_indexes.iter().map(|&index| {
168                        let pubkey = lookup_table_account
169                            .addresses
170                            .get(index as usize)
171                            .ok_or(Error::InvalidAddressLookupTableAccount)?
172                            .to_owned();
173
174                        Ok(AccountMeta {
175                            pubkey,
176                            is_writable: false,
177                            is_signer: false,
178                        })
179                    }))
180                    .collect::<Result<Vec<_>, Error>>()
181            })
182            .collect::<Result<Vec<_>, Error>>()?
183            .into_iter()
184            .flatten()
185            .collect::<Vec<_>>();
186
187        // endregion: -- Account Metas --
188
189        Ok([
190            lookup_table_account_metas,
191            static_account_metas,
192            loaded_account_metas,
193        ]
194        .concat())
195    }
196}
197
198impl VaultTransactionMessageExt for TransactionMessage {
199    fn as_transaction_message(&self) -> &TransactionMessage {
200        self
201    }
202}