account_compression/
context.rs

1use anchor_lang::{
2    prelude::{AccountInfo, AccountLoader, ProgramError},
3    solana_program::{msg, pubkey::Pubkey},
4    Discriminator as AnchorDiscriminator, Key, ToAccountInfo,
5};
6use light_account_checks::{discriminator::Discriminator, error::AccountError};
7use light_batched_merkle_tree::{
8    merkle_tree::BatchedMerkleTreeAccount, queue::BatchedQueueAccount,
9};
10use light_concurrent_merkle_tree::zero_copy::ConcurrentMerkleTreeZeroCopyMut;
11use light_hasher::Poseidon;
12use light_indexed_merkle_tree::zero_copy::IndexedMerkleTreeZeroCopyMut;
13use light_merkle_tree_metadata::TreeType;
14
15use crate::{
16    address_merkle_tree_from_bytes_zero_copy_mut,
17    errors::AccountCompressionErrorCode,
18    state_merkle_tree_from_bytes_zero_copy_mut,
19    utils::check_signer_is_registered_or_authority::{
20        manual_check_signer_is_registered_or_authority, GroupAccess,
21    },
22    AddressMerkleTreeAccount, QueueAccount, StateMerkleTreeAccount,
23};
24
25impl GroupAccess for BatchedQueueAccount<'_> {
26    fn get_owner(&self) -> Pubkey {
27        self.metadata.access_metadata.owner.into()
28    }
29
30    fn get_program_owner(&self) -> Pubkey {
31        self.metadata
32            .access_metadata
33            .program_owner
34            .to_bytes()
35            .into()
36    }
37}
38use super::RegisteredProgram;
39
40/// AccountCompressionProgramAccount
41#[derive(Debug)]
42pub enum AcpAccount<'a, 'info> {
43    Authority(&'a AccountInfo<'info>),
44    RegisteredProgramPda(&'a AccountInfo<'info>),
45    SystemProgram(&'a AccountInfo<'info>),
46    OutputQueue(BatchedQueueAccount<'info>),
47    BatchedStateTree(BatchedMerkleTreeAccount<'info>),
48    BatchedAddressTree(BatchedMerkleTreeAccount<'info>),
49    StateTree((Pubkey, ConcurrentMerkleTreeZeroCopyMut<'info, Poseidon, 26>)),
50    AddressTree(
51        (
52            Pubkey,
53            IndexedMerkleTreeZeroCopyMut<'info, Poseidon, usize, 26, 16>,
54        ),
55    ),
56    AddressQueue(Pubkey, AccountInfo<'info>),
57    V1Queue(AccountInfo<'info>),
58    Unknown(),
59}
60
61impl<'a, 'info> AcpAccount<'a, 'info> {
62    /// Merkle tree and queue accounts
63    #[inline(always)]
64    pub fn from_account_infos(
65        account_infos: &'info [AccountInfo<'info>],
66        authority: &'a AccountInfo<'info>,
67        invoked_by_program: bool,
68        // TODO: remove in separate pr because it impacts photon derivation.
69        _bump: u8,
70    ) -> std::result::Result<Vec<AcpAccount<'a, 'info>>, ProgramError> {
71        let mut vec = Vec::with_capacity(account_infos.len());
72        let mut skip = 0;
73        let derived_address = match invoked_by_program {
74            true => {
75                let account_info = &account_infos[0];
76                let data = account_info.try_borrow_data()?;
77                if RegisteredProgram::DISCRIMINATOR != &data[..8] {
78                    return Err(AccountError::InvalidDiscriminator.into());
79                }
80                if account_info.owner != &crate::ID {
81                    return Err(AccountError::AccountOwnedByWrongProgram.into());
82                }
83                let account = bytemuck::from_bytes::<RegisteredProgram>(&data[8..]);
84
85                if account.registered_program_signer_pda != *authority.key {
86                    return Err(AccountError::InvalidSigner.into());
87                }
88                skip += 1;
89                Some((
90                    account.registered_program_signer_pda,
91                    account.group_authority_pda,
92                ))
93            }
94            false => None,
95        };
96
97        account_infos.iter().skip(skip).try_for_each(
98            |account_info| -> Result<(), ProgramError> {
99                let account = AcpAccount::try_from_account_info(
100                    account_info,
101                    &AcpAccount::Authority(authority),
102                    &derived_address,
103                )?;
104                vec.push(account);
105                Ok(())
106            },
107        )?;
108        Ok(vec)
109    }
110
111    /// Try to deserialize and check account info:
112    /// 1. Owner is crate::ID
113    /// 2. match discriminator
114    /// 3. check signer is registered program or authority
115    ///    (Unless the account is a v1 queue account.
116    ///    v1 queue accounts are always used in combination
117    ///    with a v1 Merkle tree and are checked that these
118    ///    are associated to it.)
119    #[inline(always)]
120    pub(crate) fn try_from_account_info(
121        account_info: &'info AccountInfo<'info>,
122        authority: &AcpAccount<'a, 'info>,
123        registered_program_pda: &Option<(Pubkey, Pubkey)>,
124    ) -> anchor_lang::Result<AcpAccount<'a, 'info>> {
125        if crate::ID != *account_info.owner {
126            msg!("Invalid owner {:?}", account_info.owner);
127            msg!("key {:?}", account_info.key());
128            return Err(ProgramError::from(AccountError::AccountOwnedByWrongProgram).into());
129        }
130        let mut discriminator = [0u8; 8];
131        {
132            let data = account_info.try_borrow_data()?;
133            discriminator.copy_from_slice(&data[..8]);
134        }
135        match &discriminator[..] {
136            BatchedMerkleTreeAccount::LIGHT_DISCRIMINATOR_SLICE => {
137                let mut tree_type = [0u8; 8];
138                tree_type.copy_from_slice(&account_info.try_borrow_data()?[8..16]);
139                let tree_type = TreeType::from(u64::from_le_bytes(tree_type));
140                match tree_type {
141                    TreeType::AddressV2 => {
142                        let tree =
143                            BatchedMerkleTreeAccount::address_from_account_info(account_info)
144                                .map_err(ProgramError::from)?;
145                        manual_check_signer_is_registered_or_authority::<BatchedMerkleTreeAccount>(
146                            registered_program_pda,
147                            authority,
148                            &tree,
149                        )?;
150                        Ok(AcpAccount::BatchedAddressTree(tree))
151                    }
152                    TreeType::StateV2 => {
153                        let tree = BatchedMerkleTreeAccount::state_from_account_info(account_info)
154                            .map_err(ProgramError::from)?;
155                        manual_check_signer_is_registered_or_authority::<BatchedMerkleTreeAccount>(
156                            registered_program_pda,
157                            authority,
158                            &tree,
159                        )?;
160
161                        Ok(AcpAccount::BatchedStateTree(tree))
162                    }
163                    _ => Err(ProgramError::from(AccountError::BorrowAccountDataFailed).into()),
164                }
165            }
166            BatchedQueueAccount::LIGHT_DISCRIMINATOR_SLICE => {
167                let queue = BatchedQueueAccount::output_from_account_info(account_info)
168                    .map_err(ProgramError::from)?;
169
170                manual_check_signer_is_registered_or_authority::<BatchedQueueAccount>(
171                    registered_program_pda,
172                    authority,
173                    &queue,
174                )?;
175
176                Ok(AcpAccount::OutputQueue(queue))
177            }
178            StateMerkleTreeAccount::DISCRIMINATOR => {
179                {
180                    let merkle_tree =
181                        AccountLoader::<StateMerkleTreeAccount>::try_from(account_info)?;
182                    let merkle_tree = merkle_tree.load()?;
183
184                    manual_check_signer_is_registered_or_authority::<StateMerkleTreeAccount>(
185                        registered_program_pda,
186                        authority,
187                        &merkle_tree,
188                    )?;
189                }
190                let mut merkle_tree = account_info.try_borrow_mut_data()?;
191                let data_slice: &'info mut [u8] = unsafe {
192                    std::slice::from_raw_parts_mut(merkle_tree.as_mut_ptr(), merkle_tree.len())
193                };
194                Ok(AcpAccount::StateTree((
195                    account_info.key(),
196                    state_merkle_tree_from_bytes_zero_copy_mut(data_slice)?,
197                )))
198            }
199            AddressMerkleTreeAccount::DISCRIMINATOR => {
200                {
201                    let merkle_tree =
202                        AccountLoader::<AddressMerkleTreeAccount>::try_from(account_info)?;
203                    let merkle_tree = merkle_tree.load()?;
204                    manual_check_signer_is_registered_or_authority::<AddressMerkleTreeAccount>(
205                        registered_program_pda,
206                        authority,
207                        &merkle_tree,
208                    )?;
209                }
210                let mut merkle_tree = account_info.try_borrow_mut_data()?;
211                let data_slice: &'info mut [u8] = unsafe {
212                    std::slice::from_raw_parts_mut(merkle_tree.as_mut_ptr(), merkle_tree.len())
213                };
214                Ok(AcpAccount::AddressTree((
215                    account_info.key(),
216                    address_merkle_tree_from_bytes_zero_copy_mut(data_slice)?,
217                )))
218            }
219            QueueAccount::DISCRIMINATOR => {
220                msg!("queue account: {:?}", account_info.key());
221                Ok(AcpAccount::V1Queue(account_info.to_account_info()))
222            }
223            _ => Err(AccountCompressionErrorCode::InvalidAccount.into()),
224        }
225    }
226}
account_compression/context.rs

account_compression/
context.rs
