[wip]: Update initialize instruction

program/src/instruction.rs

@@ -24,12 +24,11 @@ pub enum ProgramMetadataInstruction {
     /// by the program upgrade authority.
     ///
     /// ### Accounts
-    ///  0. `[  w  ]` The metadata account to initialize.
-    ///  1. `[ w,o ]` (optional) Buffer account to copy data from.
-    ///  2. `[  s  ]` The program upgrade authority.
-    ///  3. `[     ]` Program account.
-    ///  4. `[     ]` Program data account.
-    ///  5. `[     ]` System program.
+    ///  0. `[  w  ]` Metadata account to initialize.
+    ///  1. `[  s  ]` Authority (for canonical, must match program upgrade authority).
+    ///  2. `[     ]` Program account.
+    ///  3. `[  o  ]` Program data account.
+    ///  4. `[  o  ]` System program.
     ///
     /// ### Instruction data
     ///  - `[u8; 16]`: seed

program/src/processor/initialize.rs

@@ -45,32 +45,16 @@ pub fn initialize(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramR
     };
 
     // Access accounts.
-    let [metadata, buffer, authority, program, program_data, _system_program] = accounts else {
+    let [metadata, authority, program, program_data, _system_program, _remaining @ ..] = accounts
+    else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
-    // Get data from buffer or remaining instruction data.
-    let (data, has_buffer) = {
-        let has_buffer = buffer.key() != &ID;
-        let has_remaining_data = !remaining_data.is_empty();
-        let data = match (has_remaining_data, has_buffer) {
-            (true, false) => remaining_data,
-            (false, true) => {
-                let buffer_data = unsafe { buffer.borrow_data_unchecked() };
-                if buffer_data.is_empty() || buffer_data[0] != AccountDiscriminator::Buffer as u8 {
-                    return Err(ProgramError::InvalidAccountData);
-                }
-                &buffer_data[Header::LEN..]
-            }
-            _ => return Err(ProgramError::InvalidInstructionData),
-        };
-        (data, has_buffer)
-    };
-
     // Validate authority.
     if !authority.is_signer() {
         return Err(ProgramError::MissingRequiredSignature);
     }
+
     let canonical = is_program_authority(program, program_data, authority.key())?;
 
     // Validatate metadata seeds.
@@ -84,52 +68,80 @@ pub fn initialize(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramR
         return Err(ProgramError::InvalidSeeds);
     }
 
-    // Ensure metadata account is not already initialized.
-    if !metadata.data_is_empty() {
-        return Err(ProgramError::AccountAlreadyInitialized);
-    }
+    let mut metadata_account_data = unsafe { metadata.borrow_mut_data_unchecked() };
+    let discriminator = if metadata_account_data.is_empty() {
+        None
+    } else {
+        Some(AccountDiscriminator::try_from(metadata_account_data[0])?)
+    };
 
-    // Use the buffer account to fund the metadata account, if provided.
-    if has_buffer {
-        unsafe {
-            let metadata_lamports = metadata.borrow_mut_lamports_unchecked();
-            let buffer_lamports = buffer.borrow_mut_lamports_unchecked();
-            // Move the buffer lamports to the metadata account.
-            *metadata_lamports += *buffer_lamports;
-            *buffer_lamports = 0;
+    let data_length = match discriminator {
+        Some(AccountDiscriminator::Buffer) => {
+            // When using a pre-allocated buffer, no remaining instruction data
+            // is allowed.
+            if !remaining_data.is_empty() {
+                return Err(ProgramError::InvalidAccountData);
+            }
+            metadata_account_data.len() - Header::LEN
         }
-    }
+        Some(AccountDiscriminator::Metadata) => {
+            return Err(ProgramError::AccountAlreadyInitialized)
+        }
+        None => {
+            // Ensure remaining data is provided.
+            if remaining_data.is_empty() {
+                return Err(ProgramError::InvalidInstructionData);
+            }
 
-    // Allocate and assign the metadata account.
-    let signer_bump = &[bump];
-    let signer_seeds: &[Seed] = if canonical {
-        &[
-            Seed::from(program.key()),
-            Seed::from(args.seed.as_ref()),
-            Seed::from(signer_bump),
-        ]
-    } else {
-        &[
-            Seed::from(program.key()),
-            Seed::from(authority.key()),
-            Seed::from(args.seed.as_ref()),
-            Seed::from(signer_bump),
-        ]
+            // Allocate and assign the metadata account.
+            let signer_bump = &[bump];
+            let signer_seeds: &[Seed] = if canonical {
+                &[
+                    Seed::from(program.key()),
+                    Seed::from(args.seed.as_ref()),
+                    Seed::from(signer_bump),
+                ]
+            } else {
+                &[
+                    Seed::from(program.key()),
+                    Seed::from(authority.key()),
+                    Seed::from(args.seed.as_ref()),
+                    Seed::from(signer_bump),
+                ]
+            };
+            let signer = &[Signer::from(signer_seeds)];
+
+            Allocate {
+                account: metadata,
+                space: (Header::LEN + remaining_data.len()) as u64,
+            }
+            .invoke_signed(signer)?;
+            Assign {
+                account: metadata,
+                owner: &crate::ID,
+            }
+            .invoke_signed(signer)?;
+
+            // Refresh the metadata account data reference.
+            metadata_account_data = unsafe { metadata.borrow_mut_data_unchecked() };
+
+            // Copy the instruction remaining data to the metadata account.
+            unsafe {
+                sol_memcpy(
+                    &mut metadata_account_data[Header::LEN..],
+                    remaining_data,
+                    remaining_data.len(),
+                );
+            }
+
+            remaining_data.len()
+        }
+        _ => {
+            return Err(ProgramError::InvalidAccountData);
+        }
     };
-    let signer = &[Signer::from(signer_seeds)];
-    Allocate {
-        account: metadata,
-        space: (Header::LEN + data.len()) as u64,
-    }
-    .invoke_signed(signer)?;
-    Assign {
-        account: metadata,
-        owner: &crate::ID,
-    }
-    .invoke_signed(signer)?;
 
     // Initialize the metadata account.
-    let metadata_account_data = unsafe { metadata.borrow_mut_data_unchecked() };
     let header = unsafe { Header::load_mut_unchecked(metadata_account_data) };
     header.discriminator = AccountDiscriminator::Metadata as u8;
     header.program = *program.key();
@@ -143,16 +155,12 @@ pub fn initialize(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramR
     header.encoding = Encoding::try_from(args.encoding)? as u8;
     header.compression = Compression::try_from(args.compression)? as u8;
     header.format = Format::try_from(args.format)? as u8;
-    header.data_source = DataSource::try_from(args.data_source)? as u8;
-    header.data_length = (data.len() as u32).to_le_bytes();
-    unsafe {
-        sol_memcpy(&mut metadata_account_data[Header::LEN..], data, data.len());
-    }
-
-    // Close the buffer account, if provided.
-    if has_buffer {
-        buffer.close()?;
-    }
+    header.data_source = {
+        let data_source = DataSource::try_from(args.data_source)?;
+        data_source.validate_data_length(data_length)?;
+        data_source as u8
+    };
+    header.data_length = (data_length as u32).to_le_bytes();
 
     Ok(())
 }

program/src/processor/mod.rs

@@ -27,6 +27,12 @@ fn is_program_authority(
     program_data: &AccountInfo,
     authority: &Pubkey,
 ) -> Result<bool, ProgramError> {
+    // If we don't have a program data to check against, we can't verify
+    // the program upgrade authority.
+    if program_data.key() == &crate::ID {
+        return Ok(false);
+    }
+
     // Program checks.
     let expected_program_data = {
         let data = unsafe { program.borrow_data_unchecked() };

program/src/state/data.rs

@@ -64,5 +64,5 @@ pub struct ExternalData {
 }
 
 impl ExternalData {
-    const LEN: usize = core::mem::size_of::<ExternalData>();
+    pub const LEN: usize = core::mem::size_of::<ExternalData>();
 }

program/src/state/mod.rs

@@ -1,9 +1,9 @@
 pub mod data;
 pub mod header;
 
-use data::Data;
+use data::{Data, ExternalData};
 use header::Header;
-use pinocchio::{program_error::ProgramError, pubkey::Pubkey};
+use pinocchio::{program_error::ProgramError, pubkey::Pubkey, ProgramResult};
 
 /// Struct to represent the contents of a `Metadata` account.
 pub struct Metadata<'a> {
@@ -129,6 +129,20 @@ pub enum DataSource {
     External,
 }
 
+impl DataSource {
+    #[inline(always)]
+    pub fn validate_data_length(&self, length: usize) -> ProgramResult {
+        match (self, length) {
+            (DataSource::Direct | DataSource::Url, l) if l > 0 => Ok(()),
+            (DataSource::External, ExternalData::LEN) => Ok(()),
+            _ => {
+                // TODO: use custom error (invalid data length)
+                Err(ProgramError::InvalidAccountData)
+            }
+        }
+    }
+}
+
 impl TryFrom<u8> for DataSource {
     type Error = ProgramError;
 

program/tests/initialize.rs

@@ -8,7 +8,7 @@ use solana_sdk::{account::AccountSharedData, pubkey::Pubkey, system_program};
 use spl_program_metadata::state::header::Header;
 
 #[test]
-fn test_initialize_mint() {
+fn test_initialize_canonical() {
     let authority_key = Pubkey::new_unique();
 
     let program_data_key = Pubkey::new_unique();
@@ -27,8 +27,7 @@ fn test_initialize_mint() {
     let instruction = initialize(
         &authority_key,
         &program_key,
-        &program_data_key,
-        None,
+        Some(&program_data_key),
         InitializeArgs {
             canonical: true,
             seed,
@@ -63,3 +62,61 @@ fn test_initialize_mint() {
         ],
     );
 }
+
+#[test]
+fn test_initialize_non_canonical() {
+    let authority_key = Pubkey::new_unique();
+
+    let program_data_key = Pubkey::new_unique();
+    let program_data_account = setup_program_data_account(Some(&authority_key));
+
+    let program_key = Pubkey::new_unique();
+    let program_account = setup_program_account(&program_data_key);
+
+    let mut seed = [0u8; 16];
+    seed[0..3].copy_from_slice("idl".as_bytes());
+
+    let (metadata_key, _) = Pubkey::find_program_address(
+        &[program_key.as_ref(), authority_key.as_ref(), &seed],
+        &PROGRAM_ID,
+    );
+    let metadata_account = create_empty_account(Header::LEN + 10, system_program::ID);
+
+    let instruction = initialize(
+        &authority_key,
+        &program_key,
+        None,
+        InitializeArgs {
+            canonical: false,
+            seed,
+            encoding: 0,
+            compression: 0,
+            format: 0,
+            data_source: 0,
+        },
+        Some(&[1u8; 10]),
+    )
+    .unwrap();
+
+    let mollusk = Mollusk::new(&PROGRAM_ID, "spl_program_metadata");
+    mollusk.process_and_validate_instruction_chain(
+        &[instruction],
+        &[
+            (metadata_key, metadata_account),
+            (PROGRAM_ID, AccountSharedData::default()),
+            (authority_key, AccountSharedData::default()),
+            (program_key, program_account),
+            (program_data_key, program_data_account),
+            keyed_account_for_system_program(),
+        ],
+        &[
+            Check::success(),
+            // account discriminator
+            Check::account(&metadata_key).data_slice(0, &[2]).build(),
+            // metadata data
+            Check::account(&metadata_key)
+                .data_slice(Header::LEN, &[1u8; 10])
+                .build(),
+        ],
+    );
+}

program/tests/setup/initialize.rs

@@ -10,8 +10,7 @@ use super::PROGRAM_ID;
 pub fn initialize(
     authority: &Pubkey,
     program: &Pubkey,
-    program_data: &Pubkey,
-    _buffer: Option<&Pubkey>,
+    program_data: Option<&Pubkey>,
     args: InitializeArgs,
     instruction_data: Option<&[u8]>,
 ) -> Result<Instruction, ProgramError> {
@@ -24,10 +23,9 @@ pub fn initialize(
 
     let accounts = vec![
         AccountMeta::new(metadata_key, false),
-        AccountMeta::new_readonly(PROGRAM_ID, false),
         AccountMeta::new_readonly(*authority, true),
         AccountMeta::new_readonly(*program, false),
-        AccountMeta::new_readonly(*program_data, false),
+        AccountMeta::new_readonly(*program_data.unwrap_or(&PROGRAM_ID), false),
         AccountMeta::new_readonly(system_program::ID, false),
     ];