fix(token): reject foreign-owned definitions on initialize

integration_tests/tests/token.rs

@@ -36,6 +36,10 @@ impl Ids {
         token_methods::TOKEN_ID
     }
 
+    fn foreign_token_program() -> nssa_core::program::ProgramId {
+        [0xfeed_u32; 8]
+    }
+
     fn token_definition() -> AccountId {
         AccountId::from(&PublicKey::new_from_private_key(&Keys::def_key()))
     }
@@ -63,6 +67,19 @@ impl Accounts {
         }
     }
 
+    fn token_definition_foreign_owner() -> Account {
+        Account {
+            program_owner: Ids::foreign_token_program(),
+            balance: 0_u128,
+            data: Data::from(&TokenDefinition::Fungible {
+                name: String::from("Gold"),
+                total_supply: 1_000_000_u128,
+                metadata_id: None,
+            }),
+            nonce: Nonce(0),
+        }
+    }
+
     fn holder_init() -> Account {
         Account {
             program_owner: Ids::token_program(),
@@ -167,6 +184,78 @@ fn token_new_fungible_definition() {
     );
 }
 
+#[test]
+fn token_initialize_account_succeeds_for_canonical_definition() {
+    let mut state = state_for_token_tests_without_recipient();
+
+    let instruction = token_core::Instruction::InitializeAccount;
+
+    let message = public_transaction::Message::try_new(
+        Ids::token_program(),
+        vec![Ids::token_definition(), Ids::recipient()],
+        vec![Nonce(0)],
+        instruction,
+    )
+    .unwrap();
+
+    let witness_set =
+        public_transaction::WitnessSet::for_message(&message, &[&Keys::recipient_key()]);
+
+    let tx = PublicTransaction::new(message, witness_set);
+    state.transition_from_public_transaction(&tx, 0, 0).unwrap();
+
+    assert_eq!(
+        state.get_account_by_id(Ids::token_definition()),
+        Accounts::token_definition_init()
+    );
+    assert_eq!(
+        state.get_account_by_id(Ids::recipient()),
+        Account {
+            program_owner: Ids::token_program(),
+            balance: 0_u128,
+            data: Data::from(&TokenHolding::Fungible {
+                definition_id: Ids::token_definition(),
+                balance: 0_u128,
+            }),
+            nonce: Nonce(1),
+        }
+    );
+}
+
+#[test]
+fn token_initialize_account_rejects_foreign_owned_definition() {
+    let mut state = state_for_token_tests_without_recipient();
+    state.force_insert_account(
+        Ids::token_definition(),
+        Accounts::token_definition_foreign_owner(),
+    );
+
+    let instruction = token_core::Instruction::InitializeAccount;
+
+    let message = public_transaction::Message::try_new(
+        Ids::token_program(),
+        vec![Ids::token_definition(), Ids::recipient()],
+        vec![Nonce(0)],
+        instruction,
+    )
+    .unwrap();
+
+    let witness_set =
+        public_transaction::WitnessSet::for_message(&message, &[&Keys::recipient_key()]);
+
+    let tx = PublicTransaction::new(message, witness_set);
+    assert!(state.transition_from_public_transaction(&tx, 0, 0).is_err());
+
+    assert_eq!(
+        state.get_account_by_id(Ids::token_definition()),
+        Accounts::token_definition_foreign_owner()
+    );
+    assert_eq!(
+        state.get_account_by_id(Ids::recipient()),
+        Account::default()
+    );
+}
+
 #[test]
 fn token_transfer() {
     let mut state = state_for_token_tests();
@@ -395,6 +484,44 @@ fn token_mint() {
     );
 }
 
+#[test]
+fn token_mint_rejects_foreign_owned_definition() {
+    let mut state = state_for_token_tests_without_recipient();
+    state.force_insert_account(
+        Ids::token_definition(),
+        Accounts::token_definition_foreign_owner(),
+    );
+
+    let instruction = token_core::Instruction::Mint {
+        amount_to_mint: 500_000_u128,
+    };
+
+    let message = public_transaction::Message::try_new(
+        Ids::token_program(),
+        vec![Ids::token_definition(), Ids::recipient()],
+        vec![Nonce(0), Nonce(0)],
+        instruction,
+    )
+    .unwrap();
+
+    let witness_set = public_transaction::WitnessSet::for_message(
+        &message,
+        &[&Keys::def_key(), &Keys::recipient_key()],
+    );
+
+    let tx = PublicTransaction::new(message, witness_set);
+    assert!(state.transition_from_public_transaction(&tx, 0, 0).is_err());
+
+    assert_eq!(
+        state.get_account_by_id(Ids::token_definition()),
+        Accounts::token_definition_foreign_owner()
+    );
+    assert_eq!(
+        state.get_account_by_id(Ids::recipient()),
+        Account::default()
+    );
+}
+
 #[test]
 fn token_mint_fresh_public_recipient_requires_authorization() {
     let mut state = state_for_token_tests_without_recipient();

token/methods/guest/src/bin/token.rs

@@ -1,17 +1,134 @@
-#![no_main]
+#![cfg_attr(not(test), no_main)]
 
+use nssa_core::{
+    account::AccountWithMetadata,
+    program::{AccountPostState, ProgramId},
+};
+#[cfg(not(test))]
+use nssa_core::program::{ProgramInput, ProgramOutput};
+#[cfg(test)]
 use spel_framework::prelude::*;
-use nssa_core::account::AccountWithMetadata;
+use token_core::Instruction as TokenInstruction;
 
+#[cfg(not(test))]
 risc0_zkvm::guest::entry!(main);
 
+#[cfg(not(test))]
+fn main() {
+    let (
+        ProgramInput {
+            self_program_id,
+            caller_program_id,
+            pre_states,
+            instruction,
+        },
+        instruction_words,
+    ) = nssa_core::program::read_nssa_inputs::<TokenInstruction>();
+    let pre_states_clone = pre_states.clone();
+
+    let post_states = dispatch(instruction, pre_states, self_program_id);
+
+    ProgramOutput::new(
+        self_program_id,
+        caller_program_id,
+        instruction_words,
+        pre_states_clone,
+        post_states,
+    )
+    .write();
+}
+
+#[cfg_attr(test, allow(dead_code))]
+fn dispatch(
+    instruction: TokenInstruction,
+    pre_states: Vec<AccountWithMetadata>,
+    self_program_id: ProgramId,
+) -> Vec<AccountPostState> {
+    match instruction {
+        TokenInstruction::Transfer { amount_to_transfer } => {
+            let [sender, recipient] = expect_accounts(pre_states);
+            token_program::transfer::transfer(sender, recipient, amount_to_transfer)
+        }
+        TokenInstruction::NewFungibleDefinition { name, total_supply } => {
+            let [definition_target_account, holding_target_account] =
+                expect_accounts(pre_states);
+            token_program::new_definition::new_fungible_definition(
+                definition_target_account,
+                holding_target_account,
+                name,
+                total_supply,
+            )
+        }
+        TokenInstruction::NewDefinitionWithMetadata {
+            new_definition,
+            metadata,
+        } => {
+            let [
+                definition_target_account,
+                holding_target_account,
+                metadata_target_account,
+            ] = expect_accounts(pre_states);
+            token_program::new_definition::new_definition_with_metadata(
+                definition_target_account,
+                holding_target_account,
+                metadata_target_account,
+                new_definition,
+                *metadata,
+            )
+        }
+        TokenInstruction::InitializeAccount => {
+            let [definition_account, account_to_initialize] = expect_accounts(pre_states);
+            token_program::initialize::initialize_account(
+                definition_account,
+                account_to_initialize,
+                self_program_id,
+            )
+        }
+        TokenInstruction::Burn { amount_to_burn } => {
+            let [definition_account, user_holding_account] = expect_accounts(pre_states);
+            token_program::burn::burn(
+                definition_account,
+                user_holding_account,
+                amount_to_burn,
+            )
+        }
+        TokenInstruction::Mint { amount_to_mint } => {
+            let [definition_account, user_holding_account] = expect_accounts(pre_states);
+            token_program::mint::mint(
+                definition_account,
+                user_holding_account,
+                amount_to_mint,
+                self_program_id,
+            )
+        }
+        TokenInstruction::PrintNft => {
+            let [master_account, printed_account] = expect_accounts(pre_states);
+            token_program::print_nft::print_nft(master_account, printed_account)
+        }
+    }
+}
+
+#[cfg_attr(test, allow(dead_code))]
+fn expect_accounts<const N: usize>(
+    pre_states: Vec<AccountWithMetadata>,
+) -> [AccountWithMetadata; N] {
+    let actual = pre_states.len();
+    pre_states.try_into().unwrap_or_else(|_| {
+        panic!("Account count mismatch: expected {N}, got {actual}");
+    })
+}
+
+// Kept for `idl-gen`, which parses source annotations directly.
+#[cfg(test)]
+#[allow(dead_code)]
 #[lez_program(instruction = "token_core::Instruction")]
 mod token {
     #[allow(unused_imports)]
     use super::*;
 
     /// Transfer tokens from sender to recipient.
     /// Fresh public recipients must be explicitly authorized in the same transaction.
+    #[allow(dead_code, clippy::boxed_local)]
     #[instruction]
     pub fn transfer(
         sender: AccountWithMetadata,
@@ -27,6 +144,7 @@ mod token {
 
     /// Create a new fungible token definition without metadata.
     /// Definition and holding targets must be uninitialized and authorized.
+    #[allow(dead_code, clippy::boxed_local)]
     #[instruction]
     pub fn new_fungible_definition(
         definition_target_account: AccountWithMetadata,
@@ -46,6 +164,7 @@ mod token {
 
     /// Create a new fungible or non-fungible token definition with metadata.
     /// Definition, holding, and metadata targets must be uninitialized and authorized.
+    #[allow(dead_code, clippy::boxed_local)]
     #[instruction]
     pub fn new_definition_with_metadata(
         definition_target_account: AccountWithMetadata,
@@ -67,6 +186,7 @@ mod token {
 
     /// Initialize a token holding account for a given token definition.
     /// The holding target must be uninitialized and authorized.
+    #[allow(dead_code, clippy::boxed_local)]
     #[instruction]
     pub fn initialize_account(
         definition_account: AccountWithMetadata,
@@ -76,11 +196,13 @@ mod token {
             token_program::initialize::initialize_account(
                 definition_account,
                 account_to_initialize,
+                nssa_core::program::DEFAULT_PROGRAM_ID,
             ),
         ))
     }
 
     /// Burn tokens from the holder's account.
+    #[allow(dead_code, clippy::boxed_local)]
     #[instruction]
     pub fn burn(
         definition_account: AccountWithMetadata,
@@ -96,6 +218,7 @@ mod token {
 
     /// Mint new tokens to the holder's account.
     /// Fresh public holders must be explicitly authorized in the same transaction.
+    #[allow(dead_code, clippy::boxed_local)]
     #[instruction]
     pub fn mint(
         definition_account: AccountWithMetadata,
@@ -106,11 +229,13 @@ mod token {
             definition_account,
             user_holding_account,
             amount_to_mint,
+            nssa_core::program::DEFAULT_PROGRAM_ID,
         )))
     }
 
     /// Print a new NFT from the master copy.
     /// The printed copy target must be uninitialized and authorized.
+    #[allow(dead_code, clippy::boxed_local)]
     #[instruction]
     pub fn print_nft(
         master_account: AccountWithMetadata,

token/src/initialize.rs

@@ -1,12 +1,13 @@
 use nssa_core::{
     account::{Account, AccountWithMetadata, Data},
-    program::{AccountPostState, Claim},
+    program::{AccountPostState, Claim, ProgramId},
 };
 use token_core::{TokenDefinition, TokenHolding};
 
 pub fn initialize_account(
     definition_account: AccountWithMetadata,
     account_to_initialize: AccountWithMetadata,
+    token_program_id: ProgramId,
 ) -> Vec<AccountPostState> {
     assert_eq!(
         account_to_initialize.account,
@@ -17,11 +18,11 @@ pub fn initialize_account(
         account_to_initialize.is_authorized,
         "Account to initialize must be authorized"
     );
+    assert_eq!(
+        definition_account.account.program_owner, token_program_id,
+        "Token definition must be owned by token program"
+    );
 
-    // TODO: #212 We should check that this is an account owned by the token program.
-    // This check can't be done here since the ID of the program is known only after compiling it
-    //
-    // Check definition account is valid
     let definition = TokenDefinition::try_from(&definition_account.account.data)
         .expect("Definition account must be valid");
     let holding =