diff --git a/token/program-2022/src/extension/memo_transfer/mod.rs b/token/program-2022/src/extension/memo_transfer/mod.rs index 633aaa916b7..082333c41af 100644 --- a/token/program-2022/src/extension/memo_transfer/mod.rs +++ b/token/program-2022/src/extension/memo_transfer/mod.rs @@ -25,8 +25,8 @@ impl Extension for MemoTransfer { } /// Determine if a memo is required for transfers into this account -pub fn memo_required(mut account_state: StateWithExtensionsMut) -> bool { - if let Ok(extension) = account_state.get_extension_mut::() { +pub fn memo_required(account_state: &StateWithExtensionsMut) -> bool { + if let Ok(extension) = account_state.get_extension::() { return extension.require_incoming_transfer_memos.into(); } false diff --git a/token/program-2022/src/extension/mod.rs b/token/program-2022/src/extension/mod.rs index 0912b9bd180..8892b105259 100644 --- a/token/program-2022/src/extension/mod.rs +++ b/token/program-2022/src/extension/mod.rs @@ -434,11 +434,30 @@ impl<'data, S: BaseState> StateWithExtensionsMut<'data, S> { } } - /// Unpack a portion of the TLV data as the desired type + /// Unpack a portion of the TLV data as the desired type that allows modifying the type pub fn get_extension_mut(&mut self) -> Result<&mut V, ProgramError> { self.init_or_get_extension(false) } + /// Unpack a portion of the TLV data as the desired type + pub fn get_extension(&self) -> Result<&V, ProgramError> { + if V::TYPE.get_account_type() != S::ACCOUNT_TYPE { + return Err(ProgramError::InvalidAccountData); + } + let TlvIndices { + type_start, + length_start, + value_start, + } = get_extension_indices::(self.tlv_data, false)?; + + if self.tlv_data[type_start..].len() < V::TYPE.get_tlv_len() { + return Err(ProgramError::InvalidAccountData); + } + let length = pod_from_bytes::(&self.tlv_data[length_start..value_start])?; + let value_end = value_start.saturating_add(usize::from(*length)); + pod_from_bytes::(&self.tlv_data[value_start..value_end]) + } + /// Packs base state data into the base data portion pub fn pack_base(&mut self) { S::pack_into_slice(&self.base, self.base_data); diff --git a/token/program-2022/src/processor.rs b/token/program-2022/src/processor.rs index 34809eb848c..0f1c988b91d 100644 --- a/token/program-2022/src/processor.rs +++ b/token/program-2022/src/processor.rs @@ -8,7 +8,7 @@ use { confidential_transfer::{self, ConfidentialTransferAccount}, default_account_state::{self, DefaultAccountState}, immutable_owner::ImmutableOwner, - memo_transfer::{self, MemoTransfer}, + memo_transfer::{self, memo_required}, mint_close_authority::MintCloseAuthority, reallocate, transfer_fee::{self, TransferFeeAmount, TransferFeeConfig}, @@ -373,10 +373,8 @@ impl Processor { return Err(TokenError::MintMismatch.into()); } - if let Ok(extension) = dest_account.get_extension_mut::() { - if extension.require_incoming_transfer_memos.into() { - // TODO: use get_processed_instructions syscall to check for memo - } + if memo_required(&dest_account) { + // TODO: use get_processed_instructions syscall to check for memo } source_account.base.amount = source_account