Fix instruction data validation

Author: febo

p-token/src/processor/set_authority.rs

@@ -13,7 +13,22 @@ use super::validate_owner;
 pub fn process_set_authority(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
     // Validates the instruction data.
 
-    let args = SetAuthority::try_from_bytes(instruction_data)?;
+    // SAFETY: The expected size of the instruction data is either 2 or 34 bytes:
+    //   - authority_type (1 byte)
+    //   - option + new_authority (1 byte + 32 bytes)
+    let (authority_type, new_authority) = unsafe {
+        match instruction_data.len() {
+            2 if *instruction_data.get_unchecked(1) == 0 => (
+                AuthorityType::try_from(*instruction_data.get_unchecked(0))?,
+                None,
+            ),
+            34 if *instruction_data.get_unchecked(1) == 1 => (
+                AuthorityType::try_from(*instruction_data.get_unchecked(0))?,
+                Some(&*(instruction_data.as_ptr().add(2) as *const Pubkey)),
+            ),
+            _ => return Err(ProgramError::InvalidInstructionData),
+        }
+    };
 
     // Validates the accounts.
 
@@ -30,11 +45,11 @@ pub fn process_set_authority(accounts: &[AccountInfo], instruction_data: &[u8])
             return Err(TokenError::AccountFrozen.into());
         }
 
-        match args.authority_type()? {
+        match authority_type {
             AuthorityType::AccountOwner => {
                 validate_owner(&account.owner, authority_info, remaining)?;
 
-                if let Some(authority) = args.new_authority() {
+                if let Some(authority) = new_authority {
                     account.owner = *authority;
                 } else {
                     return Err(TokenError::InvalidInstruction.into());
@@ -51,7 +66,7 @@ pub fn process_set_authority(accounts: &[AccountInfo], instruction_data: &[u8])
                 let authority = account.close_authority().unwrap_or(&account.owner);
                 validate_owner(authority, authority_info, remaining)?;
 
-                if let Some(authority) = args.new_authority() {
+                if let Some(authority) = new_authority {
                     account.set_close_authority(authority);
                 } else {
                     account.clear_close_authority();
@@ -66,15 +81,15 @@ pub fn process_set_authority(accounts: &[AccountInfo], instruction_data: &[u8])
         // `load_mut` validates that the mint is initialized.
         let mint = unsafe { load_mut::<Mint>(account_info.borrow_mut_data_unchecked())? };
 
-        match args.authority_type()? {
+        match authority_type {
             AuthorityType::MintTokens => {
                 // Once a mint's supply is fixed, it cannot be undone by setting a new
                 // mint_authority.
                 let mint_authority = mint.mint_authority().ok_or(TokenError::FixedSupply)?;
 
                 validate_owner(mint_authority, authority_info, remaining)?;
 
-                if let Some(authority) = args.new_authority() {
+                if let Some(authority) = new_authority {
                     mint.set_mint_authority(authority);
                 } else {
                     mint.clear_mint_authority();
@@ -89,7 +104,7 @@ pub fn process_set_authority(accounts: &[AccountInfo], instruction_data: &[u8])
 
                 validate_owner(freeze_authority, authority_info, remaining)?;
 
-                if let Some(authority) = args.new_authority() {
+                if let Some(authority) = new_authority {
                     mint.set_freeze_authority(authority);
                 } else {
                     mint.clear_freeze_authority();
@@ -105,40 +120,3 @@ pub fn process_set_authority(accounts: &[AccountInfo], instruction_data: &[u8])
 
     Ok(())
 }
-
-#[repr(C)]
-struct SetAuthority {
-    authority_type: u8,
-
-    new_authority: (u8, Pubkey),
-}
-
-impl SetAuthority {
-    #[inline]
-    pub fn try_from_bytes(bytes: &[u8]) -> Result<&SetAuthority, ProgramError> {
-        // The minimum expected size of the instruction data is either 2 or 34 bytes:
-        //   - authority_type (1 byte)
-        //   - option + new_authority (1 byte + 32 bytes)
-        unsafe {
-            match bytes.len() {
-                2 if *bytes.get_unchecked(1) == 0 => Ok(&*(bytes.as_ptr() as *const SetAuthority)),
-                34 => Ok(&*(bytes.as_ptr() as *const SetAuthority)),
-                _ => Err(ProgramError::InvalidInstructionData),
-            }
-        }
-    }
-
-    #[inline]
-    pub fn authority_type(&self) -> Result<AuthorityType, ProgramError> {
-        self.authority_type.try_into()
-    }
-
-    #[inline]
-    pub fn new_authority(&self) -> Option<&Pubkey> {
-        if self.new_authority.0 == 0 {
-            Option::None
-        } else {
-            Option::Some(&self.new_authority.1)
-        }
-    }
-}

p-token/src/processor/shared/initialize_mint.rs

@@ -19,7 +19,25 @@ pub fn process_initialize_mint(
 ) -> ProgramResult {
     // Validates the instruction data.
 
-    let args = InitializeMint::try_from_bytes(instruction_data)?;
+    // SAFETY: The minimum size of the instruction data is either 34 or 66 bytes:
+    //   - decimals (1 byte)
+    //   - mint_authority (32 bytes)
+    //   - option + freeze_authority (1 byte + 32 bytes)
+    let (decimals, mint_authority, freeze_authority) = unsafe {
+        match instruction_data.len() {
+            34 if *instruction_data.get_unchecked(33) == 0 => (
+                *instruction_data.get_unchecked(0),
+                &*(instruction_data.as_ptr().add(1) as *const Pubkey),
+                None,
+            ),
+            66 if *instruction_data.get_unchecked(33) == 1 => (
+                *instruction_data.get_unchecked(0),
+                &*(instruction_data.as_ptr().add(1) as *const Pubkey),
+                Some(&*(instruction_data.as_ptr().add(34) as *const Pubkey)),
+            ),
+            _ => return Err(ProgramError::InvalidInstructionData),
+        }
+    };
 
     // Validates the accounts.
 
@@ -60,50 +78,12 @@ pub fn process_initialize_mint(
     // Initialize the mint.
 
     mint.set_initialized();
-    mint.set_mint_authority(&args.mint_authority);
-    mint.decimals = args.decimals;
+    mint.set_mint_authority(mint_authority);
+    mint.decimals = decimals;
 
-    if let Some(freeze_authority) = args.freeze_authority() {
+    if let Some(freeze_authority) = freeze_authority {
         mint.set_freeze_authority(freeze_authority);
     }
 
     Ok(())
 }
-
-/// Instruction data for the `InitializeMint` instruction.
-#[repr(C)]
-struct InitializeMint {
-    pub(crate) decimals: u8,
-
-    pub(crate) mint_authority: Pubkey,
-
-    freeze_authority: (u8, Pubkey),
-}
-
-impl InitializeMint {
-    #[inline]
-    pub fn try_from_bytes(bytes: &[u8]) -> Result<&InitializeMint, ProgramError> {
-        // The minimum expected size of the instruction data is either 34 or 66 bytes:
-        //   - decimals (1 byte)
-        //   - mint_authority (32 bytes)
-        //   - option + freeze_authority (1 byte + 32 bytes)
-        unsafe {
-            match bytes.len() {
-                34 if *bytes.get_unchecked(33) == 0 => {
-                    Ok(&*(bytes.as_ptr() as *const InitializeMint))
-                }
-                66 => Ok(&*(bytes.as_ptr() as *const InitializeMint)),
-                _ => Err(ProgramError::InvalidInstructionData),
-            }
-        }
-    }
-
-    #[inline]
-    pub fn freeze_authority(&self) -> Option<&Pubkey> {
-        if self.freeze_authority.0 == 0 {
-            Option::None
-        } else {
-            Option::Some(&self.freeze_authority.1)
-        }
-    }
-}