chore: remove unwraps system program (#1803)

ananas-block · coderabbitai[bot] · web-flow · commit da6459d6d591 · 2025-06-15T14:08:45.000+01:00
* chore: remove unwraps system program

* fix: add out of bounds check for concurrent Merkle tree deserialization

* Update anchor-programs/system/src/errors.rs

Co-authored-by: coderabbitai[bot] &lt;136622811+coderabbitai[bot]@users.noreply.github.com&gt;

---------

Co-authored-by: coderabbitai[bot] &lt;136622811+coderabbitai[bot]@users.noreply.github.com&gt;
diff --git a/anchor-programs/system/src/errors.rs b/anchor-programs/system/src/errors.rs
@@ -86,4 +86,6 @@ pub enum SystemProgramError {
     CpiContextAlreadySet,
     InvalidTreeHeight,
     TooManyOutputAccounts,
+    #[msg("Failed to borrow account data")]
+    BorrowingDataFailed,
 }
diff --git a/js/compressed-token/src/types.ts b/js/compressed-token/src/types.ts
@@ -73,7 +73,6 @@ export type BatchCompressInstructionData = {
     bump: number;
 };
 
-
 export type MintToInstructionData = {
     recipients: PublicKey[];
     amounts: BN[];
diff --git a/programs/system/src/account_compression_state/address.rs b/programs/system/src/account_compression_state/address.rs
@@ -20,36 +20,23 @@ pub struct AddressMerkleTreeAccount {
 pub fn address_merkle_tree_from_bytes_zero_copy(
     data: &[u8],
 ) -> Result<IndexedMerkleTreeZeroCopy<Poseidon, usize, 26, 16>> {
-    let data = &data[8 + mem::size_of::<AddressMerkleTreeAccount>()..];
-    let merkle_tree = IndexedMerkleTreeZeroCopy::from_bytes_zero_copy(data).unwrap();
-    Ok(merkle_tree)
-}
-
-pub fn address_merkle_tree_from_bytes_zero_copy_init(
-    data: &mut [u8],
-    height: usize,
-    canopy_depth: usize,
-    changelog_capacity: usize,
-    roots_capacity: usize,
-    indexed_changelog_capacity: usize,
-) -> Result<IndexedMerkleTreeZeroCopyMut<Poseidon, usize, 26, 16>> {
-    let data = &mut data[8 + mem::size_of::<AddressMerkleTreeAccount>()..];
-    let merkle_tree = IndexedMerkleTreeZeroCopyMut::from_bytes_zero_copy_init(
-        data,
-        height,
-        canopy_depth,
-        changelog_capacity,
-        roots_capacity,
-        indexed_changelog_capacity,
-    )
-    .unwrap();
+    let required_size = 8 + mem::size_of::<AddressMerkleTreeAccount>();
+    if data.len() < required_size {
+        return Err(crate::errors::SystemProgramError::InvalidAccount.into());
+    }
+    let data = &data[required_size..];
+    let merkle_tree = IndexedMerkleTreeZeroCopy::from_bytes_zero_copy(data)?;
     Ok(merkle_tree)
 }
 
 pub fn address_merkle_tree_from_bytes_zero_copy_mut(
     data: &mut [u8],
 ) -> Result<IndexedMerkleTreeZeroCopyMut<Poseidon, usize, 26, 16>> {
-    let data = &mut data[8 + mem::size_of::<AddressMerkleTreeAccount>()..];
-    let merkle_tree = IndexedMerkleTreeZeroCopyMut::from_bytes_zero_copy_mut(data).unwrap();
+    let required_size = 8 + mem::size_of::<AddressMerkleTreeAccount>();
+    if data.len() < required_size {
+        return Err(crate::errors::SystemProgramError::InvalidAccount.into());
+    }
+    let data = &mut data[required_size..];
+    let merkle_tree = IndexedMerkleTreeZeroCopyMut::from_bytes_zero_copy_mut(data)?;
     Ok(merkle_tree)
 }
diff --git a/programs/system/src/account_compression_state/state.rs b/programs/system/src/account_compression_state/state.rs
@@ -49,37 +49,26 @@ impl StateMerkleTreeAccount {
     }
 }
 
-pub fn state_merkle_tree_from_bytes_zero_copy_init(
-    data: &mut [u8],
-    height: usize,
-    canopy_depth: usize,
-    changelog_capacity: usize,
-    roots_capacity: usize,
-) -> Result<ConcurrentMerkleTreeZeroCopyMut<Poseidon, 26>> {
-    let data = &mut data[8 + mem::size_of::<StateMerkleTreeAccount>()..];
-    let merkle_tree = ConcurrentMerkleTreeZeroCopyMut::from_bytes_zero_copy_init(
-        data,
-        height,
-        canopy_depth,
-        changelog_capacity,
-        roots_capacity,
-    )
-    .unwrap();
-    Ok(merkle_tree)
-}
-
 pub fn state_merkle_tree_from_bytes_zero_copy(
     data: &[u8],
 ) -> Result<ConcurrentMerkleTreeZeroCopy<Poseidon, 26>> {
-    let data = &data[8 + mem::size_of::<StateMerkleTreeAccount>()..];
-    let merkle_tree = ConcurrentMerkleTreeZeroCopy::from_bytes_zero_copy(data).unwrap();
+    let required_size = 8 + mem::size_of::<StateMerkleTreeAccount>();
+    if data.len() < required_size {
+        return Err(crate::errors::SystemProgramError::InvalidAccount.into());
+    }
+    let data = &data[required_size..];
+    let merkle_tree = ConcurrentMerkleTreeZeroCopy::from_bytes_zero_copy(data)?;
     Ok(merkle_tree)
 }
 
 pub fn state_merkle_tree_from_bytes_zero_copy_mut(
     data: &mut [u8],
 ) -> Result<ConcurrentMerkleTreeZeroCopyMut<Poseidon, 26>> {
-    let data = &mut data[8 + mem::size_of::<StateMerkleTreeAccount>()..];
-    let merkle_tree = ConcurrentMerkleTreeZeroCopyMut::from_bytes_zero_copy_mut(data).unwrap();
+    let required_size = 8 + mem::size_of::<StateMerkleTreeAccount>();
+    if data.len() < required_size {
+        return Err(crate::errors::SystemProgramError::InvalidAccount.into());
+    }
+    let data = &mut data[required_size..];
+    let merkle_tree = ConcurrentMerkleTreeZeroCopyMut::from_bytes_zero_copy_mut(data)?;
     Ok(merkle_tree)
 }
diff --git a/programs/system/src/accounts/remaining_account_checks.rs b/programs/system/src/accounts/remaining_account_checks.rs
@@ -68,7 +68,7 @@ pub(crate) fn try_from_account_info<'a, 'info: 'a>(
     {
         let data = account_info
             .try_borrow_data()
-            .map_err(|_| SystemProgramError::InvalidAccount)?;
+            .map_err(|_| SystemProgramError::BorrowingDataFailed)?;
 
         if data.len() < 8 {
             return Ok(AcpAccount::Unknown());
@@ -82,20 +82,18 @@ pub(crate) fn try_from_account_info<'a, 'info: 'a>(
             tree_type.copy_from_slice(
                 &account_info
                     .try_borrow_data()
-                    .map_err(|_| SystemProgramError::InvalidAccount)?[8..16],
+                    .map_err(|_| SystemProgramError::BorrowingDataFailed)?[8..16],
             );
             let tree_type = TreeType::from(u64::from_le_bytes(tree_type));
             match tree_type {
                 TreeType::AddressV2 => {
-                    let tree =
-                        BatchedMerkleTreeAccount::address_from_account_info(account_info).unwrap();
+                    let tree = BatchedMerkleTreeAccount::address_from_account_info(account_info)?;
                     let program_owner = tree.metadata.access_metadata.program_owner;
                     // for batched trees we set the fee when setting the rollover fee.
                     Ok((AcpAccount::BatchedAddressTree(tree), program_owner))
                 }
                 TreeType::StateV2 => {
-                    let tree =
-                        BatchedMerkleTreeAccount::state_from_account_info(account_info).unwrap();
+                    let tree = BatchedMerkleTreeAccount::state_from_account_info(account_info)?;
                     let program_owner = tree.metadata.access_metadata.program_owner;
                     Ok((AcpAccount::BatchedStateTree(tree), program_owner))
                 }
@@ -111,14 +109,16 @@ pub(crate) fn try_from_account_info<'a, 'info: 'a>(
             }
         }
         BatchedQueueAccount::LIGHT_DISCRIMINATOR => {
-            let queue = BatchedQueueAccount::output_from_account_info(account_info).unwrap();
+            let queue = BatchedQueueAccount::output_from_account_info(account_info)?;
             let program_owner = queue.metadata.access_metadata.program_owner;
             Ok((AcpAccount::OutputQueue(queue), program_owner))
         }
         STATE_MERKLE_TREE_ACCOUNT_DISCRIMINATOR => {
             let program_owner = {
-                check_owner(&ACCOUNT_COMPRESSION_PROGRAM_ID, account_info).unwrap();
-                let data = account_info.try_borrow_data().unwrap();
+                check_owner(&ACCOUNT_COMPRESSION_PROGRAM_ID, account_info)?;
+                let data = account_info
+                    .try_borrow_data()
+                    .map_err(|_| SystemProgramError::BorrowingDataFailed)?;
                 let merkle_tree = bytemuck::from_bytes::<StateMerkleTreeAccount>(
                     &data[8..StateMerkleTreeAccount::LEN],
                 );
@@ -144,15 +144,18 @@ pub(crate) fn try_from_account_info<'a, 'info: 'a>(
             Ok((
                 AcpAccount::StateTree((
                     (*account_info.key()).into(),
-                    state_merkle_tree_from_bytes_zero_copy_mut(data_slice).unwrap(),
+                    state_merkle_tree_from_bytes_zero_copy_mut(data_slice)
+                        .map_err(|e| SystemProgramError::ProgramError(e.into()))?,
                 )),
                 program_owner,
             ))
         }
         ADDRESS_MERKLE_TREE_ACCOUNT_DISCRIMINATOR => {
             let program_owner = {
-                check_owner(&ACCOUNT_COMPRESSION_PROGRAM_ID, account_info).unwrap();
-                let data = account_info.try_borrow_data().unwrap();
+                check_owner(&ACCOUNT_COMPRESSION_PROGRAM_ID, account_info)?;
+                let data = account_info
+                    .try_borrow_data()
+                    .map_err(|_| SystemProgramError::BorrowingDataFailed)?;
 
                 let merkle_tree = bytemuck::from_bytes::<AddressMerkleTreeAccount>(
                     &data[8..AddressMerkleTreeAccount::LEN],
@@ -170,14 +173,17 @@ pub(crate) fn try_from_account_info<'a, 'info: 'a>(
             Ok((
                 AcpAccount::AddressTree((
                     (*account_info.key()).into(),
-                    address_merkle_tree_from_bytes_zero_copy_mut(data_slice).unwrap(),
+                    address_merkle_tree_from_bytes_zero_copy_mut(data_slice)
+                        .map_err(|e| SystemProgramError::ProgramError(e.into()))?,
                 )),
                 program_owner,
             ))
         }
         QUEUE_ACCOUNT_DISCRIMINATOR => {
-            check_owner(&ACCOUNT_COMPRESSION_PROGRAM_ID, account_info).unwrap();
-            let data = account_info.try_borrow_data().unwrap();
+            check_owner(&ACCOUNT_COMPRESSION_PROGRAM_ID, account_info)?;
+            let data = account_info
+                .try_borrow_data()
+                .map_err(|_| SystemProgramError::BorrowingDataFailed)?;
             let queue = bytemuck::from_bytes::<QueueAccount>(&data[8..QueueAccount::LEN]);
 
             if queue.metadata.queue_type == QueueType::AddressV1 as u64 {
diff --git a/programs/system/src/errors.rs b/programs/system/src/errors.rs
@@ -1,3 +1,8 @@
+use light_account_checks::error::AccountError;
+use light_batched_merkle_tree::errors::BatchedMerkleTreeError;
+use light_concurrent_merkle_tree::errors::ConcurrentMerkleTreeError;
+use light_indexed_merkle_tree::errors::IndexedMerkleTreeError;
+use light_zero_copy::errors::ZeroCopyError;
 use pinocchio::program_error::ProgramError;
 use thiserror::Error;
 
@@ -107,10 +112,90 @@ pub enum SystemProgramError {
     InvalidTreeHeight,
     #[error("TooManyOutputAccounts")]
     TooManyOutputAccounts,
+    #[error("Batched Merkle tree error {0}")]
+    BatchedMerkleTreeError(#[from] BatchedMerkleTreeError),
+    #[error("Concurrent Merkle tree error {0}")]
+    ConcurrentMerkleTreeError(#[from] ConcurrentMerkleTreeError),
+    #[error("Indexed Merkle tree error {0}")]
+    IndexedMerkleTreeError(#[from] IndexedMerkleTreeError),
+    #[error("Account checks error {0}")]
+    AccountError(#[from] AccountError),
+    #[error("Zero copy error {0}")]
+    ZeroCopyError(#[from] ZeroCopyError),
+    #[error("Program error code: {0}")]
+    ProgramError(u64),
+    #[error("Borrowing data failed")]
+    BorrowingDataFailed,
+}
+
+impl From<SystemProgramError> for u32 {
+    fn from(e: SystemProgramError) -> u32 {
+        match e {
+            SystemProgramError::SumCheckFailed => 6000,
+            SystemProgramError::SignerCheckFailed => 6001,
+            SystemProgramError::CpiSignerCheckFailed => 6002,
+            SystemProgramError::ComputeInputSumFailed => 6003,
+            SystemProgramError::ComputeOutputSumFailed => 6004,
+            SystemProgramError::ComputeRpcSumFailed => 6005,
+            SystemProgramError::InvalidAddress => 6006,
+            SystemProgramError::DeriveAddressError => 6007,
+            SystemProgramError::CompressedSolPdaUndefinedForCompressSol => 6008,
+            SystemProgramError::DecompressLamportsUndefinedForCompressSol => 6009,
+            SystemProgramError::CompressedSolPdaUndefinedForDecompressSol => 6010,
+            SystemProgramError::DeCompressLamportsUndefinedForDecompressSol => 6011,
+            SystemProgramError::DecompressRecipientUndefinedForDecompressSol => 6012,
+            SystemProgramError::WriteAccessCheckFailed => 6013,
+            SystemProgramError::InvokingProgramNotProvided => 6014,
+            SystemProgramError::InvalidCapacity => 6015,
+            SystemProgramError::InvalidMerkleTreeOwner => 6016,
+            SystemProgramError::ProofIsNone => 6017,
+            SystemProgramError::ProofIsSome => 6018,
+            SystemProgramError::EmptyInputs => 6019,
+            SystemProgramError::CpiContextAccountUndefined => 6020,
+            SystemProgramError::CpiContextEmpty => 6021,
+            SystemProgramError::CpiContextMissing => 6022,
+            SystemProgramError::DecompressionRecipientDefined => 6023,
+            SystemProgramError::SolPoolPdaDefined => 6024,
+            SystemProgramError::AppendStateFailed => 6025,
+            SystemProgramError::InstructionNotCallable => 6026,
+            SystemProgramError::CpiContextFeePayerMismatch => 6027,
+            SystemProgramError::CpiContextAssociatedMerkleTreeMismatch => 6028,
+            SystemProgramError::NoInputs => 6029,
+            SystemProgramError::InputMerkleTreeIndicesNotInOrder => 6030,
+            SystemProgramError::OutputMerkleTreeIndicesNotInOrder => 6031,
+            SystemProgramError::OutputMerkleTreeNotUnique => 6032,
+            SystemProgramError::DataFieldUndefined => 6033,
+            SystemProgramError::ReadOnlyAddressAlreadyExists => 6034,
+            SystemProgramError::ReadOnlyAccountDoesNotExist => 6035,
+            SystemProgramError::HashChainInputsLenghtInconsistent => 6036,
+            SystemProgramError::InvalidAddressTreeHeight => 6037,
+            SystemProgramError::InvalidStateTreeHeight => 6038,
+            SystemProgramError::InvalidArgument => 6039,
+            SystemProgramError::InvalidAccount => 6040,
+            SystemProgramError::AddressMerkleTreeAccountDiscriminatorMismatch => 6041,
+            SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch => 6042,
+            SystemProgramError::ProofVerificationFailed => 6043,
+            SystemProgramError::InvalidAccountMode => 6044,
+            SystemProgramError::InvalidInstructionDataDiscriminator => 6045,
+            SystemProgramError::NewAddressAssignedIndexOutOfBounds => 6046,
+            SystemProgramError::AddressIsNone => 6047,
+            SystemProgramError::AddressDoesNotMatch => 6048,
+            SystemProgramError::CpiContextAlreadySet => 6049,
+            SystemProgramError::InvalidTreeHeight => 6050,
+            SystemProgramError::TooManyOutputAccounts => 6051,
+            SystemProgramError::BatchedMerkleTreeError(e) => e.into(),
+            SystemProgramError::IndexedMerkleTreeError(e) => e.into(),
+            SystemProgramError::ConcurrentMerkleTreeError(e) => e.into(),
+            SystemProgramError::AccountError(e) => e.into(),
+            SystemProgramError::ProgramError(e) => u32::try_from(e).unwrap_or(0),
+            SystemProgramError::BorrowingDataFailed => 6052,
+            SystemProgramError::ZeroCopyError(e) => e.into(),
+        }
+    }
 }
 
 impl From<SystemProgramError> for ProgramError {
     fn from(e: SystemProgramError) -> ProgramError {
-        ProgramError::Custom(e as u32 + 6000)
+        ProgramError::Custom(e.into())
     }
 }
diff --git a/programs/system/src/invoke_cpi/process_cpi_context.rs b/programs/system/src/invoke_cpi/process_cpi_context.rs
@@ -105,7 +105,8 @@ pub fn set_cpi_context<'a, 'info, T: InstructionData<'a>>(
     use borsh::{BorshDeserialize, BorshSerialize};
     let cpi_context_account = {
         let data = cpi_context_account_info.try_borrow_data()?;
-        let mut cpi_context_account = CpiContextAccount::deserialize(&mut &data[8..]).unwrap();
+        let mut cpi_context_account = CpiContextAccount::deserialize(&mut &data[8..])
+            .map_err(|_| SystemProgramError::BorrowingDataFailed)?;
         if instruction_data.cpi_context().unwrap().first_set_context {
             cpi_context_account.context.clear();
             cpi_context_account.fee_payer = fee_payer;
diff --git a/programs/system/src/lib.rs b/programs/system/src/lib.rs
@@ -91,7 +91,7 @@ pub fn invoke<'a, 'b, 'c: 'info, 'info>(
     // remove vec prefix
     let instruction_data = &instruction_data[4..];
 
-    let (inputs, _) = ZInstructionDataInvoke::zero_copy_at(instruction_data).unwrap();
+    let (inputs, _) = ZInstructionDataInvoke::zero_copy_at(instruction_data)?;
     let (ctx, remaining_accounts) = InvokeInstruction::from_account_infos(accounts)?;
 
     input_compressed_accounts_signer_check(
@@ -117,7 +117,7 @@ pub fn invoke_cpi<'a, 'b, 'c: 'info, 'info>(
     // remove vec prefix
     let instruction_data = &instruction_data[4..];
 
-    let (inputs, _) = ZInstructionDataInvokeCpi::zero_copy_at(instruction_data).unwrap();
+    let (inputs, _) = ZInstructionDataInvokeCpi::zero_copy_at(instruction_data)?;
 
     let (ctx, remaining_accounts) = InvokeCpiInstruction::from_account_infos(accounts)?;
 

Original file line number	Diff line number	Diff line change
`@@ -86,4 +86,6 @@ pub enum SystemProgramError {`
`86`	`86`	`CpiContextAlreadySet,`
`87`	`87`	`InvalidTreeHeight,`
`88`	`88`	`TooManyOutputAccounts,`
	`89`	`+ #[msg("Failed to borrow account data")]`
	`90`	`+ BorrowingDataFailed,`
`89`	`91`	`}`