program: Improve account length validation (#55)

* Stricter account length validation

* Address review comments

Author: febo

program/src/processor/allocate.rs

@@ -150,7 +150,7 @@ pub fn allocate(accounts: &mut [AccountView], instruction_data: &[u8]) -> Progra
 
     // SAFETY: single mutable borrow of the `buffer` account data. The legth of the buffer account
     // data has been checked to be at least `Buffer::LEN` and uninitialized.
-    let buffer_header = unsafe { Buffer::from_bytes_mut_unchecked(buffer.borrow_unchecked_mut()) };
+    let buffer_header = Buffer::from_bytes_mut(unsafe { buffer.borrow_unchecked_mut() })?;
     buffer_header.discriminator = AccountDiscriminator::Buffer as u8;
     buffer_header.authority = (*authority.address()).into();
 

program/src/processor/close.rs

@@ -32,16 +32,17 @@ pub fn close(accounts: &mut [AccountView]) -> ProgramResult {
     let account_data = if account.is_data_empty() {
         return Err(ProgramError::UninitializedAccount);
     } else {
+        // SAFETY: Single borrow of the `account` data.
         unsafe { account.borrow_unchecked() }
     };
 
     match AccountDiscriminator::try_from(account_data[0])? {
         AccountDiscriminator::Buffer => {
-            let buffer = unsafe { Buffer::from_bytes_unchecked(account_data) };
+            let buffer = Buffer::from_bytes(account_data)?;
             validate_authority(buffer, authority, program, program_data)?
         }
         AccountDiscriminator::Metadata => {
-            let header = validate_metadata(account)?;
+            let header = validate_metadata(account_data)?;
             validate_authority(header, authority, program, program_data)?
         }
         _ => return Err(ProgramError::InvalidAccountData),

program/src/processor/extend.rs

@@ -55,7 +55,7 @@ pub fn extend(accounts: &mut [AccountView], instruction_data: &[u8]) -> ProgramR
                 validate_authority(buffer, authority, program, program_data)?
             }
             Ok(AccountDiscriminator::Metadata) => {
-                let metadata = validate_metadata(account)?;
+                let metadata = validate_metadata(data)?;
                 validate_authority(metadata, authority, program, program_data)?
             }
             _ => return Err(ProgramError::InvalidAccountData),

program/src/processor/initialize.rs

@@ -96,11 +96,14 @@ pub fn initialize(accounts: &mut [AccountView], instruction_data: &[u8]) -> Prog
             // When using a pre-allocated buffer, no remaining instruction data
             // is allowed.
             if !remaining_data.is_empty() {
-                return Err(ProgramError::InvalidAccountData);
+                return Err(ProgramError::InvalidInstructionData);
             }
-            // A pre-allocated buffer length is at least the size of the
+            // A pre-allocated buffer length should be at least the size of the
             // `Header`.
-            metadata.data_len() - Header::LEN
+            metadata
+                .data_len()
+                .checked_sub(Header::LEN)
+                .ok_or(ProgramError::InvalidAccountData)?
         }
         Some(AccountDiscriminator::Metadata) => {
             return Err(ProgramError::AccountAlreadyInitialized)
@@ -177,7 +180,7 @@ pub fn initialize(accounts: &mut [AccountView], instruction_data: &[u8]) -> Prog
     // Initialize the metadata account.
 
     // SAFETY: there are no other active borrows to `metadata` account data and
-    // the account discriminator has been validated.
+    // the account length has been validated to be sufficient to hold a `Header`.
     let header = unsafe { Header::from_bytes_mut_unchecked(metadata.borrow_unchecked_mut()) };
 
     header.discriminator = AccountDiscriminator::Metadata as u8;

program/src/processor/mod.rs

@@ -114,8 +114,8 @@ fn is_program_authority(
 ///   be [`AccountDiscriminator::Metadata`].
 /// - The `metadata` account must be mutable (`mutable = true`).
 #[inline(always)]
-fn validate_metadata(metadata: &AccountView) -> Result<&Header, ProgramError> {
-    let header = unsafe { Header::from_bytes_unchecked(metadata.borrow_unchecked()) };
+fn validate_metadata(bytes: &[u8]) -> Result<&Header, ProgramError> {
+    let header = Header::from_bytes(bytes)?;
     if header.discriminator != AccountDiscriminator::Metadata as u8 {
         return Err(ProgramError::UninitializedAccount);
     }

program/src/processor/set_data.rs

@@ -40,27 +40,31 @@ pub fn set_data(accounts: &mut [AccountView], instruction_data: &[u8]) -> Progra
     //
     // Note that program owned and writable checks are done implicitly by writing
     // to the account.
+    {
+        // metadata
+        // - must be initialized
+        // - must be mutable
 
-    // metadata
-    // - must be initialized
-    // - must be mutable
+        // SAFETY: Scoped immutable borrow of `metadata` account data for validation.
+        let metadata_account_data = unsafe { metadata.borrow_unchecked() };
+        let header = validate_metadata(metadata_account_data)?;
 
-    let header = validate_metadata(metadata)?;
+        // authority
+        // - must be a signer
+        // - must match the authority set on the `metadata` account OR it must be the
+        //   program upgrade authority if the `metadata` account is canonical
 
-    // authority
-    // - must be a signer
-    // - must match the authority set on the `metadata` account OR it must be the
-    //   program upgrade authority if the `metadata` account is canonical
-
-    validate_authority(header, authority, program, program_data)?;
+        validate_authority(header, authority, program, program_data)?;
+    }
 
-    // Get data from buffer or remaining instruction data, if any.
     let has_buffer = buffer.address() != &crate::ID;
+    // Get data from buffer or remaining instruction data, if any.
     let data = match (optional_data, has_buffer) {
         (Some((data_source, Some(remaining_data))), false) => Some((data_source, remaining_data)),
         (Some((data_source, None)), true) => {
-            // SAFETY: singe immutable borrow of `buffer` account data.
+            // SAFETY: single immutable borrow of `buffer` account data.
             let buffer_data = unsafe { buffer.borrow_unchecked() };
+
             match AccountDiscriminator::try_from_bytes(buffer_data)? {
                 Some(AccountDiscriminator::Buffer) => {
                     Some((data_source, &buffer_data[Header::LEN..]))

program/src/processor/set_immutable.rs

@@ -1,6 +1,6 @@
 use pinocchio::{account::AccountView, error::ProgramError, ProgramResult};
 
-use crate::{error::ProgramMetadataError, state::header::Header};
+use crate::state::header::Header;
 
 use super::{validate_authority, validate_metadata};
 
@@ -22,7 +22,9 @@ pub fn set_immutable(accounts: &mut [AccountView]) -> ProgramResult {
     // - must be initialized
     // - must be mutable
 
-    let header = validate_metadata(metadata)?;
+    // SAFETY: There are no active borrows of the `metadata` account data.
+    let metadata_account_data = unsafe { metadata.borrow_unchecked_mut() };
+    let header = validate_metadata(metadata_account_data)?;
 
     // authority
     // - must be a signer
@@ -33,15 +35,10 @@ pub fn set_immutable(accounts: &mut [AccountView]) -> ProgramResult {
 
     // Make the metadata account immutable.
 
-    // SAFETY: There are no active borrows of the `metadata` account data and the
-    // account has been validated.
-    let header = unsafe { Header::from_bytes_mut_unchecked(metadata.borrow_unchecked_mut()) };
-
-    if header.mutable() {
-        header.mutable = 0;
-    } else {
-        return Err(ProgramMetadataError::ImmutableMetadataAccount.into());
-    }
+    // SAFETY: `metadata_account_data` has been validated to be initialized
+    // and mutable.
+    let header = unsafe { Header::from_bytes_mut_unchecked(metadata_account_data) };
+    header.mutable = 0;
 
     Ok(())
 }

program/src/processor/trim.rs

@@ -43,10 +43,10 @@ pub fn trim(accounts: &mut [AccountView]) -> ProgramResult {
                 account.data_len()
             }
             Ok(AccountDiscriminator::Metadata) => {
-                let metadata = validate_metadata(account)?;
-                validate_authority(metadata, authority, program, program_data)?;
+                let header = validate_metadata(data)?;
+                validate_authority(header, authority, program, program_data)?;
                 // The length of the data is never more than `10_000_000`.
-                Header::LEN + metadata.data_length() as usize
+                Header::LEN + header.data_length() as usize
             }
             _ => return Err(ProgramError::UninitializedAccount),
         }

program/src/processor/write.rs

@@ -47,27 +47,28 @@ pub fn write(accounts: &mut [AccountView], instruction_data: &[u8]) -> ProgramRe
             return Err(ProgramError::InvalidAccountData);
         }
 
-        // SAFETY: `buffer` account data is guaranteed to be a `Buffer`.
-        let buffer_header = unsafe { Buffer::from_bytes_unchecked(data) };
+        // `data` was validated to have a `Buffer` discriminator.
+        let buffer_header = Buffer::from_bytes(data)?;
 
         if Some(authority.address()) != buffer_header.authority.as_ref() {
             return Err(ProgramError::IncorrectAuthority);
         }
 
-        // Determine from where to copy the data.
+        // Determine from where to copy the data, ether from the instruction data
+        // or the source buffer account.
         let instruction_data = match args.data() {
             source_data if !source_data.is_empty() => Some(source_data),
             _ => None,
         };
 
-        let buffer_data = if source_buffer.address() != &crate::ID {
+        let source_buffer_data = if source_buffer.address() != &crate::ID {
             // SAFETY: singe immutable borrow of `source_buffer` account data.
             Some(unsafe { source_buffer.borrow_unchecked() })
         } else {
             None
         };
 
-        let source_data = match (instruction_data, buffer_data) {
+        let source_data = match (instruction_data, source_buffer_data) {
             (Some(instruction_data), None) => instruction_data,
             (None, Some(buffer_data)) => match AccountDiscriminator::try_from_bytes(buffer_data)? {
                 Some(AccountDiscriminator::Buffer) => &buffer_data[Header::LEN..],
@@ -76,7 +77,8 @@ pub fn write(accounts: &mut [AccountView], instruction_data: &[u8]) -> ProgramRe
             _ => return Err(ProgramError::InvalidInstructionData),
         };
 
-        // The length of the data to write is validated by the `try_minimum_balance`.
+        // The length of the data to write is validated by the `try_minimum_balance` and
+        // `offset` is a `u32` value and `source_data` is at most `10_000_000` bytes.
         (max(data.len(), offset + source_data.len()), source_data)
     };
 

program/src/state/mod.rs

@@ -32,19 +32,6 @@ impl<'a> Metadata<'a> {
         let data = Data::from_bytes(header.data_source()?, &bytes[Header::LEN..])?;
         Ok(Self { header, data })
     }
-
-    /*
-    /// Return a `Metadata` from the given bytes.
-    ///
-    /// # Safety
-    ///
-    /// The caller must ensure that `bytes` contains a valid representation of `Metadata`.
-    pub(crate) unsafe fn from_bytes_unchecked(bytes: &'a [u8]) -> Self {
-        let header = Header::from_bytes_unchecked(bytes);
-        let data = Data::from_bytes_unchecked(header.data_source().unwrap(), &bytes[Header::LEN..]);
-        Self { header, data }
-    }
-    */
 }
 
 /// Utility trait for an account.