Add malicious ATA methods and corresponding integration tests

- Introduced a new Cargo.toml for the malicious ATA guest program.
- Implemented malicious ATA methods in `malicious_ata.rs`, including create, transfer, and burn functions that simulate malicious behavior.
- Updated the integration test suite to include tests for the malicious ATA program, ensuring it is rejected in various value paths.
- Enhanced existing tests to accommodate the new malicious ATA program and verify that it does not affect the expected state of the AMM.
This commit is contained in:
Ricardo Guilherme Schmidt 2026-04-23 11:28:10 -03:00
parent 9444d72c60
commit e41d847667
No known key found for this signature in database
GPG Key ID: 1396EA17DE132FFE
10 changed files with 4576 additions and 16 deletions

10
Cargo.lock generated
View File

@ -1595,6 +1595,7 @@ dependencies = [
"amm_core",
"ata-methods",
"ata_core",
"malicious-ata-methods",
"nssa",
"nssa_core",
"token-methods",
@ -1758,6 +1759,15 @@ version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
[[package]]
name = "malicious-ata-methods"
version = "0.1.0"
dependencies = [
"ata_core",
"risc0-build",
"risc0-zkvm",
]
[[package]]
name = "malloc_buf"
version = "0.0.6"

View File

@ -10,12 +10,14 @@ members = [
"ata",
"ata/methods",
"integration_tests",
"integration_tests/malicious_ata_methods",
"tools/idl-gen",
]
exclude = [
"token/methods/guest",
"amm/methods/guest",
"ata/methods/guest",
"integration_tests/malicious_ata_methods/guest",
]
resolver = "2"

View File

@ -12,3 +12,4 @@ ata_core = { workspace = true }
token-methods = { path = "../token/methods" }
amm-methods = { path = "../amm/methods" }
ata-methods = { path = "../ata/methods" }
malicious-ata-methods = { path = "malicious_ata_methods" }

View File

@ -0,0 +1,14 @@
[package]
name = "malicious-ata-methods"
version = "0.1.0"
edition = "2021"
[build-dependencies]
risc0-build = "=3.0.5"
[dependencies]
risc0-zkvm = { version = "=3.0.5", features = ["std"] }
ata_core = { path = "../../ata/core" }
[package.metadata.risc0]
methods = ["guest"]

View File

@ -0,0 +1,3 @@
fn main() {
risc0_build::embed_methods();
}

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,18 @@
[package]
name = "malicious-ata-guest"
version = "0.1.0"
edition = "2021"
[workspace]
[[bin]]
name = "malicious_ata"
path = "src/bin/malicious_ata.rs"
[dependencies]
spel-framework = { git = "https://github.com/logos-co/spel.git", tag = "v0.2.0-rc.2", package = "spel-framework" }
nssa_core = { git = "https://github.com/logos-blockchain/logos-execution-zone.git", tag = "v0.2.0-rc1" }
risc0-zkvm = { version = "=3.0.5", default-features = false }
ata_core = { path = "../../../ata/core" }
serde = { version = "1.0", features = ["derive"] }
borsh = "1.5"

View File

@ -0,0 +1,78 @@
#![no_main]
use nssa_core::{
account::AccountWithMetadata,
program::{AccountPostState, ProgramId},
};
use spel_framework::prelude::*;
risc0_zkvm::guest::entry!(main);
fn unchanged_post_state(account: AccountWithMetadata) -> AccountPostState {
AccountPostState::new(account.account)
}
#[lez_program(instruction = "ata_core::Instruction")]
mod malicious_ata {
#[allow(unused_imports)]
use super::*;
/// Intentionally malicious test helper. It acknowledges ATA creation without creating
/// anything, proving callers must not trust a caller-supplied ATA program ID.
#[instruction]
pub fn create(
owner: AccountWithMetadata,
token_definition: AccountWithMetadata,
ata_account: AccountWithMetadata,
ata_program_id: ProgramId,
) -> SpelResult {
let _ = ata_program_id;
assert!(owner.is_authorized, "Owner authorization is missing");
Ok(SpelOutput::states_only(vec![
unchanged_post_state(owner),
unchanged_post_state(token_definition),
unchanged_post_state(ata_account),
]))
}
/// Intentionally malicious test helper. It returns success without debiting the sender
/// or crediting the recipient.
#[instruction]
pub fn transfer(
owner: AccountWithMetadata,
sender_ata: AccountWithMetadata,
recipient: AccountWithMetadata,
ata_program_id: ProgramId,
amount: u128,
) -> SpelResult {
let _ = (ata_program_id, amount);
assert!(owner.is_authorized, "Owner authorization is missing");
Ok(SpelOutput::states_only(vec![
unchanged_post_state(owner),
unchanged_post_state(sender_ata),
unchanged_post_state(recipient),
]))
}
/// Intentionally malicious test helper. It returns success without burning the LP position
/// or decreasing the LP definition supply.
#[instruction]
pub fn burn(
owner: AccountWithMetadata,
holder_ata: AccountWithMetadata,
token_definition: AccountWithMetadata,
ata_program_id: ProgramId,
amount: u128,
) -> SpelResult {
let _ = (ata_program_id, amount);
assert!(owner.is_authorized, "Owner authorization is missing");
Ok(SpelOutput::states_only(vec![
unchanged_post_state(owner),
unchanged_post_state(holder_ata),
unchanged_post_state(token_definition),
]))
}
}

View File

@ -0,0 +1 @@
include!(concat!(env!("OUT_DIR"), "/methods.rs"));

View File

@ -1,13 +1,17 @@
use amm_core::{
PoolDefinition, FEE_TIER_BPS_1, FEE_TIER_BPS_100, FEE_TIER_BPS_30, FEE_TIER_BPS_5,
MINIMUM_LIQUIDITY,
PoolDefinition, FEE_BPS_DENOMINATOR, FEE_TIER_BPS_1, FEE_TIER_BPS_100, FEE_TIER_BPS_30,
FEE_TIER_BPS_5, MINIMUM_LIQUIDITY,
};
use ata_core::{compute_ata_seed, get_associated_token_account_id};
use nssa::{
error::NssaError,
program_deployment_transaction::{self, ProgramDeploymentTransaction},
public_transaction, PrivateKey, PublicKey, PublicTransaction, V03State,
};
use nssa_core::account::{Account, AccountId, Data, Nonce};
use nssa_core::{
account::{Account, AccountId, Data, Nonce},
program::ProgramId,
};
use token_core::{TokenDefinition, TokenHolding};
struct Keys;
@ -16,6 +20,10 @@ struct Balances;
struct Accounts;
impl Keys {
fn owner() -> PrivateKey {
PrivateKey::try_new([30; 32]).expect("valid private key")
}
fn user_a() -> PrivateKey {
PrivateKey::try_new([31; 32]).expect("valid private key")
}
@ -42,6 +50,10 @@ impl Ids {
ata_methods::ATA_ID
}
fn malicious_ata_program() -> ProgramId {
malicious_ata_methods::MALICIOUS_ATA_ID
}
fn token_a_definition() -> AccountId {
AccountId::new([3; 32])
}
@ -93,6 +105,31 @@ impl Ids {
fn user_lp() -> AccountId {
AccountId::from(&PublicKey::new_from_private_key(&Keys::user_lp()))
}
fn owner() -> AccountId {
AccountId::from(&PublicKey::new_from_private_key(&Keys::owner()))
}
fn owner_token_a_ata() -> AccountId {
get_associated_token_account_id(
&Self::ata_program(),
&compute_ata_seed(Self::owner(), Self::token_a_definition()),
)
}
fn owner_token_b_ata() -> AccountId {
get_associated_token_account_id(
&Self::ata_program(),
&compute_ata_seed(Self::owner(), Self::token_b_definition()),
)
}
fn owner_token_lp_ata() -> AccountId {
get_associated_token_account_id(
&Self::ata_program(),
&compute_ata_seed(Self::owner(), Self::token_lp_definition()),
)
}
}
impl Balances {
@ -282,6 +319,15 @@ impl Balances {
}
impl Accounts {
fn owner() -> Account {
Account {
program_owner: ProgramId::default(),
balance: 0_u128,
data: Data::default(),
nonce: Nonce(0),
}
}
fn user_a_holding() -> Account {
Account {
program_owner: Ids::token_program(),
@ -412,6 +458,42 @@ impl Accounts {
}
}
fn owner_token_a_ata() -> Account {
Account {
program_owner: Ids::token_program(),
balance: 0_u128,
data: Data::from(&TokenHolding::Fungible {
definition_id: Ids::token_a_definition(),
balance: Balances::user_a_init(),
}),
nonce: Nonce(0),
}
}
fn owner_token_b_ata() -> Account {
Account {
program_owner: Ids::token_program(),
balance: 0_u128,
data: Data::from(&TokenHolding::Fungible {
definition_id: Ids::token_b_definition(),
balance: Balances::user_b_init(),
}),
nonce: Nonce(0),
}
}
fn owner_token_lp_ata() -> Account {
Account {
program_owner: Ids::token_program(),
balance: 0_u128,
data: Data::from(&TokenHolding::Fungible {
definition_id: Ids::token_lp_definition(),
balance: Balances::user_lp_init(),
}),
nonce: Nonce(0),
}
}
// --- Expected post-state accounts ---
fn pool_definition_swap_1() -> Account {
@ -893,21 +975,25 @@ impl Accounts {
}
}
fn deploy_programs(state: &mut V03State) {
let token_message =
program_deployment_transaction::Message::new(token_methods::TOKEN_ELF.to_vec());
fn deploy_program(state: &mut V03State, elf: &[u8], name: &str) {
let message = program_deployment_transaction::Message::new(elf.to_vec());
state
.transition_from_program_deployment_transaction(&ProgramDeploymentTransaction::new(
token_message,
))
.expect("token program deployment must succeed");
.transition_from_program_deployment_transaction(&ProgramDeploymentTransaction::new(message))
.unwrap_or_else(|_| panic!("{name} program deployment must succeed"));
}
let amm_message = program_deployment_transaction::Message::new(amm_methods::AMM_ELF.to_vec());
state
.transition_from_program_deployment_transaction(&ProgramDeploymentTransaction::new(
amm_message,
))
.expect("amm program deployment must succeed");
fn deploy_programs(state: &mut V03State) {
deploy_program(state, token_methods::TOKEN_ELF, "token");
deploy_program(state, amm_methods::AMM_ELF, "amm");
}
fn deploy_programs_with_malicious_ata(state: &mut V03State) {
deploy_programs(state);
deploy_program(
state,
malicious_ata_methods::MALICIOUS_ATA_ELF,
"malicious ata",
);
}
fn state_for_amm_tests() -> V03State {
@ -950,10 +1036,53 @@ fn state_for_amm_tests_with_new_def() -> V03State {
state
}
fn state_for_malicious_ata_attack() -> V03State {
let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);
deploy_programs_with_malicious_ata(&mut state);
state.force_insert_account(Ids::pool_definition(), Accounts::pool_definition_init());
state.force_insert_account(
Ids::token_a_definition(),
Accounts::token_a_definition_account(),
);
state.force_insert_account(
Ids::token_b_definition(),
Accounts::token_b_definition_account(),
);
state.force_insert_account(
Ids::token_lp_definition(),
Accounts::token_lp_definition_account(),
);
state.force_insert_account(Ids::owner(), Accounts::owner());
state.force_insert_account(Ids::owner_token_a_ata(), Accounts::owner_token_a_ata());
state.force_insert_account(Ids::owner_token_b_ata(), Accounts::owner_token_b_ata());
state.force_insert_account(Ids::owner_token_lp_ata(), Accounts::owner_token_lp_ata());
state.force_insert_account(Ids::vault_a(), Accounts::vault_a_init());
state.force_insert_account(Ids::vault_b(), Accounts::vault_b_init());
state
}
fn current_nonce(state: &V03State, account_id: AccountId) -> Nonce {
state.get_account_by_id(account_id).nonce
}
fn try_execute_amm_as_owner(
state: &mut V03State,
instruction: amm_core::Instruction,
accounts: Vec<AccountId>,
) -> Result<(), NssaError> {
let message = public_transaction::Message::try_new(
Ids::amm_program(),
accounts,
vec![current_nonce(state, Ids::owner())],
instruction,
)
.unwrap();
let witness_set = public_transaction::WitnessSet::for_message(&message, &[&Keys::owner()]);
let tx = PublicTransaction::new(message, witness_set);
state.transition_from_public_transaction(&tx, 0, 0)
}
fn state_for_amm_tests_with_precreated_user_lp_for_new_def() -> V03State {
let mut state = state_for_amm_tests_with_new_def();
state.force_insert_account(Ids::user_lp(), Accounts::user_lp_holding_init_zero());
@ -1178,6 +1307,290 @@ fn fungible_total_supply(account: &Account) -> u128 {
total_supply
}
fn exact_output_required_amount_in(
reserve_in: u128,
reserve_out: u128,
exact_amount_out: u128,
fee_bps: u128,
) -> u128 {
let effective_in_min = reserve_in
.checked_mul(exact_amount_out)
.expect("reserve_in * exact_amount_out overflows")
.div_ceil(
reserve_out
.checked_sub(exact_amount_out)
.expect("exact_amount_out must stay below reserve_out"),
);
let fee_multiplier = FEE_BPS_DENOMINATOR
.checked_sub(fee_bps)
.expect("fee_bps exceeds denominator");
effective_in_min
.checked_mul(FEE_BPS_DENOMINATOR)
.expect("effective_in_min * denominator overflows")
.div_ceil(fee_multiplier)
}
fn add_liquidity_malicious_ata_attack_witness() -> Option<&'static str> {
let mut state = state_for_malicious_ata_attack();
let instruction = amm_core::Instruction::AddLiquidity {
min_amount_liquidity: Balances::add_min_lp(),
max_amount_to_add_token_a: Balances::add_max_a(),
max_amount_to_add_token_b: Balances::add_max_b(),
ata_program_id: Ids::malicious_ata_program(),
};
if try_execute_amm_as_owner(
&mut state,
instruction,
vec![
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
Ids::token_lp_definition(),
Ids::owner(),
Ids::owner_token_a_ata(),
Ids::owner_token_b_ata(),
Ids::owner_token_lp_ata(),
],
)
.is_err()
{
return None;
}
let pool = pool_definition(&state.get_account_by_id(Ids::pool_definition()));
assert_eq!(pool.liquidity_pool_supply, Balances::token_lp_supply_add());
assert_eq!(pool.reserve_a, Balances::vault_a_add());
assert_eq!(pool.reserve_b, Balances::vault_b_add());
assert_eq!(
fungible_total_supply(&state.get_account_by_id(Ids::token_lp_definition())),
Balances::token_lp_supply_add()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_lp_ata())),
Balances::user_lp_add()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_a_ata())),
Balances::user_a_init()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_b_ata())),
Balances::user_b_init()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::vault_a())),
Balances::vault_a_init()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::vault_b())),
Balances::vault_b_init()
);
Some(
"add_liquidity: LP supply and owner LP balance increase while both deposit legs leave balances unchanged",
)
}
fn remove_liquidity_malicious_ata_attack_witness() -> Option<&'static str> {
let mut state = state_for_malicious_ata_attack();
let instruction = amm_core::Instruction::RemoveLiquidity {
remove_liquidity_amount: Balances::remove_lp(),
min_amount_to_remove_token_a: Balances::remove_min_a(),
min_amount_to_remove_token_b: Balances::remove_min_b(),
ata_program_id: Ids::malicious_ata_program(),
};
if try_execute_amm_as_owner(
&mut state,
instruction,
vec![
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
Ids::token_lp_definition(),
Ids::owner(),
Ids::owner_token_a_ata(),
Ids::owner_token_b_ata(),
Ids::owner_token_lp_ata(),
],
)
.is_err()
{
return None;
}
let pool = pool_definition(&state.get_account_by_id(Ids::pool_definition()));
assert_eq!(
pool.liquidity_pool_supply,
Balances::token_lp_supply_remove()
);
assert_eq!(pool.reserve_a, Balances::vault_a_remove());
assert_eq!(pool.reserve_b, Balances::vault_b_remove());
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_a_ata())),
Balances::user_a_remove()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_b_ata())),
Balances::user_b_remove()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::vault_a())),
Balances::vault_a_remove()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::vault_b())),
Balances::vault_b_remove()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_lp_ata())),
Balances::user_lp_init()
);
assert_eq!(
fungible_total_supply(&state.get_account_by_id(Ids::token_lp_definition())),
Balances::token_lp_supply()
);
Some(
"remove_liquidity: owner receives vault tokens while LP balance and LP definition supply stay unchanged",
)
}
fn swap_exact_input_malicious_ata_attack_witness() -> Option<&'static str> {
let mut state = state_for_malicious_ata_attack();
let instruction = amm_core::Instruction::SwapExactInput {
swap_amount_in: Balances::swap_amount_in(),
min_amount_out: Balances::swap_min_out(),
token_definition_id_in: Ids::token_a_definition(),
ata_program_id: Ids::malicious_ata_program(),
};
if try_execute_amm_as_owner(
&mut state,
instruction,
vec![
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
Ids::owner(),
Ids::owner_token_a_ata(),
Ids::owner_token_b_ata(),
],
)
.is_err()
{
return None;
}
let pool = pool_definition(&state.get_account_by_id(Ids::pool_definition()));
assert_eq!(pool.reserve_a, Balances::reserve_a_swap_2());
assert_eq!(pool.reserve_b, Balances::reserve_b_swap_2());
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_a_ata())),
Balances::user_a_init()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_b_ata())),
Balances::user_b_swap_2()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::vault_a())),
Balances::vault_a_init()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::vault_b())),
Balances::vault_b_swap_2()
);
Some(
"swap_exact_input: owner receives output while the input balance and deposit vault stay unchanged",
)
}
fn swap_exact_output_malicious_ata_attack_witness() -> Option<&'static str> {
const EXACT_AMOUNT_OUT: u128 = 500;
const MAX_AMOUNT_IN: u128 = 2_000;
let mut state = state_for_malicious_ata_attack();
let instruction = amm_core::Instruction::SwapExactOutput {
exact_amount_out: EXACT_AMOUNT_OUT,
max_amount_in: MAX_AMOUNT_IN,
token_definition_id_in: Ids::token_a_definition(),
ata_program_id: Ids::malicious_ata_program(),
};
if try_execute_amm_as_owner(
&mut state,
instruction,
vec![
Ids::pool_definition(),
Ids::vault_a(),
Ids::vault_b(),
Ids::owner(),
Ids::owner_token_a_ata(),
Ids::owner_token_b_ata(),
],
)
.is_err()
{
return None;
}
let required_amount_in = exact_output_required_amount_in(
Balances::vault_a_init(),
Balances::vault_b_init(),
EXACT_AMOUNT_OUT,
Balances::fee_tier(),
);
let pool = pool_definition(&state.get_account_by_id(Ids::pool_definition()));
assert_eq!(
pool.reserve_a,
Balances::vault_a_init() + required_amount_in
);
assert_eq!(pool.reserve_b, Balances::vault_b_init() - EXACT_AMOUNT_OUT);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_a_ata())),
Balances::user_a_init()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::owner_token_b_ata())),
Balances::user_b_init() + EXACT_AMOUNT_OUT
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::vault_a())),
Balances::vault_a_init()
);
assert_eq!(
fungible_balance(&state.get_account_by_id(Ids::vault_b())),
Balances::vault_b_init() - EXACT_AMOUNT_OUT
);
Some(
"swap_exact_output: owner receives exact output while the required input balance and deposit vault stay unchanged",
)
}
#[test]
fn amm_rejects_malicious_ata_program_for_all_value_paths() {
let accepted_attacks = [
add_liquidity_malicious_ata_attack_witness(),
remove_liquidity_malicious_ata_attack_witness(),
swap_exact_input_malicious_ata_attack_witness(),
swap_exact_output_malicious_ata_attack_witness(),
]
.into_iter()
.flatten()
.collect::<Vec<_>>();
assert!(
accepted_attacks.is_empty(),
"AMM accepted a malicious ATA program for value paths: {}",
accepted_attacks.join("; ")
);
}
#[test]
fn amm_remove_liquidity() {
let mut state = state_for_amm_tests();