safe_token_2022/extension/
reallocate.rs

1use {
2    crate::{
3        error::TokenError,
4        extension::{
5            set_account_type, AccountType, BaseStateWithExtensions, ExtensionType,
6            StateWithExtensions,
7        },
8        processor::Processor,
9        state::Account,
10    },
11    solana_program::{
12        account_info::{next_account_info, AccountInfo},
13        entrypoint::ProgramResult,
14        msg,
15        program::invoke,
16        pubkey::Pubkey,
17        system_instruction,
18        sysvar::{rent::Rent, Sysvar},
19    },
20};
21
22/// Processes a [Reallocate](enum.TokenInstruction.html) instruction
23pub fn process_reallocate(
24    program_id: &Pubkey,
25    accounts: &[AccountInfo],
26    new_extension_types: Vec<ExtensionType>,
27) -> ProgramResult {
28    let account_info_iter = &mut accounts.iter();
29    let token_account_info = next_account_info(account_info_iter)?;
30    let payer_info = next_account_info(account_info_iter)?;
31    let system_program_info = next_account_info(account_info_iter)?;
32    let authority_info = next_account_info(account_info_iter)?;
33    let authority_info_data_len = authority_info.data_len();
34
35    // check that account is the right type and validate owner
36    let mut current_extension_types = {
37        let token_account = token_account_info.data.borrow();
38        let account = StateWithExtensions::<Account>::unpack(&token_account)?;
39        Processor::validate_owner(
40            program_id,
41            &account.base.owner,
42            authority_info,
43            authority_info_data_len,
44            account_info_iter.as_slice(),
45        )?;
46        account.get_extension_types()?
47    };
48
49    // check that all desired extensions are for the right account type
50    if new_extension_types
51        .iter()
52        .any(|extension_type| extension_type.get_account_type() != AccountType::Account)
53    {
54        return Err(TokenError::InvalidState.into());
55    }
56    // ExtensionType::get_account_len() dedupes types, so just a dumb concatenation is fine here
57    current_extension_types.extend_from_slice(&new_extension_types);
58    let needed_account_len = ExtensionType::get_account_len::<Account>(&current_extension_types);
59
60    // if account is already large enough, return early
61    if token_account_info.data_len() >= needed_account_len {
62        return Ok(());
63    }
64
65    // reallocate
66    msg!(
67        "account needs realloc, +{:?} bytes",
68        needed_account_len - token_account_info.data_len()
69    );
70    token_account_info.realloc(needed_account_len, false)?;
71
72    // if additional lamports needed to remain rent-exempt, transfer them
73    let rent = Rent::get()?;
74    let new_minimum_balance = rent.minimum_balance(needed_account_len);
75    let lamports_diff = new_minimum_balance.saturating_sub(token_account_info.lamports());
76    invoke(
77        &system_instruction::transfer(payer_info.key, token_account_info.key, lamports_diff),
78        &[
79            payer_info.clone(),
80            token_account_info.clone(),
81            system_program_info.clone(),
82        ],
83    )?;
84
85    // unpack to set account_type, if needed
86    let mut token_account = token_account_info.data.borrow_mut();
87    set_account_type::<Account>(&mut token_account)?;
88
89    Ok(())
90}