diff --git a/p-token/src/entrypoint.rs b/p-token/src/entrypoint.rs index 2a483822..9a0f4bd2 100644 --- a/p-token/src/entrypoint.rs +++ b/p-token/src/entrypoint.rs @@ -1,25 +1,235 @@ use { crate::processor::*, + core::{ + mem::{size_of, transmute, MaybeUninit}, + slice::from_raw_parts, + }, pinocchio::{ account_info::AccountInfo, - no_allocator, nostd_panic_handler, program_entrypoint, + entrypoint::{deserialize, NON_DUP_MARKER}, + hint::likely, + log::sol_log, + no_allocator, nostd_panic_handler, program_error::{ProgramError, ToStr}, - pubkey::Pubkey, - ProgramResult, + ProgramResult, MAX_TX_ACCOUNTS, SUCCESS, + }, + pinocchio_token_interface::{ + error::TokenError, + instruction::TokenInstruction, + state::{account::Account, mint::Mint, Transmutable}, }, - pinocchio_token_interface::error::TokenError, }; -program_entrypoint!(process_instruction); // Do not allocate memory. no_allocator!(); // Use the no_std panic handler. nostd_panic_handler!(); +/// Custom program entrypoint to give priority to `transfer` and +/// `transfer_checked` instructions. +/// +/// The entrypoint prioritizes the transfer instruction by validating +/// account data lengths and instruction data. When it can reliably +/// determine that the instruction is a transfer, it will invoke the +/// processor directly. +#[no_mangle] +#[allow(clippy::arithmetic_side_effects)] +pub unsafe extern "C" fn entrypoint(input: *mut u8) -> u64 { + // Constants that apply to both `transfer` and `transfer_checked`. + + /// Offset for the first account. + const ACCOUNT1_HEADER_OFFSET: usize = 0x0008; + + /// Offset for the first account data length. This is + /// expected to be a token account (165 bytes). + const ACCOUNT1_DATA_LEN: usize = 0x0058; + + /// Offset for the second account. + const ACCOUNT2_HEADER_OFFSET: usize = 0x2910; + + /// Offset for the second account data length. This is + /// expected to be a token account for `transfer` (165 bytes) + /// or a mint account for `transfer_checked` (82 bytes). + const ACCOUNT2_DATA_LEN: usize = 0x2960; + + // Constants that apply to `transfer_checked` (instruction 12). + + /// Offset for the third account. + const IX12_ACCOUNT3_HEADER_OFFSET: usize = 0x51c8; + + /// Offset for the third account data length. This is + /// expected to be a token account (165 bytes). + const IX12_ACCOUNT3_DATA_LEN: usize = 0x5218; + + /// Offset for the fourth account. + const IX12_ACCOUNT4_HEADER_OFFSET: usize = 0x7ad0; + + /// Offset for the fourth account data length. + /// + /// This is expected to be an account with variable data + /// length. + const IX12_ACCOUNT4_DATA_LEN: usize = 0x7b20; + + /// Expected offset for the instruction data in the case the + /// fourth (authority) account has zero data. + /// + /// This value is adjusted before it is used. + const IX12_EXPECTED_INSTRUCTION_DATA_LEN_OFFSET: usize = 0xa330; + + // Constants that apply to `transfer` (instruction 3). + + /// Offset for the third account. + /// + /// Note that this assumes that both first and second accounts + /// have zero data, which is being validated before the offset + /// is used. + const IX3_ACCOUNT3_HEADER_OFFSET: usize = 0x5218; + + /// Offset for the third account data length. + /// + /// This is expected to be an account with variable data + /// length. + const IX3_ACCOUNT3_DATA_LEN: usize = 0x5268; + + /// Expected offset for the instruction data in the case the + /// third (authority) account has zero data. + /// + /// This value is adjusted before it is used. + const IX3_INSTRUCTION_DATA_LEN_OFFSET: usize = 0x7a78; + + /// Align an address to the next multiple of 8. + #[inline(always)] + fn align(input: u64) -> u64 { + (input + 7) & (!7) + } + + // Fast path for `transfer_checked`. + // + // It expects 4 accounts: + // 1. source: must be a token account (165 length) + // 2. mint: must be a mint account (82 length) + // 3. destination: must be a token account (165 length) + // 4. authority: can be any account (variable length) + // + // Instruction data is expected to be at least 9 bytes + // and discriminator equal to 12. + if *input == 4 + && (*input.add(ACCOUNT1_DATA_LEN).cast::() == Account::LEN as u64) + && (*input.add(ACCOUNT2_HEADER_OFFSET) == NON_DUP_MARKER) + && (*input.add(ACCOUNT2_DATA_LEN).cast::() == Mint::LEN as u64) + && (*input.add(IX12_ACCOUNT3_HEADER_OFFSET) == NON_DUP_MARKER) + && (*input.add(IX12_ACCOUNT3_DATA_LEN).cast::() == Account::LEN as u64) + && (*input.add(IX12_ACCOUNT4_HEADER_OFFSET) == NON_DUP_MARKER) + { + // The `authority` account can have variable data length. + let account_4_data_len_aligned = + align(*input.add(IX12_ACCOUNT4_DATA_LEN).cast::()) as usize; + let offset = IX12_EXPECTED_INSTRUCTION_DATA_LEN_OFFSET + account_4_data_len_aligned; + + // Check that we have enough instruction data. + // + // Expected: instruction discriminator (u8) + amount (u64) + decimals (u8) + if input.add(offset).cast::().read() >= 10 { + let discriminator = input.add(offset + size_of::()).cast::().read(); + + // Check for transfer discriminator. + if likely(discriminator == TokenInstruction::TransferChecked as u8) { + // instruction data length (u64) + discriminator (u8) + let instruction_data = unsafe { from_raw_parts(input.add(offset + 9), 9) }; + + let accounts = unsafe { + [ + transmute::<*mut u8, AccountInfo>(input.add(ACCOUNT1_HEADER_OFFSET)), + transmute::<*mut u8, AccountInfo>(input.add(ACCOUNT2_HEADER_OFFSET)), + transmute::<*mut u8, AccountInfo>(input.add(IX12_ACCOUNT3_HEADER_OFFSET)), + transmute::<*mut u8, AccountInfo>(input.add(IX12_ACCOUNT4_HEADER_OFFSET)), + ] + }; + + #[cfg(feature = "logging")] + pinocchio::msg!("Instruction: TransferChecked"); + + return match process_transfer_checked(&accounts, instruction_data) { + Ok(()) => SUCCESS, + Err(error) => { + log_error(&error); + error.into() + } + }; + } + } + } + // Fast path for `transfer`. + // + // It expects 3 accounts: + // 1. source: must be a token account (165 length) + // 2. destination: must be a token account (165 length) + // 3. authority: can be any account (variable length) + // + // Instruction data is expected to be at least 8 bytes + // and discriminator equal to 3. + else if *input == 3 + && (*input.add(ACCOUNT1_DATA_LEN).cast::() == Account::LEN as u64) + && (*input.add(ACCOUNT2_HEADER_OFFSET) == NON_DUP_MARKER) + && (*input.add(ACCOUNT2_DATA_LEN).cast::() == Account::LEN as u64) + && (*input.add(IX3_ACCOUNT3_HEADER_OFFSET) == NON_DUP_MARKER) + { + // The `authority` account can have variable data length. + let account_3_data_len_aligned = + align(*input.add(IX3_ACCOUNT3_DATA_LEN).cast::()) as usize; + let offset = IX3_INSTRUCTION_DATA_LEN_OFFSET + account_3_data_len_aligned; + + // Check that we have enough instruction data. + if likely(input.add(offset).cast::().read() >= 9) { + let discriminator = input.add(offset + size_of::()).cast::().read(); + + // Check for transfer discriminator. + if likely(discriminator == TokenInstruction::Transfer as u8) { + let instruction_data = + unsafe { from_raw_parts(input.add(offset + 9), size_of::()) }; + + let accounts = unsafe { + [ + transmute::<*mut u8, AccountInfo>(input.add(ACCOUNT1_HEADER_OFFSET)), + transmute::<*mut u8, AccountInfo>(input.add(ACCOUNT2_HEADER_OFFSET)), + transmute::<*mut u8, AccountInfo>(input.add(IX3_ACCOUNT3_HEADER_OFFSET)), + ] + }; + + #[cfg(feature = "logging")] + pinocchio::msg!("Instruction: Transfer"); + + return match process_transfer(&accounts, instruction_data) { + Ok(()) => SUCCESS, + Err(error) => { + log_error(&error); + error.into() + } + }; + } + } + } + + // Entrypoint for the remaining instructions. + + const UNINIT: MaybeUninit = MaybeUninit::::uninit(); + let mut accounts = [UNINIT; { MAX_TX_ACCOUNTS }]; + + let (_, count, instruction_data) = deserialize(input, &mut accounts); + + match process_instruction( + from_raw_parts(accounts.as_ptr() as _, count), + instruction_data, + ) { + Ok(()) => SUCCESS, + Err(error) => error.into(), + } +} + /// Log an error. #[cold] fn log_error(error: &ProgramError) { - pinocchio::log::sol_log(error.to_str::()); + sol_log(error.to_str::()); } /// Process an instruction. @@ -30,11 +240,7 @@ fn log_error(error: &ProgramError) { /// instructions, since it is not sound to have a "batch" instruction inside /// another "batch" instruction. #[inline(always)] -pub fn process_instruction( - _program_id: &Pubkey, - accounts: &[AccountInfo], - instruction_data: &[u8], -) -> ProgramResult { +pub fn process_instruction(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult { let [discriminator, remaining @ ..] = instruction_data else { return Err(TokenError::InvalidInstruction.into()); }; @@ -138,12 +344,12 @@ pub(crate) fn inner_process_instruction( process_burn_checked(accounts, instruction_data) } - // 16 - InitializeAccount2 - 16 => { + // 17 - SyncNative + 17 => { #[cfg(feature = "logging")] - pinocchio::msg!("Instruction: InitializeAccount2"); + pinocchio::msg!("Instruction: SyncNative"); - process_initialize_account2(accounts, instruction_data) + process_sync_native(accounts) } // 18 - InitializeAccount3 18 => { @@ -159,6 +365,13 @@ pub(crate) fn inner_process_instruction( process_initialize_mint2(accounts, instruction_data) } + // 22 - InitializeImmutableOwner + 22 => { + #[cfg(feature = "logging")] + pinocchio::msg!("Instruction: InitializeImmutableOwner"); + + process_initialize_immutable_owner(accounts) + } d => inner_process_remaining_instruction(accounts, instruction_data, d), } } @@ -231,12 +444,12 @@ fn inner_process_remaining_instruction( process_mint_to_checked(accounts, instruction_data) } - // 17 - SyncNative - 17 => { + // 16 - InitializeAccount2 + 16 => { #[cfg(feature = "logging")] - pinocchio::msg!("Instruction: SyncNative"); + pinocchio::msg!("Instruction: InitializeAccount2"); - process_sync_native(accounts) + process_initialize_account2(accounts, instruction_data) } // 19 - InitializeMultisig2 19 => { @@ -252,13 +465,6 @@ fn inner_process_remaining_instruction( process_get_account_data_size(accounts) } - // 22 - InitializeImmutableOwner - 22 => { - #[cfg(feature = "logging")] - pinocchio::msg!("Instruction: InitializeImmutableOwner"); - - process_initialize_immutable_owner(accounts) - } // 23 - AmountToUiAmount 23 => { #[cfg(feature = "logging")]