From af81719414c7b9f534003f87bf6326c4d926fe39 Mon Sep 17 00:00:00 2001 From: moudyellaz Date: Thu, 2 Apr 2026 21:17:21 +0200 Subject: [PATCH] feat: add flash swap integration tests --- nssa/src/state.rs | 316 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 316 insertions(+) diff --git a/nssa/src/state.rs b/nssa/src/state.rs index fa76c2e6..85d0a053 100644 --- a/nssa/src/state.rs +++ b/nssa/src/state.rs @@ -3479,4 +3479,320 @@ pub mod tests { let state_from_bytes: V03State = borsh::from_slice(&bytes).unwrap(); assert_eq!(state, state_from_bytes); } + + // ── Flash Swap integration tests ────────────────────────────────────────── + + /// Mirror of the guest `FlashSwapInstruction` enum so we can serialise + /// instructions on the host side. + #[derive(serde::Serialize, serde::Deserialize)] + enum FlashSwapInstruction { + Initiate { + token_program_id: ProgramId, + callback_program_id: ProgramId, + amount_out: u128, + callback_instruction_data: Vec, + vault_after_transfer: AccountWithMetadata, + receiver_after_transfer: AccountWithMetadata, + vault_after_callback: AccountWithMetadata, + }, + InvariantCheck { + min_vault_balance: u128, + }, + } + + /// Mirror of the guest `CallbackInstruction`. + #[derive(serde::Serialize, serde::Deserialize)] + struct CallbackInstruction { + return_funds: bool, + token_program_id: ProgramId, + amount: u128, + vault_after_return: Option, + receiver_after_return: Option, + } + + /// Build a flash-swap `PublicTransaction` ready for execution. + /// + /// `vault_id` / `receiver_id` are the two accounts passed to the initiator. + /// `instruction` is the already-built `FlashSwapInstruction`. + fn build_flash_swap_tx( + initiator: &Program, + vault_id: AccountId, + receiver_id: AccountId, + instruction: FlashSwapInstruction, + ) -> PublicTransaction { + let message = public_transaction::Message::try_new( + initiator.id(), + vec![vault_id, receiver_id], + vec![], // no signers — vault is PDA-authorised + instruction, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + PublicTransaction::new(message, witness_set) + } + + #[test] + fn flash_swap_successful() { + let initiator = Program::flash_swap_initiator(); + let callback = Program::flash_swap_callback(); + let token = Program::authenticated_transfer_program(); + + let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0u8; 32]))); + let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1u8; 32]))); + + let initial_balance: u128 = 1000; + let amount_out: u128 = 100; + + let vault_account = Account { + program_owner: token.id(), + balance: initial_balance, + ..Account::default() + }; + let receiver_account = Account { + program_owner: token.id(), + balance: 0, + ..Account::default() + }; + + let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs(); + state.force_insert_account(vault_id, vault_account); + state.force_insert_account(receiver_id, receiver_account); + + // Pre-simulated intermediate states: + // After transfer (vault→receiver, amount_out): + let vault_after_transfer = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: initial_balance - amount_out, ..Account::default() }, + false, vault_id, + ); + let receiver_after_transfer = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: amount_out, ..Account::default() }, + false, receiver_id, + ); + + // After callback returns funds (receiver→vault, amount_out): + let vault_after_callback = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: initial_balance, ..Account::default() }, + false, vault_id, + ); + + // Callback instruction: return funds + let cb_instruction = CallbackInstruction { + return_funds: true, + token_program_id: token.id(), + amount: amount_out, + vault_after_return: Some(AccountWithMetadata::new( + Account { program_owner: token.id(), balance: initial_balance, ..Account::default() }, + false, vault_id, + )), + receiver_after_return: Some(AccountWithMetadata::new( + Account { program_owner: token.id(), balance: 0, ..Account::default() }, + false, receiver_id, + )), + }; + let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); + + let instruction = FlashSwapInstruction::Initiate { + token_program_id: token.id(), + callback_program_id: callback.id(), + amount_out, + callback_instruction_data: cb_data, + vault_after_transfer, + receiver_after_transfer, + vault_after_callback, + }; + + let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); + let result = state.transition_from_public_transaction(&tx, 1, 0); + assert!(result.is_ok(), "flash swap should succeed: {result:?}"); + + // Vault balance restored, receiver back to 0 + assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance); + assert_eq!(state.get_account_by_id(receiver_id).balance, 0); + } + + #[test] + fn flash_swap_callback_keeps_funds_rollback() { + let initiator = Program::flash_swap_initiator(); + let callback = Program::flash_swap_callback(); + let token = Program::authenticated_transfer_program(); + + let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0u8; 32]))); + let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1u8; 32]))); + + let initial_balance: u128 = 1000; + let amount_out: u128 = 100; + + let vault_account = Account { + program_owner: token.id(), + balance: initial_balance, + ..Account::default() + }; + let receiver_account = Account { + program_owner: token.id(), + balance: 0, + ..Account::default() + }; + + let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs(); + state.force_insert_account(vault_id, vault_account); + state.force_insert_account(receiver_id, receiver_account); + + // Pre-simulated intermediate states (same as successful case for steps 1-2): + let vault_after_transfer = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: initial_balance - amount_out, ..Account::default() }, + false, vault_id, + ); + let receiver_after_transfer = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: amount_out, ..Account::default() }, + false, receiver_id, + ); + + // After callback that does NOT return funds — vault stays drained: + let vault_after_callback = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: initial_balance - amount_out, ..Account::default() }, + false, vault_id, + ); + + // Callback instruction: do NOT return funds + let cb_instruction = CallbackInstruction { + return_funds: false, + token_program_id: token.id(), + amount: amount_out, + vault_after_return: None, + receiver_after_return: None, + }; + let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); + + let instruction = FlashSwapInstruction::Initiate { + token_program_id: token.id(), + callback_program_id: callback.id(), + amount_out, + callback_instruction_data: cb_data, + vault_after_transfer, + receiver_after_transfer, + vault_after_callback, + }; + + let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); + let result = state.transition_from_public_transaction(&tx, 1, 0); + + // Invariant check fails → entire tx rolls back + assert!(result.is_err(), "flash swap should fail when callback keeps funds"); + + // State unchanged (rollback) + assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance); + assert_eq!(state.get_account_by_id(receiver_id).balance, 0); + } + + #[test] + fn flash_swap_self_call_targets_correct_program() { + // Zero-amount flash swap: the invariant self-call still runs and succeeds + // because vault balance doesn't decrease. + let initiator = Program::flash_swap_initiator(); + let callback = Program::flash_swap_callback(); + let token = Program::authenticated_transfer_program(); + + let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0u8; 32]))); + let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1u8; 32]))); + + let initial_balance: u128 = 1000; + + let vault_account = Account { + program_owner: token.id(), + balance: initial_balance, + ..Account::default() + }; + let receiver_account = Account { + program_owner: token.id(), + balance: 0, + ..Account::default() + }; + + let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs(); + state.force_insert_account(vault_id, vault_account); + state.force_insert_account(receiver_id, receiver_account); + + // Zero-amount transfer: states remain unchanged after transfer + let vault_after_transfer = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: initial_balance, ..Account::default() }, + false, vault_id, + ); + let receiver_after_transfer = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: 0, ..Account::default() }, + false, receiver_id, + ); + // Callback with zero amount, return_funds=true (no-op effectively) + let vault_after_callback = AccountWithMetadata::new( + Account { program_owner: token.id(), balance: initial_balance, ..Account::default() }, + false, vault_id, + ); + + let cb_instruction = CallbackInstruction { + return_funds: true, + token_program_id: token.id(), + amount: 0, + vault_after_return: Some(AccountWithMetadata::new( + Account { program_owner: token.id(), balance: initial_balance, ..Account::default() }, + false, vault_id, + )), + receiver_after_return: Some(AccountWithMetadata::new( + Account { program_owner: token.id(), balance: 0, ..Account::default() }, + false, receiver_id, + )), + }; + let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); + + let instruction = FlashSwapInstruction::Initiate { + token_program_id: token.id(), + callback_program_id: callback.id(), + amount_out: 0, + callback_instruction_data: cb_data, + vault_after_transfer, + receiver_after_transfer, + vault_after_callback, + }; + + let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); + let result = state.transition_from_public_transaction(&tx, 1, 0); + assert!(result.is_ok(), "zero-amount flash swap should succeed: {result:?}"); + } + + #[test] + fn flash_swap_standalone_invariant_check_rejected() { + // Calling InvariantCheck directly (not as a chained self-call) should fail + // because caller_program_id will be None. + let initiator = Program::flash_swap_initiator(); + let token = Program::authenticated_transfer_program(); + + let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0u8; 32]))); + + let vault_account = Account { + program_owner: token.id(), + balance: 1000, + ..Account::default() + }; + + let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs(); + state.force_insert_account(vault_id, vault_account); + + let instruction = FlashSwapInstruction::InvariantCheck { + min_vault_balance: 1000, + }; + + let message = public_transaction::Message::try_new( + initiator.id(), + vec![vault_id], + vec![], + instruction, + ) + .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(), + "standalone InvariantCheck should be rejected (caller_program_id is None)" + ); + } }