token-2022: Add TLV extension state infrastructure and initial extension structs

programs/token-2022/src/state/account_state.rs

@@ -1,3 +1,5 @@
+use solana_program_error::ProgramError;
+
 #[repr(u8)]
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub enum AccountState {
@@ -14,13 +16,13 @@ pub enum AccountState {
     Frozen,
 }
 
-impl From<u8> for AccountState {
-    fn from(value: u8) -> Self {
+impl TryFrom<u8> for AccountState {
+    type Error = ProgramError;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
         match value {
-            0 => AccountState::Uninitialized,
-            1 => AccountState::Initialized,
-            2 => AccountState::Frozen,
-            _ => panic!("invalid account state value: {value}"),
+            0..=2 => Ok(unsafe { core::mem::transmute::<u8, AccountState>(value) }),
+            _ => Err(ProgramError::InvalidAccountData),
         }
     }
 }

programs/token-2022/src/state/account_type.rs

@@ -0,0 +1,7 @@
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum AccountType {
+    Uninitialized,
+    Mint,
+    Account,
+}

programs/token-2022/src/state/extension/default_account_state.rs

@@ -0,0 +1,38 @@
+use {
+    super::{sealed, ExtensionPod, ExtensionType, ExtensionValue},
+    crate::state::AccountState,
+    solana_program_error::ProgramError,
+};
+
+/// Default account state extension data (1 byte).
+///
+/// When set on a mint, all new token accounts are initialized
+/// with this state.
+#[repr(C)]
+#[derive(Clone, Copy, Debug, Default, PartialEq)]
+pub struct DefaultAccountStateExtension {
+    state: u8,
+}
+
+impl DefaultAccountStateExtension {
+    pub const LEN: usize = core::mem::size_of::<DefaultAccountStateExtension>();
+
+    #[inline(always)]
+    pub fn state(&self) -> Result<AccountState, ProgramError> {
+        AccountState::try_from(self.state)
+    }
+
+    #[inline(always)]
+    pub fn set_state(&mut self, state: AccountState) {
+        self.state = state as u8;
+    }
+}
+
+// SAFETY: `DefaultAccountStateExtension` is repr(C), contains only `u8`,
+// has no padding, and all bit patterns are valid.
+impl sealed::SealedExtensionPod for DefaultAccountStateExtension {}
+unsafe impl ExtensionPod for DefaultAccountStateExtension {}
+
+impl ExtensionValue for DefaultAccountStateExtension {
+    const TYPE: ExtensionType = ExtensionType::DefaultAccountState;
+}

programs/token-2022/src/state/extension/mod.rs

@@ -0,0 +1,279 @@
+pub mod default_account_state;
+pub mod permanent_delegate;
+mod state;
+pub mod transfer_hook;
+pub mod transfer_hook_account;
+
+use {
+    super::{AccountType, Mint, Multisig},
+    solana_program_error::ProgramError,
+};
+pub use {
+    default_account_state::DefaultAccountStateExtension,
+    permanent_delegate::PermanentDelegateExtension,
+    state::{StateWithExtensions, StateWithExtensionsMut},
+    transfer_hook::TransferHookExtension,
+    transfer_hook_account::TransferHookAccountExtension,
+};
+
+/// Maximum number of distinct extension types (excluding `Uninitialized`).
+///
+/// Useful for pre-allocating the output buffer when calling
+/// `collect_extension_types_from_tlv` or the wrapper `get_extension_types`
+/// methods.
+pub const MAX_EXTENSIONS: usize = 28;
+
+const BASE_ACCOUNT_LEN: usize = 165;
+const MINT_PADDING_LEN: usize = BASE_ACCOUNT_LEN - Mint::BASE_LEN;
+const ZERO_MINT_PADDING: [u8; MINT_PADDING_LEN] = [0u8; MINT_PADDING_LEN];
+const ACCOUNT_TYPE_INDEX: usize = BASE_ACCOUNT_LEN;
+const TLV_START_INDEX: usize = ACCOUNT_TYPE_INDEX + 1;
+const TLV_HEADER_LEN: usize = 4;
+
+/// Token-2022 error discriminants used by this state module.
+///
+/// Keep these values aligned with SPL Token-2022's `TokenError`.
+#[repr(u32)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum TokenError {
+    /// Extension not found in account data.
+    ExtensionNotFound = 48,
+}
+
+impl From<TokenError> for ProgramError {
+    #[inline(always)]
+    fn from(error: TokenError) -> Self {
+        ProgramError::Custom(error as u32)
+    }
+}
+
+/// SPL Token-2022 `TokenError::ExtensionNotFound` discriminant.
+pub const EXTENSION_NOT_FOUND_ERROR_CODE: u32 = TokenError::ExtensionNotFound as u32;
+
+mod sealed {
+    pub trait SealedExtensionPod {}
+}
+
+/// Marker trait for plain extension payload values that are safe to
+/// reinterpret from bytes.
+///
+/// # Safety
+///
+/// Implementers must be plain data with no padding and no invalid bit-patterns.
+pub unsafe trait ExtensionPod: sealed::SealedExtensionPod + 'static {}
+
+#[cold]
+pub const fn extension_not_found_error() -> ProgramError {
+    ProgramError::Custom(TokenError::ExtensionNotFound as u32)
+}
+
+#[inline(always)]
+pub fn is_extension_not_found_error(error: &ProgramError) -> bool {
+    matches!(
+        error,
+        ProgramError::Custom(code) if *code == EXTENSION_NOT_FOUND_ERROR_CODE
+    )
+}
+
+#[repr(u16)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+#[non_exhaustive]
+pub enum ExtensionType {
+    Uninitialized = 0,
+    TransferFeeConfig = 1,
+    TransferFeeAmount = 2,
+    MintCloseAuthority = 3,
+    ConfidentialTransferMint = 4,
+    ConfidentialTransferAccount = 5,
+    DefaultAccountState = 6,
+    ImmutableOwner = 7,
+    MemoTransfer = 8,
+    NonTransferable = 9,
+    InterestBearingConfig = 10,
+    CpiGuard = 11,
+    PermanentDelegate = 12,
+    NonTransferableAccount = 13,
+    TransferHook = 14,
+    TransferHookAccount = 15,
+    ConfidentialTransferFeeConfig = 16,
+    ConfidentialTransferFeeAmount = 17,
+    MetadataPointer = 18,
+    TokenMetadata = 19,
+    GroupPointer = 20,
+    TokenGroup = 21,
+    GroupMemberPointer = 22,
+    TokenGroupMember = 23,
+    ConfidentialMintBurn = 24,
+    ScaledUiAmount = 25,
+    Pausable = 26,
+    PausableAccount = 27,
+}
+
+/// Marker for typed extension values that can be decoded from TLV entries.
+pub trait ExtensionValue: ExtensionPod {
+    const TYPE: ExtensionType;
+}
+
+/// Trait for supported token-2022 base account types that can host TLV
+/// extensions.
+pub trait ExtensionBaseState: Sized {
+    const BASE_LEN: usize;
+    const ACCOUNT_TYPE: AccountType;
+
+    fn validate_extensions_data(data: &[u8]) -> Result<(), ProgramError>;
+
+    /// # Safety
+    ///
+    /// The caller must ensure the provided bytes contain a valid base state.
+    unsafe fn from_bytes_unchecked(data: &[u8]) -> &Self;
+
+    /// # Safety
+    ///
+    /// The caller must ensure the provided bytes contain a valid mutable base
+    /// state.
+    unsafe fn from_bytes_unchecked_mut(data: &mut [u8]) -> &mut Self;
+}
+
+#[inline(always)]
+const fn extension_account_type(extension_type: ExtensionType) -> AccountType {
+    match extension_type {
+        ExtensionType::Uninitialized => AccountType::Uninitialized,
+        ExtensionType::TransferFeeConfig
+        | ExtensionType::MintCloseAuthority
+        | ExtensionType::ConfidentialTransferMint
+        | ExtensionType::DefaultAccountState
+        | ExtensionType::NonTransferable
+        | ExtensionType::InterestBearingConfig
+        | ExtensionType::PermanentDelegate
+        | ExtensionType::TransferHook
+        | ExtensionType::ConfidentialTransferFeeConfig
+        | ExtensionType::MetadataPointer
+        | ExtensionType::TokenMetadata
+        | ExtensionType::GroupPointer
+        | ExtensionType::TokenGroup
+        | ExtensionType::GroupMemberPointer
+        | ExtensionType::TokenGroupMember
+        | ExtensionType::ConfidentialMintBurn
+        | ExtensionType::ScaledUiAmount
+        | ExtensionType::Pausable => AccountType::Mint,
+        ExtensionType::TransferFeeAmount
+        | ExtensionType::ConfidentialTransferAccount
+        | ExtensionType::ImmutableOwner
+        | ExtensionType::MemoTransfer
+        | ExtensionType::CpiGuard
+        | ExtensionType::NonTransferableAccount
+        | ExtensionType::TransferHookAccount
+        | ExtensionType::ConfidentialTransferFeeAmount
+        | ExtensionType::PausableAccount => AccountType::Account,
+    }
+}
+
+#[inline(always)]
+fn extension_type_from_u16(extension_type: u16) -> Result<ExtensionType, ProgramError> {
+    match extension_type {
+        0..=27 => Ok(unsafe { core::mem::transmute::<u16, ExtensionType>(extension_type) }),
+        _ => Err(ProgramError::InvalidAccountData),
+    }
+}
+
+#[inline(always)]
+const fn adjust_len_for_multisig(account_len: usize) -> usize {
+    if account_len == Multisig::LEN {
+        account_len.saturating_add(core::mem::size_of::<ExtensionType>())
+    } else {
+        account_len
+    }
+}
+
+#[inline(always)]
+fn validate_extension_account_type(
+    extension_type: ExtensionType,
+    expected: AccountType,
+) -> Result<(), ProgramError> {
+    if extension_account_type(extension_type) != expected {
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    Ok(())
+}
+
+#[inline]
+fn validate_mint_extensions_data(data: &[u8]) -> Result<(), ProgramError> {
+    if data.len() <= ACCOUNT_TYPE_INDEX {
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    if data[Mint::BASE_LEN..BASE_ACCOUNT_LEN] != ZERO_MINT_PADDING {
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    if data[ACCOUNT_TYPE_INDEX] != AccountType::Mint as u8 {
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    Ok(())
+}
+
+#[inline]
+fn validate_token_extensions_data(data: &[u8]) -> Result<(), ProgramError> {
+    if data.len() <= ACCOUNT_TYPE_INDEX {
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    if data[ACCOUNT_TYPE_INDEX] != AccountType::Account as u8 {
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    Ok(())
+}
+
+/// Returns the fixed byte length of the given extension's value payload,
+/// or `None` if the extension type is not yet supported.
+///
+/// Only a subset of extension types have known sizes registered here.
+/// Unsupported types will cause [`try_calculate_account_len`] to fail.
+#[inline(always)]
+pub const fn extension_value_len(extension_type: ExtensionType) -> Option<usize> {
+    match extension_type {
+        ExtensionType::DefaultAccountState => Some(DefaultAccountStateExtension::LEN),
+        ExtensionType::PermanentDelegate => Some(PermanentDelegateExtension::LEN),
+        ExtensionType::TransferHook => Some(TransferHookExtension::LEN),
+        ExtensionType::TransferHookAccount => Some(TransferHookAccountExtension::LEN),
+        _ => None,
+    }
+}
+
+#[inline]
+pub fn try_calculate_account_len<B: ExtensionBaseState>(
+    extension_types: &[ExtensionType],
+) -> Result<usize, ProgramError> {
+    if extension_types.is_empty() {
+        return Ok(B::BASE_LEN);
+    }
+
+    let mut total_len = TLV_START_INDEX;
+    let mut i = 0;
+
+    while i < extension_types.len() {
+        let extension_type = extension_types[i];
+        validate_extension_account_type(extension_type, B::ACCOUNT_TYPE)?;
+
+        let mut j = 0;
+        while j < i {
+            if extension_types[j] == extension_type {
+                return Err(ProgramError::InvalidInstructionData);
+            }
+            j += 1;
+        }
+
+        let value_len =
+            extension_value_len(extension_type).ok_or(ProgramError::InvalidInstructionData)?;
+
+        total_len = total_len
+            .checked_add(TLV_HEADER_LEN + value_len)
+            .ok_or(ProgramError::InvalidInstructionData)?;
+        i += 1;
+    }
+
+    Ok(adjust_len_for_multisig(total_len))
+}

programs/token-2022/src/state/extension/permanent_delegate.rs

@@ -0,0 +1,37 @@
+use {
+    super::{sealed, ExtensionPod, ExtensionType, ExtensionValue},
+    solana_address::Address,
+};
+
+/// Permanent delegate extension data for mints (32 bytes).
+///
+/// When set on a mint, the delegate has unrestricted transfer and
+/// burn authority over all token accounts for the mint.
+#[repr(C)]
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct PermanentDelegateExtension {
+    delegate: Address,
+}
+
+impl PermanentDelegateExtension {
+    pub const LEN: usize = core::mem::size_of::<PermanentDelegateExtension>();
+
+    #[inline(always)]
+    pub fn delegate(&self) -> &Address {
+        &self.delegate
+    }
+
+    #[inline(always)]
+    pub fn set_delegate(&mut self, delegate: &Address) {
+        self.delegate = delegate.clone();
+    }
+}
+
+// SAFETY: `PermanentDelegateExtension` is repr(C), contains only `Address`
+// (`[u8; 32]`), has no padding, and all bit patterns are valid.
+impl sealed::SealedExtensionPod for PermanentDelegateExtension {}
+unsafe impl ExtensionPod for PermanentDelegateExtension {}
+
+impl ExtensionValue for PermanentDelegateExtension {
+    const TYPE: ExtensionType = ExtensionType::PermanentDelegate;
+}