From fb8491593571061a1767aceae0112372cd8998e9 Mon Sep 17 00:00:00 2001 From: Marvin Jones Date: Wed, 15 Jul 2026 15:00:05 -0400 Subject: [PATCH 1/2] fix(state_machine): reject public transactions that silently drop a declared account MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ValidatedStateDiff::from_public_transaction never checked that the accounts touched in a program's output matched the caller-declared message.account_ids — it just folded whatever pairs the program returned into the diff. A program (or a macro-generated dispatcher wrapping one) that silently drops an account from both sides of its own output together stays internally consistent (pre_states.len() == post_states.len()) and passes validate_execution's existing checks, so the dropped account simply vanishes with no error. Add a check after the chained-call loop: every account_id in message.account_ids must appear in the final state_diff, or the transaction is rejected with the new DeclaredAccountMissingFromOutput error. Add a dropped_account test-guest program that reproduces the exact shape of the bug (two pre_states in, one consistent (pre, post) pair out) and a regression test proving the transaction is now rejected. Verified the test fails with Ok(()) when the check is removed, and passes once it's restored. --- lee/state_machine/src/error.rs | 3 ++ lee/state_machine/src/lib.rs | 8 +++ lee/state_machine/src/state/tests/mod.rs | 1 + .../src/state/tests/public_program_rules.rs | 52 +++++++++++++++++++ .../src/validated_state_diff/mod.rs | 15 ++++++ .../guest/src/bin/dropped_account.rs | 39 ++++++++++++++ 6 files changed, 118 insertions(+) create mode 100644 lee/state_machine/test_methods/guest/src/bin/dropped_account.rs 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(); +} From 09af47b862830f567c8e844e08d36374ed920cb1 Mon Sep 17 00:00:00 2001 From: Marvin Jones Date: Wed, 15 Jul 2026 15:20:04 -0400 Subject: [PATCH 2/2] style: tighten comments and apply nightly rustfmt Condense the explanatory comments on the new account-accounting check and the dropped_account test guest, and fix imprecise wording (the account is dropped from both pre_state and post_states together, not just pre_state). Also applies cargo +nightly fmt's wrapping of the new error message. --- lee/state_machine/src/error.rs | 4 +++- lee/state_machine/src/validated_state_diff/mod.rs | 6 +----- .../test_methods/guest/src/bin/dropped_account.rs | 8 ++------ 3 files changed, 6 insertions(+), 12 deletions(-) diff --git a/lee/state_machine/src/error.rs b/lee/state_machine/src/error.rs index 2d0221bc..94e66e8b 100644 --- a/lee/state_machine/src/error.rs +++ b/lee/state_machine/src/error.rs @@ -132,7 +132,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")] + #[error( + "Account {account_id} was declared in the transaction but is missing from the program output" + )] DeclaredAccountMissingFromOutput { account_id: AccountId }, } diff --git a/lee/state_machine/src/validated_state_diff/mod.rs b/lee/state_machine/src/validated_state_diff/mod.rs index a740e1ef..e4a92d63 100644 --- a/lee/state_machine/src/validated_state_diff/mod.rs +++ b/lee/state_machine/src/validated_state_diff/mod.rs @@ -294,11 +294,7 @@ 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. + // diff. for account_id in &message.account_ids { ensure!( state_diff.contains_key(account_id), 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 index e79dbc3b..348eefa8 100644 --- a/lee/state_machine/test_methods/guest/src/bin/dropped_account.rs +++ b/lee/state_machine/test_methods/guest/src/bin/dropped_account.rs @@ -5,12 +5,8 @@ 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. +/// Differs from `missing_output` because the `pre_state` and `post_states` lengths match. We +/// simply drop the account from both before returning them as part of the program's output. fn main() { let ( ProgramInput {