Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion artifacts/ata-idl.json
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,12 @@
"init": false
}
],
"args": []
"args": [
{
"name": "token_program_id",
"type": "program_id"
}
]
},
{
"name": "transfer",
Expand All @@ -49,6 +54,10 @@
}
],
"args": [
{
"name": "token_program_id",
"type": "program_id"
},
{
"name": "amount",
"type": "u128"
Expand Down Expand Up @@ -78,6 +87,10 @@
}
],
"args": [
{
"name": "token_program_id",
"type": "program_id"
},
{
"name": "amount",
"type": "u128"
Expand Down
46 changes: 32 additions & 14 deletions ata/core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,17 @@ 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):
/// - Owner account
/// - 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.
Expand All @@ -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.
Expand All @@ -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()
Expand All @@ -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,
Expand Down
16 changes: 14 additions & 2 deletions ata/methods/guest/src/bin/ata.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand All @@ -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 {
Comment thread
3esmit marked this conversation as resolved.
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]
Expand All @@ -38,6 +44,7 @@ mod ata {
owner: AccountWithMetadata,
sender_ata: AccountWithMetadata,
recipient: AccountWithMetadata,
token_program_id: ProgramId,
amount: u128,
) -> SpelResult {
Comment thread
3esmit marked this conversation as resolved.
let (post_states, chained_calls) =
Expand All @@ -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) =
Expand All @@ -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))
Expand Down
23 changes: 20 additions & 3 deletions ata/src/burn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<AccountPostState>, Vec<ChainedCall>) {
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()),
Expand Down
21 changes: 20 additions & 1 deletion ata/src/create.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<AccountPostState>, Vec<ChainedCall>) {
// 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),
Expand Down
Loading
Loading