diff --git a/artifacts/ata-idl.json b/artifacts/ata-idl.json index 53818d0..11e99cb 100644 --- a/artifacts/ata-idl.json +++ b/artifacts/ata-idl.json @@ -24,7 +24,12 @@ "init": false } ], - "args": [] + "args": [ + { + "name": "token_program_id", + "type": "program_id" + } + ] }, { "name": "transfer", @@ -49,6 +54,10 @@ } ], "args": [ + { + "name": "token_program_id", + "type": "program_id" + }, { "name": "amount", "type": "u128" @@ -78,6 +87,10 @@ } ], "args": [ + { + "name": "token_program_id", + "type": "program_id" + }, { "name": "amount", "type": "u128" diff --git a/ata/core/src/lib.rs b/ata/core/src/lib.rs index 0256c57..f4a8ac0 100644 --- a/ata/core/src/lib.rs +++ b/ata/core/src/lib.rs @@ -7,7 +7,7 @@ use serde::{Deserialize, Serialize}; #[derive(Serialize, Deserialize)] pub enum Instruction { - /// Create the Associated Token Account for (owner, definition). + /// Create the Associated Token Account for (token program, owner, definition). /// Idempotent: no-op if the account already exists. /// /// Required accounts (3): @@ -15,8 +15,9 @@ pub enum Instruction { /// - Token definition account /// - Associated token account (default/uninitialized, or already initialized) /// - /// `token_program_id` is derived from `token_definition.account.program_owner`. - Create, + /// `token_program_id` is explicit so callers can support multiple token programs without + /// letting account metadata choose downstream code. + Create { token_program_id: ProgramId }, /// Transfer tokens FROM owner's ATA to a recipient token holding account. /// Uses ATA PDA seeds to authorize the chained Token::Transfer call. @@ -29,8 +30,12 @@ pub enum Instruction { /// - owned by the same token program as the sender ATA, /// - and point at the same token definition as the sender. /// - /// `token_program_id` is derived from `sender_ata.account.program_owner`. - Transfer { amount: u128 }, + /// `token_program_id` is explicit so callers can support multiple token programs without + /// letting account metadata choose downstream code. + Transfer { + token_program_id: ProgramId, + amount: u128, + }, /// Burn tokens FROM owner's ATA. /// Uses PDA seeds to authorize the ATA in the chained Token::Burn call. @@ -40,15 +45,27 @@ pub enum Instruction { /// - Owner's ATA (the holding to burn from) /// - Token definition account /// - /// `token_program_id` is derived from `holder_ata.account.program_owner`. - Burn { amount: u128 }, + /// `token_program_id` is explicit so callers can support multiple token programs without + /// letting account metadata choose downstream code. + Burn { + token_program_id: ProgramId, + amount: u128, + }, } -pub fn compute_ata_seed(owner_id: AccountId, definition_id: AccountId) -> PdaSeed { +pub fn compute_ata_seed( + token_program_id: ProgramId, + owner_id: AccountId, + definition_id: AccountId, +) -> PdaSeed { use risc0_zkvm::sha::{Impl, Sha256}; - let mut bytes = [0u8; 64]; - bytes[0..32].copy_from_slice(&owner_id.to_bytes()); - bytes[32..64].copy_from_slice(&definition_id.to_bytes()); + let mut bytes = [0u8; 96]; + for (index, word) in token_program_id.iter().enumerate() { + let offset = index * 4; + bytes[offset..offset + 4].copy_from_slice(&word.to_le_bytes()); + } + bytes[32..64].copy_from_slice(&owner_id.to_bytes()); + bytes[64..96].copy_from_slice(&definition_id.to_bytes()); PdaSeed::new( Impl::hash_bytes(&bytes) .as_bytes() @@ -61,15 +78,16 @@ pub fn get_associated_token_account_id(ata_program_id: &ProgramId, seed: &PdaSee AccountId::for_public_pda(ata_program_id, seed) } -/// Verify the ATA's address matches `(ata_program_id, owner, definition)` and return -/// the [`PdaSeed`] for use in chained calls. +/// Verify the ATA's address matches `(ata_program_id, token_program_id, owner, definition)` and +/// return the [`PdaSeed`] for use in chained calls. pub fn verify_ata_and_get_seed( ata_account: &AccountWithMetadata, owner: &AccountWithMetadata, + token_program_id: ProgramId, definition_id: AccountId, ata_program_id: ProgramId, ) -> PdaSeed { - let seed = compute_ata_seed(owner.account_id, definition_id); + let seed = compute_ata_seed(token_program_id, owner.account_id, definition_id); let expected_id = get_associated_token_account_id(&ata_program_id, &seed); assert_eq!( ata_account.account_id, expected_id, diff --git a/ata/methods/guest/src/bin/ata.rs b/ata/methods/guest/src/bin/ata.rs index 46ccdb8..7d892b2 100644 --- a/ata/methods/guest/src/bin/ata.rs +++ b/ata/methods/guest/src/bin/ata.rs @@ -2,7 +2,7 @@ use spel_framework::prelude::*; use spel_framework::context::ProgramContext; -use nssa_core::account::AccountWithMetadata; +use nssa_core::{account::AccountWithMetadata, program::ProgramId}; risc0_zkvm::guest::entry!(main); @@ -11,25 +11,31 @@ mod ata { #[allow(unused_imports)] use super::*; - /// Create the Associated Token Account for (owner, definition). + /// Create the Associated Token Account for (token program, owner, definition). /// Idempotent: no-op if the account already exists. + /// The token program is selected explicitly by `token_program_id`; the token definition and + /// any existing ATA occupant must be owned by that program. #[instruction] pub fn create( ctx: ProgramContext, owner: AccountWithMetadata, token_definition: AccountWithMetadata, ata_account: AccountWithMetadata, + token_program_id: ProgramId, ) -> SpelResult { let (post_states, chained_calls) = ata_program::create::create_associated_token_account( owner, token_definition, ata_account, ctx.self_program_id, + token_program_id, ); Ok(spel_framework::SpelOutput::execute(post_states, chained_calls)) } /// Transfer tokens FROM owner's ATA to a recipient token holding account. + /// The token program is selected explicitly by `token_program_id`; the sender ATA and recipient + /// holding must be owned by that program. /// The recipient holding must already be initialized, be owned by the same token program /// as the sender ATA, and point at the same token definition as the sender. #[instruction] @@ -38,6 +44,7 @@ mod ata { owner: AccountWithMetadata, sender_ata: AccountWithMetadata, recipient: AccountWithMetadata, + token_program_id: ProgramId, amount: u128, ) -> SpelResult { let (post_states, chained_calls) = @@ -46,18 +53,22 @@ mod ata { sender_ata, recipient, ctx.self_program_id, + token_program_id, amount, ); Ok(spel_framework::SpelOutput::execute(post_states, chained_calls)) } /// Burn tokens FROM owner's ATA. + /// The token program is selected explicitly by `token_program_id`; the holder ATA and token + /// definition must be owned by that program. #[instruction] pub fn burn( ctx: ProgramContext, owner: AccountWithMetadata, holder_ata: AccountWithMetadata, token_definition: AccountWithMetadata, + token_program_id: ProgramId, amount: u128, ) -> SpelResult { let (post_states, chained_calls) = @@ -66,6 +77,7 @@ mod ata { holder_ata, token_definition, ctx.self_program_id, + token_program_id, amount, ); Ok(spel_framework::SpelOutput::execute(post_states, chained_calls)) diff --git a/ata/src/burn.rs b/ata/src/burn.rs index 4940fde..3386e6d 100644 --- a/ata/src/burn.rs +++ b/ata/src/burn.rs @@ -9,15 +9,32 @@ pub fn burn_from_associated_token_account( holder_ata: AccountWithMetadata, token_definition: AccountWithMetadata, ata_program_id: ProgramId, + token_program_id: ProgramId, amount: u128, ) -> (Vec, Vec) { - let token_program_id = holder_ata.account.program_owner; assert!(owner.is_authorized, "Owner authorization is missing"); + assert_eq!( + holder_ata.account.program_owner, token_program_id, + "Holder ATA must be owned by expected token program" + ); + assert_eq!( + token_definition.account.program_owner, token_program_id, + "Token definition must be owned by expected token program" + ); let definition_id = TokenHolding::try_from(&holder_ata.account.data) .expect("Holder ATA must hold a valid token") .definition_id(); - let seed = - ata_core::verify_ata_and_get_seed(&holder_ata, &owner, definition_id, ata_program_id); + assert_eq!( + definition_id, token_definition.account_id, + "Holder ATA token definition does not match" + ); + let seed = ata_core::verify_ata_and_get_seed( + &holder_ata, + &owner, + token_program_id, + definition_id, + ata_program_id, + ); let post_states = vec![ AccountPostState::new(owner.account.clone()), diff --git a/ata/src/create.rs b/ata/src/create.rs index 060f0c7..6239b11 100644 --- a/ata/src/create.rs +++ b/ata/src/create.rs @@ -2,27 +2,46 @@ use nssa_core::{ account::{Account, AccountWithMetadata}, program::{AccountPostState, ChainedCall, Claim, ProgramId}, }; +use token_core::{TokenDefinition, TokenHolding}; pub fn create_associated_token_account( owner: AccountWithMetadata, token_definition: AccountWithMetadata, ata_account: AccountWithMetadata, ata_program_id: ProgramId, + token_program_id: ProgramId, ) -> (Vec, Vec) { // No explicit owner authorization check is needed here: ATA creation is idempotent, so the // call itself may proceed without `owner.is_authorized`. If the owner account is still // default, the returned post-state will still carry `Claim::Authorized` so the runtime can // claim that owner account when needed. - let token_program_id = token_definition.account.program_owner; + assert_eq!( + token_definition.account.program_owner, token_program_id, + "Token definition must be owned by expected token program" + ); + let _definition = TokenDefinition::try_from(&token_definition.account.data) + .expect("Token definition must be valid"); let seed = ata_core::verify_ata_and_get_seed( &ata_account, &owner, + token_program_id, token_definition.account_id, ata_program_id, ); // Idempotent: already initialized → no-op if ata_account.account != Account::default() { + assert_eq!( + ata_account.account.program_owner, token_program_id, + "Existing ATA must be owned by expected token program" + ); + let holding = TokenHolding::try_from(&ata_account.account.data) + .expect("Existing ATA must hold a valid token"); + assert_eq!( + holding.definition_id(), + token_definition.account_id, + "Existing ATA token definition does not match" + ); return ( vec![ AccountPostState::new_claimed_if_default(owner.account.clone(), Claim::Authorized), diff --git a/ata/src/tests.rs b/ata/src/tests.rs index 8ee4df9..595cfdd 100644 --- a/ata/src/tests.rs +++ b/ata/src/tests.rs @@ -7,6 +7,7 @@ use token_core::{TokenDefinition, TokenHolding}; const ATA_PROGRAM_ID: nssa_core::program::ProgramId = [1u32; 8]; const TOKEN_PROGRAM_ID: nssa_core::program::ProgramId = [2u32; 8]; +const OTHER_TOKEN_PROGRAM_ID: nssa_core::program::ProgramId = [3u32; 8]; fn owner_id() -> AccountId { AccountId::new([0x01u8; 32]) @@ -19,7 +20,7 @@ fn definition_id() -> AccountId { fn ata_id() -> AccountId { get_associated_token_account_id( &ATA_PROGRAM_ID, - &compute_ata_seed(owner_id(), definition_id()), + &compute_ata_seed(TOKEN_PROGRAM_ID, owner_id(), definition_id()), ) } @@ -79,6 +80,7 @@ fn create_emits_chained_call_for_uninitialized_ata() { definition_account(), uninitialized_ata_account(), ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, ); assert_eq!(post_states.len(), 3); @@ -91,7 +93,11 @@ fn create_emits_chained_call_for_uninitialized_ata() { vec![definition_account(), authorized_ata], &token_core::Instruction::InitializeAccount, ) - .with_pda_seeds(vec![compute_ata_seed(owner_id(), definition_id())]); + .with_pda_seeds(vec![compute_ata_seed( + TOKEN_PROGRAM_ID, + owner_id(), + definition_id(), + )]); assert_eq!(chained_calls, vec![expected_call]); } @@ -103,6 +109,7 @@ fn create_is_idempotent_for_initialized_ata() { definition_account(), initialized_ata_account(), ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, ); assert_eq!(post_states.len(), 3); @@ -126,27 +133,41 @@ fn create_panics_on_wrong_ata_address() { definition_account(), wrong_ata, ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, ); } #[test] fn get_associated_token_account_id_is_deterministic() { - let seed = compute_ata_seed(owner_id(), definition_id()); + let seed = compute_ata_seed(TOKEN_PROGRAM_ID, owner_id(), definition_id()); let id1 = get_associated_token_account_id(&ATA_PROGRAM_ID, &seed); let id2 = get_associated_token_account_id(&ATA_PROGRAM_ID, &seed); assert_eq!(id1, id2); } +#[test] +fn get_associated_token_account_id_differs_by_token_program() { + let id1 = get_associated_token_account_id( + &ATA_PROGRAM_ID, + &compute_ata_seed(TOKEN_PROGRAM_ID, owner_id(), definition_id()), + ); + let id2 = get_associated_token_account_id( + &ATA_PROGRAM_ID, + &compute_ata_seed(OTHER_TOKEN_PROGRAM_ID, owner_id(), definition_id()), + ); + assert_ne!(id1, id2); +} + #[test] fn get_associated_token_account_id_differs_by_owner() { let other_owner = AccountId::new([0x99u8; 32]); let id1 = get_associated_token_account_id( &ATA_PROGRAM_ID, - &compute_ata_seed(owner_id(), definition_id()), + &compute_ata_seed(TOKEN_PROGRAM_ID, owner_id(), definition_id()), ); let id2 = get_associated_token_account_id( &ATA_PROGRAM_ID, - &compute_ata_seed(other_owner, definition_id()), + &compute_ata_seed(TOKEN_PROGRAM_ID, other_owner, definition_id()), ); assert_ne!(id1, id2); } @@ -156,13 +177,63 @@ fn get_associated_token_account_id_differs_by_definition() { let other_def = AccountId::new([0x99u8; 32]); let id1 = get_associated_token_account_id( &ATA_PROGRAM_ID, - &compute_ata_seed(owner_id(), definition_id()), + &compute_ata_seed(TOKEN_PROGRAM_ID, owner_id(), definition_id()), + ); + let id2 = get_associated_token_account_id( + &ATA_PROGRAM_ID, + &compute_ata_seed(TOKEN_PROGRAM_ID, owner_id(), other_def), ); - let id2 = - get_associated_token_account_id(&ATA_PROGRAM_ID, &compute_ata_seed(owner_id(), other_def)); assert_ne!(id1, id2); } +#[test] +#[should_panic(expected = "Token definition must be owned by expected token program")] +fn create_panics_when_definition_is_owned_by_unexpected_token_program() { + let mut definition = definition_account(); + definition.account.program_owner = OTHER_TOKEN_PROGRAM_ID; + + crate::create::create_associated_token_account( + owner_account(), + definition, + uninitialized_ata_account(), + ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, + ); +} + +#[test] +#[should_panic(expected = "Existing ATA must be owned by expected token program")] +fn create_panics_when_existing_ata_is_owned_by_unexpected_token_program() { + let mut ata = initialized_ata_account(); + ata.account.program_owner = OTHER_TOKEN_PROGRAM_ID; + + crate::create::create_associated_token_account( + owner_account(), + definition_account(), + ata, + ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, + ); +} + +#[test] +#[should_panic(expected = "Existing ATA token definition does not match")] +fn create_panics_when_existing_ata_definition_mismatches_requested_definition() { + let mut ata = initialized_ata_account(); + ata.account.data = Data::from(&TokenHolding::Fungible { + definition_id: AccountId::new([0xAAu8; 32]), + balance: 100, + }); + + crate::create::create_associated_token_account( + owner_account(), + definition_account(), + ata, + ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, + ); +} + fn recipient_id() -> AccountId { AccountId::new([0x03u8; 32]) } @@ -190,6 +261,7 @@ fn transfer_emits_chained_call_for_initialized_recipient() { initialized_ata_account(), initialized_recipient_account(), ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, 25, ); @@ -205,7 +277,11 @@ fn transfer_emits_chained_call_for_initialized_recipient() { amount_to_transfer: 25, }, ) - .with_pda_seeds(vec![compute_ata_seed(owner_id(), definition_id())]); + .with_pda_seeds(vec![compute_ata_seed( + TOKEN_PROGRAM_ID, + owner_id(), + definition_id(), + )]); assert_eq!(chained_calls, vec![expected_call]); } @@ -221,6 +297,7 @@ fn transfer_panics_when_owner_not_authorized() { initialized_ata_account(), initialized_recipient_account(), ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, 1, ); } @@ -239,6 +316,23 @@ fn transfer_panics_when_recipient_is_default() { initialized_ata_account(), default_recipient, ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, + 1, + ); +} + +#[test] +#[should_panic(expected = "Sender ATA must be owned by expected token program")] +fn transfer_panics_when_sender_ata_is_owned_by_unexpected_token_program() { + let mut sender = initialized_ata_account(); + sender.account.program_owner = OTHER_TOKEN_PROGRAM_ID; + + crate::transfer::transfer_from_associated_token_account( + owner_account(), + sender, + initialized_recipient_account(), + ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, 1, ); } @@ -254,6 +348,7 @@ fn transfer_panics_when_recipient_is_foreign_owned() { initialized_ata_account(), foreign_recipient, ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, 1, ); } @@ -269,6 +364,7 @@ fn transfer_panics_when_recipient_data_is_malformed() { initialized_ata_account(), malformed_recipient, ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, 1, ); } @@ -287,6 +383,85 @@ fn transfer_panics_when_recipient_definition_mismatches_sender() { initialized_ata_account(), mismatched_recipient, ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, + 1, + ); +} + +#[test] +fn burn_emits_chained_call_for_initialized_ata() { + let (post_states, chained_calls) = crate::burn::burn_from_associated_token_account( + owner_account(), + initialized_ata_account(), + definition_account(), + ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, + 25, + ); + + assert_eq!(post_states.len(), 3); + assert_eq!(chained_calls.len(), 1); + + let mut holder_auth = initialized_ata_account(); + holder_auth.is_authorized = true; + let expected_call = ChainedCall::new( + TOKEN_PROGRAM_ID, + vec![definition_account(), holder_auth], + &token_core::Instruction::Burn { amount_to_burn: 25 }, + ) + .with_pda_seeds(vec![compute_ata_seed( + TOKEN_PROGRAM_ID, + owner_id(), + definition_id(), + )]); + + assert_eq!(chained_calls, vec![expected_call]); +} + +#[test] +#[should_panic(expected = "Holder ATA must be owned by expected token program")] +fn burn_panics_when_holder_ata_is_owned_by_unexpected_token_program() { + let mut holder = initialized_ata_account(); + holder.account.program_owner = OTHER_TOKEN_PROGRAM_ID; + + crate::burn::burn_from_associated_token_account( + owner_account(), + holder, + definition_account(), + ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, + 1, + ); +} + +#[test] +#[should_panic(expected = "Token definition must be owned by expected token program")] +fn burn_panics_when_definition_is_owned_by_unexpected_token_program() { + let mut definition = definition_account(); + definition.account.program_owner = OTHER_TOKEN_PROGRAM_ID; + + crate::burn::burn_from_associated_token_account( + owner_account(), + initialized_ata_account(), + definition, + ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, + 1, + ); +} + +#[test] +#[should_panic(expected = "Holder ATA token definition does not match")] +fn burn_panics_when_holder_definition_mismatches_supplied_definition() { + let mut definition = definition_account(); + definition.account_id = AccountId::new([0xBBu8; 32]); + + crate::burn::burn_from_associated_token_account( + owner_account(), + initialized_ata_account(), + definition, + ATA_PROGRAM_ID, + TOKEN_PROGRAM_ID, 1, ); } diff --git a/ata/src/transfer.rs b/ata/src/transfer.rs index 3fa1668..c816f26 100644 --- a/ata/src/transfer.rs +++ b/ata/src/transfer.rs @@ -9,16 +9,21 @@ pub fn transfer_from_associated_token_account( sender_ata: AccountWithMetadata, recipient: AccountWithMetadata, ata_program_id: ProgramId, + token_program_id: ProgramId, amount: u128, ) -> (Vec, Vec) { - let token_program_id = sender_ata.account.program_owner; assert!(owner.is_authorized, "Owner authorization is missing"); + assert_eq!( + sender_ata.account.program_owner, token_program_id, + "Sender ATA must be owned by expected token program" + ); let sender_definition_id = TokenHolding::try_from(&sender_ata.account.data) .expect("Sender ATA must hold a valid token") .definition_id(); let sender_seed = ata_core::verify_ata_and_get_seed( &sender_ata, &owner, + token_program_id, sender_definition_id, ata_program_id, ); diff --git a/integration_tests/tests/ata.rs b/integration_tests/tests/ata.rs index df13983..d88bd3b 100644 --- a/integration_tests/tests/ata.rs +++ b/integration_tests/tests/ata.rs @@ -58,12 +58,20 @@ impl Ids { } fn owner_ata() -> AccountId { - let seed = compute_ata_seed(Self::owner(), Self::token_definition()); + let seed = compute_ata_seed( + Self::token_program(), + Self::owner(), + Self::token_definition(), + ); get_associated_token_account_id(&Self::ata_program(), &seed) } fn recipient_ata() -> AccountId { - let seed = compute_ata_seed(Self::recipient(), Self::token_definition()); + let seed = compute_ata_seed( + Self::token_program(), + Self::recipient(), + Self::token_definition(), + ); get_associated_token_account_id(&Self::ata_program(), &seed) } } @@ -105,6 +113,19 @@ impl Accounts { nonce: Nonce(0), } } + + fn foreign_owned_token_definition() -> Account { + Account { + program_owner: [99; 8], + balance: 0_u128, + data: Data::from(&TokenDefinition::Fungible { + name: String::from("Foreign Gold"), + total_supply: 1_000_000_u128, + metadata_id: None, + }), + nonce: Nonce(0), + } + } } fn deploy_programs(state: &mut V03State) { @@ -144,7 +165,9 @@ fn ata_create() { deploy_programs(&mut state); state.force_insert_account(Ids::token_definition(), Accounts::token_definition_init()); - let instruction = ata_core::Instruction::Create; + let instruction = ata_core::Instruction::Create { + token_program_id: Ids::token_program(), + }; let message = public_transaction::Message::try_new( Ids::ata_program(), @@ -177,7 +200,9 @@ fn ata_create() { fn ata_create_is_idempotent() { let mut state = state_for_ata_tests(); - let instruction = ata_core::Instruction::Create; + let instruction = ata_core::Instruction::Create { + token_program_id: Ids::token_program(), + }; let message = public_transaction::Message::try_new( Ids::ata_program(), @@ -207,11 +232,104 @@ fn ata_create_is_idempotent() { ); } +#[test] +fn ata_create_rejects_definition_owned_by_unexpected_token_program() { + let mut state = V03State::new_with_genesis_accounts(&[], vec![], 0); + deploy_programs(&mut state); + state.force_insert_account( + Ids::token_definition(), + Accounts::foreign_owned_token_definition(), + ); + + let instruction = ata_core::Instruction::Create { + token_program_id: Ids::token_program(), + }; + + let message = public_transaction::Message::try_new( + Ids::ata_program(), + vec![Ids::owner(), Ids::token_definition(), Ids::owner_ata()], + vec![Nonce(0)], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&Keys::owner_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::owner_ata()), + Account::default() + ); +} + +#[test] +fn ata_create_rejects_existing_ata_owned_by_unexpected_token_program() { + let mut state = V03State::new_with_genesis_accounts(&[], vec![], 0); + deploy_programs(&mut state); + state.force_insert_account(Ids::token_definition(), Accounts::token_definition_init()); + + let mut foreign_ata = Accounts::owner_ata_init(); + foreign_ata.program_owner = [99; 8]; + state.force_insert_account(Ids::owner_ata(), foreign_ata.clone()); + + let instruction = ata_core::Instruction::Create { + token_program_id: Ids::token_program(), + }; + + let message = public_transaction::Message::try_new( + Ids::ata_program(), + vec![Ids::owner(), Ids::token_definition(), Ids::owner_ata()], + vec![Nonce(0)], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&Keys::owner_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::owner_ata()), foreign_ata); +} + +#[test] +fn ata_create_rejects_existing_ata_with_mismatched_definition() { + let mut state = V03State::new_with_genesis_accounts(&[], vec![], 0); + deploy_programs(&mut state); + state.force_insert_account(Ids::token_definition(), Accounts::token_definition_init()); + + let mut mismatched_ata = Accounts::owner_ata_init(); + mismatched_ata.data = Data::from(&TokenHolding::Fungible { + definition_id: Ids::recipient(), + balance: 1_000_000_u128, + }); + state.force_insert_account(Ids::owner_ata(), mismatched_ata.clone()); + + let instruction = ata_core::Instruction::Create { + token_program_id: Ids::token_program(), + }; + + let message = public_transaction::Message::try_new( + Ids::ata_program(), + vec![Ids::owner(), Ids::token_definition(), Ids::owner_ata()], + vec![Nonce(0)], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&Keys::owner_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::owner_ata()), mismatched_ata); +} + #[test] fn ata_transfer() { let mut state = state_for_ata_tests_with_precreated_recipient_ata(); let instruction = ata_core::Instruction::Transfer { + token_program_id: Ids::token_program(), amount: 400_000_u128, }; @@ -259,7 +377,10 @@ fn ata_transfer() { fn ata_transfer_rejects_default_recipient() { let mut state = state_for_ata_tests(); - let instruction = ata_core::Instruction::Transfer { amount: 1_u128 }; + let instruction = ata_core::Instruction::Transfer { + token_program_id: Ids::token_program(), + amount: 1_u128, + }; let message = public_transaction::Message::try_new( Ids::ata_program(), @@ -303,7 +424,10 @@ fn ata_transfer_rejects_mismatched_definition_recipient() { }; state.force_insert_account(Ids::recipient_ata(), mismatched_recipient.clone()); - let instruction = ata_core::Instruction::Transfer { amount: 1_u128 }; + let instruction = ata_core::Instruction::Transfer { + token_program_id: Ids::token_program(), + amount: 1_u128, + }; let message = public_transaction::Message::try_new( Ids::ata_program(), @@ -333,6 +457,7 @@ fn ata_burn() { let mut state = state_for_ata_tests(); let instruction = ata_core::Instruction::Burn { + token_program_id: Ids::token_program(), amount: 300_000_u128, }; @@ -391,7 +516,7 @@ fn ata_create_from_private_owner() { let owner_id = AccountId::from(&owner_npk); // ATA derived from the private owner - let seed = compute_ata_seed(owner_id, Ids::token_definition()); + let seed = compute_ata_seed(Ids::token_program(), owner_id, Ids::token_definition()); let owner_ata_id = get_associated_token_account_id(&Ids::ata_program(), &seed); // Pre-states: private uninitialized owner (mask=2), public token definition (mask=0), public @@ -404,7 +529,9 @@ fn ata_create_from_private_owner() { ); let ata_pre = AccountWithMetadata::new(Account::default(), false, owner_ata_id); - let instruction = ata_core::Instruction::Create; + let instruction = ata_core::Instruction::Create { + token_program_id: Ids::token_program(), + }; let instruction_data = Program::serialize_instruction(instruction).unwrap(); // Ephemeral key for encrypting the private owner's post-state