diff --git a/lee/state_machine/src/error.rs b/lee/state_machine/src/error.rs index 2f073746..2d0221bc 100644 --- a/lee/state_machine/src/error.rs +++ b/lee/state_machine/src/error.rs @@ -131,6 +131,9 @@ pub enum InvalidProgramBehaviorError { #[error("Called program {program_id:?} which is not listed in dependencies")] UndeclaredProgramDependency { program_id: ProgramId }, + + #[error("Account {account_id} was declared in the transaction but is missing from the program output")] + DeclaredAccountMissingFromOutput { account_id: AccountId }, } #[cfg(test)] diff --git a/lee/state_machine/src/lib.rs b/lee/state_machine/src/lib.rs index f8cc034a..bd1f940c 100644 --- a/lee/state_machine/src/lib.rs +++ b/lee/state_machine/src/lib.rs @@ -77,6 +77,14 @@ mod test_methods { ) } + #[must_use] + pub const fn dropped_account() -> Program { + Program::new_unchecked( + test_methods::DROPPED_ACCOUNT_ID, + Cow::Borrowed(test_methods::DROPPED_ACCOUNT_ELF), + ) + } + #[must_use] pub const fn program_owner_changer() -> Program { Program::new_unchecked( diff --git a/lee/state_machine/src/state/tests/mod.rs b/lee/state_machine/src/state/tests/mod.rs index 55a99ead..caf555b5 100644 --- a/lee/state_machine/src/state/tests/mod.rs +++ b/lee/state_machine/src/state/tests/mod.rs @@ -48,6 +48,7 @@ impl V03State { self.insert_program(crate::test_methods::nonce_changer()); self.insert_program(crate::test_methods::extra_output()); self.insert_program(crate::test_methods::missing_output()); + self.insert_program(crate::test_methods::dropped_account()); self.insert_program(crate::test_methods::program_owner_changer()); self.insert_program(crate::test_methods::data_changer()); self.insert_program(crate::test_methods::minter()); diff --git a/lee/state_machine/src/state/tests/public_program_rules.rs b/lee/state_machine/src/state/tests/public_program_rules.rs index 1c67aeff..236bddcf 100644 --- a/lee/state_machine/src/state/tests/public_program_rules.rs +++ b/lee/state_machine/src/state/tests/public_program_rules.rs @@ -79,6 +79,58 @@ fn program_should_fail_with_missing_output_accounts() { )); } +/// A program can drop an entire account from its own output — both its `pre_state` and +/// `post_state` together, not just one side — while staying internally consistent +/// (`pre_states.len() == post_states.len()` within its own report, so `validate_execution`'s +/// length check alone can't catch it). This must still be rejected: every account the caller +/// declared in the transaction must appear somewhere in the final diff. +#[test] +fn program_should_fail_if_it_drops_a_declared_account() { + // Both accounts need a non-default program_owner: an account left at DEFAULT_PROGRAM_ID with + // non-default data would itself violate the (separate, pre-existing) "claim before mutating a + // default-owned account" rule the moment it's echoed back — unrelated to what this test + // targets. `with_public_account_balances` leaves program_owner at DEFAULT_PROGRAM_ID, so use + // `with_public_accounts` to set it explicitly instead. + let mut state = V03State::new() + .with_public_accounts([ + ( + AccountId::new([1; 32]), + Account { + program_owner: crate::test_methods::dropped_account().id(), + balance: 100, + ..Account::default() + }, + ), + ( + AccountId::new([2; 32]), + Account { + program_owner: crate::test_methods::dropped_account().id(), + balance: 0, + ..Account::default() + }, + ), + ]) + .with_test_programs(); + let account_ids = vec![AccountId::new([1; 32]), AccountId::new([2; 32])]; + let program_id = crate::test_methods::dropped_account().id(); + let message = + public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!( + matches!( + result, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::DeclaredAccountMissingFromOutput { account_id } + )) if account_id == AccountId::new([2; 32]) + ), + "expected DeclaredAccountMissingFromOutput for the dropped account, got {result:?}" + ); +} + #[test] fn program_should_fail_if_modifies_program_owner_with_only_non_default_program_owner() { let initial_data = [( diff --git a/lee/state_machine/src/validated_state_diff/mod.rs b/lee/state_machine/src/validated_state_diff/mod.rs index 8fae4fee..a740e1ef 100644 --- a/lee/state_machine/src/validated_state_diff/mod.rs +++ b/lee/state_machine/src/validated_state_diff/mod.rs @@ -293,6 +293,21 @@ impl ValidatedStateDiff { ); } + // Every account the caller declared as part of the transaction must appear in the final + // diff. A well-behaved program's output always echoes back every account it was handed + // (unchanged if it only reads it) — this catches a program (or a macro-generated + // dispatcher wrapping one) that silently drops an account from its own output instead. + // Chained calls may still introduce accounts beyond this original list; this only checks + // the reverse direction. + for account_id in &message.account_ids { + ensure!( + state_diff.contains_key(account_id), + InvalidProgramBehaviorError::DeclaredAccountMissingFromOutput { + account_id: *account_id + } + ); + } + Ok(Self(StateDiff { signer_account_ids, public_diff: state_diff, diff --git a/lee/state_machine/test_methods/guest/src/bin/dropped_account.rs b/lee/state_machine/test_methods/guest/src/bin/dropped_account.rs new file mode 100644 index 00000000..e79dbc3b --- /dev/null +++ b/lee/state_machine/test_methods/guest/src/bin/dropped_account.rs @@ -0,0 +1,39 @@ +use lee_core::program::{AccountPostState, ProgramInput, ProgramOutput, read_lee_inputs}; + +type Instruction = (); + +/// Silently drops the second account entirely from its own output: given two `pre_states`, it +/// returns only one `(pre, post)` pair, echoing the first account back unchanged. +/// +/// Unlike `missing_output` (whose own `pre_states.len() != post_states.len()`, tripping +/// `validate_execution`'s internal length check directly), this program's own output is +/// internally consistent — one `pre_state` and one matching `post_state` — it just reports fewer +/// accounts than it was handed. This mirrors a well-behaved-looking dispatcher that filters an +/// account out of both sides of its output together (e.g. a stale-signer-nonce workaround that's +/// too broad), rather than an obviously malformed program. +fn main() { + let ( + ProgramInput { + self_program_id, + caller_program_id, + pre_states, + .. + }, + instruction_words, + ) = read_lee_inputs::(); + + let Ok([pre1, _pre2]) = <[_; 2]>::try_from(pre_states) else { + return; + }; + + let account_pre1 = pre1.account.clone(); + + ProgramOutput::new( + self_program_id, + caller_program_id, + instruction_words, + vec![pre1], + vec![AccountPostState::new(account_pre1)], + ) + .write(); +}