Merge 09af47b862830f567c8e844e08d36374ed920cb1 into 041cf68cd63acd9c4c2d57492a0a0590396c27de

lee/state_machine/src/error.rs

@@ -131,6 +131,11 @@ 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)]

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(

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());

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 = [(

lee/state_machine/src/validated_state_diff/mod.rs

@@ -293,6 +293,17 @@ impl ValidatedStateDiff {
             );
         }
 
+        // Every account the caller declared as part of the transaction must appear in the final
+        // diff.
+        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,

lee/state_machine/test_methods/guest/src/bin/dropped_account.rs

@@ -0,0 +1,35 @@
+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.
+///
+/// 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 {
+            self_program_id,
+            caller_program_id,
+            pre_states,
+            ..
+        },
+        instruction_words,
+    ) = read_lee_inputs::<Instruction>();
+
+    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();
+}