spl_token_2022/extension/transfer_fee/
processor.rs

1use {
2    crate::{
3        check_program_account,
4        error::TokenError,
5        extension::{
6            transfer_fee::{
7                instruction::TransferFeeInstruction, TransferFee, TransferFeeAmount,
8                TransferFeeConfig, MAX_FEE_BASIS_POINTS,
9            },
10            BaseStateWithExtensions, BaseStateWithExtensionsMut, PodStateWithExtensions,
11            PodStateWithExtensionsMut,
12        },
13        pod::{PodAccount, PodMint},
14        processor::Processor,
15    },
16    solana_program::{
17        account_info::{next_account_info, AccountInfo},
18        clock::Clock,
19        entrypoint::ProgramResult,
20        msg,
21        program_option::COption,
22        pubkey::Pubkey,
23        sysvar::Sysvar,
24    },
25    std::convert::TryInto,
26};
27
28fn process_initialize_transfer_fee_config(
29    accounts: &[AccountInfo],
30    transfer_fee_config_authority: COption<Pubkey>,
31    withdraw_withheld_authority: COption<Pubkey>,
32    transfer_fee_basis_points: u16,
33    maximum_fee: u64,
34) -> ProgramResult {
35    let account_info_iter = &mut accounts.iter();
36    let mint_account_info = next_account_info(account_info_iter)?;
37
38    let mut mint_data = mint_account_info.data.borrow_mut();
39    let mut mint = PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut mint_data)?;
40    let extension = mint.init_extension::<TransferFeeConfig>(true)?;
41    extension.transfer_fee_config_authority = transfer_fee_config_authority.try_into()?;
42    extension.withdraw_withheld_authority = withdraw_withheld_authority.try_into()?;
43    extension.withheld_amount = 0u64.into();
44
45    if transfer_fee_basis_points > MAX_FEE_BASIS_POINTS {
46        return Err(TokenError::TransferFeeExceedsMaximum.into());
47    }
48    // To be safe, set newer and older transfer fees to the same thing on init,
49    // but only newer will actually be used
50    let epoch = Clock::get()?.epoch;
51    let transfer_fee = TransferFee {
52        epoch: epoch.into(),
53        transfer_fee_basis_points: transfer_fee_basis_points.into(),
54        maximum_fee: maximum_fee.into(),
55    };
56    extension.older_transfer_fee = transfer_fee;
57    extension.newer_transfer_fee = transfer_fee;
58
59    Ok(())
60}
61
62fn process_set_transfer_fee(
63    program_id: &Pubkey,
64    accounts: &[AccountInfo],
65    transfer_fee_basis_points: u16,
66    maximum_fee: u64,
67) -> ProgramResult {
68    let account_info_iter = &mut accounts.iter();
69    let mint_account_info = next_account_info(account_info_iter)?;
70    let authority_info = next_account_info(account_info_iter)?;
71    let authority_info_data_len = authority_info.data_len();
72
73    let mut mint_data = mint_account_info.data.borrow_mut();
74    let mut mint = PodStateWithExtensionsMut::<PodMint>::unpack(&mut mint_data)?;
75    let extension = mint.get_extension_mut::<TransferFeeConfig>()?;
76
77    let transfer_fee_config_authority =
78        Option::<Pubkey>::from(extension.transfer_fee_config_authority)
79            .ok_or(TokenError::NoAuthorityExists)?;
80    Processor::validate_owner(
81        program_id,
82        &transfer_fee_config_authority,
83        authority_info,
84        authority_info_data_len,
85        account_info_iter.as_slice(),
86    )?;
87
88    if transfer_fee_basis_points > MAX_FEE_BASIS_POINTS {
89        return Err(TokenError::TransferFeeExceedsMaximum.into());
90    }
91
92    // When setting the transfer fee, we have two situations:
93    // * newer transfer fee epoch <= current epoch: newer transfer fee is the active
94    //   one, so overwrite older transfer fee with newer, then overwrite newer
95    //   transfer fee
96    // * newer transfer fee epoch >= next epoch: it was never used, so just
97    //   overwrite next transfer fee
98    let epoch = Clock::get()?.epoch;
99    if u64::from(extension.newer_transfer_fee.epoch) <= epoch {
100        extension.older_transfer_fee = extension.newer_transfer_fee;
101    }
102    // set two epochs ahead to avoid rug pulls at the end of an epoch
103    let newer_fee_start_epoch = epoch.saturating_add(2);
104    let transfer_fee = TransferFee {
105        epoch: newer_fee_start_epoch.into(),
106        transfer_fee_basis_points: transfer_fee_basis_points.into(),
107        maximum_fee: maximum_fee.into(),
108    };
109    extension.newer_transfer_fee = transfer_fee;
110
111    Ok(())
112}
113
114fn process_withdraw_withheld_tokens_from_mint(
115    program_id: &Pubkey,
116    accounts: &[AccountInfo],
117) -> ProgramResult {
118    let account_info_iter = &mut accounts.iter();
119    let mint_account_info = next_account_info(account_info_iter)?;
120    let destination_account_info = next_account_info(account_info_iter)?;
121    let authority_info = next_account_info(account_info_iter)?;
122    let authority_info_data_len = authority_info.data_len();
123
124    // unnecessary check, but helps for clarity
125    check_program_account(mint_account_info.owner)?;
126
127    let mut mint_data = mint_account_info.data.borrow_mut();
128    let mut mint = PodStateWithExtensionsMut::<PodMint>::unpack(&mut mint_data)?;
129    let extension = mint.get_extension_mut::<TransferFeeConfig>()?;
130
131    let withdraw_withheld_authority = Option::<Pubkey>::from(extension.withdraw_withheld_authority)
132        .ok_or(TokenError::NoAuthorityExists)?;
133    Processor::validate_owner(
134        program_id,
135        &withdraw_withheld_authority,
136        authority_info,
137        authority_info_data_len,
138        account_info_iter.as_slice(),
139    )?;
140
141    let mut destination_account_data = destination_account_info.data.borrow_mut();
142    let destination_account =
143        PodStateWithExtensionsMut::<PodAccount>::unpack(&mut destination_account_data)?;
144    if destination_account.base.mint != *mint_account_info.key {
145        return Err(TokenError::MintMismatch.into());
146    }
147    if destination_account.base.is_frozen() {
148        return Err(TokenError::AccountFrozen.into());
149    }
150    let withheld_amount = u64::from(extension.withheld_amount);
151    extension.withheld_amount = 0.into();
152    destination_account.base.amount = u64::from(destination_account.base.amount)
153        .checked_add(withheld_amount)
154        .ok_or(TokenError::Overflow)?
155        .into();
156
157    Ok(())
158}
159
160fn harvest_from_account<'b>(
161    mint_key: &'b Pubkey,
162    token_account_info: &'b AccountInfo<'_>,
163) -> Result<u64, TokenError> {
164    let mut token_account_data = token_account_info.data.borrow_mut();
165    let mut token_account =
166        PodStateWithExtensionsMut::<PodAccount>::unpack(&mut token_account_data)
167            .map_err(|_| TokenError::InvalidState)?;
168    if token_account.base.mint != *mint_key {
169        return Err(TokenError::MintMismatch);
170    }
171    check_program_account(token_account_info.owner).map_err(|_| TokenError::InvalidState)?;
172    let token_account_extension = token_account
173        .get_extension_mut::<TransferFeeAmount>()
174        .map_err(|_| TokenError::InvalidState)?;
175    let account_withheld_amount = u64::from(token_account_extension.withheld_amount);
176    token_account_extension.withheld_amount = 0.into();
177    Ok(account_withheld_amount)
178}
179
180fn process_harvest_withheld_tokens_to_mint(accounts: &[AccountInfo]) -> ProgramResult {
181    let account_info_iter = &mut accounts.iter();
182    let mint_account_info = next_account_info(account_info_iter)?;
183    let token_account_infos = account_info_iter.as_slice();
184
185    let mut mint_data = mint_account_info.data.borrow_mut();
186    let mut mint = PodStateWithExtensionsMut::<PodMint>::unpack(&mut mint_data)?;
187    let mint_extension = mint.get_extension_mut::<TransferFeeConfig>()?;
188
189    for token_account_info in token_account_infos {
190        match harvest_from_account(mint_account_info.key, token_account_info) {
191            Ok(amount) => {
192                let mint_withheld_amount = u64::from(mint_extension.withheld_amount);
193                mint_extension.withheld_amount = mint_withheld_amount
194                    .checked_add(amount)
195                    .ok_or(TokenError::Overflow)?
196                    .into();
197            }
198            Err(e) => {
199                msg!("Error harvesting from {}: {}", token_account_info.key, e);
200            }
201        }
202    }
203    Ok(())
204}
205
206fn process_withdraw_withheld_tokens_from_accounts(
207    program_id: &Pubkey,
208    accounts: &[AccountInfo],
209    num_token_accounts: u8,
210) -> ProgramResult {
211    let account_info_iter = &mut accounts.iter();
212    let mint_account_info = next_account_info(account_info_iter)?;
213    let destination_account_info = next_account_info(account_info_iter)?;
214    let authority_info = next_account_info(account_info_iter)?;
215    let authority_info_data_len = authority_info.data_len();
216    let account_infos = account_info_iter.as_slice();
217    let num_signers = account_infos
218        .len()
219        .saturating_sub(num_token_accounts as usize);
220
221    // unnecessary check, but helps for clarity
222    check_program_account(mint_account_info.owner)?;
223
224    let mint_data = mint_account_info.data.borrow();
225    let mint = PodStateWithExtensions::<PodMint>::unpack(&mint_data)?;
226    let extension = mint.get_extension::<TransferFeeConfig>()?;
227
228    let withdraw_withheld_authority = Option::<Pubkey>::from(extension.withdraw_withheld_authority)
229        .ok_or(TokenError::NoAuthorityExists)?;
230    Processor::validate_owner(
231        program_id,
232        &withdraw_withheld_authority,
233        authority_info,
234        authority_info_data_len,
235        &account_infos[..num_signers],
236    )?;
237
238    let mut destination_account_data = destination_account_info.data.borrow_mut();
239    let mut destination_account =
240        PodStateWithExtensionsMut::<PodAccount>::unpack(&mut destination_account_data)?;
241    if destination_account.base.mint != *mint_account_info.key {
242        return Err(TokenError::MintMismatch.into());
243    }
244    if destination_account.base.is_frozen() {
245        return Err(TokenError::AccountFrozen.into());
246    }
247    for account_info in &account_infos[num_signers..] {
248        // self-harvest, can't double-borrow the underlying data
249        if account_info.key == destination_account_info.key {
250            let token_account_extension = destination_account
251                .get_extension_mut::<TransferFeeAmount>()
252                .map_err(|_| TokenError::InvalidState)?;
253            let account_withheld_amount = u64::from(token_account_extension.withheld_amount);
254            token_account_extension.withheld_amount = 0.into();
255            destination_account.base.amount = u64::from(destination_account.base.amount)
256                .checked_add(account_withheld_amount)
257                .ok_or(TokenError::Overflow)?
258                .into();
259        } else {
260            match harvest_from_account(mint_account_info.key, account_info) {
261                Ok(amount) => {
262                    destination_account.base.amount = u64::from(destination_account.base.amount)
263                        .checked_add(amount)
264                        .ok_or(TokenError::Overflow)?
265                        .into();
266                }
267                Err(e) => {
268                    msg!("Error harvesting from {}: {}", account_info.key, e);
269                }
270            }
271        }
272    }
273
274    Ok(())
275}
276
277pub(crate) fn process_instruction(
278    program_id: &Pubkey,
279    accounts: &[AccountInfo],
280    input: &[u8],
281) -> ProgramResult {
282    let instruction = TransferFeeInstruction::unpack(input)?;
283    check_program_account(program_id)?;
284
285    match instruction {
286        TransferFeeInstruction::InitializeTransferFeeConfig {
287            transfer_fee_config_authority,
288            withdraw_withheld_authority,
289            transfer_fee_basis_points,
290            maximum_fee,
291        } => process_initialize_transfer_fee_config(
292            accounts,
293            transfer_fee_config_authority,
294            withdraw_withheld_authority,
295            transfer_fee_basis_points,
296            maximum_fee,
297        ),
298        TransferFeeInstruction::TransferCheckedWithFee {
299            amount,
300            decimals,
301            fee,
302        } => {
303            msg!("TransferFeeInstruction: TransferCheckedWithFee");
304            Processor::process_transfer(program_id, accounts, amount, Some(decimals), Some(fee))
305        }
306        TransferFeeInstruction::WithdrawWithheldTokensFromMint => {
307            msg!("TransferFeeInstruction: WithdrawWithheldTokensFromMint");
308            process_withdraw_withheld_tokens_from_mint(program_id, accounts)
309        }
310        TransferFeeInstruction::WithdrawWithheldTokensFromAccounts { num_token_accounts } => {
311            msg!("TransferFeeInstruction: WithdrawWithheldTokensFromAccounts");
312            process_withdraw_withheld_tokens_from_accounts(program_id, accounts, num_token_accounts)
313        }
314        TransferFeeInstruction::HarvestWithheldTokensToMint => {
315            msg!("TransferFeeInstruction: HarvestWithheldTokensToMint");
316            process_harvest_withheld_tokens_to_mint(accounts)
317        }
318        TransferFeeInstruction::SetTransferFee {
319            transfer_fee_basis_points,
320            maximum_fee,
321        } => {
322            msg!("TransferFeeInstruction: SetTransferFee");
323            process_set_transfer_fee(program_id, accounts, transfer_fee_basis_points, maximum_fee)
324        }
325    }
326}
spl_token_2022/extension/transfer_fee/processor.rs

spl_token_2022/extension/transfer_fee/
processor.rs
