test: verify malicious self_program_id is rejected in public execution

nssa/src/program.rs

@@ -338,6 +338,13 @@ mod tests {
             Self::new(FLASH_SWAP_CALLBACK_ELF.to_vec())
                 .expect("flash_swap_callback must be a valid Risc0 program")
         }
+
+        #[must_use]
+        pub fn malicious_self_program_id() -> Self {
+            use test_program_methods::MALICIOUS_SELF_PROGRAM_ID_ELF;
+            Self::new(MALICIOUS_SELF_PROGRAM_ID_ELF.to_vec())
+                .expect("malicious_self_program_id must be a valid Risc0 program")
+        }
     }
 
     #[test]

nssa/src/state.rs

@@ -384,6 +384,7 @@ pub mod tests {
             self.insert_program(Program::validity_window());
             self.insert_program(Program::flash_swap_initiator());
             self.insert_program(Program::flash_swap_callback());
+            self.insert_program(Program::malicious_self_program_id());
             self
         }
 
@@ -3714,4 +3715,30 @@ pub mod tests {
             "standalone InvariantCheck should be rejected (caller_program_id is None)"
         );
     }
+
+    #[test]
+    fn malicious_self_program_id_rejected_in_public_execution() {
+        let program = Program::malicious_self_program_id();
+        let acc_id = AccountId::from(program.id());
+        let account = Account::default();
+
+        let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs();
+        state.force_insert_account(acc_id, account);
+
+        let message = public_transaction::Message::try_new(
+            program.id(),
+            vec![acc_id],
+            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!(
+            result.is_err(),
+            "program with wrong self_program_id in output should be rejected"
+        );
+    }
 }

test_program_methods/guest/src/bin/malicious_self_program_id.rs

@@ -0,0 +1,31 @@
+use nssa_core::program::{
+    AccountPostState, DEFAULT_PROGRAM_ID, ProgramInput, ProgramOutput, read_nssa_inputs,
+};
+
+type Instruction = ();
+
+fn main() {
+    let (
+        ProgramInput {
+            self_program_id: _, // ignore the correct ID
+            caller_program_id: _,
+            pre_states,
+            instruction: (),
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
+
+    let post_states = pre_states
+        .iter()
+        .map(|a| AccountPostState::new(a.account.clone()))
+        .collect();
+
+    // Deliberately output wrong self_program_id
+    ProgramOutput::new(
+        DEFAULT_PROGRAM_ID, // WRONG: should be self_program_id
+        instruction_words,
+        pre_states,
+        post_states,
+    )
+    .write();
+}