safe_token_2022/extension/cpi_guard/
processor.rs

1use {
2    crate::{
3        check_program_account,
4        error::TokenError,
5        extension::{
6            cpi_guard::{in_cpi, instruction::CpiGuardInstruction, CpiGuard},
7            StateWithExtensionsMut,
8        },
9        instruction::decode_instruction_type,
10        processor::Processor,
11        state::Account,
12    },
13    solana_program::{
14        account_info::{next_account_info, AccountInfo},
15        entrypoint::ProgramResult,
16        msg,
17        pubkey::Pubkey,
18    },
19};
20
21/// Toggle the CpiGuard extension, initializing the extension if not already present.
22fn process_toggle_cpi_guard(
23    program_id: &Pubkey,
24    accounts: &[AccountInfo],
25    enable: bool,
26) -> ProgramResult {
27    let account_info_iter = &mut accounts.iter();
28    let token_account_info = next_account_info(account_info_iter)?;
29    let owner_info = next_account_info(account_info_iter)?;
30    let owner_info_data_len = owner_info.data_len();
31
32    let mut account_data = token_account_info.data.borrow_mut();
33    let mut account = StateWithExtensionsMut::<Account>::unpack(&mut account_data)?;
34
35    Processor::validate_owner(
36        program_id,
37        &account.base.owner,
38        owner_info,
39        owner_info_data_len,
40        account_info_iter.as_slice(),
41    )?;
42
43    if in_cpi() {
44        return Err(TokenError::CpiGuardSettingsLocked.into());
45    }
46
47    let extension = if let Ok(extension) = account.get_extension_mut::<CpiGuard>() {
48        extension
49    } else {
50        account.init_extension::<CpiGuard>(true)?
51    };
52    extension.lock_cpi = enable.into();
53    Ok(())
54}
55
56pub(crate) fn process_instruction(
57    program_id: &Pubkey,
58    accounts: &[AccountInfo],
59    input: &[u8],
60) -> ProgramResult {
61    check_program_account(program_id)?;
62
63    match decode_instruction_type(input)? {
64        CpiGuardInstruction::Enable => {
65            msg!("CpiGuardInstruction::Enable");
66            process_toggle_cpi_guard(program_id, accounts, true /* enable */)
67        }
68        CpiGuardInstruction::Disable => {
69            msg!("CpiGuardInstruction::Disable");
70            process_toggle_cpi_guard(program_id, accounts, false /* disable */)
71        }
72    }
73}