refactor: Improve authority validation on close

program/src/processor/close.rs

@@ -1,6 +1,6 @@
 use pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult};
 
-use crate::state::AccountDiscriminator;
+use crate::state::{buffer::Buffer, AccountDiscriminator};
 
 use super::{validate_authority, validate_metadata};
 
@@ -15,6 +15,10 @@ pub fn close(accounts: &[AccountInfo]) -> ProgramResult {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
+    if !authority.is_signer() {
+        return Err(ProgramError::MissingRequiredSignature);
+    }
+
     let account_data = if account.data_is_empty() {
         return Err(ProgramError::UninitializedAccount);
     } else {
@@ -25,19 +29,20 @@ pub fn close(accounts: &[AccountInfo]) -> ProgramResult {
     //  - account: program owned is implicitly checked since we are writing
     //    to the account
 
-    match AccountDiscriminator::try_from(account_data[0])? {
-        AccountDiscriminator::Buffer => {
-            // The authority of a buffer account must be the buffer account
-            // itself.
-            if !(account.key() == authority.key() && authority.is_signer()) {
-                return Err(ProgramError::MissingRequiredSignature);
+    // We only need to validate the authority if it is not a keypair buffer,
+    // since we already validated that the authority is a signer.
+    if account.key() != authority.key() {
+        match AccountDiscriminator::try_from(account_data[0])? {
+            AccountDiscriminator::Buffer => {
+                let buffer = unsafe { Buffer::load_unchecked(account_data) };
+                validate_authority(buffer, authority, program, program_data)?
             }
+            AccountDiscriminator::Metadata => {
+                let header = validate_metadata(account)?;
+                validate_authority(header, authority, program, program_data)?
+            }
+            _ => return Err(ProgramError::InvalidAccountData),
         }
-        AccountDiscriminator::Metadata => {
-            let header = validate_metadata(account)?;
-            validate_authority(header, authority, program, program_data)?
-        }
-        _ => return Err(ProgramError::InvalidAccountData),
     }
 
     // Move the lamports to the destination account.

program/src/processor/mod.rs

@@ -1,6 +1,6 @@
 use pinocchio::{account_info::AccountInfo, program_error::ProgramError, pubkey::Pubkey};
 
-use crate::state::{header::Header, AccountDiscriminator};
+use crate::state::{header::Header, Account, AccountDiscriminator};
 
 pub mod allocate;
 pub mod close;
@@ -113,8 +113,8 @@ fn validate_metadata(metadata: &AccountInfo) -> Result<&Header, ProgramError> {
 /// - [explicit] The `authority` account must match the authority set on the `metadata` account OR
 ///   it must be the program upgrade authority if the `metadata` account is canonical (see `is_program_authority`).
 #[inline(always)]
-fn validate_authority(
-    metadata_header: &Header,
+fn validate_authority<T: Account>(
+    account: &T,
     authority: &AccountInfo,
     program: &AccountInfo,
     program_data: &AccountInfo,
@@ -123,19 +123,21 @@ fn validate_authority(
     if !authority.is_signer() {
         return Err(ProgramError::MissingRequiredSignature);
     }
+
     // The authority is the set authority.
-    let explicitly_authorized = match metadata_header.authority.as_ref() {
+    let explicitly_authorized = match account.get_authority() {
         Some(metadata_authority) => metadata_authority == authority.key(),
         None => false,
     };
+
     // The authority is the program upgrade authority for canonical metadata accounts.
     let authorized = explicitly_authorized
-        || (metadata_header.canonical()
-            && program.key() == &metadata_header.program
+        || (account.is_canonical(program.key())
             && is_program_authority(program, program_data, authority.key())?);
+
     if !authorized {
-        return Err(ProgramError::IncorrectAuthority);
+        Err(ProgramError::IncorrectAuthority)
+    } else {
+        Ok(())
     }
-
-    Ok(())
 }

program/src/state/buffer.rs

@@ -1,6 +1,6 @@
 use pinocchio::{program_error::ProgramError, pubkey::Pubkey};
 
-use super::{AccountDiscriminator, ZeroableOption, SEED_LEN};
+use super::{Account, AccountDiscriminator, ZeroableOption, SEED_LEN};
 
 /// Buffer account header.
 #[repr(C)]
@@ -65,3 +65,13 @@ impl Buffer {
         &mut *(bytes.as_mut_ptr() as *mut Self)
     }
 }
+
+impl Account for Buffer {
+    fn get_authority(&self) -> Option<&Pubkey> {
+        self.authority.as_ref()
+    }
+
+    fn is_canonical(&self, program: &Pubkey) -> bool {
+        self.canonical() && self.program.as_ref() == Some(program)
+    }
+}

program/src/state/header.rs

@@ -1,7 +1,8 @@
 use pinocchio::{program_error::ProgramError, pubkey::Pubkey};
 
 use super::{
-    AccountDiscriminator, Compression, DataSource, Encoding, Format, ZeroableOption, SEED_LEN,
+    Account, AccountDiscriminator, Compression, DataSource, Encoding, Format, ZeroableOption,
+    SEED_LEN,
 };
 
 /// Metadata account header.
@@ -106,3 +107,13 @@ impl Header {
         &mut *(bytes.as_mut_ptr() as *mut Self)
     }
 }
+
+impl Account for Header {
+    fn get_authority(&self) -> Option<&Pubkey> {
+        self.authority.as_ref()
+    }
+
+    fn is_canonical(&self, program: &Pubkey) -> bool {
+        self.canonical() && self.program == *program
+    }
+}

program/src/state/mod.rs

@@ -32,6 +32,16 @@ impl<'a> Metadata<'a> {
     }
 }
 
+/// Utility trait for an account.
+pub(crate) trait Account {
+    /// Returns the account authority, if there is one.
+    fn get_authority(&self) -> Option<&Pubkey>;
+
+    /// Indicates whether the PDA represents a canonical PDA for
+    /// the given program.
+    fn is_canonical(&self, program: &Pubkey) -> bool;
+}
+
 /// Account discriminators.
 #[repr(u8)]
 #[derive(Clone, Copy, Debug)]