diff --git a/token/program-2022/src/extension/mod.rs b/token/program-2022/src/extension/mod.rs index 659d50b6116..2ddac9e0a0a 100644 --- a/token/program-2022/src/extension/mod.rs +++ b/token/program-2022/src/extension/mod.rs @@ -1217,35 +1217,26 @@ pub fn alloc_and_serialize( overwrite: bool, ) -> Result<(), ProgramError> { let previous_account_len = account_info.try_data_len()?; - let (new_account_len, extension_already_exists) = { + let new_account_len = { let data = account_info.try_borrow_data()?; let state = StateWithExtensions::::unpack(&data)?; - let new_account_len = state.try_get_new_account_len::()?; - let extension_already_exists = state.get_extension_bytes::().is_ok(); - (new_account_len, extension_already_exists) + state.try_get_new_account_len::()? }; - - if extension_already_exists { - if !overwrite { - return Err(TokenError::ExtensionAlreadyInitialized.into()); - } else { - // Overwrite the extension - let mut buffer = account_info.try_borrow_mut_data()?; - let mut state = StateWithExtensionsMut::::unpack(&mut buffer)?; - let extension = state.get_extension_mut::()?; - *extension = *new_extension; - } - } else { - // Realloc the account, then write the new extension + + // Realloc the account first, if needed + if new_account_len > previous_account_len { account_info.realloc(new_account_len, false)?; - let mut buffer = account_info.try_borrow_mut_data()?; - if previous_account_len <= BASE_ACCOUNT_LENGTH { - set_account_type::(*buffer)?; - } - let mut state = StateWithExtensionsMut::::unpack(&mut buffer)?; - let extension = state.init_extension::(false)?; - *extension = *new_extension; } + let mut buffer = account_info.try_borrow_mut_data()?; + if previous_account_len <= BASE_ACCOUNT_LENGTH { + set_account_type::(*buffer)?; + } + let mut state = StateWithExtensionsMut::::unpack(&mut buffer)?; + + // Write the extension + let extension = state.init_extension::(overwrite)?; + *extension = *new_extension; + Ok(()) }