counter(native): avoid panic, validate owner, checked_add, fail unknown instructions

basics/counter/native/program/src/lib.rs

@@ -20,38 +20,59 @@ use solana_program::entrypoint;
 entrypoint!(process_instruction);
 
 pub fn process_instruction(
-    _program_id: &Pubkey,
+    program_id: &Pubkey,
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let (instruction_discriminant, instruction_data_inner) = instruction_data.split_at(1);
-    match instruction_discriminant[0] {
+    let (instruction_discriminant, instruction_data_inner) = instruction_data
+        .split_first()
+        .ok_or(ProgramError::InvalidInstructionData)?;
+
+    match instruction_discriminant {
         0 => {
             msg!("Instruction: Increment");
-            process_increment_counter(accounts, instruction_data_inner)?;
+            process_increment_counter(program_id, accounts, instruction_data_inner)?;
         }
         _ => {
-            msg!("Error: unknown instruction")
+            msg!("Error: unknown instruction");
+            return Err(ProgramError::InvalidInstructionData);
         }
     }
+
     Ok(())
 }
 
 pub fn process_increment_counter(
+    program_id: &Pubkey,
     accounts: &[AccountInfo],
     _instruction_data: &[u8],
 ) -> Result<(), ProgramError> {
     let account_info_iter = &mut accounts.iter();
 
     let counter_account = next_account_info(account_info_iter)?;
-    assert!(
-        counter_account.is_writable,
-        "Counter account must be writable"
-    );
-
-    let mut counter = Counter::try_from_slice(&counter_account.try_borrow_mut_data()?)?;
-    counter.count += 1;
-    counter.serialize(&mut *counter_account.data.borrow_mut())?;
+
+    // Never panic in on-chain programs: panics consume all compute units and are hard to debug.
+    if !counter_account.is_writable {
+        msg!("Counter account must be writable");
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    // Ownership check prevents unexpected state corruption when a wrong account is passed.
+    if counter_account.owner != program_id {
+        msg!("Counter account is not owned by this program");
+        return Err(ProgramError::IncorrectProgramId);
+    }
+
+    // Use a single mutable borrow to avoid runtime borrow conflicts.
+    let mut data = counter_account.try_borrow_mut_data()?;
+    let mut counter = Counter::try_from_slice(&data)?;
+
+    counter.count = counter
+        .count
+        .checked_add(1)
+        .ok_or(ProgramError::ArithmeticOverflow)?;
+
+    counter.serialize(&mut &mut data[..])?;
 
     msg!("Counter state incremented to {:?}", counter.count);
     Ok(())

basics/counter/native/program/tests/test.rs

@@ -53,9 +53,83 @@ fn test_counter() {
         svm.latest_blockhash(),
     );
 
-    let _ = svm.send_transaction(tx).is_ok();
+    assert!(svm.send_transaction(tx).is_ok());
 
     let counter_account_data = svm.get_account(&counter_account.pubkey()).unwrap().data;
     let counter = Counter::try_from_slice(&counter_account_data).unwrap();
     assert_eq!(counter.count, 1);
 }
+
+#[test]
+fn test_unknown_instruction_fails() {
+    let program_id = Pubkey::new_unique();
+    let program_bytes = include_bytes!("../../tests/fixtures/counter_solana_native.so");
+
+    let mut svm = LiteSVM::new();
+    svm.add_program(program_id, program_bytes).unwrap();
+
+    let payer = Keypair::new();
+    svm.airdrop(&payer.pubkey(), LAMPORTS_PER_SOL).unwrap();
+
+    let ix = Instruction {
+        program_id,
+        accounts: vec![],
+        data: vec![99],
+    };
+
+    let tx = Transaction::new_signed_with_payer(
+        &[ix],
+        Some(&payer.pubkey()),
+        &[payer],
+        svm.latest_blockhash(),
+    );
+
+    assert!(svm.send_transaction(tx).is_err());
+}
+
+#[test]
+fn test_readonly_counter_fails() {
+    let program_id = Pubkey::new_unique();
+    let program_bytes = include_bytes!("../../tests/fixtures/counter_solana_native.so");
+
+    let mut svm = LiteSVM::new();
+    svm.add_program(program_id, program_bytes).unwrap();
+
+    let payer = Keypair::new();
+    let counter_account = Keypair::new();
+
+    svm.airdrop(&payer.pubkey(), LAMPORTS_PER_SOL * 10).unwrap();
+
+    let counter_account_size = std::mem::size_of::<Counter>();
+    let create_ix = create_account(
+        &payer.pubkey(),
+        &counter_account.pubkey(),
+        Rent::default().minimum_balance(counter_account_size),
+        counter_account_size as u64,
+        &program_id,
+    );
+
+    let tx = Transaction::new_signed_with_payer(
+        &[create_ix],
+        Some(&payer.pubkey()),
+        &[&payer, &counter_account],
+        svm.latest_blockhash(),
+    );
+    assert!(svm.send_transaction(tx).is_ok());
+
+    // Mark counter account as readonly; program should return an error (not panic).
+    let ix = Instruction {
+        program_id,
+        accounts: vec![AccountMeta::new_readonly(counter_account.pubkey(), false)],
+        data: vec![0],
+    };
+
+    let tx = Transaction::new_signed_with_payer(
+        &[ix],
+        Some(&payer.pubkey()),
+        &[payer],
+        svm.latest_blockhash(),
+    );
+
+    assert!(svm.send_transaction(tx).is_err());
+}