diff --git a/basics/counter/native/program/src/lib.rs b/basics/counter/native/program/src/lib.rs index f8bded9dc..e50985268 100644 --- a/basics/counter/native/program/src/lib.rs +++ b/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(()) diff --git a/basics/counter/native/program/tests/test.rs b/basics/counter/native/program/tests/test.rs index 474631327..1f1bf6ca1 100644 --- a/basics/counter/native/program/tests/test.rs +++ b/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::(); + 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()); +}