Add safety comments (#8)

* Remove spl-token test cases

* Update CU values

* Add safety comments

* Use git dependencies

Author: febo

p-token/Cargo.toml

@@ -20,7 +20,6 @@ test-sbf = []
 [dependencies]
 pinocchio = { workspace = true }
 pinocchio-log = { workspace = true }
-pinocchio-pubkey = { workspace = true }
 token-interface = { version = "^0", path = "../interface" }
 
 [dev-dependencies]

p-token/src/processor/amount_to_ui_amount.rs

@@ -23,7 +23,8 @@ pub fn process_amount_to_ui_amount(
 
     let mint_info = accounts.first().ok_or(ProgramError::NotEnoughAccountKeys)?;
     check_account_owner(mint_info)?;
-    // SAFETY: there is a single borrow to the `Mint` account.
+    // SAFETY: single immutable borrow to `mint_info` account data and
+    // `load` validates that the mint is initialized.
     let mint = unsafe {
         load::<Mint>(mint_info.borrow_data_unchecked()).map_err(|_| TokenError::InvalidMint)?
     };

p-token/src/processor/burn_checked.rs

@@ -4,20 +4,20 @@ use super::shared;
 
 #[inline(always)]
 pub fn process_burn_checked(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let (amount, decimals) = instruction_data.split_at(core::mem::size_of::<u64>());
-    let amount = u64::from_le_bytes(
-        amount
-            .try_into()
-            .map_err(|_error| ProgramError::InvalidInstructionData)?,
-    );
+    // expected u64 (8) + u8 (1)
+    let (amount, decimals) = if instruction_data.len() == 9 {
+        let (amount, decimals) = instruction_data.split_at(core::mem::size_of::<u64>());
+        (
+            u64::from_le_bytes(
+                amount
+                    .try_into()
+                    .map_err(|_error| ProgramError::InvalidInstructionData)?,
+            ),
+            decimals.first(),
+        )
+    } else {
+        return Err(ProgramError::InvalidInstructionData);
+    };
 
-    shared::burn::process_burn(
-        accounts,
-        amount,
-        Some(
-            *decimals
-                .first()
-                .ok_or(ProgramError::InvalidInstructionData)?,
-        ),
-    )
+    shared::burn::process_burn(accounts, amount, decimals.copied())
 }

p-token/src/processor/close_account.rs

@@ -3,14 +3,16 @@ use pinocchio::{
 };
 use token_interface::{
     error::TokenError,
-    state::{account::Account, load_mut},
+    state::{account::Account, load},
 };
 
 use super::validate_owner;
 
-/// Incinerator address.
-const INCINERATOR_ID: Pubkey =
-    pinocchio_pubkey::pubkey!("1nc1nerator11111111111111111111111111111111");
+/// Incinerator (`1nc1nerator11111111111111111111111111111111`) address.
+const INCINERATOR_ID: Pubkey = [
+    0, 51, 144, 114, 141, 52, 17, 96, 121, 189, 201, 17, 191, 255, 0, 219, 212, 77, 46, 205, 204,
+    247, 156, 166, 225, 0, 56, 225, 0, 0, 0, 0,
+];
 
 #[inline(always)]
 pub fn process_close_account(accounts: &[AccountInfo]) -> ProgramResult {
@@ -24,26 +26,30 @@ pub fn process_close_account(accounts: &[AccountInfo]) -> ProgramResult {
     // raw pointer.
     if source_account_info == destination_account_info {
         return Err(ProgramError::InvalidAccountData);
-    }
-
-    let source_account =
-        unsafe { load_mut::<Account>(source_account_info.borrow_mut_data_unchecked())? };
-
-    if !source_account.is_native() && source_account.amount() != 0 {
-        return Err(TokenError::NonNativeHasBalance.into());
-    }
-
-    let authority = source_account
-        .close_authority()
-        .unwrap_or(&source_account.owner);
-
-    if !source_account.is_owned_by_system_program_or_incinerator() {
-        validate_owner(authority, authority_info, remaining)?;
-    } else if destination_account_info.key() != &INCINERATOR_ID {
-        return Err(ProgramError::InvalidAccountData);
+    } else {
+        // SAFETY: scoped immutable borrow to `source_account_info` account data and
+        // `load` validates that the account is initialized.
+        let source_account =
+            unsafe { load::<Account>(source_account_info.borrow_data_unchecked())? };
+
+        if !source_account.is_native() && source_account.amount() != 0 {
+            return Err(TokenError::NonNativeHasBalance.into());
+        }
+
+        let authority = source_account
+            .close_authority()
+            .unwrap_or(&source_account.owner);
+
+        if !source_account.is_owned_by_system_program_or_incinerator() {
+            validate_owner(authority, authority_info, remaining)?;
+        } else if destination_account_info.key() != &INCINERATOR_ID {
+            return Err(ProgramError::InvalidAccountData);
+        }
     }
 
     let destination_starting_lamports = destination_account_info.lamports();
+    // SAFETY: single mutable borrow to `destination_account_info` lamports and
+    // there are no "active" borrows of `source_account_info` account data.
     unsafe {
         // Moves the lamports to the destination account.
         *destination_account_info.borrow_mut_lamports_unchecked() = destination_starting_lamports

p-token/src/processor/get_account_data_size.rs

@@ -17,6 +17,8 @@ pub fn process_get_account_data_size(accounts: &[AccountInfo]) -> ProgramResult
     // Make sure the mint is valid.
     check_account_owner(mint_info)?;
 
+    // SAFETY: single immutable borrow to `mint_info` account data and
+    // `load` validates that the mint is initialized.
     let _ = unsafe {
         load::<Mint>(mint_info.borrow_data_unchecked()).map_err(|_| TokenError::InvalidMint)?
     };

p-token/src/processor/initialize_account2.rs

@@ -1,4 +1,9 @@
-use pinocchio::{account_info::AccountInfo, pubkey::Pubkey, ProgramResult};
+use pinocchio::{
+    account_info::AccountInfo,
+    program_error::ProgramError,
+    pubkey::{Pubkey, PUBKEY_BYTES},
+    ProgramResult,
+};
 
 use super::shared;
 
@@ -7,6 +12,13 @@ pub fn process_initialize_account2(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let owner = unsafe { &*(instruction_data.as_ptr() as *const Pubkey) };
+    // SAFETY: validate `instruction_data` length.
+    let owner = unsafe {
+        if instruction_data.len() != PUBKEY_BYTES {
+            return Err(ProgramError::InvalidInstructionData);
+        } else {
+            &*(instruction_data.as_ptr() as *const Pubkey)
+        }
+    };
     shared::initialize_account::process_initialize_account(accounts, Some(owner), true)
 }

p-token/src/processor/initialize_account3.rs

@@ -1,4 +1,9 @@
-use pinocchio::{account_info::AccountInfo, pubkey::Pubkey, ProgramResult};
+use pinocchio::{
+    account_info::AccountInfo,
+    program_error::ProgramError,
+    pubkey::{Pubkey, PUBKEY_BYTES},
+    ProgramResult,
+};
 
 use super::shared;
 
@@ -7,6 +12,13 @@ pub fn process_initialize_account3(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let owner = unsafe { &*(instruction_data.as_ptr() as *const Pubkey) };
+    // SAFETY: validate `instruction_data` length.
+    let owner = unsafe {
+        if instruction_data.len() != PUBKEY_BYTES {
+            return Err(ProgramError::InvalidInstructionData);
+        } else {
+            &*(instruction_data.as_ptr() as *const Pubkey)
+        }
+    };
     shared::initialize_account::process_initialize_account(accounts, Some(owner), false)
 }

p-token/src/processor/initialize_immutable_owner.rs

@@ -8,6 +8,7 @@ use token_interface::{
 pub fn process_initialize_immutable_owner(accounts: &[AccountInfo]) -> ProgramResult {
     let token_account_info = accounts.first().ok_or(ProgramError::NotEnoughAccountKeys)?;
 
+    // SAFETY: single immutable borrow to `token_account_info` account data.
     let account = unsafe { load_unchecked::<Account>(token_account_info.borrow_data_unchecked())? };
 
     if account.is_initialized() {

p-token/src/processor/initialize_mint.rs

@@ -35,6 +35,7 @@ pub fn process_initialize_mint(
         (mint_info, None)
     };
 
+    // SAFETY: single mutable borrow to `mint_info` account data.
     let mint = unsafe { load_mut_unchecked::<Mint>(mint_info.borrow_mut_data_unchecked())? };
 
     if mint.is_initialized() {
@@ -44,7 +45,9 @@ pub fn process_initialize_mint(
     // Check rent-exempt status of the mint account.
 
     let is_exempt = if let Some(rent_sysvar_info) = rent_sysvar_info {
-        let rent = unsafe { Rent::from_bytes(rent_sysvar_info.borrow_data_unchecked()) };
+        // SAFETY: single immutable borrow to `rent_sysvar_info`; account ID and length are
+        // checked by `from_account_info_unchecked`.
+        let rent = unsafe { Rent::from_account_info_unchecked(rent_sysvar_info)? };
         rent.is_exempt(mint_info.lamports(), size_of::<Mint>())
     } else {
         Rent::get()?.is_exempt(mint_info.lamports(), size_of::<Mint>())
@@ -81,7 +84,7 @@ impl InitializeMint<'_> {
         // - decimals (1 byte)
         // - mint_authority (32 bytes)
         // - option + freeze_authority (1 byte + 32 bytes)
-        if bytes.len() < 34 {
+        if bytes.len() < 34 || (bytes[33] == 1 && bytes.len() < 66) {
             return Err(ProgramError::InvalidInstructionData);
         }
 
@@ -93,16 +96,19 @@ impl InitializeMint<'_> {
 
     #[inline]
     pub fn decimals(&self) -> u8 {
+        // SAFETY: the `bytes` length was validated in `try_from_bytes`.
         unsafe { *self.raw }
     }
 
     #[inline]
     pub fn mint_authority(&self) -> &Pubkey {
+        // SAFETY: the `bytes` length was validated in `try_from_bytes`.
         unsafe { &*(self.raw.add(1) as *const Pubkey) }
     }
 
     #[inline]
     pub fn freeze_authority(&self) -> Option<&Pubkey> {
+        // SAFETY: the `bytes` length was validated in `try_from_bytes`.
         unsafe {
             if *self.raw.add(33) == 0 {
                 Option::None

p-token/src/processor/mint_to_checked.rs

@@ -4,17 +4,20 @@ use super::shared;
 
 #[inline(always)]
 pub fn process_mint_to_checked(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let (amount, decimals) = instruction_data.split_at(core::mem::size_of::<u64>());
+    // expected u64 (8) + u8 (1)
+    let (amount, decimals) = if instruction_data.len() == 9 {
+        let (amount, decimals) = instruction_data.split_at(core::mem::size_of::<u64>());
+        (
+            u64::from_le_bytes(
+                amount
+                    .try_into()
+                    .map_err(|_error| ProgramError::InvalidInstructionData)?,
+            ),
+            decimals.first(),
+        )
+    } else {
+        return Err(ProgramError::InvalidInstructionData);
+    };
 
-    let amount = u64::from_le_bytes(
-        amount
-            .try_into()
-            .map_err(|_error| ProgramError::InvalidInstructionData)?,
-    );
-
-    shared::mint_to::process_mint_to(
-        accounts,
-        amount,
-        Some(*decimals.first().ok_or(ProgramError::InvalidAccountData)?),
-    )
+    shared::mint_to::process_mint_to(accounts, amount, decimals.copied())
 }