spl_token_2022/extension/
mod.rs

1//! Extensions available to token mints and accounts
2
3#[cfg(feature = "serde-traits")]
4use serde::{Deserialize, Serialize};
5use {
6    crate::{
7        error::TokenError,
8        extension::{
9            confidential_mint_burn::ConfidentialMintBurn,
10            confidential_transfer::{ConfidentialTransferAccount, ConfidentialTransferMint},
11            confidential_transfer_fee::{
12                ConfidentialTransferFeeAmount, ConfidentialTransferFeeConfig,
13            },
14            cpi_guard::CpiGuard,
15            default_account_state::DefaultAccountState,
16            group_member_pointer::GroupMemberPointer,
17            group_pointer::GroupPointer,
18            immutable_owner::ImmutableOwner,
19            interest_bearing_mint::InterestBearingConfig,
20            memo_transfer::MemoTransfer,
21            metadata_pointer::MetadataPointer,
22            mint_close_authority::MintCloseAuthority,
23            non_transferable::{NonTransferable, NonTransferableAccount},
24            pausable::{PausableAccount, PausableConfig},
25            permanent_delegate::PermanentDelegate,
26            scaled_ui_amount::ScaledUiAmountConfig,
27            transfer_fee::{TransferFeeAmount, TransferFeeConfig},
28            transfer_hook::{TransferHook, TransferHookAccount},
29        },
30        pod::{PodAccount, PodMint},
31        state::{Account, Mint, Multisig, PackedSizeOf},
32    },
33    bytemuck::{Pod, Zeroable},
34    num_enum::{IntoPrimitive, TryFromPrimitive},
35    solana_program::{
36        account_info::AccountInfo,
37        program_error::ProgramError,
38        program_pack::{IsInitialized, Pack},
39    },
40    spl_pod::{
41        bytemuck::{pod_from_bytes, pod_from_bytes_mut, pod_get_packed_len},
42        primitives::PodU16,
43    },
44    spl_token_group_interface::state::{TokenGroup, TokenGroupMember},
45    spl_type_length_value::variable_len_pack::VariableLenPack,
46    std::{
47        cmp::Ordering,
48        convert::{TryFrom, TryInto},
49        mem::size_of,
50    },
51};
52
53/// Confidential Transfer extension
54pub mod confidential_transfer;
55/// Confidential Transfer Fee extension
56pub mod confidential_transfer_fee;
57/// CPI Guard extension
58pub mod cpi_guard;
59/// Default Account State extension
60pub mod default_account_state;
61/// Group Member Pointer extension
62pub mod group_member_pointer;
63/// Group Pointer extension
64pub mod group_pointer;
65/// Immutable Owner extension
66pub mod immutable_owner;
67/// Interest-Bearing Mint extension
68pub mod interest_bearing_mint;
69/// Memo Transfer extension
70pub mod memo_transfer;
71/// Metadata Pointer extension
72pub mod metadata_pointer;
73/// Mint Close Authority extension
74pub mod mint_close_authority;
75/// Non Transferable extension
76pub mod non_transferable;
77/// Pausable extension
78pub mod pausable;
79/// Permanent Delegate extension
80pub mod permanent_delegate;
81/// Utility to reallocate token accounts
82pub mod reallocate;
83/// Scaled UI Amount extension
84pub mod scaled_ui_amount;
85/// Token-group extension
86pub mod token_group;
87/// Token-metadata extension
88pub mod token_metadata;
89/// Transfer Fee extension
90pub mod transfer_fee;
91/// Transfer Hook extension
92pub mod transfer_hook;
93
94/// Confidential mint-burn extension
95pub mod confidential_mint_burn;
96
97/// Length in TLV structure
98#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
99#[repr(transparent)]
100pub struct Length(PodU16);
101impl From<Length> for usize {
102    fn from(n: Length) -> Self {
103        Self::from(u16::from(n.0))
104    }
105}
106impl TryFrom<usize> for Length {
107    type Error = ProgramError;
108    fn try_from(n: usize) -> Result<Self, Self::Error> {
109        u16::try_from(n)
110            .map(|v| Self(PodU16::from(v)))
111            .map_err(|_| ProgramError::AccountDataTooSmall)
112    }
113}
114
115/// Helper function to get the current `TlvIndices` from the current spot
116fn get_tlv_indices(type_start: usize) -> TlvIndices {
117    let length_start = type_start.saturating_add(size_of::<ExtensionType>());
118    let value_start = length_start.saturating_add(pod_get_packed_len::<Length>());
119    TlvIndices {
120        type_start,
121        length_start,
122        value_start,
123    }
124}
125
126/// Helper function to tack on the size of an extension bytes if an account with
127/// extensions is exactly the size of a multisig
128const fn adjust_len_for_multisig(account_len: usize) -> usize {
129    if account_len == Multisig::LEN {
130        account_len.saturating_add(size_of::<ExtensionType>())
131    } else {
132        account_len
133    }
134}
135
136/// Helper function to calculate exactly how many bytes a value will take up,
137/// given the value's length
138const fn add_type_and_length_to_len(value_len: usize) -> usize {
139    value_len
140        .saturating_add(size_of::<ExtensionType>())
141        .saturating_add(pod_get_packed_len::<Length>())
142}
143
144/// Helper struct for returning the indices of the type, length, and value in
145/// a TLV entry
146#[derive(Debug)]
147struct TlvIndices {
148    pub type_start: usize,
149    pub length_start: usize,
150    pub value_start: usize,
151}
152fn get_extension_indices<V: Extension>(
153    tlv_data: &[u8],
154    init: bool,
155) -> Result<TlvIndices, ProgramError> {
156    let mut start_index = 0;
157    let v_account_type = V::TYPE.get_account_type();
158    while start_index < tlv_data.len() {
159        let tlv_indices = get_tlv_indices(start_index);
160        if tlv_data.len() < tlv_indices.value_start {
161            return Err(ProgramError::InvalidAccountData);
162        }
163        let extension_type =
164            ExtensionType::try_from(&tlv_data[tlv_indices.type_start..tlv_indices.length_start])?;
165        let account_type = extension_type.get_account_type();
166        if extension_type == V::TYPE {
167            // found an instance of the extension that we're initializing, return!
168            return Ok(tlv_indices);
169        // got to an empty spot, init here, or error if we're searching, since
170        // nothing is written after an Uninitialized spot
171        } else if extension_type == ExtensionType::Uninitialized {
172            if init {
173                return Ok(tlv_indices);
174            } else {
175                return Err(TokenError::ExtensionNotFound.into());
176            }
177        } else if v_account_type != account_type {
178            return Err(TokenError::ExtensionTypeMismatch.into());
179        } else {
180            let length = pod_from_bytes::<Length>(
181                &tlv_data[tlv_indices.length_start..tlv_indices.value_start],
182            )?;
183            let value_end_index = tlv_indices.value_start.saturating_add(usize::from(*length));
184            start_index = value_end_index;
185        }
186    }
187    Err(ProgramError::InvalidAccountData)
188}
189
190/// Basic information about the TLV buffer, collected from iterating through all
191/// entries
192#[derive(Debug, PartialEq)]
193struct TlvDataInfo {
194    /// The extension types written in the TLV buffer
195    extension_types: Vec<ExtensionType>,
196    /// The total number bytes allocated for all TLV entries.
197    ///
198    /// Each TLV entry's allocated bytes comprises two bytes for the `type`, two
199    /// bytes for the `length`, and `length` number of bytes for the `value`.
200    used_len: usize,
201}
202
203/// Fetches basic information about the TLV buffer by iterating through all
204/// TLV entries.
205fn get_tlv_data_info(tlv_data: &[u8]) -> Result<TlvDataInfo, ProgramError> {
206    let mut extension_types = vec![];
207    let mut start_index = 0;
208    while start_index < tlv_data.len() {
209        let tlv_indices = get_tlv_indices(start_index);
210        if tlv_data.len() < tlv_indices.length_start {
211            // There aren't enough bytes to store the next type, which means we
212            // got to the end. The last byte could be used during a realloc!
213            return Ok(TlvDataInfo {
214                extension_types,
215                used_len: tlv_indices.type_start,
216            });
217        }
218        let extension_type =
219            ExtensionType::try_from(&tlv_data[tlv_indices.type_start..tlv_indices.length_start])?;
220        if extension_type == ExtensionType::Uninitialized {
221            return Ok(TlvDataInfo {
222                extension_types,
223                used_len: tlv_indices.type_start,
224            });
225        } else {
226            if tlv_data.len() < tlv_indices.value_start {
227                // not enough bytes to store the length, malformed
228                return Err(ProgramError::InvalidAccountData);
229            }
230            extension_types.push(extension_type);
231            let length = pod_from_bytes::<Length>(
232                &tlv_data[tlv_indices.length_start..tlv_indices.value_start],
233            )?;
234
235            let value_end_index = tlv_indices.value_start.saturating_add(usize::from(*length));
236            if value_end_index > tlv_data.len() {
237                // value blows past the size of the slice, malformed
238                return Err(ProgramError::InvalidAccountData);
239            }
240            start_index = value_end_index;
241        }
242    }
243    Ok(TlvDataInfo {
244        extension_types,
245        used_len: start_index,
246    })
247}
248
249fn get_first_extension_type(tlv_data: &[u8]) -> Result<Option<ExtensionType>, ProgramError> {
250    if tlv_data.is_empty() {
251        Ok(None)
252    } else {
253        let tlv_indices = get_tlv_indices(0);
254        if tlv_data.len() <= tlv_indices.length_start {
255            return Ok(None);
256        }
257        let extension_type =
258            ExtensionType::try_from(&tlv_data[tlv_indices.type_start..tlv_indices.length_start])?;
259        if extension_type == ExtensionType::Uninitialized {
260            Ok(None)
261        } else {
262            Ok(Some(extension_type))
263        }
264    }
265}
266
267fn check_min_len_and_not_multisig(input: &[u8], minimum_len: usize) -> Result<(), ProgramError> {
268    if input.len() == Multisig::LEN || input.len() < minimum_len {
269        Err(ProgramError::InvalidAccountData)
270    } else {
271        Ok(())
272    }
273}
274
275fn check_account_type<S: BaseState>(account_type: AccountType) -> Result<(), ProgramError> {
276    if account_type != S::ACCOUNT_TYPE {
277        Err(ProgramError::InvalidAccountData)
278    } else {
279        Ok(())
280    }
281}
282
283/// Any account with extensions must be at least `Account::LEN`.  Both mints and
284/// accounts can have extensions
285/// A mint with extensions that takes it past 165 could be indiscernible from an
286/// Account with an extension, even if we add the account type. For example,
287/// let's say we have:
288///
289/// ```text
290/// Account: 165 bytes... + [2, 0, 3, 0, 100, ....]
291///                          ^     ^       ^     ^
292///                     acct type  extension length data...
293///
294/// Mint: 82 bytes... + 83 bytes of other extension data
295///     + [2, 0, 3, 0, 100, ....]
296///      (data in extension just happens to look like this)
297/// ```
298///
299/// With this approach, we only start writing the TLV data after `Account::LEN`,
300/// which means we always know that the account type is going to be right after
301/// that. We do a special case checking for a Multisig length, because those
302/// aren't extensible under any circumstances.
303const BASE_ACCOUNT_LENGTH: usize = Account::LEN;
304/// Helper that tacks on the `AccountType` length, which gives the minimum for
305/// any account with extensions
306const BASE_ACCOUNT_AND_TYPE_LENGTH: usize = BASE_ACCOUNT_LENGTH + size_of::<AccountType>();
307
308fn type_and_tlv_indices<S: BaseState>(
309    rest_input: &[u8],
310) -> Result<Option<(usize, usize)>, ProgramError> {
311    if rest_input.is_empty() {
312        Ok(None)
313    } else {
314        let account_type_index = BASE_ACCOUNT_LENGTH.saturating_sub(S::SIZE_OF);
315        // check padding is all zeroes
316        let tlv_start_index = account_type_index.saturating_add(size_of::<AccountType>());
317        if rest_input.len() <= tlv_start_index {
318            return Err(ProgramError::InvalidAccountData);
319        }
320        if rest_input[..account_type_index] != vec![0; account_type_index] {
321            Err(ProgramError::InvalidAccountData)
322        } else {
323            Ok(Some((account_type_index, tlv_start_index)))
324        }
325    }
326}
327
328/// Checks a base buffer to verify if it is an Account without having to
329/// completely deserialize it
330fn is_initialized_account(input: &[u8]) -> Result<bool, ProgramError> {
331    const ACCOUNT_INITIALIZED_INDEX: usize = 108; // See state.rs#L99
332
333    if input.len() != BASE_ACCOUNT_LENGTH {
334        return Err(ProgramError::InvalidAccountData);
335    }
336    Ok(input[ACCOUNT_INITIALIZED_INDEX] != 0)
337}
338
339fn get_extension_bytes<S: BaseState, V: Extension>(tlv_data: &[u8]) -> Result<&[u8], ProgramError> {
340    if V::TYPE.get_account_type() != S::ACCOUNT_TYPE {
341        return Err(ProgramError::InvalidAccountData);
342    }
343    let TlvIndices {
344        type_start: _,
345        length_start,
346        value_start,
347    } = get_extension_indices::<V>(tlv_data, false)?;
348    // get_extension_indices has checked that tlv_data is long enough to include
349    // these indices
350    let length = pod_from_bytes::<Length>(&tlv_data[length_start..value_start])?;
351    let value_end = value_start.saturating_add(usize::from(*length));
352    if tlv_data.len() < value_end {
353        return Err(ProgramError::InvalidAccountData);
354    }
355    Ok(&tlv_data[value_start..value_end])
356}
357
358fn get_extension_bytes_mut<S: BaseState, V: Extension>(
359    tlv_data: &mut [u8],
360) -> Result<&mut [u8], ProgramError> {
361    if V::TYPE.get_account_type() != S::ACCOUNT_TYPE {
362        return Err(ProgramError::InvalidAccountData);
363    }
364    let TlvIndices {
365        type_start: _,
366        length_start,
367        value_start,
368    } = get_extension_indices::<V>(tlv_data, false)?;
369    // get_extension_indices has checked that tlv_data is long enough to include
370    // these indices
371    let length = pod_from_bytes::<Length>(&tlv_data[length_start..value_start])?;
372    let value_end = value_start.saturating_add(usize::from(*length));
373    if tlv_data.len() < value_end {
374        return Err(ProgramError::InvalidAccountData);
375    }
376    Ok(&mut tlv_data[value_start..value_end])
377}
378
379/// Calculate the new expected size if the state allocates the given number
380/// of bytes for the given extension type.
381///
382/// Provides the correct answer regardless if the extension is already present
383/// in the TLV data.
384fn try_get_new_account_len_for_extension_len<S: BaseState, V: Extension>(
385    tlv_data: &[u8],
386    new_extension_len: usize,
387) -> Result<usize, ProgramError> {
388    // get the new length used by the extension
389    let new_extension_tlv_len = add_type_and_length_to_len(new_extension_len);
390    let tlv_info = get_tlv_data_info(tlv_data)?;
391    // If we're adding an extension, then we must have at least BASE_ACCOUNT_LENGTH
392    // and account type
393    let current_len = tlv_info
394        .used_len
395        .saturating_add(BASE_ACCOUNT_AND_TYPE_LENGTH);
396    // get the current length used by the extension
397    let current_extension_len = get_extension_bytes::<S, V>(tlv_data)
398        .map(|x| add_type_and_length_to_len(x.len()))
399        .unwrap_or(0);
400    let new_len = current_len
401        .saturating_sub(current_extension_len)
402        .saturating_add(new_extension_tlv_len);
403    Ok(adjust_len_for_multisig(new_len))
404}
405
406/// Trait for base state with extension
407pub trait BaseStateWithExtensions<S: BaseState> {
408    /// Get the buffer containing all extension data
409    fn get_tlv_data(&self) -> &[u8];
410
411    /// Fetch the bytes for a TLV entry
412    fn get_extension_bytes<V: Extension>(&self) -> Result<&[u8], ProgramError> {
413        get_extension_bytes::<S, V>(self.get_tlv_data())
414    }
415
416    /// Unpack a portion of the TLV data as the desired type
417    fn get_extension<V: Extension + Pod>(&self) -> Result<&V, ProgramError> {
418        pod_from_bytes::<V>(self.get_extension_bytes::<V>()?)
419    }
420
421    /// Unpacks a portion of the TLV data as the desired variable-length type
422    fn get_variable_len_extension<V: Extension + VariableLenPack>(
423        &self,
424    ) -> Result<V, ProgramError> {
425        let data = get_extension_bytes::<S, V>(self.get_tlv_data())?;
426        V::unpack_from_slice(data)
427    }
428
429    /// Iterates through the TLV entries, returning only the types
430    fn get_extension_types(&self) -> Result<Vec<ExtensionType>, ProgramError> {
431        get_tlv_data_info(self.get_tlv_data()).map(|x| x.extension_types)
432    }
433
434    /// Get just the first extension type, useful to track mixed initialization
435    fn get_first_extension_type(&self) -> Result<Option<ExtensionType>, ProgramError> {
436        get_first_extension_type(self.get_tlv_data())
437    }
438
439    /// Get the total number of bytes used by TLV entries and the base type
440    fn try_get_account_len(&self) -> Result<usize, ProgramError> {
441        let tlv_info = get_tlv_data_info(self.get_tlv_data())?;
442        if tlv_info.extension_types.is_empty() {
443            Ok(S::SIZE_OF)
444        } else {
445            let total_len = tlv_info
446                .used_len
447                .saturating_add(BASE_ACCOUNT_AND_TYPE_LENGTH);
448            Ok(adjust_len_for_multisig(total_len))
449        }
450    }
451    /// Calculate the new expected size if the state allocates the given
452    /// fixed-length extension instance.
453    /// If the state already has the extension, the resulting account length
454    /// will be unchanged.
455    fn try_get_new_account_len<V: Extension + Pod>(&self) -> Result<usize, ProgramError> {
456        try_get_new_account_len_for_extension_len::<S, V>(
457            self.get_tlv_data(),
458            pod_get_packed_len::<V>(),
459        )
460    }
461
462    /// Calculate the new expected size if the state allocates the given
463    /// variable-length extension instance.
464    fn try_get_new_account_len_for_variable_len_extension<V: Extension + VariableLenPack>(
465        &self,
466        new_extension: &V,
467    ) -> Result<usize, ProgramError> {
468        try_get_new_account_len_for_extension_len::<S, V>(
469            self.get_tlv_data(),
470            new_extension.get_packed_len()?,
471        )
472    }
473}
474
475/// Encapsulates owned immutable base state data (mint or account) with possible
476/// extensions
477#[derive(Clone, Debug, PartialEq)]
478pub struct StateWithExtensionsOwned<S: BaseState> {
479    /// Unpacked base data
480    pub base: S,
481    /// Raw TLV data, deserialized on demand
482    tlv_data: Vec<u8>,
483}
484impl<S: BaseState + Pack> StateWithExtensionsOwned<S> {
485    /// Unpack base state, leaving the extension data as a slice
486    ///
487    /// Fails if the base state is not initialized.
488    pub fn unpack(mut input: Vec<u8>) -> Result<Self, ProgramError> {
489        check_min_len_and_not_multisig(&input, S::SIZE_OF)?;
490        let mut rest = input.split_off(S::SIZE_OF);
491        let base = S::unpack(&input)?;
492        if let Some((account_type_index, tlv_start_index)) = type_and_tlv_indices::<S>(&rest)? {
493            // type_and_tlv_indices() checks that returned indexes are within range
494            let account_type = AccountType::try_from(rest[account_type_index])
495                .map_err(|_| ProgramError::InvalidAccountData)?;
496            check_account_type::<S>(account_type)?;
497            let tlv_data = rest.split_off(tlv_start_index);
498            Ok(Self { base, tlv_data })
499        } else {
500            Ok(Self {
501                base,
502                tlv_data: vec![],
503            })
504        }
505    }
506}
507
508impl<S: BaseState> BaseStateWithExtensions<S> for StateWithExtensionsOwned<S> {
509    fn get_tlv_data(&self) -> &[u8] {
510        &self.tlv_data
511    }
512}
513
514/// Encapsulates immutable base state data (mint or account) with possible
515/// extensions
516#[derive(Debug, PartialEq)]
517pub struct StateWithExtensions<'data, S: BaseState + Pack> {
518    /// Unpacked base data
519    pub base: S,
520    /// Slice of data containing all TLV data, deserialized on demand
521    tlv_data: &'data [u8],
522}
523impl<'data, S: BaseState + Pack> StateWithExtensions<'data, S> {
524    /// Unpack base state, leaving the extension data as a slice
525    ///
526    /// Fails if the base state is not initialized.
527    pub fn unpack(input: &'data [u8]) -> Result<Self, ProgramError> {
528        check_min_len_and_not_multisig(input, S::SIZE_OF)?;
529        let (base_data, rest) = input.split_at(S::SIZE_OF);
530        let base = S::unpack(base_data)?;
531        let tlv_data = unpack_tlv_data::<S>(rest)?;
532        Ok(Self { base, tlv_data })
533    }
534}
535impl<'a, S: BaseState + Pack> BaseStateWithExtensions<S> for StateWithExtensions<'a, S> {
536    fn get_tlv_data(&self) -> &[u8] {
537        self.tlv_data
538    }
539}
540
541/// Encapsulates immutable base state data (mint or account) with possible
542/// extensions, where the base state is Pod for zero-copy serde.
543#[derive(Debug, PartialEq)]
544pub struct PodStateWithExtensions<'data, S: BaseState + Pod> {
545    /// Unpacked base data
546    pub base: &'data S,
547    /// Slice of data containing all TLV data, deserialized on demand
548    tlv_data: &'data [u8],
549}
550impl<'data, S: BaseState + Pod> PodStateWithExtensions<'data, S> {
551    /// Unpack base state, leaving the extension data as a slice
552    ///
553    /// Fails if the base state is not initialized.
554    pub fn unpack(input: &'data [u8]) -> Result<Self, ProgramError> {
555        check_min_len_and_not_multisig(input, S::SIZE_OF)?;
556        let (base_data, rest) = input.split_at(S::SIZE_OF);
557        let base = pod_from_bytes::<S>(base_data)?;
558        if !base.is_initialized() {
559            Err(ProgramError::UninitializedAccount)
560        } else {
561            let tlv_data = unpack_tlv_data::<S>(rest)?;
562            Ok(Self { base, tlv_data })
563        }
564    }
565}
566impl<'a, S: BaseState + Pod> BaseStateWithExtensions<S> for PodStateWithExtensions<'a, S> {
567    fn get_tlv_data(&self) -> &[u8] {
568        self.tlv_data
569    }
570}
571
572/// Trait for mutable base state with extension
573pub trait BaseStateWithExtensionsMut<S: BaseState>: BaseStateWithExtensions<S> {
574    /// Get the underlying TLV data as mutable
575    fn get_tlv_data_mut(&mut self) -> &mut [u8];
576
577    /// Get the underlying account type as mutable
578    fn get_account_type_mut(&mut self) -> &mut [u8];
579
580    /// Unpack a portion of the TLV data as the base mutable bytes
581    fn get_extension_bytes_mut<V: Extension>(&mut self) -> Result<&mut [u8], ProgramError> {
582        get_extension_bytes_mut::<S, V>(self.get_tlv_data_mut())
583    }
584
585    /// Unpack a portion of the TLV data as the desired type that allows
586    /// modifying the type
587    fn get_extension_mut<V: Extension + Pod>(&mut self) -> Result<&mut V, ProgramError> {
588        pod_from_bytes_mut::<V>(self.get_extension_bytes_mut::<V>()?)
589    }
590
591    /// Packs a variable-length extension into its appropriate data segment.
592    /// Fails if space hasn't already been allocated for the given extension
593    fn pack_variable_len_extension<V: Extension + VariableLenPack>(
594        &mut self,
595        extension: &V,
596    ) -> Result<(), ProgramError> {
597        let data = self.get_extension_bytes_mut::<V>()?;
598        // NOTE: Do *not* use `pack`, since the length check will cause
599        // reallocations to smaller sizes to fail
600        extension.pack_into_slice(data)
601    }
602
603    /// Packs the default extension data into an open slot if not already found
604    /// in the data buffer. If extension is already found in the buffer, it
605    /// overwrites the existing extension with the default state if
606    /// `overwrite` is set. If extension found, but `overwrite` is not set,
607    /// it returns error.
608    fn init_extension<V: Extension + Pod + Default>(
609        &mut self,
610        overwrite: bool,
611    ) -> Result<&mut V, ProgramError> {
612        let length = pod_get_packed_len::<V>();
613        let buffer = self.alloc::<V>(length, overwrite)?;
614        let extension_ref = pod_from_bytes_mut::<V>(buffer)?;
615        *extension_ref = V::default();
616        Ok(extension_ref)
617    }
618
619    /// Reallocate and overwrite the TLV entry for the given variable-length
620    /// extension.
621    ///
622    /// Returns an error if the extension is not present, or if there is not
623    /// enough space in the buffer.
624    fn realloc_variable_len_extension<V: Extension + VariableLenPack>(
625        &mut self,
626        new_extension: &V,
627    ) -> Result<(), ProgramError> {
628        let data = self.realloc::<V>(new_extension.get_packed_len()?)?;
629        new_extension.pack_into_slice(data)
630    }
631
632    /// Reallocate the TLV entry for the given extension to the given number of
633    /// bytes.
634    ///
635    /// If the new length is smaller, it will compact the rest of the buffer and
636    /// zero out the difference at the end. If it's larger, it will move the
637    /// rest of the buffer data and zero out the new data.
638    ///
639    /// Returns an error if the extension is not present, or if this is not
640    /// enough space in the buffer.
641    fn realloc<V: Extension + VariableLenPack>(
642        &mut self,
643        length: usize,
644    ) -> Result<&mut [u8], ProgramError> {
645        let tlv_data = self.get_tlv_data_mut();
646        let TlvIndices {
647            type_start: _,
648            length_start,
649            value_start,
650        } = get_extension_indices::<V>(tlv_data, false)?;
651        let tlv_len = get_tlv_data_info(tlv_data).map(|x| x.used_len)?;
652        let data_len = tlv_data.len();
653
654        let length_ref = pod_from_bytes_mut::<Length>(&mut tlv_data[length_start..value_start])?;
655        let old_length = usize::from(*length_ref);
656
657        // Length check to avoid a panic later in `copy_within`
658        if old_length < length {
659            let new_tlv_len = tlv_len.saturating_add(length.saturating_sub(old_length));
660            if new_tlv_len > data_len {
661                return Err(ProgramError::InvalidAccountData);
662            }
663        }
664
665        // write new length after the check, to avoid getting into a bad situation
666        // if trying to recover from an error
667        *length_ref = Length::try_from(length)?;
668
669        let old_value_end = value_start.saturating_add(old_length);
670        let new_value_end = value_start.saturating_add(length);
671        tlv_data.copy_within(old_value_end..tlv_len, new_value_end);
672        match old_length.cmp(&length) {
673            Ordering::Greater => {
674                // realloc to smaller, zero out the end
675                let new_tlv_len = tlv_len.saturating_sub(old_length.saturating_sub(length));
676                tlv_data[new_tlv_len..tlv_len].fill(0);
677            }
678            Ordering::Less => {
679                // realloc to bigger, zero out the new bytes
680                tlv_data[old_value_end..new_value_end].fill(0);
681            }
682            Ordering::Equal => {} // nothing needed!
683        }
684
685        Ok(&mut tlv_data[value_start..new_value_end])
686    }
687
688    /// Allocate the given number of bytes for the given variable-length
689    /// extension and write its contents into the TLV buffer.
690    ///
691    /// This can only be used for variable-sized types, such as `String` or
692    /// `Vec`. `Pod` types must use `init_extension`
693    fn init_variable_len_extension<V: Extension + VariableLenPack>(
694        &mut self,
695        extension: &V,
696        overwrite: bool,
697    ) -> Result<(), ProgramError> {
698        let data = self.alloc::<V>(extension.get_packed_len()?, overwrite)?;
699        extension.pack_into_slice(data)
700    }
701
702    /// Allocate some space for the extension in the TLV data
703    fn alloc<V: Extension>(
704        &mut self,
705        length: usize,
706        overwrite: bool,
707    ) -> Result<&mut [u8], ProgramError> {
708        if V::TYPE.get_account_type() != S::ACCOUNT_TYPE {
709            return Err(ProgramError::InvalidAccountData);
710        }
711        let tlv_data = self.get_tlv_data_mut();
712        let TlvIndices {
713            type_start,
714            length_start,
715            value_start,
716        } = get_extension_indices::<V>(tlv_data, true)?;
717
718        if tlv_data[type_start..].len() < add_type_and_length_to_len(length) {
719            return Err(ProgramError::InvalidAccountData);
720        }
721        let extension_type = ExtensionType::try_from(&tlv_data[type_start..length_start])?;
722
723        if extension_type == ExtensionType::Uninitialized || overwrite {
724            // write extension type
725            let extension_type_array: [u8; 2] = V::TYPE.into();
726            let extension_type_ref = &mut tlv_data[type_start..length_start];
727            extension_type_ref.copy_from_slice(&extension_type_array);
728            // write length
729            let length_ref =
730                pod_from_bytes_mut::<Length>(&mut tlv_data[length_start..value_start])?;
731
732            // check that the length is the same if we're doing an alloc
733            // with overwrite, otherwise a realloc should be done
734            if overwrite && extension_type == V::TYPE && usize::from(*length_ref) != length {
735                return Err(TokenError::InvalidLengthForAlloc.into());
736            }
737
738            *length_ref = Length::try_from(length)?;
739
740            let value_end = value_start.saturating_add(length);
741            Ok(&mut tlv_data[value_start..value_end])
742        } else {
743            // extension is already initialized, but no overwrite permission
744            Err(TokenError::ExtensionAlreadyInitialized.into())
745        }
746    }
747
748    /// If `extension_type` is an Account-associated `ExtensionType` that
749    /// requires initialization on `InitializeAccount`, this method packs
750    /// the default relevant `Extension` of an `ExtensionType` into an open
751    /// slot if not already found in the data buffer, otherwise overwrites
752    /// the existing extension with the default state. For all other
753    /// `ExtensionType`s, this is a no-op.
754    fn init_account_extension_from_type(
755        &mut self,
756        extension_type: ExtensionType,
757    ) -> Result<(), ProgramError> {
758        if extension_type.get_account_type() != AccountType::Account {
759            return Ok(());
760        }
761        match extension_type {
762            ExtensionType::TransferFeeAmount => {
763                self.init_extension::<TransferFeeAmount>(true).map(|_| ())
764            }
765            ExtensionType::ImmutableOwner => {
766                self.init_extension::<ImmutableOwner>(true).map(|_| ())
767            }
768            ExtensionType::NonTransferableAccount => self
769                .init_extension::<NonTransferableAccount>(true)
770                .map(|_| ()),
771            ExtensionType::TransferHookAccount => {
772                self.init_extension::<TransferHookAccount>(true).map(|_| ())
773            }
774            // ConfidentialTransfers are currently opt-in only, so this is a no-op for extra safety
775            // on InitializeAccount
776            ExtensionType::ConfidentialTransferAccount => Ok(()),
777            ExtensionType::PausableAccount => {
778                self.init_extension::<PausableAccount>(true).map(|_| ())
779            }
780            #[cfg(test)]
781            ExtensionType::AccountPaddingTest => {
782                self.init_extension::<AccountPaddingTest>(true).map(|_| ())
783            }
784            _ => unreachable!(),
785        }
786    }
787
788    /// Write the account type into the buffer, done during the base
789    /// state initialization
790    /// Noops if there is no room for an extension in the account, needed for
791    /// pure base mints / accounts.
792    fn init_account_type(&mut self) -> Result<(), ProgramError> {
793        let first_extension_type = self.get_first_extension_type()?;
794        let account_type = self.get_account_type_mut();
795        if !account_type.is_empty() {
796            if let Some(extension_type) = first_extension_type {
797                let account_type = extension_type.get_account_type();
798                if account_type != S::ACCOUNT_TYPE {
799                    return Err(TokenError::ExtensionBaseMismatch.into());
800                }
801            }
802            account_type[0] = S::ACCOUNT_TYPE.into();
803        }
804        Ok(())
805    }
806
807    /// Check that the account type on the account (if initialized) matches the
808    /// account type for any extensions initialized on the TLV data
809    fn check_account_type_matches_extension_type(&self) -> Result<(), ProgramError> {
810        if let Some(extension_type) = self.get_first_extension_type()? {
811            let account_type = extension_type.get_account_type();
812            if account_type != S::ACCOUNT_TYPE {
813                return Err(TokenError::ExtensionBaseMismatch.into());
814            }
815        }
816        Ok(())
817    }
818}
819
820/// Encapsulates mutable base state data (mint or account) with possible
821/// extensions
822#[derive(Debug, PartialEq)]
823pub struct StateWithExtensionsMut<'data, S: BaseState> {
824    /// Unpacked base data
825    pub base: S,
826    /// Raw base data
827    base_data: &'data mut [u8],
828    /// Writable account type
829    account_type: &'data mut [u8],
830    /// Slice of data containing all TLV data, deserialized on demand
831    tlv_data: &'data mut [u8],
832}
833impl<'data, S: BaseState + Pack> StateWithExtensionsMut<'data, S> {
834    /// Unpack base state, leaving the extension data as a mutable slice
835    ///
836    /// Fails if the base state is not initialized.
837    pub fn unpack(input: &'data mut [u8]) -> Result<Self, ProgramError> {
838        check_min_len_and_not_multisig(input, S::SIZE_OF)?;
839        let (base_data, rest) = input.split_at_mut(S::SIZE_OF);
840        let base = S::unpack(base_data)?;
841        let (account_type, tlv_data) = unpack_type_and_tlv_data_mut::<S>(rest)?;
842        Ok(Self {
843            base,
844            base_data,
845            account_type,
846            tlv_data,
847        })
848    }
849
850    /// Unpack an uninitialized base state, leaving the extension data as a
851    /// mutable slice
852    ///
853    /// Fails if the base state has already been initialized.
854    pub fn unpack_uninitialized(input: &'data mut [u8]) -> Result<Self, ProgramError> {
855        check_min_len_and_not_multisig(input, S::SIZE_OF)?;
856        let (base_data, rest) = input.split_at_mut(S::SIZE_OF);
857        let base = S::unpack_unchecked(base_data)?;
858        if base.is_initialized() {
859            return Err(TokenError::AlreadyInUse.into());
860        }
861        let (account_type, tlv_data) = unpack_uninitialized_type_and_tlv_data_mut::<S>(rest)?;
862        let state = Self {
863            base,
864            base_data,
865            account_type,
866            tlv_data,
867        };
868        state.check_account_type_matches_extension_type()?;
869        Ok(state)
870    }
871
872    /// Packs base state data into the base data portion
873    pub fn pack_base(&mut self) {
874        S::pack_into_slice(&self.base, self.base_data);
875    }
876}
877impl<'a, S: BaseState> BaseStateWithExtensions<S> for StateWithExtensionsMut<'a, S> {
878    fn get_tlv_data(&self) -> &[u8] {
879        self.tlv_data
880    }
881}
882impl<'a, S: BaseState> BaseStateWithExtensionsMut<S> for StateWithExtensionsMut<'a, S> {
883    fn get_tlv_data_mut(&mut self) -> &mut [u8] {
884        self.tlv_data
885    }
886    fn get_account_type_mut(&mut self) -> &mut [u8] {
887        self.account_type
888    }
889}
890
891/// Encapsulates mutable base state data (mint or account) with possible
892/// extensions, where the base state is Pod for zero-copy serde.
893#[derive(Debug, PartialEq)]
894pub struct PodStateWithExtensionsMut<'data, S: BaseState> {
895    /// Unpacked base data
896    pub base: &'data mut S,
897    /// Writable account type
898    account_type: &'data mut [u8],
899    /// Slice of data containing all TLV data, deserialized on demand
900    tlv_data: &'data mut [u8],
901}
902impl<'data, S: BaseState + Pod> PodStateWithExtensionsMut<'data, S> {
903    /// Unpack base state, leaving the extension data as a mutable slice
904    ///
905    /// Fails if the base state is not initialized.
906    pub fn unpack(input: &'data mut [u8]) -> Result<Self, ProgramError> {
907        check_min_len_and_not_multisig(input, S::SIZE_OF)?;
908        let (base_data, rest) = input.split_at_mut(S::SIZE_OF);
909        let base = pod_from_bytes_mut::<S>(base_data)?;
910        if !base.is_initialized() {
911            Err(ProgramError::UninitializedAccount)
912        } else {
913            let (account_type, tlv_data) = unpack_type_and_tlv_data_mut::<S>(rest)?;
914            Ok(Self {
915                base,
916                account_type,
917                tlv_data,
918            })
919        }
920    }
921
922    /// Unpack an uninitialized base state, leaving the extension data as a
923    /// mutable slice
924    ///
925    /// Fails if the base state has already been initialized.
926    pub fn unpack_uninitialized(input: &'data mut [u8]) -> Result<Self, ProgramError> {
927        check_min_len_and_not_multisig(input, S::SIZE_OF)?;
928        let (base_data, rest) = input.split_at_mut(S::SIZE_OF);
929        let base = pod_from_bytes_mut::<S>(base_data)?;
930        if base.is_initialized() {
931            return Err(TokenError::AlreadyInUse.into());
932        }
933        let (account_type, tlv_data) = unpack_uninitialized_type_and_tlv_data_mut::<S>(rest)?;
934        let state = Self {
935            base,
936            account_type,
937            tlv_data,
938        };
939        state.check_account_type_matches_extension_type()?;
940        Ok(state)
941    }
942}
943
944impl<'a, S: BaseState> BaseStateWithExtensions<S> for PodStateWithExtensionsMut<'a, S> {
945    fn get_tlv_data(&self) -> &[u8] {
946        self.tlv_data
947    }
948}
949impl<'a, S: BaseState> BaseStateWithExtensionsMut<S> for PodStateWithExtensionsMut<'a, S> {
950    fn get_tlv_data_mut(&mut self) -> &mut [u8] {
951        self.tlv_data
952    }
953    fn get_account_type_mut(&mut self) -> &mut [u8] {
954        self.account_type
955    }
956}
957
958fn unpack_tlv_data<S: BaseState>(rest: &[u8]) -> Result<&[u8], ProgramError> {
959    if let Some((account_type_index, tlv_start_index)) = type_and_tlv_indices::<S>(rest)? {
960        // type_and_tlv_indices() checks that returned indexes are within range
961        let account_type = AccountType::try_from(rest[account_type_index])
962            .map_err(|_| ProgramError::InvalidAccountData)?;
963        check_account_type::<S>(account_type)?;
964        Ok(&rest[tlv_start_index..])
965    } else {
966        Ok(&[])
967    }
968}
969
970fn unpack_type_and_tlv_data_with_check_mut<
971    S: BaseState,
972    F: Fn(AccountType) -> Result<(), ProgramError>,
973>(
974    rest: &mut [u8],
975    check_fn: F,
976) -> Result<(&mut [u8], &mut [u8]), ProgramError> {
977    if let Some((account_type_index, tlv_start_index)) = type_and_tlv_indices::<S>(rest)? {
978        // type_and_tlv_indices() checks that returned indexes are within range
979        let account_type = AccountType::try_from(rest[account_type_index])
980            .map_err(|_| ProgramError::InvalidAccountData)?;
981        check_fn(account_type)?;
982        let (account_type, tlv_data) = rest.split_at_mut(tlv_start_index);
983        Ok((
984            &mut account_type[account_type_index..tlv_start_index],
985            tlv_data,
986        ))
987    } else {
988        Ok((&mut [], &mut []))
989    }
990}
991
992fn unpack_type_and_tlv_data_mut<S: BaseState>(
993    rest: &mut [u8],
994) -> Result<(&mut [u8], &mut [u8]), ProgramError> {
995    unpack_type_and_tlv_data_with_check_mut::<S, _>(rest, check_account_type::<S>)
996}
997
998fn unpack_uninitialized_type_and_tlv_data_mut<S: BaseState>(
999    rest: &mut [u8],
1000) -> Result<(&mut [u8], &mut [u8]), ProgramError> {
1001    unpack_type_and_tlv_data_with_check_mut::<S, _>(rest, |account_type| {
1002        if account_type != AccountType::Uninitialized {
1003            Err(ProgramError::InvalidAccountData)
1004        } else {
1005            Ok(())
1006        }
1007    })
1008}
1009
1010/// If `AccountType` is uninitialized, set it to the `BaseState`'s
1011/// `ACCOUNT_TYPE`; if `AccountType` is already set, check is set correctly for
1012/// `BaseState`. This method assumes that the `base_data` has already been
1013/// packed with data of the desired type.
1014pub fn set_account_type<S: BaseState>(input: &mut [u8]) -> Result<(), ProgramError> {
1015    check_min_len_and_not_multisig(input, S::SIZE_OF)?;
1016    let (base_data, rest) = input.split_at_mut(S::SIZE_OF);
1017    if S::ACCOUNT_TYPE == AccountType::Account && !is_initialized_account(base_data)? {
1018        return Err(ProgramError::InvalidAccountData);
1019    }
1020    if let Some((account_type_index, _tlv_start_index)) = type_and_tlv_indices::<S>(rest)? {
1021        let mut account_type = AccountType::try_from(rest[account_type_index])
1022            .map_err(|_| ProgramError::InvalidAccountData)?;
1023        if account_type == AccountType::Uninitialized {
1024            rest[account_type_index] = S::ACCOUNT_TYPE.into();
1025            account_type = S::ACCOUNT_TYPE;
1026        }
1027        check_account_type::<S>(account_type)?;
1028        Ok(())
1029    } else {
1030        Err(ProgramError::InvalidAccountData)
1031    }
1032}
1033
1034/// Different kinds of accounts. Note that `Mint`, `Account`, and `Multisig`
1035/// types are determined exclusively by the size of the account, and are not
1036/// included in the account data. `AccountType` is only included if extensions
1037/// have been initialized.
1038#[repr(u8)]
1039#[derive(Clone, Copy, Debug, PartialEq, TryFromPrimitive, IntoPrimitive)]
1040pub enum AccountType {
1041    /// Marker for 0 data
1042    Uninitialized,
1043    /// Mint account with additional extensions
1044    Mint,
1045    /// Token holding account with additional extensions
1046    Account,
1047}
1048impl Default for AccountType {
1049    fn default() -> Self {
1050        Self::Uninitialized
1051    }
1052}
1053
1054/// Extensions that can be applied to mints or accounts.  Mint extensions must
1055/// only be applied to mint accounts, and account extensions must only be
1056/// applied to token holding accounts.
1057#[repr(u16)]
1058#[cfg_attr(feature = "serde-traits", derive(Serialize, Deserialize))]
1059#[cfg_attr(feature = "serde-traits", serde(rename_all = "camelCase"))]
1060#[derive(Clone, Copy, Debug, PartialEq, TryFromPrimitive, IntoPrimitive)]
1061pub enum ExtensionType {
1062    /// Used as padding if the account size would otherwise be 355, same as a
1063    /// multisig
1064    Uninitialized,
1065    /// Includes transfer fee rate info and accompanying authorities to withdraw
1066    /// and set the fee
1067    TransferFeeConfig,
1068    /// Includes withheld transfer fees
1069    TransferFeeAmount,
1070    /// Includes an optional mint close authority
1071    MintCloseAuthority,
1072    /// Auditor configuration for confidential transfers
1073    ConfidentialTransferMint,
1074    /// State for confidential transfers
1075    ConfidentialTransferAccount,
1076    /// Specifies the default Account::state for new Accounts
1077    DefaultAccountState,
1078    /// Indicates that the Account owner authority cannot be changed
1079    ImmutableOwner,
1080    /// Require inbound transfers to have memo
1081    MemoTransfer,
1082    /// Indicates that the tokens from this mint can't be transferred
1083    NonTransferable,
1084    /// Tokens accrue interest over time,
1085    InterestBearingConfig,
1086    /// Locks privileged token operations from happening via CPI
1087    CpiGuard,
1088    /// Includes an optional permanent delegate
1089    PermanentDelegate,
1090    /// Indicates that the tokens in this account belong to a non-transferable
1091    /// mint
1092    NonTransferableAccount,
1093    /// Mint requires a CPI to a program implementing the "transfer hook"
1094    /// interface
1095    TransferHook,
1096    /// Indicates that the tokens in this account belong to a mint with a
1097    /// transfer hook
1098    TransferHookAccount,
1099    /// Includes encrypted withheld fees and the encryption public that they are
1100    /// encrypted under
1101    ConfidentialTransferFeeConfig,
1102    /// Includes confidential withheld transfer fees
1103    ConfidentialTransferFeeAmount,
1104    /// Mint contains a pointer to another account (or the same account) that
1105    /// holds metadata
1106    MetadataPointer,
1107    /// Mint contains token-metadata
1108    TokenMetadata,
1109    /// Mint contains a pointer to another account (or the same account) that
1110    /// holds group configurations
1111    GroupPointer,
1112    /// Mint contains token group configurations
1113    TokenGroup,
1114    /// Mint contains a pointer to another account (or the same account) that
1115    /// holds group member configurations
1116    GroupMemberPointer,
1117    /// Mint contains token group member configurations
1118    TokenGroupMember,
1119    /// Mint allowing the minting and burning of confidential tokens
1120    ConfidentialMintBurn,
1121    /// Tokens whose UI amount is scaled by a given amount
1122    ScaledUiAmount,
1123    /// Tokens where minting / burning / transferring can be paused
1124    Pausable,
1125    /// Indicates that the account belongs to a pausable mint
1126    PausableAccount,
1127
1128    /// Test variable-length mint extension
1129    #[cfg(test)]
1130    VariableLenMintTest = u16::MAX - 2,
1131    /// Padding extension used to make an account exactly Multisig::LEN, used
1132    /// for testing
1133    #[cfg(test)]
1134    AccountPaddingTest,
1135    /// Padding extension used to make a mint exactly Multisig::LEN, used for
1136    /// testing
1137    #[cfg(test)]
1138    MintPaddingTest,
1139}
1140impl TryFrom<&[u8]> for ExtensionType {
1141    type Error = ProgramError;
1142    fn try_from(a: &[u8]) -> Result<Self, Self::Error> {
1143        Self::try_from(u16::from_le_bytes(
1144            a.try_into().map_err(|_| ProgramError::InvalidAccountData)?,
1145        ))
1146        .map_err(|_| ProgramError::InvalidAccountData)
1147    }
1148}
1149impl From<ExtensionType> for [u8; 2] {
1150    fn from(a: ExtensionType) -> Self {
1151        u16::from(a).to_le_bytes()
1152    }
1153}
1154impl ExtensionType {
1155    /// Returns true if the given extension type is sized
1156    ///
1157    /// Most extension types should be sized, so any variable-length extension
1158    /// types should be added here by hand
1159    const fn sized(&self) -> bool {
1160        match self {
1161            ExtensionType::TokenMetadata => false,
1162            #[cfg(test)]
1163            ExtensionType::VariableLenMintTest => false,
1164            _ => true,
1165        }
1166    }
1167
1168    /// Get the data length of the type associated with the enum
1169    ///
1170    /// Fails if the extension type has a variable length
1171    fn try_get_type_len(&self) -> Result<usize, ProgramError> {
1172        if !self.sized() {
1173            return Err(ProgramError::InvalidArgument);
1174        }
1175        Ok(match self {
1176            ExtensionType::Uninitialized => 0,
1177            ExtensionType::TransferFeeConfig => pod_get_packed_len::<TransferFeeConfig>(),
1178            ExtensionType::TransferFeeAmount => pod_get_packed_len::<TransferFeeAmount>(),
1179            ExtensionType::MintCloseAuthority => pod_get_packed_len::<MintCloseAuthority>(),
1180            ExtensionType::ImmutableOwner => pod_get_packed_len::<ImmutableOwner>(),
1181            ExtensionType::ConfidentialTransferMint => {
1182                pod_get_packed_len::<ConfidentialTransferMint>()
1183            }
1184            ExtensionType::ConfidentialTransferAccount => {
1185                pod_get_packed_len::<ConfidentialTransferAccount>()
1186            }
1187            ExtensionType::DefaultAccountState => pod_get_packed_len::<DefaultAccountState>(),
1188            ExtensionType::MemoTransfer => pod_get_packed_len::<MemoTransfer>(),
1189            ExtensionType::NonTransferable => pod_get_packed_len::<NonTransferable>(),
1190            ExtensionType::InterestBearingConfig => pod_get_packed_len::<InterestBearingConfig>(),
1191            ExtensionType::CpiGuard => pod_get_packed_len::<CpiGuard>(),
1192            ExtensionType::PermanentDelegate => pod_get_packed_len::<PermanentDelegate>(),
1193            ExtensionType::NonTransferableAccount => pod_get_packed_len::<NonTransferableAccount>(),
1194            ExtensionType::TransferHook => pod_get_packed_len::<TransferHook>(),
1195            ExtensionType::TransferHookAccount => pod_get_packed_len::<TransferHookAccount>(),
1196            ExtensionType::ConfidentialTransferFeeConfig => {
1197                pod_get_packed_len::<ConfidentialTransferFeeConfig>()
1198            }
1199            ExtensionType::ConfidentialTransferFeeAmount => {
1200                pod_get_packed_len::<ConfidentialTransferFeeAmount>()
1201            }
1202            ExtensionType::MetadataPointer => pod_get_packed_len::<MetadataPointer>(),
1203            ExtensionType::TokenMetadata => unreachable!(),
1204            ExtensionType::GroupPointer => pod_get_packed_len::<GroupPointer>(),
1205            ExtensionType::TokenGroup => pod_get_packed_len::<TokenGroup>(),
1206            ExtensionType::GroupMemberPointer => pod_get_packed_len::<GroupMemberPointer>(),
1207            ExtensionType::TokenGroupMember => pod_get_packed_len::<TokenGroupMember>(),
1208            ExtensionType::ConfidentialMintBurn => pod_get_packed_len::<ConfidentialMintBurn>(),
1209            ExtensionType::ScaledUiAmount => pod_get_packed_len::<ScaledUiAmountConfig>(),
1210            ExtensionType::Pausable => pod_get_packed_len::<PausableConfig>(),
1211            ExtensionType::PausableAccount => pod_get_packed_len::<PausableAccount>(),
1212            #[cfg(test)]
1213            ExtensionType::AccountPaddingTest => pod_get_packed_len::<AccountPaddingTest>(),
1214            #[cfg(test)]
1215            ExtensionType::MintPaddingTest => pod_get_packed_len::<MintPaddingTest>(),
1216            #[cfg(test)]
1217            ExtensionType::VariableLenMintTest => unreachable!(),
1218        })
1219    }
1220
1221    /// Get the TLV length for an `ExtensionType`
1222    ///
1223    /// Fails if the extension type has a variable length
1224    fn try_get_tlv_len(&self) -> Result<usize, ProgramError> {
1225        Ok(add_type_and_length_to_len(self.try_get_type_len()?))
1226    }
1227
1228    /// Get the TLV length for a set of `ExtensionType`s
1229    ///
1230    /// Fails if any of the extension types has a variable length
1231    fn try_get_total_tlv_len(extension_types: &[Self]) -> Result<usize, ProgramError> {
1232        // dedupe extensions
1233        let mut extensions = vec![];
1234        for extension_type in extension_types {
1235            if !extensions.contains(&extension_type) {
1236                extensions.push(extension_type);
1237            }
1238        }
1239        extensions.iter().map(|e| e.try_get_tlv_len()).sum()
1240    }
1241
1242    /// Get the required account data length for the given `ExtensionType`s
1243    ///
1244    /// Fails if any of the extension types has a variable length
1245    pub fn try_calculate_account_len<S: BaseState>(
1246        extension_types: &[Self],
1247    ) -> Result<usize, ProgramError> {
1248        if extension_types.is_empty() {
1249            Ok(S::SIZE_OF)
1250        } else {
1251            let extension_size = Self::try_get_total_tlv_len(extension_types)?;
1252            let total_len = extension_size.saturating_add(BASE_ACCOUNT_AND_TYPE_LENGTH);
1253            Ok(adjust_len_for_multisig(total_len))
1254        }
1255    }
1256
1257    /// Get the associated account type
1258    pub fn get_account_type(&self) -> AccountType {
1259        match self {
1260            ExtensionType::Uninitialized => AccountType::Uninitialized,
1261            ExtensionType::TransferFeeConfig
1262            | ExtensionType::MintCloseAuthority
1263            | ExtensionType::ConfidentialTransferMint
1264            | ExtensionType::DefaultAccountState
1265            | ExtensionType::NonTransferable
1266            | ExtensionType::InterestBearingConfig
1267            | ExtensionType::PermanentDelegate
1268            | ExtensionType::TransferHook
1269            | ExtensionType::ConfidentialTransferFeeConfig
1270            | ExtensionType::MetadataPointer
1271            | ExtensionType::TokenMetadata
1272            | ExtensionType::GroupPointer
1273            | ExtensionType::TokenGroup
1274            | ExtensionType::GroupMemberPointer
1275            | ExtensionType::ConfidentialMintBurn
1276            | ExtensionType::TokenGroupMember
1277            | ExtensionType::ScaledUiAmount
1278            | ExtensionType::Pausable => AccountType::Mint,
1279            ExtensionType::ImmutableOwner
1280            | ExtensionType::TransferFeeAmount
1281            | ExtensionType::ConfidentialTransferAccount
1282            | ExtensionType::MemoTransfer
1283            | ExtensionType::NonTransferableAccount
1284            | ExtensionType::TransferHookAccount
1285            | ExtensionType::CpiGuard
1286            | ExtensionType::ConfidentialTransferFeeAmount
1287            | ExtensionType::PausableAccount => AccountType::Account,
1288            #[cfg(test)]
1289            ExtensionType::VariableLenMintTest => AccountType::Mint,
1290            #[cfg(test)]
1291            ExtensionType::AccountPaddingTest => AccountType::Account,
1292            #[cfg(test)]
1293            ExtensionType::MintPaddingTest => AccountType::Mint,
1294        }
1295    }
1296
1297    /// Based on a set of `AccountType::Mint` `ExtensionType`s, get the list of
1298    /// `AccountType::Account` `ExtensionType`s required on `InitializeAccount`
1299    pub fn get_required_init_account_extensions(mint_extension_types: &[Self]) -> Vec<Self> {
1300        let mut account_extension_types = vec![];
1301        for extension_type in mint_extension_types {
1302            match extension_type {
1303                ExtensionType::TransferFeeConfig => {
1304                    account_extension_types.push(ExtensionType::TransferFeeAmount);
1305                }
1306                ExtensionType::NonTransferable => {
1307                    account_extension_types.push(ExtensionType::NonTransferableAccount);
1308                    account_extension_types.push(ExtensionType::ImmutableOwner);
1309                }
1310                ExtensionType::TransferHook => {
1311                    account_extension_types.push(ExtensionType::TransferHookAccount);
1312                }
1313                ExtensionType::Pausable => {
1314                    account_extension_types.push(ExtensionType::PausableAccount);
1315                }
1316                #[cfg(test)]
1317                ExtensionType::MintPaddingTest => {
1318                    account_extension_types.push(ExtensionType::AccountPaddingTest);
1319                }
1320                _ => {}
1321            }
1322        }
1323        account_extension_types
1324    }
1325
1326    /// Check for invalid combination of mint extensions
1327    pub fn check_for_invalid_mint_extension_combinations(
1328        mint_extension_types: &[Self],
1329    ) -> Result<(), TokenError> {
1330        let mut transfer_fee_config = false;
1331        let mut confidential_transfer_mint = false;
1332        let mut confidential_transfer_fee_config = false;
1333        let mut confidential_mint_burn = false;
1334        let mut interest_bearing = false;
1335        let mut scaled_ui_amount = false;
1336
1337        for extension_type in mint_extension_types {
1338            match extension_type {
1339                ExtensionType::TransferFeeConfig => transfer_fee_config = true,
1340                ExtensionType::ConfidentialTransferMint => confidential_transfer_mint = true,
1341                ExtensionType::ConfidentialTransferFeeConfig => {
1342                    confidential_transfer_fee_config = true
1343                }
1344                ExtensionType::ConfidentialMintBurn => confidential_mint_burn = true,
1345                ExtensionType::InterestBearingConfig => interest_bearing = true,
1346                ExtensionType::ScaledUiAmount => scaled_ui_amount = true,
1347                _ => (),
1348            }
1349        }
1350
1351        if confidential_transfer_fee_config && !(transfer_fee_config && confidential_transfer_mint)
1352        {
1353            return Err(TokenError::InvalidExtensionCombination);
1354        }
1355
1356        if transfer_fee_config && confidential_transfer_mint && !confidential_transfer_fee_config {
1357            return Err(TokenError::InvalidExtensionCombination);
1358        }
1359
1360        if confidential_mint_burn && !confidential_transfer_mint {
1361            return Err(TokenError::InvalidExtensionCombination);
1362        }
1363
1364        if scaled_ui_amount && interest_bearing {
1365            return Err(TokenError::InvalidExtensionCombination);
1366        }
1367
1368        Ok(())
1369    }
1370}
1371
1372/// Trait for base states, specifying the associated enum
1373pub trait BaseState: PackedSizeOf + IsInitialized {
1374    /// Associated extension type enum, checked at the start of TLV entries
1375    const ACCOUNT_TYPE: AccountType;
1376}
1377impl BaseState for Account {
1378    const ACCOUNT_TYPE: AccountType = AccountType::Account;
1379}
1380impl BaseState for Mint {
1381    const ACCOUNT_TYPE: AccountType = AccountType::Mint;
1382}
1383impl BaseState for PodAccount {
1384    const ACCOUNT_TYPE: AccountType = AccountType::Account;
1385}
1386impl BaseState for PodMint {
1387    const ACCOUNT_TYPE: AccountType = AccountType::Mint;
1388}
1389
1390/// Trait to be implemented by all extension states, specifying which extension
1391/// and account type they are associated with
1392pub trait Extension {
1393    /// Associated extension type enum, checked at the start of TLV entries
1394    const TYPE: ExtensionType;
1395}
1396
1397/// Padding a mint account to be exactly `Multisig::LEN`.
1398/// We need to pad 185 bytes, since `Multisig::LEN = 355`, `Account::LEN = 165`,
1399/// `size_of::<AccountType>() = 1`, `size_of::<ExtensionType>() = 2`,
1400/// `size_of::<Length>() = 2`.
1401///
1402/// ```
1403/// assert_eq!(355 - 165 - 1 - 2 - 2, 185);
1404/// ```
1405#[cfg(test)]
1406#[repr(C)]
1407#[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
1408pub struct MintPaddingTest {
1409    /// Largest value under 185 that implements Pod
1410    pub padding1: [u8; 128],
1411    /// Largest value under 57 that implements Pod
1412    pub padding2: [u8; 48],
1413    /// Exact value needed to finish the padding
1414    pub padding3: [u8; 9],
1415}
1416#[cfg(test)]
1417impl Extension for MintPaddingTest {
1418    const TYPE: ExtensionType = ExtensionType::MintPaddingTest;
1419}
1420#[cfg(test)]
1421impl Default for MintPaddingTest {
1422    fn default() -> Self {
1423        Self {
1424            padding1: [1; 128],
1425            padding2: [2; 48],
1426            padding3: [3; 9],
1427        }
1428    }
1429}
1430/// Account version of the `MintPadding`
1431#[cfg(test)]
1432#[repr(C)]
1433#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
1434pub struct AccountPaddingTest(MintPaddingTest);
1435#[cfg(test)]
1436impl Extension for AccountPaddingTest {
1437    const TYPE: ExtensionType = ExtensionType::AccountPaddingTest;
1438}
1439
1440/// Packs a fixed-length extension into a TLV space
1441///
1442/// This function reallocates the account as needed to accommodate for the
1443/// change in space.
1444///
1445/// If the extension already exists, it will overwrite the existing extension
1446/// if `overwrite` is `true`, otherwise it will return an error.
1447///
1448/// If the extension does not exist, it will reallocate the account and write
1449/// the extension into the TLV buffer.
1450///
1451/// NOTE: Since this function deals with fixed-size extensions, it does not
1452/// handle _decreasing_ the size of an account's data buffer, like the function
1453/// `alloc_and_serialize_variable_len_extension` does.
1454pub(crate) fn alloc_and_serialize<S: BaseState + Pod, V: Default + Extension + Pod>(
1455    account_info: &AccountInfo,
1456    new_extension: &V,
1457    overwrite: bool,
1458) -> Result<(), ProgramError> {
1459    let previous_account_len = account_info.try_data_len()?;
1460    let new_account_len = {
1461        let data = account_info.try_borrow_data()?;
1462        let state = PodStateWithExtensions::<S>::unpack(&data)?;
1463        state.try_get_new_account_len::<V>()?
1464    };
1465
1466    // Realloc the account first, if needed
1467    if new_account_len > previous_account_len {
1468        account_info.realloc(new_account_len, false)?;
1469    }
1470    let mut buffer = account_info.try_borrow_mut_data()?;
1471    if previous_account_len <= BASE_ACCOUNT_LENGTH {
1472        set_account_type::<S>(*buffer)?;
1473    }
1474    let mut state = PodStateWithExtensionsMut::<S>::unpack(&mut buffer)?;
1475
1476    // Write the extension
1477    let extension = state.init_extension::<V>(overwrite)?;
1478    *extension = *new_extension;
1479
1480    Ok(())
1481}
1482
1483/// Packs a variable-length extension into a TLV space
1484///
1485/// This function reallocates the account as needed to accommodate for the
1486/// change in space, then reallocates in the TLV buffer, and finally writes the
1487/// bytes.
1488///
1489/// NOTE: Unlike the `reallocate` instruction, this function will reduce the
1490/// size of an account if it has too many bytes allocated for the given value.
1491pub(crate) fn alloc_and_serialize_variable_len_extension<
1492    S: BaseState + Pod,
1493    V: Extension + VariableLenPack,
1494>(
1495    account_info: &AccountInfo,
1496    new_extension: &V,
1497    overwrite: bool,
1498) -> Result<(), ProgramError> {
1499    let previous_account_len = account_info.try_data_len()?;
1500    let (new_account_len, extension_already_exists) = {
1501        let data = account_info.try_borrow_data()?;
1502        let state = PodStateWithExtensions::<S>::unpack(&data)?;
1503        let new_account_len =
1504            state.try_get_new_account_len_for_variable_len_extension(new_extension)?;
1505        let extension_already_exists = state.get_extension_bytes::<V>().is_ok();
1506        (new_account_len, extension_already_exists)
1507    };
1508
1509    if extension_already_exists && !overwrite {
1510        return Err(TokenError::ExtensionAlreadyInitialized.into());
1511    }
1512
1513    if previous_account_len < new_account_len {
1514        // account size increased, so realloc the account, then the TLV entry, then
1515        // write data
1516        account_info.realloc(new_account_len, false)?;
1517        let mut buffer = account_info.try_borrow_mut_data()?;
1518        if extension_already_exists {
1519            let mut state = PodStateWithExtensionsMut::<S>::unpack(&mut buffer)?;
1520            state.realloc_variable_len_extension(new_extension)?;
1521        } else {
1522            if previous_account_len <= BASE_ACCOUNT_LENGTH {
1523                set_account_type::<S>(*buffer)?;
1524            }
1525            // now alloc in the TLV buffer and write the data
1526            let mut state = PodStateWithExtensionsMut::<S>::unpack(&mut buffer)?;
1527            state.init_variable_len_extension(new_extension, false)?;
1528        }
1529    } else {
1530        // do it backwards otherwise, write the state, realloc TLV, then the account
1531        let mut buffer = account_info.try_borrow_mut_data()?;
1532        let mut state = PodStateWithExtensionsMut::<S>::unpack(&mut buffer)?;
1533        if extension_already_exists {
1534            state.realloc_variable_len_extension(new_extension)?;
1535        } else {
1536            // this situation can happen if we have an overallocated buffer
1537            state.init_variable_len_extension(new_extension, false)?;
1538        }
1539
1540        let removed_bytes = previous_account_len
1541            .checked_sub(new_account_len)
1542            .ok_or(ProgramError::AccountDataTooSmall)?;
1543        if removed_bytes > 0 {
1544            // this is probably fine, but be safe and avoid invalidating references
1545            drop(buffer);
1546            account_info.realloc(new_account_len, false)?;
1547        }
1548    }
1549    Ok(())
1550}
1551
1552#[cfg(test)]
1553mod test {
1554    use {
1555        super::*,
1556        crate::{
1557            pod::test::{TEST_POD_ACCOUNT, TEST_POD_MINT},
1558            state::test::{TEST_ACCOUNT_SLICE, TEST_MINT_SLICE},
1559        },
1560        bytemuck::Pod,
1561        solana_program::{
1562            account_info::{Account as GetAccount, IntoAccountInfo},
1563            clock::Epoch,
1564            entrypoint::MAX_PERMITTED_DATA_INCREASE,
1565            pubkey::Pubkey,
1566        },
1567        spl_pod::{
1568            bytemuck::pod_bytes_of,
1569            optional_keys::OptionalNonZeroPubkey,
1570            primitives::{PodBool, PodU64},
1571        },
1572        transfer_fee::test::test_transfer_fee_config,
1573    };
1574
1575    /// Test fixed-length struct
1576    #[repr(C)]
1577    #[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
1578    struct FixedLenMintTest {
1579        data: [u8; 8],
1580    }
1581    impl Extension for FixedLenMintTest {
1582        const TYPE: ExtensionType = ExtensionType::MintPaddingTest;
1583    }
1584
1585    /// Test variable-length struct
1586    #[derive(Clone, Debug, PartialEq)]
1587    struct VariableLenMintTest {
1588        data: Vec<u8>,
1589    }
1590    impl Extension for VariableLenMintTest {
1591        const TYPE: ExtensionType = ExtensionType::VariableLenMintTest;
1592    }
1593    impl VariableLenPack for VariableLenMintTest {
1594        fn pack_into_slice(&self, dst: &mut [u8]) -> Result<(), ProgramError> {
1595            let data_start = size_of::<u64>();
1596            let end = data_start + self.data.len();
1597            if dst.len() < end {
1598                Err(ProgramError::InvalidAccountData)
1599            } else {
1600                dst[..data_start].copy_from_slice(&self.data.len().to_le_bytes());
1601                dst[data_start..end].copy_from_slice(&self.data);
1602                Ok(())
1603            }
1604        }
1605        fn unpack_from_slice(src: &[u8]) -> Result<Self, ProgramError> {
1606            let data_start = size_of::<u64>();
1607            let length = u64::from_le_bytes(src[..data_start].try_into().unwrap()) as usize;
1608            if src[data_start..data_start + length].len() != length {
1609                return Err(ProgramError::InvalidAccountData);
1610            }
1611            let data = Vec::from(&src[data_start..data_start + length]);
1612            Ok(Self { data })
1613        }
1614        fn get_packed_len(&self) -> Result<usize, ProgramError> {
1615            Ok(size_of::<u64>().saturating_add(self.data.len()))
1616        }
1617    }
1618
1619    const MINT_WITH_EXTENSION: &[u8] = &[
1620        1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1621        1, 1, 1, 1, 1, 1, 42, 0, 0, 0, 0, 0, 0, 0, 7, 1, 1, 0, 0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
1622        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, // base mint
1623        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
1624        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
1625        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // padding
1626        1, // account type
1627        3, 0, // extension type
1628        32, 0, // length
1629        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1630        1, 1, // data
1631    ];
1632
1633    const ACCOUNT_WITH_EXTENSION: &[u8] = &[
1634        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1635        1, 1, // mint
1636        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
1637        2, 2, // owner
1638        3, 0, 0, 0, 0, 0, 0, 0, // amount
1639        1, 0, 0, 0, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
1640        4, 4, 4, 4, 4, 4, // delegate
1641        2, // account state
1642        1, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, // is native
1643        6, 0, 0, 0, 0, 0, 0, 0, // delegated amount
1644        1, 0, 0, 0, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
1645        7, 7, 7, 7, 7, 7, // close authority
1646        2, // account type
1647        15, 0, // extension type
1648        1, 0, // length
1649        1, // data
1650    ];
1651
1652    #[test]
1653    fn unpack_opaque_buffer() {
1654        // Mint
1655        let state = PodStateWithExtensions::<PodMint>::unpack(MINT_WITH_EXTENSION).unwrap();
1656        assert_eq!(state.base, &TEST_POD_MINT);
1657        let extension = state.get_extension::<MintCloseAuthority>().unwrap();
1658        let close_authority =
1659            OptionalNonZeroPubkey::try_from(Some(Pubkey::new_from_array([1; 32]))).unwrap();
1660        assert_eq!(extension.close_authority, close_authority);
1661        assert_eq!(
1662            state.get_extension::<TransferFeeConfig>(),
1663            Err(ProgramError::InvalidAccountData)
1664        );
1665        assert_eq!(
1666            PodStateWithExtensions::<PodAccount>::unpack(MINT_WITH_EXTENSION),
1667            Err(ProgramError::UninitializedAccount)
1668        );
1669
1670        let state = PodStateWithExtensions::<PodMint>::unpack(TEST_MINT_SLICE).unwrap();
1671        assert_eq!(state.base, &TEST_POD_MINT);
1672
1673        let mut test_mint = TEST_MINT_SLICE.to_vec();
1674        let state = PodStateWithExtensionsMut::<PodMint>::unpack(&mut test_mint).unwrap();
1675        assert_eq!(state.base, &TEST_POD_MINT);
1676
1677        // Account
1678        let state = PodStateWithExtensions::<PodAccount>::unpack(ACCOUNT_WITH_EXTENSION).unwrap();
1679        assert_eq!(state.base, &TEST_POD_ACCOUNT);
1680        let extension = state.get_extension::<TransferHookAccount>().unwrap();
1681        let transferring = PodBool::from(true);
1682        assert_eq!(extension.transferring, transferring);
1683        assert_eq!(
1684            PodStateWithExtensions::<PodMint>::unpack(ACCOUNT_WITH_EXTENSION),
1685            Err(ProgramError::InvalidAccountData)
1686        );
1687
1688        let state = PodStateWithExtensions::<PodAccount>::unpack(TEST_ACCOUNT_SLICE).unwrap();
1689        assert_eq!(state.base, &TEST_POD_ACCOUNT);
1690
1691        let mut test_account = TEST_ACCOUNT_SLICE.to_vec();
1692        let state = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut test_account).unwrap();
1693        assert_eq!(state.base, &TEST_POD_ACCOUNT);
1694    }
1695
1696    #[test]
1697    fn mint_fail_unpack_opaque_buffer() {
1698        // input buffer too small
1699        let mut buffer = vec![0, 3];
1700        assert_eq!(
1701            PodStateWithExtensions::<PodMint>::unpack(&buffer),
1702            Err(ProgramError::InvalidAccountData)
1703        );
1704        assert_eq!(
1705            PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer),
1706            Err(ProgramError::InvalidAccountData)
1707        );
1708        assert_eq!(
1709            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer),
1710            Err(ProgramError::InvalidAccountData)
1711        );
1712
1713        // tweak the account type
1714        let mut buffer = MINT_WITH_EXTENSION.to_vec();
1715        buffer[BASE_ACCOUNT_LENGTH] = 3;
1716        assert_eq!(
1717            PodStateWithExtensions::<PodMint>::unpack(&buffer),
1718            Err(ProgramError::InvalidAccountData)
1719        );
1720
1721        // clear the mint initialized byte
1722        let mut buffer = MINT_WITH_EXTENSION.to_vec();
1723        buffer[45] = 0;
1724        assert_eq!(
1725            PodStateWithExtensions::<PodMint>::unpack(&buffer),
1726            Err(ProgramError::UninitializedAccount)
1727        );
1728
1729        // tweak the padding
1730        let mut buffer = MINT_WITH_EXTENSION.to_vec();
1731        buffer[PodMint::SIZE_OF] = 100;
1732        assert_eq!(
1733            PodStateWithExtensions::<PodMint>::unpack(&buffer),
1734            Err(ProgramError::InvalidAccountData)
1735        );
1736
1737        // tweak the extension type
1738        let mut buffer = MINT_WITH_EXTENSION.to_vec();
1739        buffer[BASE_ACCOUNT_LENGTH + 1] = 2;
1740        let state = PodStateWithExtensions::<PodMint>::unpack(&buffer).unwrap();
1741        assert_eq!(
1742            state.get_extension::<TransferFeeConfig>(),
1743            Err(ProgramError::Custom(
1744                TokenError::ExtensionTypeMismatch as u32
1745            ))
1746        );
1747
1748        // tweak the length, too big
1749        let mut buffer = MINT_WITH_EXTENSION.to_vec();
1750        buffer[BASE_ACCOUNT_LENGTH + 3] = 100;
1751        let state = PodStateWithExtensions::<PodMint>::unpack(&buffer).unwrap();
1752        assert_eq!(
1753            state.get_extension::<TransferFeeConfig>(),
1754            Err(ProgramError::InvalidAccountData)
1755        );
1756
1757        // tweak the length, too small
1758        let mut buffer = MINT_WITH_EXTENSION.to_vec();
1759        buffer[BASE_ACCOUNT_LENGTH + 3] = 10;
1760        let state = PodStateWithExtensions::<PodMint>::unpack(&buffer).unwrap();
1761        assert_eq!(
1762            state.get_extension::<TransferFeeConfig>(),
1763            Err(ProgramError::InvalidAccountData)
1764        );
1765
1766        // data buffer is too small
1767        let buffer = &MINT_WITH_EXTENSION[..MINT_WITH_EXTENSION.len() - 1];
1768        let state = PodStateWithExtensions::<PodMint>::unpack(buffer).unwrap();
1769        assert_eq!(
1770            state.get_extension::<MintCloseAuthority>(),
1771            Err(ProgramError::InvalidAccountData)
1772        );
1773    }
1774
1775    #[test]
1776    fn account_fail_unpack_opaque_buffer() {
1777        // input buffer too small
1778        let mut buffer = vec![0, 3];
1779        assert_eq!(
1780            PodStateWithExtensions::<PodAccount>::unpack(&buffer),
1781            Err(ProgramError::InvalidAccountData)
1782        );
1783        assert_eq!(
1784            PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer),
1785            Err(ProgramError::InvalidAccountData)
1786        );
1787        assert_eq!(
1788            PodStateWithExtensionsMut::<PodAccount>::unpack_uninitialized(&mut buffer),
1789            Err(ProgramError::InvalidAccountData)
1790        );
1791
1792        // input buffer invalid
1793        // all 5's - not a valid `AccountState`
1794        let mut buffer = vec![5; BASE_ACCOUNT_LENGTH];
1795        assert_eq!(
1796            PodStateWithExtensions::<PodAccount>::unpack(&buffer),
1797            Err(ProgramError::UninitializedAccount)
1798        );
1799        assert_eq!(
1800            PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer),
1801            Err(ProgramError::UninitializedAccount)
1802        );
1803
1804        // tweak the account type
1805        let mut buffer = ACCOUNT_WITH_EXTENSION.to_vec();
1806        buffer[BASE_ACCOUNT_LENGTH] = 3;
1807        assert_eq!(
1808            PodStateWithExtensions::<PodAccount>::unpack(&buffer),
1809            Err(ProgramError::InvalidAccountData)
1810        );
1811
1812        // clear the state byte
1813        let mut buffer = ACCOUNT_WITH_EXTENSION.to_vec();
1814        buffer[108] = 0;
1815        assert_eq!(
1816            PodStateWithExtensions::<PodAccount>::unpack(&buffer),
1817            Err(ProgramError::UninitializedAccount)
1818        );
1819
1820        // tweak the extension type
1821        let mut buffer = ACCOUNT_WITH_EXTENSION.to_vec();
1822        buffer[BASE_ACCOUNT_LENGTH + 1] = 12;
1823        let state = PodStateWithExtensions::<PodAccount>::unpack(&buffer).unwrap();
1824        assert_eq!(
1825            state.get_extension::<TransferHookAccount>(),
1826            Err(ProgramError::Custom(
1827                TokenError::ExtensionTypeMismatch as u32
1828            ))
1829        );
1830
1831        // tweak the length, too big
1832        let mut buffer = ACCOUNT_WITH_EXTENSION.to_vec();
1833        buffer[BASE_ACCOUNT_LENGTH + 3] = 100;
1834        let state = PodStateWithExtensions::<PodAccount>::unpack(&buffer).unwrap();
1835        assert_eq!(
1836            state.get_extension::<TransferHookAccount>(),
1837            Err(ProgramError::InvalidAccountData)
1838        );
1839
1840        // tweak the length, too small
1841        let mut buffer = ACCOUNT_WITH_EXTENSION.to_vec();
1842        buffer[BASE_ACCOUNT_LENGTH + 3] = 10;
1843        let state = PodStateWithExtensions::<PodAccount>::unpack(&buffer).unwrap();
1844        assert_eq!(
1845            state.get_extension::<TransferHookAccount>(),
1846            Err(ProgramError::InvalidAccountData)
1847        );
1848
1849        // data buffer is too small
1850        let buffer = &ACCOUNT_WITH_EXTENSION[..ACCOUNT_WITH_EXTENSION.len() - 1];
1851        let state = PodStateWithExtensions::<PodAccount>::unpack(buffer).unwrap();
1852        assert_eq!(
1853            state.get_extension::<TransferHookAccount>(),
1854            Err(ProgramError::InvalidAccountData)
1855        );
1856    }
1857
1858    #[test]
1859    fn get_extension_types_with_opaque_buffer() {
1860        // incorrect due to the length
1861        assert_eq!(
1862            get_tlv_data_info(&[1, 0, 1, 1]).unwrap_err(),
1863            ProgramError::InvalidAccountData,
1864        );
1865        // incorrect due to the huge enum number
1866        assert_eq!(
1867            get_tlv_data_info(&[0, 1, 0, 0]).unwrap_err(),
1868            ProgramError::InvalidAccountData,
1869        );
1870        // correct due to the good enum number and zero length
1871        assert_eq!(
1872            get_tlv_data_info(&[1, 0, 0, 0]).unwrap(),
1873            TlvDataInfo {
1874                extension_types: vec![ExtensionType::try_from(1).unwrap()],
1875                used_len: add_type_and_length_to_len(0),
1876            }
1877        );
1878        // correct since it's just uninitialized data at the end
1879        assert_eq!(
1880            get_tlv_data_info(&[0, 0]).unwrap(),
1881            TlvDataInfo {
1882                extension_types: vec![],
1883                used_len: 0
1884            }
1885        );
1886    }
1887
1888    #[test]
1889    fn mint_with_extension_pack_unpack() {
1890        let mint_size = ExtensionType::try_calculate_account_len::<PodMint>(&[
1891            ExtensionType::MintCloseAuthority,
1892            ExtensionType::TransferFeeConfig,
1893        ])
1894        .unwrap();
1895        let mut buffer = vec![0; mint_size];
1896
1897        // fail unpack
1898        assert_eq!(
1899            PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer),
1900            Err(ProgramError::UninitializedAccount),
1901        );
1902
1903        let mut state =
1904            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
1905        // fail init account extension
1906        assert_eq!(
1907            state.init_extension::<TransferFeeAmount>(true),
1908            Err(ProgramError::InvalidAccountData),
1909        );
1910
1911        // success write extension
1912        let close_authority =
1913            OptionalNonZeroPubkey::try_from(Some(Pubkey::new_from_array([1; 32]))).unwrap();
1914        let extension = state.init_extension::<MintCloseAuthority>(true).unwrap();
1915        extension.close_authority = close_authority;
1916        assert_eq!(
1917            &state.get_extension_types().unwrap(),
1918            &[ExtensionType::MintCloseAuthority]
1919        );
1920
1921        // fail init extension when already initialized
1922        assert_eq!(
1923            state.init_extension::<MintCloseAuthority>(false),
1924            Err(ProgramError::Custom(
1925                TokenError::ExtensionAlreadyInitialized as u32
1926            ))
1927        );
1928
1929        // fail unpack as account, a mint extension was written
1930        assert_eq!(
1931            PodStateWithExtensionsMut::<PodAccount>::unpack_uninitialized(&mut buffer),
1932            Err(ProgramError::Custom(
1933                TokenError::ExtensionBaseMismatch as u32
1934            ))
1935        );
1936
1937        // fail unpack again, still no base data
1938        assert_eq!(
1939            PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer.clone()),
1940            Err(ProgramError::UninitializedAccount),
1941        );
1942
1943        // write base mint
1944        let mut state =
1945            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
1946        *state.base = TEST_POD_MINT;
1947        state.init_account_type().unwrap();
1948
1949        // check raw buffer
1950        let mut expect = TEST_MINT_SLICE.to_vec();
1951        expect.extend_from_slice(&[0; BASE_ACCOUNT_LENGTH - PodMint::SIZE_OF]); // padding
1952        expect.push(AccountType::Mint.into());
1953        expect.extend_from_slice(&(ExtensionType::MintCloseAuthority as u16).to_le_bytes());
1954        expect
1955            .extend_from_slice(&(pod_get_packed_len::<MintCloseAuthority>() as u16).to_le_bytes());
1956        expect.extend_from_slice(&[1; 32]); // data
1957        expect.extend_from_slice(&[0; size_of::<ExtensionType>()]);
1958        expect.extend_from_slice(&[0; size_of::<Length>()]);
1959        expect.extend_from_slice(&[0; size_of::<TransferFeeConfig>()]);
1960        assert_eq!(expect, buffer);
1961
1962        // unpack uninitialized will now fail because the PodMint is now initialized
1963        assert_eq!(
1964            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer.clone()),
1965            Err(TokenError::AlreadyInUse.into()),
1966        );
1967
1968        // check unpacking
1969        let mut state = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap();
1970
1971        // update base
1972        *state.base = TEST_POD_MINT;
1973        state.base.supply = (u64::from(state.base.supply) + 100).into();
1974
1975        // check unpacking
1976        let unpacked_extension = state.get_extension_mut::<MintCloseAuthority>().unwrap();
1977        assert_eq!(*unpacked_extension, MintCloseAuthority { close_authority });
1978
1979        // update extension
1980        let close_authority = OptionalNonZeroPubkey::try_from(None).unwrap();
1981        unpacked_extension.close_authority = close_authority;
1982
1983        // check updates are propagated
1984        let base = *state.base;
1985        let state = PodStateWithExtensions::<PodMint>::unpack(&buffer).unwrap();
1986        assert_eq!(state.base, &base);
1987        let unpacked_extension = state.get_extension::<MintCloseAuthority>().unwrap();
1988        assert_eq!(*unpacked_extension, MintCloseAuthority { close_authority });
1989
1990        // check raw buffer
1991        let mut expect = vec![];
1992        expect.extend_from_slice(bytemuck::bytes_of(&base));
1993        expect.extend_from_slice(&[0; BASE_ACCOUNT_LENGTH - PodMint::SIZE_OF]); // padding
1994        expect.push(AccountType::Mint.into());
1995        expect.extend_from_slice(&(ExtensionType::MintCloseAuthority as u16).to_le_bytes());
1996        expect
1997            .extend_from_slice(&(pod_get_packed_len::<MintCloseAuthority>() as u16).to_le_bytes());
1998        expect.extend_from_slice(&[0; 32]);
1999        expect.extend_from_slice(&[0; size_of::<ExtensionType>()]);
2000        expect.extend_from_slice(&[0; size_of::<Length>()]);
2001        expect.extend_from_slice(&[0; size_of::<TransferFeeConfig>()]);
2002        assert_eq!(expect, buffer);
2003
2004        // fail unpack as an account
2005        assert_eq!(
2006            PodStateWithExtensions::<PodAccount>::unpack(&buffer),
2007            Err(ProgramError::UninitializedAccount),
2008        );
2009
2010        let mut state = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap();
2011        // init one more extension
2012        let mint_transfer_fee = test_transfer_fee_config();
2013        let new_extension = state.init_extension::<TransferFeeConfig>(true).unwrap();
2014        new_extension.transfer_fee_config_authority =
2015            mint_transfer_fee.transfer_fee_config_authority;
2016        new_extension.withdraw_withheld_authority = mint_transfer_fee.withdraw_withheld_authority;
2017        new_extension.withheld_amount = mint_transfer_fee.withheld_amount;
2018        new_extension.older_transfer_fee = mint_transfer_fee.older_transfer_fee;
2019        new_extension.newer_transfer_fee = mint_transfer_fee.newer_transfer_fee;
2020
2021        assert_eq!(
2022            &state.get_extension_types().unwrap(),
2023            &[
2024                ExtensionType::MintCloseAuthority,
2025                ExtensionType::TransferFeeConfig
2026            ]
2027        );
2028
2029        // check raw buffer
2030        let mut expect = vec![];
2031        expect.extend_from_slice(pod_bytes_of(&base));
2032        expect.extend_from_slice(&[0; BASE_ACCOUNT_LENGTH - PodMint::SIZE_OF]); // padding
2033        expect.push(AccountType::Mint.into());
2034        expect.extend_from_slice(&(ExtensionType::MintCloseAuthority as u16).to_le_bytes());
2035        expect
2036            .extend_from_slice(&(pod_get_packed_len::<MintCloseAuthority>() as u16).to_le_bytes());
2037        expect.extend_from_slice(&[0; 32]); // data
2038        expect.extend_from_slice(&(ExtensionType::TransferFeeConfig as u16).to_le_bytes());
2039        expect.extend_from_slice(&(pod_get_packed_len::<TransferFeeConfig>() as u16).to_le_bytes());
2040        expect.extend_from_slice(pod_bytes_of(&mint_transfer_fee));
2041        assert_eq!(expect, buffer);
2042
2043        // fail to init one more extension that does not fit
2044        let mut state = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap();
2045        assert_eq!(
2046            state.init_extension::<MintPaddingTest>(true),
2047            Err(ProgramError::InvalidAccountData),
2048        );
2049    }
2050
2051    #[test]
2052    fn mint_extension_any_order() {
2053        let mint_size = ExtensionType::try_calculate_account_len::<PodMint>(&[
2054            ExtensionType::MintCloseAuthority,
2055            ExtensionType::TransferFeeConfig,
2056        ])
2057        .unwrap();
2058        let mut buffer = vec![0; mint_size];
2059
2060        let mut state =
2061            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2062        // write extensions
2063        let close_authority =
2064            OptionalNonZeroPubkey::try_from(Some(Pubkey::new_from_array([1; 32]))).unwrap();
2065        let extension = state.init_extension::<MintCloseAuthority>(true).unwrap();
2066        extension.close_authority = close_authority;
2067
2068        let mint_transfer_fee = test_transfer_fee_config();
2069        let extension = state.init_extension::<TransferFeeConfig>(true).unwrap();
2070        extension.transfer_fee_config_authority = mint_transfer_fee.transfer_fee_config_authority;
2071        extension.withdraw_withheld_authority = mint_transfer_fee.withdraw_withheld_authority;
2072        extension.withheld_amount = mint_transfer_fee.withheld_amount;
2073        extension.older_transfer_fee = mint_transfer_fee.older_transfer_fee;
2074        extension.newer_transfer_fee = mint_transfer_fee.newer_transfer_fee;
2075
2076        assert_eq!(
2077            &state.get_extension_types().unwrap(),
2078            &[
2079                ExtensionType::MintCloseAuthority,
2080                ExtensionType::TransferFeeConfig
2081            ]
2082        );
2083
2084        // write base mint
2085        let mut state =
2086            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2087        *state.base = TEST_POD_MINT;
2088        state.init_account_type().unwrap();
2089
2090        let mut other_buffer = vec![0; mint_size];
2091        let mut state =
2092            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut other_buffer).unwrap();
2093
2094        // write base mint
2095        *state.base = TEST_POD_MINT;
2096        state.init_account_type().unwrap();
2097
2098        // write extensions in a different order
2099        let mint_transfer_fee = test_transfer_fee_config();
2100        let extension = state.init_extension::<TransferFeeConfig>(true).unwrap();
2101        extension.transfer_fee_config_authority = mint_transfer_fee.transfer_fee_config_authority;
2102        extension.withdraw_withheld_authority = mint_transfer_fee.withdraw_withheld_authority;
2103        extension.withheld_amount = mint_transfer_fee.withheld_amount;
2104        extension.older_transfer_fee = mint_transfer_fee.older_transfer_fee;
2105        extension.newer_transfer_fee = mint_transfer_fee.newer_transfer_fee;
2106
2107        let close_authority =
2108            OptionalNonZeroPubkey::try_from(Some(Pubkey::new_from_array([1; 32]))).unwrap();
2109        let extension = state.init_extension::<MintCloseAuthority>(true).unwrap();
2110        extension.close_authority = close_authority;
2111
2112        assert_eq!(
2113            &state.get_extension_types().unwrap(),
2114            &[
2115                ExtensionType::TransferFeeConfig,
2116                ExtensionType::MintCloseAuthority
2117            ]
2118        );
2119
2120        // buffers are NOT the same because written in a different order
2121        assert_ne!(buffer, other_buffer);
2122        let state = PodStateWithExtensions::<PodMint>::unpack(&buffer).unwrap();
2123        let other_state = PodStateWithExtensions::<PodMint>::unpack(&other_buffer).unwrap();
2124
2125        // BUT mint and extensions are the same
2126        assert_eq!(
2127            state.get_extension::<TransferFeeConfig>().unwrap(),
2128            other_state.get_extension::<TransferFeeConfig>().unwrap()
2129        );
2130        assert_eq!(
2131            state.get_extension::<MintCloseAuthority>().unwrap(),
2132            other_state.get_extension::<MintCloseAuthority>().unwrap()
2133        );
2134        assert_eq!(state.base, other_state.base);
2135    }
2136
2137    #[test]
2138    fn mint_with_multisig_len() {
2139        let mut buffer = vec![0; Multisig::LEN];
2140        assert_eq!(
2141            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer),
2142            Err(ProgramError::InvalidAccountData),
2143        );
2144        let mint_size =
2145            ExtensionType::try_calculate_account_len::<PodMint>(&[ExtensionType::MintPaddingTest])
2146                .unwrap();
2147        assert_eq!(mint_size, Multisig::LEN + size_of::<ExtensionType>());
2148        let mut buffer = vec![0; mint_size];
2149
2150        // write base mint
2151        let mut state =
2152            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2153        *state.base = TEST_POD_MINT;
2154        state.init_account_type().unwrap();
2155
2156        // write padding
2157        let extension = state.init_extension::<MintPaddingTest>(true).unwrap();
2158        extension.padding1 = [1; 128];
2159        extension.padding2 = [1; 48];
2160        extension.padding3 = [1; 9];
2161
2162        assert_eq!(
2163            &state.get_extension_types().unwrap(),
2164            &[ExtensionType::MintPaddingTest]
2165        );
2166
2167        // check raw buffer
2168        let mut expect = TEST_MINT_SLICE.to_vec();
2169        expect.extend_from_slice(&[0; BASE_ACCOUNT_LENGTH - PodMint::SIZE_OF]); // padding
2170        expect.push(AccountType::Mint.into());
2171        expect.extend_from_slice(&(ExtensionType::MintPaddingTest as u16).to_le_bytes());
2172        expect.extend_from_slice(&(pod_get_packed_len::<MintPaddingTest>() as u16).to_le_bytes());
2173        expect.extend_from_slice(&vec![1; pod_get_packed_len::<MintPaddingTest>()]);
2174        expect.extend_from_slice(&(ExtensionType::Uninitialized as u16).to_le_bytes());
2175        assert_eq!(expect, buffer);
2176    }
2177
2178    #[test]
2179    fn account_with_extension_pack_unpack() {
2180        let account_size = ExtensionType::try_calculate_account_len::<PodAccount>(&[
2181            ExtensionType::TransferFeeAmount,
2182        ])
2183        .unwrap();
2184        let mut buffer = vec![0; account_size];
2185
2186        // fail unpack
2187        assert_eq!(
2188            PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer),
2189            Err(ProgramError::UninitializedAccount),
2190        );
2191
2192        let mut state =
2193            PodStateWithExtensionsMut::<PodAccount>::unpack_uninitialized(&mut buffer).unwrap();
2194        // fail init mint extension
2195        assert_eq!(
2196            state.init_extension::<TransferFeeConfig>(true),
2197            Err(ProgramError::InvalidAccountData),
2198        );
2199        // success write extension
2200        let withheld_amount = PodU64::from(u64::MAX);
2201        let extension = state.init_extension::<TransferFeeAmount>(true).unwrap();
2202        extension.withheld_amount = withheld_amount;
2203
2204        assert_eq!(
2205            &state.get_extension_types().unwrap(),
2206            &[ExtensionType::TransferFeeAmount]
2207        );
2208
2209        // fail unpack again, still no base data
2210        assert_eq!(
2211            PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer.clone()),
2212            Err(ProgramError::UninitializedAccount),
2213        );
2214
2215        // write base account
2216        let mut state =
2217            PodStateWithExtensionsMut::<PodAccount>::unpack_uninitialized(&mut buffer).unwrap();
2218        *state.base = TEST_POD_ACCOUNT;
2219        state.init_account_type().unwrap();
2220        let base = *state.base;
2221
2222        // check raw buffer
2223        let mut expect = TEST_ACCOUNT_SLICE.to_vec();
2224        expect.push(AccountType::Account.into());
2225        expect.extend_from_slice(&(ExtensionType::TransferFeeAmount as u16).to_le_bytes());
2226        expect.extend_from_slice(&(pod_get_packed_len::<TransferFeeAmount>() as u16).to_le_bytes());
2227        expect.extend_from_slice(&u64::from(withheld_amount).to_le_bytes());
2228        assert_eq!(expect, buffer);
2229
2230        // check unpacking
2231        let mut state = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer).unwrap();
2232        assert_eq!(state.base, &base);
2233        assert_eq!(
2234            &state.get_extension_types().unwrap(),
2235            &[ExtensionType::TransferFeeAmount]
2236        );
2237
2238        // update base
2239        *state.base = TEST_POD_ACCOUNT;
2240        state.base.amount = (u64::from(state.base.amount) + 100).into();
2241
2242        // check unpacking
2243        let unpacked_extension = state.get_extension_mut::<TransferFeeAmount>().unwrap();
2244        assert_eq!(*unpacked_extension, TransferFeeAmount { withheld_amount });
2245
2246        // update extension
2247        let withheld_amount = PodU64::from(u32::MAX as u64);
2248        unpacked_extension.withheld_amount = withheld_amount;
2249
2250        // check updates are propagated
2251        let base = *state.base;
2252        let state = PodStateWithExtensions::<PodAccount>::unpack(&buffer).unwrap();
2253        assert_eq!(state.base, &base);
2254        let unpacked_extension = state.get_extension::<TransferFeeAmount>().unwrap();
2255        assert_eq!(*unpacked_extension, TransferFeeAmount { withheld_amount });
2256
2257        // check raw buffer
2258        let mut expect = vec![];
2259        expect.extend_from_slice(pod_bytes_of(&base));
2260        expect.push(AccountType::Account.into());
2261        expect.extend_from_slice(&(ExtensionType::TransferFeeAmount as u16).to_le_bytes());
2262        expect.extend_from_slice(&(pod_get_packed_len::<TransferFeeAmount>() as u16).to_le_bytes());
2263        expect.extend_from_slice(&u64::from(withheld_amount).to_le_bytes());
2264        assert_eq!(expect, buffer);
2265
2266        // fail unpack as a mint
2267        assert_eq!(
2268            PodStateWithExtensions::<PodMint>::unpack(&buffer),
2269            Err(ProgramError::InvalidAccountData),
2270        );
2271    }
2272
2273    #[test]
2274    fn account_with_multisig_len() {
2275        let mut buffer = vec![0; Multisig::LEN];
2276        assert_eq!(
2277            PodStateWithExtensionsMut::<PodAccount>::unpack_uninitialized(&mut buffer),
2278            Err(ProgramError::InvalidAccountData),
2279        );
2280        let account_size = ExtensionType::try_calculate_account_len::<PodAccount>(&[
2281            ExtensionType::AccountPaddingTest,
2282        ])
2283        .unwrap();
2284        assert_eq!(account_size, Multisig::LEN + size_of::<ExtensionType>());
2285        let mut buffer = vec![0; account_size];
2286
2287        // write base account
2288        let mut state =
2289            PodStateWithExtensionsMut::<PodAccount>::unpack_uninitialized(&mut buffer).unwrap();
2290        *state.base = TEST_POD_ACCOUNT;
2291        state.init_account_type().unwrap();
2292
2293        // write padding
2294        let extension = state.init_extension::<AccountPaddingTest>(true).unwrap();
2295        extension.0.padding1 = [2; 128];
2296        extension.0.padding2 = [2; 48];
2297        extension.0.padding3 = [2; 9];
2298
2299        assert_eq!(
2300            &state.get_extension_types().unwrap(),
2301            &[ExtensionType::AccountPaddingTest]
2302        );
2303
2304        // check raw buffer
2305        let mut expect = TEST_ACCOUNT_SLICE.to_vec();
2306        expect.push(AccountType::Account.into());
2307        expect.extend_from_slice(&(ExtensionType::AccountPaddingTest as u16).to_le_bytes());
2308        expect
2309            .extend_from_slice(&(pod_get_packed_len::<AccountPaddingTest>() as u16).to_le_bytes());
2310        expect.extend_from_slice(&vec![2; pod_get_packed_len::<AccountPaddingTest>()]);
2311        expect.extend_from_slice(&(ExtensionType::Uninitialized as u16).to_le_bytes());
2312        assert_eq!(expect, buffer);
2313    }
2314
2315    #[test]
2316    fn test_set_account_type() {
2317        // account with buffer big enough for AccountType and Extension
2318        let mut buffer = TEST_ACCOUNT_SLICE.to_vec();
2319        let needed_len = ExtensionType::try_calculate_account_len::<PodAccount>(&[
2320            ExtensionType::ImmutableOwner,
2321        ])
2322        .unwrap()
2323            - buffer.len();
2324        buffer.append(&mut vec![0; needed_len]);
2325        let err = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer).unwrap_err();
2326        assert_eq!(err, ProgramError::InvalidAccountData);
2327        set_account_type::<PodAccount>(&mut buffer).unwrap();
2328        // unpack is viable after manual set_account_type
2329        let mut state = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer).unwrap();
2330        assert_eq!(state.base, &TEST_POD_ACCOUNT);
2331        assert_eq!(state.account_type[0], AccountType::Account as u8);
2332        state.init_extension::<ImmutableOwner>(true).unwrap(); // just confirming initialization works
2333
2334        // account with buffer big enough for AccountType only
2335        let mut buffer = TEST_ACCOUNT_SLICE.to_vec();
2336        buffer.append(&mut vec![0; 2]);
2337        let err = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer).unwrap_err();
2338        assert_eq!(err, ProgramError::InvalidAccountData);
2339        set_account_type::<PodAccount>(&mut buffer).unwrap();
2340        // unpack is viable after manual set_account_type
2341        let state = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer).unwrap();
2342        assert_eq!(state.base, &TEST_POD_ACCOUNT);
2343        assert_eq!(state.account_type[0], AccountType::Account as u8);
2344
2345        // account with AccountType already set => noop
2346        let mut buffer = TEST_ACCOUNT_SLICE.to_vec();
2347        buffer.append(&mut vec![2, 0]);
2348        let _ = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer).unwrap();
2349        set_account_type::<PodAccount>(&mut buffer).unwrap();
2350        let state = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer).unwrap();
2351        assert_eq!(state.base, &TEST_POD_ACCOUNT);
2352        assert_eq!(state.account_type[0], AccountType::Account as u8);
2353
2354        // account with wrong AccountType fails
2355        let mut buffer = TEST_ACCOUNT_SLICE.to_vec();
2356        buffer.append(&mut vec![1, 0]);
2357        let err = PodStateWithExtensionsMut::<PodAccount>::unpack(&mut buffer).unwrap_err();
2358        assert_eq!(err, ProgramError::InvalidAccountData);
2359        let err = set_account_type::<PodAccount>(&mut buffer).unwrap_err();
2360        assert_eq!(err, ProgramError::InvalidAccountData);
2361
2362        // mint with buffer big enough for AccountType and Extension
2363        let mut buffer = TEST_MINT_SLICE.to_vec();
2364        let needed_len = ExtensionType::try_calculate_account_len::<PodMint>(&[
2365            ExtensionType::MintCloseAuthority,
2366        ])
2367        .unwrap()
2368            - buffer.len();
2369        buffer.append(&mut vec![0; needed_len]);
2370        let err = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap_err();
2371        assert_eq!(err, ProgramError::InvalidAccountData);
2372        set_account_type::<PodMint>(&mut buffer).unwrap();
2373        // unpack is viable after manual set_account_type
2374        let mut state = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap();
2375        assert_eq!(state.base, &TEST_POD_MINT);
2376        assert_eq!(state.account_type[0], AccountType::Mint as u8);
2377        state.init_extension::<MintCloseAuthority>(true).unwrap();
2378
2379        // mint with buffer big enough for AccountType only
2380        let mut buffer = TEST_MINT_SLICE.to_vec();
2381        buffer.append(&mut vec![0; PodAccount::SIZE_OF - PodMint::SIZE_OF]);
2382        buffer.append(&mut vec![0; 2]);
2383        let err = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap_err();
2384        assert_eq!(err, ProgramError::InvalidAccountData);
2385        set_account_type::<PodMint>(&mut buffer).unwrap();
2386        // unpack is viable after manual set_account_type
2387        let state = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap();
2388        assert_eq!(state.base, &TEST_POD_MINT);
2389        assert_eq!(state.account_type[0], AccountType::Mint as u8);
2390
2391        // mint with AccountType already set => noop
2392        let mut buffer = TEST_MINT_SLICE.to_vec();
2393        buffer.append(&mut vec![0; PodAccount::SIZE_OF - PodMint::SIZE_OF]);
2394        buffer.append(&mut vec![1, 0]);
2395        set_account_type::<PodMint>(&mut buffer).unwrap();
2396        let state = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap();
2397        assert_eq!(state.base, &TEST_POD_MINT);
2398        assert_eq!(state.account_type[0], AccountType::Mint as u8);
2399
2400        // mint with wrong AccountType fails
2401        let mut buffer = TEST_MINT_SLICE.to_vec();
2402        buffer.append(&mut vec![0; PodAccount::SIZE_OF - PodMint::SIZE_OF]);
2403        buffer.append(&mut vec![2, 0]);
2404        let err = PodStateWithExtensionsMut::<PodMint>::unpack(&mut buffer).unwrap_err();
2405        assert_eq!(err, ProgramError::InvalidAccountData);
2406        let err = set_account_type::<PodMint>(&mut buffer).unwrap_err();
2407        assert_eq!(err, ProgramError::InvalidAccountData);
2408    }
2409
2410    #[test]
2411    fn test_set_account_type_wrongly() {
2412        // try to set PodAccount account_type to PodMint
2413        let mut buffer = TEST_ACCOUNT_SLICE.to_vec();
2414        buffer.append(&mut vec![0; 2]);
2415        let err = set_account_type::<PodMint>(&mut buffer).unwrap_err();
2416        assert_eq!(err, ProgramError::InvalidAccountData);
2417
2418        // try to set PodMint account_type to PodAccount
2419        let mut buffer = TEST_MINT_SLICE.to_vec();
2420        buffer.append(&mut vec![0; PodAccount::SIZE_OF - PodMint::SIZE_OF]);
2421        buffer.append(&mut vec![0; 2]);
2422        let err = set_account_type::<PodAccount>(&mut buffer).unwrap_err();
2423        assert_eq!(err, ProgramError::InvalidAccountData);
2424    }
2425
2426    #[test]
2427    fn test_get_required_init_account_extensions() {
2428        // Some mint extensions with no required account extensions
2429        let mint_extensions = vec![
2430            ExtensionType::MintCloseAuthority,
2431            ExtensionType::Uninitialized,
2432        ];
2433        assert_eq!(
2434            ExtensionType::get_required_init_account_extensions(&mint_extensions),
2435            vec![]
2436        );
2437
2438        // One mint extension with required account extension, one without
2439        let mint_extensions = vec![
2440            ExtensionType::TransferFeeConfig,
2441            ExtensionType::MintCloseAuthority,
2442        ];
2443        assert_eq!(
2444            ExtensionType::get_required_init_account_extensions(&mint_extensions),
2445            vec![ExtensionType::TransferFeeAmount]
2446        );
2447
2448        // Some mint extensions both with required account extensions
2449        let mint_extensions = vec![
2450            ExtensionType::TransferFeeConfig,
2451            ExtensionType::MintPaddingTest,
2452        ];
2453        assert_eq!(
2454            ExtensionType::get_required_init_account_extensions(&mint_extensions),
2455            vec![
2456                ExtensionType::TransferFeeAmount,
2457                ExtensionType::AccountPaddingTest
2458            ]
2459        );
2460
2461        // Demonstrate that method does not dedupe inputs or outputs
2462        let mint_extensions = vec![
2463            ExtensionType::TransferFeeConfig,
2464            ExtensionType::TransferFeeConfig,
2465        ];
2466        assert_eq!(
2467            ExtensionType::get_required_init_account_extensions(&mint_extensions),
2468            vec![
2469                ExtensionType::TransferFeeAmount,
2470                ExtensionType::TransferFeeAmount
2471            ]
2472        );
2473    }
2474
2475    #[test]
2476    fn mint_without_extensions() {
2477        let space = ExtensionType::try_calculate_account_len::<PodMint>(&[]).unwrap();
2478        let mut buffer = vec![0; space];
2479        assert_eq!(
2480            PodStateWithExtensionsMut::<PodAccount>::unpack_uninitialized(&mut buffer),
2481            Err(ProgramError::InvalidAccountData),
2482        );
2483
2484        // write base account
2485        let mut state =
2486            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2487        *state.base = TEST_POD_MINT;
2488        state.init_account_type().unwrap();
2489
2490        // fail init extension
2491        assert_eq!(
2492            state.init_extension::<TransferFeeConfig>(true),
2493            Err(ProgramError::InvalidAccountData),
2494        );
2495
2496        assert_eq!(TEST_MINT_SLICE, buffer);
2497    }
2498
2499    #[test]
2500    fn test_init_nonzero_default() {
2501        let mint_size =
2502            ExtensionType::try_calculate_account_len::<PodMint>(&[ExtensionType::MintPaddingTest])
2503                .unwrap();
2504        let mut buffer = vec![0; mint_size];
2505        let mut state =
2506            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2507        *state.base = TEST_POD_MINT;
2508        state.init_account_type().unwrap();
2509        let extension = state.init_extension::<MintPaddingTest>(true).unwrap();
2510        assert_eq!(extension.padding1, [1; 128]);
2511        assert_eq!(extension.padding2, [2; 48]);
2512        assert_eq!(extension.padding3, [3; 9]);
2513    }
2514
2515    #[test]
2516    fn test_init_buffer_too_small() {
2517        let mint_size = ExtensionType::try_calculate_account_len::<PodMint>(&[
2518            ExtensionType::MintCloseAuthority,
2519        ])
2520        .unwrap();
2521        let mut buffer = vec![0; mint_size - 1];
2522        let mut state =
2523            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2524        let err = state
2525            .init_extension::<MintCloseAuthority>(true)
2526            .unwrap_err();
2527        assert_eq!(err, ProgramError::InvalidAccountData);
2528
2529        state.tlv_data[0] = 3;
2530        state.tlv_data[2] = 32;
2531        let err = state.get_extension_mut::<MintCloseAuthority>().unwrap_err();
2532        assert_eq!(err, ProgramError::InvalidAccountData);
2533
2534        let mut buffer = vec![0; PodMint::SIZE_OF + 2];
2535        let err =
2536            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap_err();
2537        assert_eq!(err, ProgramError::InvalidAccountData);
2538
2539        // OK since there are two bytes for the type, which is `Uninitialized`
2540        let mut buffer = vec![0; BASE_ACCOUNT_LENGTH + 3];
2541        let mut state =
2542            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2543        let err = state.get_extension_mut::<MintCloseAuthority>().unwrap_err();
2544        assert_eq!(err, ProgramError::InvalidAccountData);
2545
2546        assert_eq!(state.get_extension_types().unwrap(), vec![]);
2547
2548        // OK, there aren't two bytes for the type, but that's fine
2549        let mut buffer = vec![0; BASE_ACCOUNT_LENGTH + 2];
2550        let state =
2551            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2552        assert_eq!(state.get_extension_types().unwrap(), []);
2553    }
2554
2555    #[test]
2556    fn test_extension_with_no_data() {
2557        let account_size = ExtensionType::try_calculate_account_len::<PodAccount>(&[
2558            ExtensionType::ImmutableOwner,
2559        ])
2560        .unwrap();
2561        let mut buffer = vec![0; account_size];
2562        let mut state =
2563            PodStateWithExtensionsMut::<PodAccount>::unpack_uninitialized(&mut buffer).unwrap();
2564        *state.base = TEST_POD_ACCOUNT;
2565        state.init_account_type().unwrap();
2566
2567        let err = state.get_extension::<ImmutableOwner>().unwrap_err();
2568        assert_eq!(
2569            err,
2570            ProgramError::Custom(TokenError::ExtensionNotFound as u32)
2571        );
2572
2573        state.init_extension::<ImmutableOwner>(true).unwrap();
2574        assert_eq!(
2575            get_first_extension_type(state.tlv_data).unwrap(),
2576            Some(ExtensionType::ImmutableOwner)
2577        );
2578        assert_eq!(
2579            get_tlv_data_info(state.tlv_data).unwrap(),
2580            TlvDataInfo {
2581                extension_types: vec![ExtensionType::ImmutableOwner],
2582                used_len: add_type_and_length_to_len(0)
2583            }
2584        );
2585    }
2586
2587    #[test]
2588    fn fail_account_len_with_metadata() {
2589        assert_eq!(
2590            ExtensionType::try_calculate_account_len::<PodMint>(&[
2591                ExtensionType::MintCloseAuthority,
2592                ExtensionType::VariableLenMintTest,
2593                ExtensionType::TransferFeeConfig,
2594            ])
2595            .unwrap_err(),
2596            ProgramError::InvalidArgument
2597        );
2598    }
2599
2600    #[test]
2601    fn alloc() {
2602        let variable_len = VariableLenMintTest { data: vec![1] };
2603        let alloc_size = variable_len.get_packed_len().unwrap();
2604        let account_size =
2605            BASE_ACCOUNT_LENGTH + size_of::<AccountType>() + add_type_and_length_to_len(alloc_size);
2606        let mut buffer = vec![0; account_size];
2607        let mut state =
2608            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2609        state
2610            .init_variable_len_extension(&variable_len, false)
2611            .unwrap();
2612
2613        // can't double alloc
2614        assert_eq!(
2615            state
2616                .init_variable_len_extension(&variable_len, false)
2617                .unwrap_err(),
2618            TokenError::ExtensionAlreadyInitialized.into()
2619        );
2620
2621        // unless overwrite is set
2622        state
2623            .init_variable_len_extension(&variable_len, true)
2624            .unwrap();
2625
2626        // can't change the size during overwrite though
2627        assert_eq!(
2628            state
2629                .init_variable_len_extension(&VariableLenMintTest { data: vec![] }, true)
2630                .unwrap_err(),
2631            TokenError::InvalidLengthForAlloc.into()
2632        );
2633
2634        // try to write too far, fail earlier
2635        assert_eq!(
2636            state
2637                .init_variable_len_extension(&VariableLenMintTest { data: vec![1, 2] }, true)
2638                .unwrap_err(),
2639            ProgramError::InvalidAccountData
2640        );
2641    }
2642
2643    #[test]
2644    fn realloc() {
2645        let small_variable_len = VariableLenMintTest {
2646            data: vec![1, 2, 3],
2647        };
2648        let base_variable_len = VariableLenMintTest {
2649            data: vec![1, 2, 3, 4],
2650        };
2651        let big_variable_len = VariableLenMintTest {
2652            data: vec![1, 2, 3, 4, 5],
2653        };
2654        let too_big_variable_len = VariableLenMintTest {
2655            data: vec![1, 2, 3, 4, 5, 6],
2656        };
2657        let account_size =
2658            ExtensionType::try_calculate_account_len::<PodMint>(&[ExtensionType::MetadataPointer])
2659                .unwrap()
2660                + add_type_and_length_to_len(big_variable_len.get_packed_len().unwrap());
2661        let mut buffer = vec![0; account_size];
2662        let mut state =
2663            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2664
2665        // alloc both types
2666        state
2667            .init_variable_len_extension(&base_variable_len, false)
2668            .unwrap();
2669        let max_pubkey =
2670            OptionalNonZeroPubkey::try_from(Some(Pubkey::new_from_array([255; 32]))).unwrap();
2671        let extension = state.init_extension::<MetadataPointer>(false).unwrap();
2672        extension.authority = max_pubkey;
2673        extension.metadata_address = max_pubkey;
2674
2675        // realloc first entry to larger
2676        state
2677            .realloc_variable_len_extension(&big_variable_len)
2678            .unwrap();
2679        let extension = state
2680            .get_variable_len_extension::<VariableLenMintTest>()
2681            .unwrap();
2682        assert_eq!(extension, big_variable_len);
2683        let extension = state.get_extension::<MetadataPointer>().unwrap();
2684        assert_eq!(extension.authority, max_pubkey);
2685        assert_eq!(extension.metadata_address, max_pubkey);
2686
2687        // realloc to smaller
2688        state
2689            .realloc_variable_len_extension(&small_variable_len)
2690            .unwrap();
2691        let extension = state
2692            .get_variable_len_extension::<VariableLenMintTest>()
2693            .unwrap();
2694        assert_eq!(extension, small_variable_len);
2695        let extension = state.get_extension::<MetadataPointer>().unwrap();
2696        assert_eq!(extension.authority, max_pubkey);
2697        assert_eq!(extension.metadata_address, max_pubkey);
2698        let diff = big_variable_len.get_packed_len().unwrap()
2699            - small_variable_len.get_packed_len().unwrap();
2700        assert_eq!(&buffer[account_size - diff..account_size], vec![0; diff]);
2701
2702        // unpack again since we dropped the last `state`
2703        let mut state =
2704            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2705        // realloc too much, fails
2706        assert_eq!(
2707            state
2708                .realloc_variable_len_extension(&too_big_variable_len)
2709                .unwrap_err(),
2710            ProgramError::InvalidAccountData,
2711        );
2712    }
2713
2714    #[test]
2715    fn account_len() {
2716        let small_variable_len = VariableLenMintTest {
2717            data: vec![20, 30, 40],
2718        };
2719        let variable_len = VariableLenMintTest {
2720            data: vec![20, 30, 40, 50],
2721        };
2722        let big_variable_len = VariableLenMintTest {
2723            data: vec![20, 30, 40, 50, 60],
2724        };
2725        let value_len = variable_len.get_packed_len().unwrap();
2726        let account_size =
2727            BASE_ACCOUNT_LENGTH + size_of::<AccountType>() + add_type_and_length_to_len(value_len);
2728        let mut buffer = vec![0; account_size];
2729        let mut state =
2730            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2731
2732        // With a new extension, new length must include padding, 1 byte for
2733        // account type, 2 bytes for type, 2 for length
2734        let current_len = state.try_get_account_len().unwrap();
2735        assert_eq!(current_len, PodMint::SIZE_OF);
2736        let new_len = state
2737            .try_get_new_account_len_for_variable_len_extension::<VariableLenMintTest>(
2738                &variable_len,
2739            )
2740            .unwrap();
2741        assert_eq!(
2742            new_len,
2743            BASE_ACCOUNT_AND_TYPE_LENGTH.saturating_add(add_type_and_length_to_len(value_len))
2744        );
2745
2746        state
2747            .init_variable_len_extension::<VariableLenMintTest>(&variable_len, false)
2748            .unwrap();
2749        let current_len = state.try_get_account_len().unwrap();
2750        assert_eq!(current_len, new_len);
2751
2752        // Reduce the extension size
2753        let new_len = state
2754            .try_get_new_account_len_for_variable_len_extension::<VariableLenMintTest>(
2755                &small_variable_len,
2756            )
2757            .unwrap();
2758        assert_eq!(current_len.checked_sub(new_len).unwrap(), 1);
2759
2760        // Increase the extension size
2761        let new_len = state
2762            .try_get_new_account_len_for_variable_len_extension::<VariableLenMintTest>(
2763                &big_variable_len,
2764            )
2765            .unwrap();
2766        assert_eq!(new_len.checked_sub(current_len).unwrap(), 1);
2767
2768        // Maintain the extension size
2769        let new_len = state
2770            .try_get_new_account_len_for_variable_len_extension::<VariableLenMintTest>(
2771                &variable_len,
2772            )
2773            .unwrap();
2774        assert_eq!(new_len, current_len);
2775    }
2776
2777    /// Test helper for mimicking the data layout an on-chain `AccountInfo`,
2778    /// which permits "reallocs" as the Solana runtime does it
2779    struct SolanaAccountData {
2780        data: Vec<u8>,
2781        lamports: u64,
2782        owner: Pubkey,
2783    }
2784    impl SolanaAccountData {
2785        /// Create a new fake solana account data. The underlying vector is
2786        /// overallocated to mimic the runtime
2787        fn new(account_data: &[u8]) -> Self {
2788            let mut data = vec![];
2789            data.extend_from_slice(&(account_data.len() as u64).to_le_bytes());
2790            data.extend_from_slice(account_data);
2791            data.extend_from_slice(&[0; MAX_PERMITTED_DATA_INCREASE]);
2792            Self {
2793                data,
2794                lamports: 10,
2795                owner: Pubkey::new_unique(),
2796            }
2797        }
2798
2799        /// Data lops off the first 8 bytes, since those store the size of the
2800        /// account for the Solana runtime
2801        fn data(&self) -> &[u8] {
2802            let start = size_of::<u64>();
2803            let len = self.len();
2804            &self.data[start..start + len]
2805        }
2806
2807        /// Gets the runtime length of the account data
2808        fn len(&self) -> usize {
2809            self.data
2810                .get(..size_of::<u64>())
2811                .and_then(|slice| slice.try_into().ok())
2812                .map(u64::from_le_bytes)
2813                .unwrap() as usize
2814        }
2815    }
2816    impl GetAccount for SolanaAccountData {
2817        fn get(&mut self) -> (&mut u64, &mut [u8], &Pubkey, bool, Epoch) {
2818            // need to pull out the data here to avoid a double-mutable borrow
2819            let start = size_of::<u64>();
2820            let len = self.len();
2821            (
2822                &mut self.lamports,
2823                &mut self.data[start..start + len],
2824                &self.owner,
2825                false,
2826                Epoch::default(),
2827            )
2828        }
2829    }
2830
2831    #[test]
2832    fn alloc_new_fixed_len_tlv_in_account_info_from_base_size() {
2833        let fixed_len = FixedLenMintTest {
2834            data: [1, 2, 3, 4, 5, 6, 7, 8],
2835        };
2836        let value_len = pod_get_packed_len::<FixedLenMintTest>();
2837        let base_account_size = PodMint::SIZE_OF;
2838        let mut buffer = vec![0; base_account_size];
2839        let state =
2840            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2841        *state.base = TEST_POD_MINT;
2842
2843        let mut data = SolanaAccountData::new(&buffer);
2844        let key = Pubkey::new_unique();
2845        let account_info = (&key, &mut data).into_account_info();
2846
2847        alloc_and_serialize::<PodMint, _>(&account_info, &fixed_len, false).unwrap();
2848        let new_account_len = BASE_ACCOUNT_AND_TYPE_LENGTH + add_type_and_length_to_len(value_len);
2849        assert_eq!(data.len(), new_account_len);
2850        let state = PodStateWithExtensions::<PodMint>::unpack(data.data()).unwrap();
2851        assert_eq!(
2852            state.get_extension::<FixedLenMintTest>().unwrap(),
2853            &fixed_len,
2854        );
2855
2856        // alloc again succeeds with "overwrite"
2857        let account_info = (&key, &mut data).into_account_info();
2858        alloc_and_serialize::<PodMint, _>(&account_info, &fixed_len, true).unwrap();
2859
2860        // alloc again fails without "overwrite"
2861        let account_info = (&key, &mut data).into_account_info();
2862        assert_eq!(
2863            alloc_and_serialize::<PodMint, _>(&account_info, &fixed_len, false).unwrap_err(),
2864            TokenError::ExtensionAlreadyInitialized.into()
2865        );
2866    }
2867
2868    #[test]
2869    fn alloc_new_variable_len_tlv_in_account_info_from_base_size() {
2870        let variable_len = VariableLenMintTest { data: vec![20, 99] };
2871        let value_len = variable_len.get_packed_len().unwrap();
2872        let base_account_size = PodMint::SIZE_OF;
2873        let mut buffer = vec![0; base_account_size];
2874        let state =
2875            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2876        *state.base = TEST_POD_MINT;
2877
2878        let mut data = SolanaAccountData::new(&buffer);
2879        let key = Pubkey::new_unique();
2880        let account_info = (&key, &mut data).into_account_info();
2881
2882        alloc_and_serialize_variable_len_extension::<PodMint, _>(
2883            &account_info,
2884            &variable_len,
2885            false,
2886        )
2887        .unwrap();
2888        let new_account_len = BASE_ACCOUNT_AND_TYPE_LENGTH + add_type_and_length_to_len(value_len);
2889        assert_eq!(data.len(), new_account_len);
2890        let state = PodStateWithExtensions::<PodMint>::unpack(data.data()).unwrap();
2891        assert_eq!(
2892            state
2893                .get_variable_len_extension::<VariableLenMintTest>()
2894                .unwrap(),
2895            variable_len
2896        );
2897
2898        // alloc again succeeds with "overwrite"
2899        let account_info = (&key, &mut data).into_account_info();
2900        alloc_and_serialize_variable_len_extension::<PodMint, _>(
2901            &account_info,
2902            &variable_len,
2903            true,
2904        )
2905        .unwrap();
2906
2907        // alloc again fails without "overwrite"
2908        let account_info = (&key, &mut data).into_account_info();
2909        assert_eq!(
2910            alloc_and_serialize_variable_len_extension::<PodMint, _>(
2911                &account_info,
2912                &variable_len,
2913                false,
2914            )
2915            .unwrap_err(),
2916            TokenError::ExtensionAlreadyInitialized.into()
2917        );
2918    }
2919
2920    #[test]
2921    fn alloc_new_fixed_len_tlv_in_account_info_from_extended_size() {
2922        let fixed_len = FixedLenMintTest {
2923            data: [1, 2, 3, 4, 5, 6, 7, 8],
2924        };
2925        let value_len = pod_get_packed_len::<FixedLenMintTest>();
2926        let account_size =
2927            ExtensionType::try_calculate_account_len::<PodMint>(&[ExtensionType::GroupPointer])
2928                .unwrap()
2929                + add_type_and_length_to_len(value_len);
2930        let mut buffer = vec![0; account_size];
2931        let mut state =
2932            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2933        *state.base = TEST_POD_MINT;
2934        state.init_account_type().unwrap();
2935
2936        let test_key =
2937            OptionalNonZeroPubkey::try_from(Some(Pubkey::new_from_array([20; 32]))).unwrap();
2938        let extension = state.init_extension::<GroupPointer>(false).unwrap();
2939        extension.authority = test_key;
2940        extension.group_address = test_key;
2941
2942        let mut data = SolanaAccountData::new(&buffer);
2943        let key = Pubkey::new_unique();
2944        let account_info = (&key, &mut data).into_account_info();
2945
2946        alloc_and_serialize::<PodMint, _>(&account_info, &fixed_len, false).unwrap();
2947        let new_account_len = BASE_ACCOUNT_AND_TYPE_LENGTH
2948            + add_type_and_length_to_len(value_len)
2949            + add_type_and_length_to_len(size_of::<GroupPointer>());
2950        assert_eq!(data.len(), new_account_len);
2951        let state = PodStateWithExtensions::<PodMint>::unpack(data.data()).unwrap();
2952        assert_eq!(
2953            state.get_extension::<FixedLenMintTest>().unwrap(),
2954            &fixed_len,
2955        );
2956        let extension = state.get_extension::<GroupPointer>().unwrap();
2957        assert_eq!(extension.authority, test_key);
2958        assert_eq!(extension.group_address, test_key);
2959
2960        // alloc again succeeds with "overwrite"
2961        let account_info = (&key, &mut data).into_account_info();
2962        alloc_and_serialize::<PodMint, _>(&account_info, &fixed_len, true).unwrap();
2963
2964        // alloc again fails without "overwrite"
2965        let account_info = (&key, &mut data).into_account_info();
2966        assert_eq!(
2967            alloc_and_serialize::<PodMint, _>(&account_info, &fixed_len, false).unwrap_err(),
2968            TokenError::ExtensionAlreadyInitialized.into()
2969        );
2970    }
2971
2972    #[test]
2973    fn alloc_new_variable_len_tlv_in_account_info_from_extended_size() {
2974        let variable_len = VariableLenMintTest { data: vec![42, 6] };
2975        let value_len = variable_len.get_packed_len().unwrap();
2976        let account_size =
2977            ExtensionType::try_calculate_account_len::<PodMint>(&[ExtensionType::MetadataPointer])
2978                .unwrap()
2979                + add_type_and_length_to_len(value_len);
2980        let mut buffer = vec![0; account_size];
2981        let mut state =
2982            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
2983        *state.base = TEST_POD_MINT;
2984        state.init_account_type().unwrap();
2985
2986        let test_key =
2987            OptionalNonZeroPubkey::try_from(Some(Pubkey::new_from_array([20; 32]))).unwrap();
2988        let extension = state.init_extension::<MetadataPointer>(false).unwrap();
2989        extension.authority = test_key;
2990        extension.metadata_address = test_key;
2991
2992        let mut data = SolanaAccountData::new(&buffer);
2993        let key = Pubkey::new_unique();
2994        let account_info = (&key, &mut data).into_account_info();
2995
2996        alloc_and_serialize_variable_len_extension::<PodMint, _>(
2997            &account_info,
2998            &variable_len,
2999            false,
3000        )
3001        .unwrap();
3002        let new_account_len = BASE_ACCOUNT_AND_TYPE_LENGTH
3003            + add_type_and_length_to_len(value_len)
3004            + add_type_and_length_to_len(size_of::<MetadataPointer>());
3005        assert_eq!(data.len(), new_account_len);
3006        let state = PodStateWithExtensions::<PodMint>::unpack(data.data()).unwrap();
3007        assert_eq!(
3008            state
3009                .get_variable_len_extension::<VariableLenMintTest>()
3010                .unwrap(),
3011            variable_len
3012        );
3013        let extension = state.get_extension::<MetadataPointer>().unwrap();
3014        assert_eq!(extension.authority, test_key);
3015        assert_eq!(extension.metadata_address, test_key);
3016
3017        // alloc again succeeds with "overwrite"
3018        let account_info = (&key, &mut data).into_account_info();
3019        alloc_and_serialize_variable_len_extension::<PodMint, _>(
3020            &account_info,
3021            &variable_len,
3022            true,
3023        )
3024        .unwrap();
3025
3026        // alloc again fails without "overwrite"
3027        let account_info = (&key, &mut data).into_account_info();
3028        assert_eq!(
3029            alloc_and_serialize_variable_len_extension::<PodMint, _>(
3030                &account_info,
3031                &variable_len,
3032                false,
3033            )
3034            .unwrap_err(),
3035            TokenError::ExtensionAlreadyInitialized.into()
3036        );
3037    }
3038
3039    #[test]
3040    fn realloc_variable_len_tlv_in_account_info() {
3041        let variable_len = VariableLenMintTest {
3042            data: vec![1, 2, 3, 4, 5],
3043        };
3044        let alloc_size = variable_len.get_packed_len().unwrap();
3045        let account_size =
3046            ExtensionType::try_calculate_account_len::<PodMint>(&[ExtensionType::MetadataPointer])
3047                .unwrap()
3048                + add_type_and_length_to_len(alloc_size);
3049        let mut buffer = vec![0; account_size];
3050        let mut state =
3051            PodStateWithExtensionsMut::<PodMint>::unpack_uninitialized(&mut buffer).unwrap();
3052        *state.base = TEST_POD_MINT;
3053        state.init_account_type().unwrap();
3054
3055        // alloc both types
3056        state
3057            .init_variable_len_extension(&variable_len, false)
3058            .unwrap();
3059        let max_pubkey =
3060            OptionalNonZeroPubkey::try_from(Some(Pubkey::new_from_array([255; 32]))).unwrap();
3061        let extension = state.init_extension::<MetadataPointer>(false).unwrap();
3062        extension.authority = max_pubkey;
3063        extension.metadata_address = max_pubkey;
3064
3065        // reallocate to smaller, make sure existing extension is fine
3066        let mut data = SolanaAccountData::new(&buffer);
3067        let key = Pubkey::new_unique();
3068        let account_info = (&key, &mut data).into_account_info();
3069        let variable_len = VariableLenMintTest { data: vec![1, 2] };
3070        alloc_and_serialize_variable_len_extension::<PodMint, _>(
3071            &account_info,
3072            &variable_len,
3073            true,
3074        )
3075        .unwrap();
3076
3077        let state = PodStateWithExtensions::<PodMint>::unpack(data.data()).unwrap();
3078        let extension = state.get_extension::<MetadataPointer>().unwrap();
3079        assert_eq!(extension.authority, max_pubkey);
3080        assert_eq!(extension.metadata_address, max_pubkey);
3081        let extension = state
3082            .get_variable_len_extension::<VariableLenMintTest>()
3083            .unwrap();
3084        assert_eq!(extension, variable_len);
3085        assert_eq!(data.len(), state.try_get_account_len().unwrap());
3086
3087        // reallocate to larger
3088        let account_info = (&key, &mut data).into_account_info();
3089        let variable_len = VariableLenMintTest {
3090            data: vec![1, 2, 3, 4, 5, 6, 7],
3091        };
3092        alloc_and_serialize_variable_len_extension::<PodMint, _>(
3093            &account_info,
3094            &variable_len,
3095            true,
3096        )
3097        .unwrap();
3098
3099        let state = PodStateWithExtensions::<PodMint>::unpack(data.data()).unwrap();
3100        let extension = state.get_extension::<MetadataPointer>().unwrap();
3101        assert_eq!(extension.authority, max_pubkey);
3102        assert_eq!(extension.metadata_address, max_pubkey);
3103        let extension = state
3104            .get_variable_len_extension::<VariableLenMintTest>()
3105            .unwrap();
3106        assert_eq!(extension, variable_len);
3107        assert_eq!(data.len(), state.try_get_account_len().unwrap());
3108
3109        // reallocate to same
3110        let account_info = (&key, &mut data).into_account_info();
3111        let variable_len = VariableLenMintTest {
3112            data: vec![7, 6, 5, 4, 3, 2, 1],
3113        };
3114        alloc_and_serialize_variable_len_extension::<PodMint, _>(
3115            &account_info,
3116            &variable_len,
3117            true,
3118        )
3119        .unwrap();
3120
3121        let state = PodStateWithExtensions::<PodMint>::unpack(data.data()).unwrap();
3122        let extension = state.get_extension::<MetadataPointer>().unwrap();
3123        assert_eq!(extension.authority, max_pubkey);
3124        assert_eq!(extension.metadata_address, max_pubkey);
3125        let extension = state
3126            .get_variable_len_extension::<VariableLenMintTest>()
3127            .unwrap();
3128        assert_eq!(extension, variable_len);
3129        assert_eq!(data.len(), state.try_get_account_len().unwrap());
3130    }
3131}