1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

#[cfg(feature = "serde-traits")]
use serde::{Deserialize, Serialize};
use {
    crate::{
        extension::{
            BaseState, BaseStateWithExtensions, Extension, ExtensionType, StateWithExtensionsMut,
        },
        state::Account,
    },
    bytemuck::{Pod, Zeroable},
    solana_program::{account_info::AccountInfo, program_error::ProgramError, pubkey::Pubkey},
    spl_pod::{optional_keys::OptionalNonZeroPubkey, primitives::PodBool},
};

/// Instructions for the TransferHook extension
pub mod instruction;
/// Instruction processor for the TransferHook extension
pub mod processor;

/// Transfer hook extension data for mints.
#[repr(C)]
#[cfg_attr(feature = "serde-traits", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde-traits", serde(rename_all = "camelCase"))]
#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
pub struct TransferHook {
    /// Authority that can set the transfer hook program id
    pub authority: OptionalNonZeroPubkey,
    /// Program that authorizes the transfer
    pub program_id: OptionalNonZeroPubkey,
}

/// Indicates that the tokens from this account belong to a mint with a transfer
/// hook
#[cfg_attr(feature = "serde-traits", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde-traits", serde(rename_all = "camelCase"))]
#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
#[repr(transparent)]
pub struct TransferHookAccount {
    /// Flag to indicate that the account is in the middle of a transfer
    pub transferring: PodBool,
}

impl Extension for TransferHook {
    const TYPE: ExtensionType = ExtensionType::TransferHook;
}

impl Extension for TransferHookAccount {
    const TYPE: ExtensionType = ExtensionType::TransferHookAccount;
}

/// Attempts to get the transfer hook program id from the TLV data, returning
/// None if the extension is not found
pub fn get_program_id<S: BaseState, BSE: BaseStateWithExtensions<S>>(
    state: &BSE,
) -> Option<Pubkey> {
    state
        .get_extension::<TransferHook>()
        .ok()
        .and_then(|e| Option::<Pubkey>::from(e.program_id))
}

/// Helper function to set the transferring flag before calling into transfer
/// hook
pub fn set_transferring(account: &mut StateWithExtensionsMut<Account>) -> Result<(), ProgramError> {
    let account_extension = account.get_extension_mut::<TransferHookAccount>()?;
    account_extension.transferring = true.into();
    Ok(())
}

/// Helper function to unset the transferring flag after a transfer
pub fn unset_transferring(account_info: &AccountInfo) -> Result<(), ProgramError> {
    let mut account_data = account_info.data.borrow_mut();
    let mut account = StateWithExtensionsMut::<Account>::unpack(&mut account_data)?;
    let account_extension = account.get_extension_mut::<TransferHookAccount>()?;
    account_extension.transferring = false.into();
    Ok(())
}