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

safe_token_2022/extension/transfer_fee/
processor.rs
