diff --git a/program-tests/compressed-token-test/tests/transfer2/mod.rs b/program-tests/compressed-token-test/tests/transfer2/mod.rs index 3eea908aeb..8f30a2ccff 100644 --- a/program-tests/compressed-token-test/tests/transfer2/mod.rs +++ b/program-tests/compressed-token-test/tests/transfer2/mod.rs @@ -6,4 +6,5 @@ pub mod no_system_program_cpi_failing; pub mod random; pub mod shared; pub mod spl_ctoken; +pub mod transfer2_with_ata; pub mod transfer_failing; diff --git a/program-tests/compressed-token-test/tests/transfer2/transfer2_with_ata.rs b/program-tests/compressed-token-test/tests/transfer2/transfer2_with_ata.rs new file mode 100644 index 0000000000..2d3a1fbc2c --- /dev/null +++ b/program-tests/compressed-token-test/tests/transfer2/transfer2_with_ata.rs @@ -0,0 +1,1317 @@ +//! Tests for Transfer2WithAta instruction. +//! +//! Transfer2WithAta enables decompress/transfer operations on compressed tokens +//! where ALL inputs have owner = ATA pubkey (compress_to_pubkey mode). +//! +//! Test coverage: +//! +//! OWNER MODE SUCCESS CASES: +//! 1. Single ATA-owned compressed token decompressed +//! 2. Multiple ATA-owned compressed tokens decompressed in single call +//! +//! OWNER MODE FAILURE CASES: +//! 1. Wrong owner signer (not wallet that owns the ATA) +//! 2. Wrong mint passed +//! 3. wallet_idx correct key but not signer +//! 4. False ATA derivation (wrong bump) +//! 5. Non-matching ATA in accounts +//! 6. Mixed ownership (some ATA-owned, some wallet-owned) +//! +//! DELEGATE MODE SUCCESS CASES: +//! 1. Single delegated ATA-owned token - delegate signs +//! 2. Multiple delegated ATA-owned tokens - delegate signs +//! +//! DELEGATE MODE FAILURE CASES (SECURITY CRITICAL): +//! 1. Delegate provided but doesn't sign +//! 2. Wrong delegate signs (different from input's delegate field) +//! 3. Inputs have different delegates (must all match) +//! 4. Input has no delegate set (delegate mode requires delegate on all inputs) +//! 5. Owner signs instead of delegate (when delegate mode is specified) +//! 6. Mixed: some inputs have delegate, some don't +//! 7. Delegate index tampered (out of bounds) +//! 8. Delegate mode with mismatched delegate pubkey in accounts vs instruction + +use light_client::indexer::Indexer; +use light_compressed_token_sdk::ctoken::{ + derive_ctoken_ata, CompressibleParams, CreateAssociatedTokenAccount, +}; +use light_ctoken_types::{ + instructions::{extensions::compressible::CompressToPubkey, mint_action::Recipient}, + state::TokenDataVersion, +}; +use light_program_test::{LightProgramTest, ProgramTestConfig}; +use light_test_utils::{airdrop_lamports, Rpc}; +use light_token_client::{ + actions::{create_mint, mint_action, mint_to_compressed}, + instructions::{ + mint_action::{MintActionParams, MintActionType}, + transfer2::{create_generic_transfer2_instruction, Transfer2InstructionType}, + transfer2_with_ata::{create_decompress_ata_instruction_rpc, DecompressAtaInput}, + }, +}; +use serial_test::serial; +use solana_sdk::{signature::Keypair, signer::Signer}; + +// ============================================================================ +// Test Context Setup +// ============================================================================ + +struct Transfer2WithAtaTestContext { + rpc: LightProgramTest, + payer: Keypair, + owner_wallet: Keypair, + mint: solana_sdk::pubkey::Pubkey, + mint_seed: Keypair, + mint_authority: Keypair, + ata: solana_sdk::pubkey::Pubkey, + ata_bump: u8, +} + +async fn setup_transfer2_with_ata_test( +) -> Result> { + let mut rpc = LightProgramTest::new(ProgramTestConfig::new_v2(false, None)).await?; + let payer = rpc.get_payer().insecure_clone(); + let owner_wallet = Keypair::new(); + let mint_authority = Keypair::new(); + let mint_seed = Keypair::new(); + + // Airdrop to owner + airdrop_lamports(&mut rpc, &owner_wallet.pubkey(), 10_000_000_000).await?; + + Ok(Transfer2WithAtaTestContext { + rpc, + payer, + owner_wallet, + mint: solana_sdk::pubkey::Pubkey::default(), // Will be set after mint creation + mint_seed, + mint_authority, + ata: solana_sdk::pubkey::Pubkey::default(), // Will be set after ATA creation + ata_bump: 0, + }) +} + +impl Transfer2WithAtaTestContext { + /// Create compressed mint + async fn create_mint(&mut self) -> Result<(), Box> { + let (mint, _) = + light_compressed_token_sdk::compressed_token::create_compressed_mint::find_spl_mint_address( + &self.mint_seed.pubkey(), + ); + + create_mint( + &mut self.rpc, + &self.mint_seed, + 6, // decimals + &self.mint_authority, + None, + None, + &self.payer, + ) + .await?; + + self.mint = mint; + let (ata, bump) = derive_ctoken_ata(&self.owner_wallet.pubkey(), &self.mint); + self.ata = ata; + self.ata_bump = bump; + + Ok(()) + } + + /// Create ATA with compress_to_pubkey enabled + async fn create_ata_with_compress_to_pubkey( + &mut self, + ) -> Result<(), Box> { + let compressible_config = self + .rpc + .test_accounts + .funding_pool_config + .compressible_config_pda; + let rent_sponsor = self.rpc.test_accounts.funding_pool_config.rent_sponsor_pda; + + // Create compress_to_pubkey with ATA seeds + let compress_to_pubkey = CompressToPubkey { + bump: self.ata_bump, + program_id: light_ctoken_types::COMPRESSED_TOKEN_PROGRAM_ID, + seeds: vec![ + self.owner_wallet.pubkey().to_bytes().to_vec(), + light_ctoken_types::COMPRESSED_TOKEN_PROGRAM_ID.to_vec(), + self.mint.to_bytes().to_vec(), + ], + }; + + let compressible_params = CompressibleParams { + compressible_config, + rent_sponsor, + pre_pay_num_epochs: 10, + lamports_per_write: None, + compress_to_account_pubkey: Some(compress_to_pubkey), + token_account_version: TokenDataVersion::ShaFlat, + }; + + let create_ata_ix = CreateAssociatedTokenAccount::new( + self.payer.pubkey(), + self.owner_wallet.pubkey(), + self.mint, + compressible_params, + ) + .instruction()?; + + self.rpc + .create_and_send_transaction(&[create_ata_ix], &self.payer.pubkey(), &[&self.payer]) + .await?; + + Ok(()) + } + + /// Mint tokens to ATA then compress (creating ATA-owned compressed tokens) + async fn mint_and_compress_to_ata( + &mut self, + amount: u64, + ) -> Result<(), Box> { + let address_tree = self.rpc.get_address_tree_v2().tree; + let compressed_mint_address = light_compressed_token_sdk::compressed_token::create_compressed_mint::derive_compressed_mint_address( + &self.mint_seed.pubkey(), + &address_tree, + ); + + // Mint to the ATA + mint_action( + &mut self.rpc, + MintActionParams { + compressed_mint_address, + mint_seed: self.mint_seed.pubkey(), + authority: self.mint_authority.pubkey(), + payer: self.payer.pubkey(), + actions: vec![MintActionType::MintToCToken { + account: self.ata, + amount, + }], + new_mint: None, + }, + &self.mint_authority, + &self.payer, + None, + ) + .await?; + + // Now compress the tokens in the ATA + // This creates compressed tokens with owner = ATA (compress_to_pubkey mode) + let output_queue = self + .rpc + .get_random_state_tree_info() + .unwrap() + .get_output_pubkey() + .unwrap(); + + let compress_ix = create_generic_transfer2_instruction( + &mut self.rpc, + vec![Transfer2InstructionType::CompressAndClose( + light_token_client::instructions::transfer2::CompressAndCloseInput { + solana_ctoken_account: self.ata, + authority: self.owner_wallet.pubkey(), + output_queue, + destination: None, + is_compressible: true, + }, + )], + self.payer.pubkey(), + false, + ) + .await?; + + self.rpc + .create_and_send_transaction( + &[compress_ix], + &self.payer.pubkey(), + &[&self.payer, &self.owner_wallet], + ) + .await?; + + Ok(()) + } + + /// Mint tokens directly to wallet (wallet-owned compressed tokens) + async fn mint_to_wallet(&mut self, amount: u64) -> Result<(), Box> { + let recipients = vec![Recipient::new(self.owner_wallet.pubkey(), amount)]; + + mint_to_compressed( + &mut self.rpc, + self.mint, + recipients, + TokenDataVersion::ShaFlat, + &self.mint_authority, + &self.payer, + ) + .await?; + + Ok(()) + } + + /// Get compressed token accounts owned by the ATA + async fn get_ata_owned_compressed_accounts( + &self, + ) -> Result, Box> + { + let accounts = self + .rpc + .indexer() + .unwrap() + .get_compressed_token_accounts_by_owner(&self.ata, None, None) + .await? + .value + .items; + Ok(accounts) + } + + /// Get compressed token accounts owned by the wallet + async fn get_wallet_owned_compressed_accounts( + &self, + ) -> Result, Box> + { + let accounts = self + .rpc + .indexer() + .unwrap() + .get_compressed_token_accounts_by_owner(&self.owner_wallet.pubkey(), None, None) + .await? + .value + .items; + Ok(accounts) + } +} + +// ============================================================================ +// SUCCESS TESTS +// ============================================================================ + +/// Test: Successfully decompress a single ATA-owned compressed token +#[tokio::test] +#[serial] +async fn test_transfer2_with_ata_single_input_success() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + + // Re-create the ATA (it was closed by compress_and_close) + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + // Get ATA-owned compressed accounts + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!( + !compressed_accounts.is_empty(), + "Should have ATA-owned compressed accounts" + ); + + // Verify owner is the ATA + for acc in &compressed_accounts { + assert_eq!( + acc.token.owner, ctx.ata, + "Compressed token owner should be ATA" + ); + } + + // Decompress using Transfer2WithAta + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts, + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, // Full balance + use_delegate: false, + }; + + let ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + ctx.rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await + .unwrap(); + + // Verify ATA now has the tokens + let ata_account = ctx.rpc.get_account(ctx.ata).await.unwrap().unwrap(); + assert!(!ata_account.data.is_empty(), "ATA should exist with tokens"); +} + +/// Test: Successfully decompress multiple ATA-owned compressed tokens in single call +#[tokio::test] +#[serial] +async fn test_transfer2_with_ata_multiple_inputs_success() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + + // Create ATA and mint+compress multiple times to create multiple compressed accounts + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(500).await.unwrap(); + + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(300).await.unwrap(); + + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + // Get all ATA-owned compressed accounts + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!( + compressed_accounts.len() >= 2, + "Should have multiple ATA-owned compressed accounts, got {}", + compressed_accounts.len() + ); + + // Decompress all using Transfer2WithAta + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + ctx.rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await + .unwrap(); + + // Verify all tokens are now in ATA + let ata_account = ctx.rpc.get_account(ctx.ata).await.unwrap().unwrap(); + assert!(!ata_account.data.is_empty(), "ATA should exist with tokens"); +} + +// ============================================================================ +// FAILURE TESTS +// ============================================================================ + +/// Test: Fail when wrong owner wallet signs (not the wallet that owns the ATA) +#[tokio::test] +#[serial] +async fn test_transfer2_with_ata_wrong_owner_signer_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + // Try with a different signer (wrong wallet) + let wrong_wallet = Keypair::new(); + airdrop_lamports(&mut ctx.rpc, &wrong_wallet.pubkey(), 1_000_000_000) + .await + .unwrap(); + + // Build instruction manually with wrong wallet + // The SDK will reject this because ATA derivation won't match + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts, + owner_wallet: wrong_wallet.pubkey(), // Wrong wallet + mint: ctx.mint, + destination_ata: ctx.ata, // This won't match derivation from wrong_wallet + decompress_amount: None, + use_delegate: false, + }; + + let result = + create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()).await; + assert!( + result.is_err(), + "Should fail when ATA doesn't match owner_wallet derivation" + ); +} + +/// Test: Fail when wrong mint is passed +#[tokio::test] +#[serial] +async fn test_transfer2_with_ata_wrong_mint_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + // Try with wrong mint + let wrong_mint = solana_sdk::pubkey::Pubkey::new_unique(); + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts, + owner_wallet: ctx.owner_wallet.pubkey(), + mint: wrong_mint, // Wrong mint + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let result = + create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()).await; + assert!( + result.is_err(), + "Should fail when mint doesn't match ATA derivation" + ); +} + +/// Test: Fail when wallet_idx is correct key but not a signer +#[tokio::test] +#[serial] +async fn test_transfer2_with_ata_wallet_not_signer_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts, + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // Sign only with payer, NOT with owner_wallet + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer]) // Missing owner_wallet! + .await; + + assert!( + result.is_err(), + "Should fail when owner_wallet doesn't sign" + ); +} + +/// Test: Fail when non-matching ATA is passed in accounts +#[tokio::test] +#[serial] +async fn test_transfer2_with_ata_wrong_ata_account_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + // Use wrong ATA address + let wrong_ata = solana_sdk::pubkey::Pubkey::new_unique(); + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts, + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: wrong_ata, // Wrong ATA! + decompress_amount: None, + use_delegate: false, + }; + + let result = + create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()).await; + assert!( + result.is_err(), + "Should fail when destination_ata doesn't match derivation" + ); +} + +/// Test: Fail when mixed ownership (some ATA-owned, some wallet-owned) +/// This tests the on-chain check that ALL inputs must have owner = ATA +#[tokio::test] +#[serial] +async fn test_transfer2_with_ata_mixed_ownership_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + + // Create ATA-owned compressed tokens + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(500).await.unwrap(); + + let ata_owned_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + + // Also mint some tokens directly to wallet (wallet-owned compressed tokens) + ctx.mint_to_wallet(300).await.unwrap(); + + // Get wallet-owned compressed accounts + let wallet_owned_accounts = ctx.get_wallet_owned_compressed_accounts().await.unwrap(); + + // Mix the accounts (if we have both types) + if ata_owned_accounts.is_empty() || wallet_owned_accounts.is_empty() { + println!("Warning: Could not get both account types for mixed ownership test"); + return; + } + + // Try to pass mixed ownership - this should fail client-side + let mut mixed_accounts = ata_owned_accounts; + mixed_accounts.extend(wallet_owned_accounts); + + // Re-create ATA for destination + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let input = DecompressAtaInput { + compressed_token_accounts: mixed_accounts, + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + // Should fail because not all inputs have owner = ATA + let result = + create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()).await; + assert!( + result.is_err(), + "Should fail when mixing ATA-owned and wallet-owned inputs" + ); +} + +/// Test: ATTACK - Modify bump to invalid value (breaks ATA derivation) +#[tokio::test] +#[serial] +async fn test_attack_wrong_bump_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + // Build instruction correctly first + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Modify bump (2nd to last byte - suffix is [wallet, mint, ata, bump, delegate]) + let data_len = ix.data.len(); + let correct_bump = ix.data[data_len - 2]; + ix.data[data_len - 2] = correct_bump.wrapping_add(1); // Invalid bump + + // This should fail on-chain + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!(result.is_err(), "Should fail with wrong ATA bump"); +} + +// ============================================================================ +// DELEGATE MODE TESTS +// ============================================================================ +// +// Key security properties for delegate mode: +// 1. use_delegate=true requires inputs to have delegate set +// 2. Delegate (derived from inputs) MUST sign +// 3. All inputs MUST have same delegate +// +// NOTE: Success test for delegate mode blocked - needs ATA-owned tokens with delegates + +/// Test: use_delegate=true but inputs have no delegate (variant 1) +#[tokio::test] +#[serial] +async fn test_delegate_mode_no_delegate_in_inputs_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // Set use_delegate=true + let data_len = ix.data.len(); + ix.data[data_len - 1] = 1; + + // Fails: inputs have no delegate + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "Must fail when use_delegate=true but inputs have no delegate" + ); +} + +/// Test: SDK rejects use_delegate=true when inputs have no delegate +#[tokio::test] +#[serial] +async fn test_sdk_rejects_use_delegate_without_delegate() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + // Inputs have no delegate + for acc in &compressed_accounts { + assert!(acc.token.delegate.is_none()); + } + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts, + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: true, // SDK should reject this + }; + + let result = + create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()).await; + assert!( + result.is_err(), + "SDK must reject use_delegate=true when inputs have no delegate" + ); +} + +/// Test: Any non-zero use_delegate value triggers delegate mode +#[tokio::test] +#[serial] +async fn test_any_nonzero_use_delegate_triggers_delegate_mode() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts, + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // Set use_delegate to large value (255) - should still be treated as true + let data_len = ix.data.len(); + ix.data[data_len - 1] = 255; + + // Fails: delegate mode triggered but inputs have no delegate + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "Any non-zero use_delegate should trigger delegate mode" + ); +} + +/// Test: ATTACK - Modify wallet_index to point to wrong account +/// This tests that ATA derivation check catches wrong wallet +#[tokio::test] +#[serial] +async fn test_attack_wrong_wallet_index_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Change wallet_index to point to a different account + let data_len = ix.data.len(); + // wallet_index is at data_len - 5 + ix.data[data_len - 5] = 0; // Point to first account (likely fee_payer or system program) + + // This should fail because ATA derivation won't match + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when wallet_index points to wrong account" + ); +} + +/// Test: ATTACK - Modify mint_index to point to wrong account +/// This tests that ATA derivation check catches wrong mint +#[tokio::test] +#[serial] +async fn test_attack_wrong_mint_index_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Change mint_index to point to a different account + let data_len = ix.data.len(); + // mint_index is at data_len - 4 + ix.data[data_len - 4] = 0; // Point to first account + + // This should fail because ATA derivation won't match + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when mint_index points to wrong account" + ); +} + +/// Test: ATTACK - Modify ata_index to point to wrong account +/// This tests that ATA derivation check catches wrong ATA +#[tokio::test] +#[serial] +async fn test_attack_wrong_ata_index_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Change ata_index to point to a different account + let data_len = ix.data.len(); + // ata_index is at data_len - 3 + ix.data[data_len - 3] = 0; // Point to first account + + // This should fail because ATA derivation won't match + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when ata_index points to wrong account" + ); +} + +/// Test: ATTACK - Duplicate indices (wallet_idx == mint_idx) +/// This could cause ATA derivation to use same key twice, breaking the derivation +#[tokio::test] +#[serial] +async fn test_attack_duplicate_wallet_mint_index_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Set mint_index to same value as wallet_index + let data_len = ix.data.len(); + let wallet_idx = ix.data[data_len - 5]; + ix.data[data_len - 4] = wallet_idx; // mint_idx = wallet_idx + + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when wallet_idx == mint_idx (ATA derivation would be wrong)" + ); +} + +/// Test: ATTACK - Duplicate indices (wallet_idx == ata_idx) +/// This tests that we can't use same account for wallet and ATA +#[tokio::test] +#[serial] +async fn test_attack_duplicate_wallet_ata_index_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Set ata_index to same value as wallet_index + let data_len = ix.data.len(); + let wallet_idx = ix.data[data_len - 5]; + ix.data[data_len - 3] = wallet_idx; // ata_idx = wallet_idx + + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when wallet_idx == ata_idx (account confusion)" + ); +} + +/// Test: ATTACK - ata_index pointing to system account slot (< 7) +/// The packed_index calculation uses saturating_sub(7), so ata_index < 7 would give packed_index = 0 +/// This could potentially match an input with owner = 0 +#[tokio::test] +#[serial] +async fn test_attack_ata_index_in_system_slot_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Set ata_index to 0 (system account slot) + // This would make ata_packed_index = 0.saturating_sub(7) = 0 + let data_len = ix.data.len(); + ix.data[data_len - 3] = 0; // ata_idx = 0 (likely system program or fee payer) + + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when ata_index is in system slot (packed_index underflow)" + ); +} + +/// Test: ATTACK - All indices point to same account +/// Maximum confusion attack +#[tokio::test] +#[serial] +async fn test_attack_all_indices_same_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Set all indices to same value + let data_len = ix.data.len(); + let single_idx = ix.data[data_len - 5]; // wallet_idx (keep same) + ix.data[data_len - 4] = single_idx; // mint_idx = wallet_idx + ix.data[data_len - 3] = single_idx; // ata_idx = wallet_idx + + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when all indices point to same account" + ); +} + +/// Test: ATTACK - Zero byte for all suffix indices +/// Tests handling of all-zero suffix +#[tokio::test] +#[serial] +async fn test_attack_all_zero_suffix_indices_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Set all indices to 0 (except delegate which stays 255) + let data_len = ix.data.len(); + ix.data[data_len - 5] = 0; // wallet_idx = 0 + ix.data[data_len - 4] = 0; // mint_idx = 0 + ix.data[data_len - 3] = 0; // ata_idx = 0 + ix.data[data_len - 2] = 0; // bump = 0 + + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when suffix has all zero indices" + ); +} + +/// Test: Empty instruction data attack (too short) +#[tokio::test] +#[serial] +async fn test_instruction_data_too_short_fails() { + let mut ctx = setup_transfer2_with_ata_test().await.unwrap(); + ctx.create_mint().await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ctx.mint_and_compress_to_ata(1000).await.unwrap(); + ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let compressed_accounts = ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert!(!compressed_accounts.is_empty()); + + let input = DecompressAtaInput { + compressed_token_accounts: compressed_accounts.clone(), + owner_wallet: ctx.owner_wallet.pubkey(), + mint: ctx.mint, + destination_ata: ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let mut ix = create_decompress_ata_instruction_rpc(&mut ctx.rpc, input, ctx.payer.pubkey()) + .await + .unwrap(); + + // ATTACK: Truncate instruction data to be too short (less than 5 bytes suffix) + ix.data.truncate(3); // Way too short + + // This should fail with InvalidInstructionData + let result = ctx + .rpc + .create_and_send_transaction(&[ix], &ctx.payer.pubkey(), &[&ctx.payer, &ctx.owner_wallet]) + .await; + + assert!( + result.is_err(), + "SECURITY: Must fail when instruction data is too short" + ); +} + +// ============================================================================ +// CU BENCHMARKS - Compare Transfer2WithAta vs Regular Transfer2 +// ============================================================================ + +/// Benchmark: Compare CU usage between Transfer2WithAta and regular Transfer2 decompression +#[tokio::test] +#[serial] +async fn test_transfer2_with_ata_cu_benchmark() { + use light_program_test::utils::simulate_cu_multi; + use light_token_client::instructions::transfer2::{ + create_generic_transfer2_instruction, DecompressInput, Transfer2InstructionType, + }; + + println!("\n========================================"); + println!("Transfer2WithAta vs Transfer2 CU Benchmark"); + println!("========================================\n"); + + // Test configurations: (num_inputs, amount_per_input) + let test_configs = [(1, 1000u64), (2, 500u64)]; + + for (num_inputs, amount_per_input) in test_configs { + println!( + "--- {} input(s), {} tokens each ---", + num_inputs, amount_per_input + ); + + // Setup for Transfer2WithAta + let mut ata_ctx = setup_transfer2_with_ata_test().await.unwrap(); + ata_ctx.create_mint().await.unwrap(); + + // Create multiple ATA-owned compressed accounts + for _ in 0..num_inputs { + ata_ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + ata_ctx + .mint_and_compress_to_ata(amount_per_input) + .await + .unwrap(); + } + + // Re-create ATA for destination + ata_ctx.create_ata_with_compress_to_pubkey().await.unwrap(); + + let ata_compressed_accounts = ata_ctx.get_ata_owned_compressed_accounts().await.unwrap(); + assert_eq!( + ata_compressed_accounts.len(), + num_inputs, + "Should have {} ATA-owned compressed accounts", + num_inputs + ); + + // Build Transfer2WithAta instruction + let ata_input = DecompressAtaInput { + compressed_token_accounts: ata_compressed_accounts, + owner_wallet: ata_ctx.owner_wallet.pubkey(), + mint: ata_ctx.mint, + destination_ata: ata_ctx.ata, + decompress_amount: None, + use_delegate: false, + }; + + let ata_ix = create_decompress_ata_instruction_rpc( + &mut ata_ctx.rpc, + ata_input, + ata_ctx.payer.pubkey(), + ) + .await + .unwrap(); + + // Measure Transfer2WithAta CU + let ata_cu = simulate_cu_multi( + &mut ata_ctx.rpc, + &ata_ctx.payer, + &ata_ix, + &[&ata_ctx.owner_wallet], + ) + .await; + + // Setup for regular Transfer2 decompression + let mut transfer2_ctx = setup_transfer2_with_ata_test().await.unwrap(); + transfer2_ctx.create_mint().await.unwrap(); + + // For regular transfer2, mint directly to wallet (wallet-owned compressed tokens) + let total_amount = amount_per_input * num_inputs as u64; + for _ in 0..num_inputs { + transfer2_ctx + .mint_to_wallet(amount_per_input) + .await + .unwrap(); + } + + // Create CToken ATA for decompression destination (regular non-compressible ATA) + let (ctoken_ata, bump) = + derive_ctoken_ata(&transfer2_ctx.owner_wallet.pubkey(), &transfer2_ctx.mint); + let create_ata_ix = light_compressed_token_sdk::ctoken::CreateAssociatedTokenAccount { + idempotent: false, + bump, + payer: transfer2_ctx.payer.pubkey(), + owner: transfer2_ctx.owner_wallet.pubkey(), + mint: transfer2_ctx.mint, + associated_token_account: ctoken_ata, + compressible: None, + } + .instruction() + .unwrap(); + + transfer2_ctx + .rpc + .create_and_send_transaction( + &[create_ata_ix], + &transfer2_ctx.payer.pubkey(), + &[&transfer2_ctx.payer], + ) + .await + .unwrap(); + + // Get wallet-owned compressed accounts + let wallet_compressed_accounts = transfer2_ctx + .get_wallet_owned_compressed_accounts() + .await + .unwrap(); + assert_eq!( + wallet_compressed_accounts.len(), + num_inputs, + "Should have {} wallet-owned compressed accounts", + num_inputs + ); + + // Build regular Transfer2 decompression instruction + let decompress_input = DecompressInput { + compressed_token_account: wallet_compressed_accounts, + decompress_amount: total_amount, + solana_token_account: ctoken_ata, + amount: total_amount, + pool_index: None, + }; + + let transfer2_ix = create_generic_transfer2_instruction( + &mut transfer2_ctx.rpc, + vec![Transfer2InstructionType::Decompress(decompress_input)], + transfer2_ctx.payer.pubkey(), + false, + ) + .await + .unwrap(); + + // Measure regular Transfer2 CU + let transfer2_cu = simulate_cu_multi( + &mut transfer2_ctx.rpc, + &transfer2_ctx.payer, + &transfer2_ix, + &[&transfer2_ctx.owner_wallet], + ) + .await; + + // Calculate and print results + let cu_diff = ata_cu as i64 - transfer2_cu as i64; + let percent_diff = if transfer2_cu > 0 { + (cu_diff as f64 / transfer2_cu as f64) * 100.0 + } else { + 0.0 + }; + + println!("{} input(s) with ata : {} cu", num_inputs, ata_cu); + println!("{} input(s) transfer2 : {} cu", num_inputs, transfer2_cu); + println!( + "{} input(s) difference : {:+} cu ({:+.1}%)", + num_inputs, cu_diff, percent_diff + ); + println!(); + } + + println!("========================================"); + println!("Benchmark complete"); + println!("========================================\n"); +} diff --git a/programs/compressed-token/program/src/create_associated_token_account.rs b/programs/compressed-token/program/src/create_associated_token_account.rs index dc19e6c5cd..4141d34fd6 100644 --- a/programs/compressed-token/program/src/create_associated_token_account.rs +++ b/programs/compressed-token/program/src/create_associated_token_account.rs @@ -163,12 +163,17 @@ fn process_compressible_config<'info>( return Err(anchor_compressed_token::ErrorCode::OneEpochPrefundingNotAllowed.into()); } - if compressible_config_ix_data + // For ATAs with compress_to_pubkey, validate the provided seeds match the ATA derivation + if let Some(compress_to_pubkey) = compressible_config_ix_data .compress_to_account_pubkey - .is_some() + .as_ref() { - msg!("Associated token accounts must not compress to pubkey"); - return Err(ProgramError::InvalidInstructionData); + compress_to_pubkey + .check_seeds(associated_token_account.key()) + .map_err(|_| { + msg!("compress_to_pubkey seeds do not match ATA derivation"); + ProgramError::InvalidInstructionData + })?; } let compressible_config_account = next_config_account(iter)?; diff --git a/programs/compressed-token/program/src/lib.rs b/programs/compressed-token/program/src/lib.rs index 77630ca7f3..cb99a96f66 100644 --- a/programs/compressed-token/program/src/lib.rs +++ b/programs/compressed-token/program/src/lib.rs @@ -16,6 +16,7 @@ pub mod extensions; pub mod mint_action; pub mod shared; pub mod transfer2; +pub mod transfer2_with_ata; pub mod withdraw_funding_pool; // Reexport the wrapped anchor program. @@ -34,6 +35,7 @@ use withdraw_funding_pool::process_withdraw_funding_pool; use crate::{ convert_account_infos::convert_account_infos, mint_action::processor::process_mint_action, + transfer2_with_ata::process_transfer2_with_ata, }; pub const LIGHT_CPI_SIGNER: CpiSigner = @@ -78,6 +80,9 @@ pub enum InstructionType { CreateAssociatedTokenAccount2 = 106, /// Create associated token account with owner and mint as accounts (idempotent) CreateAssociatedTokenAccount2Idempotent = 107, + /// Transfer2 wrapper for ATA-owned compressed tokens (compress_to_pubkey mode) + /// Allows any Transfer2 operation where some inputs have owner = ATA pubkey + Transfer2WithAta = 108, Other, } @@ -96,6 +101,7 @@ impl From for InstructionType { 105 => InstructionType::WithdrawFundingPool, 106 => InstructionType::CreateAssociatedTokenAccount2, 107 => InstructionType::CreateAssociatedTokenAccount2Idempotent, + 108 => InstructionType::Transfer2WithAta, _ => InstructionType::Other, // anchor instructions } } @@ -163,6 +169,10 @@ pub fn process_instruction( msg!("CreateAssociatedTokenAccount2Idempotent"); process_create_associated_token_account2_idempotent(accounts, &instruction_data[1..])?; } + InstructionType::Transfer2WithAta => { + msg!("Transfer2WithAta"); + process_transfer2_with_ata(accounts, &instruction_data[1..])?; + } // anchor instructions have no discriminator conflicts with InstructionType // TODO: add test for discriminator conflict _ => { diff --git a/programs/compressed-token/program/src/transfer2_with_ata.rs b/programs/compressed-token/program/src/transfer2_with_ata.rs new file mode 100644 index 0000000000..30a55e5102 --- /dev/null +++ b/programs/compressed-token/program/src/transfer2_with_ata.rs @@ -0,0 +1,143 @@ +use anchor_lang::solana_program::program_error::ProgramError; +use light_account_checks::checks::check_signer; +use light_ctoken_types::instructions::transfer2::CompressedTokenInstructionDataTransfer2; +use light_program_profiler::profile; +use light_zero_copy::traits::ZeroCopyAt; +use pinocchio::{ + account_info::AccountInfo, + instruction::{AccountMeta, Instruction, Seed, Signer}, +}; +use spl_pod::solana_msg::msg; + +use crate::{shared::cpi::slice_invoke_signed, LIGHT_CPI_SIGNER}; + +const SYSTEM_ACCOUNTS_OFFSET: usize = 7; + +/// Process transfer2 with ATA as compressed-token owner. +#[profile] +pub fn process_transfer2_with_ata( + accounts: &[AccountInfo], + instruction_data: &[u8], +) -> Result<(), ProgramError> { + // Parse suffix: [transfer2_data...] ++ [wallet_idx, mint_idx, ata_idx, bump, use_delegate] + let len = instruction_data.len(); + if len < 5 { + return Err(ProgramError::InvalidInstructionData); + } + let (transfer2_data, suffix) = instruction_data.split_at(len - 5); + let (wallet_index, mint_index, ata_index) = + (suffix[0] as usize, suffix[1] as usize, suffix[2] as usize); + let (ata_bump, use_delegate) = (suffix[3], suffix[4] != 0); + + let owner_wallet = accounts + .get(wallet_index) + .ok_or(ProgramError::NotEnoughAccountKeys)?; + let mint = accounts + .get(mint_index) + .ok_or(ProgramError::NotEnoughAccountKeys)?; + let ata_account = accounts + .get(ata_index) + .ok_or(ProgramError::NotEnoughAccountKeys)?; + + // CHECK: ATA is derived correctly + let seeds: [&[u8]; 3] = [ + owner_wallet.key().as_ref(), + LIGHT_CPI_SIGNER.program_id.as_ref(), + mint.key().as_ref(), + ]; + let derived_ata = + pinocchio_pubkey::derive_address(&seeds, Some(ata_bump), &LIGHT_CPI_SIGNER.program_id); + if *ata_account.key() != derived_ata { + msg!("ATA derivation mismatch"); + return Err(ProgramError::InvalidAccountData); + } + + let (parsed, _) = CompressedTokenInstructionDataTransfer2::zero_copy_at(transfer2_data) + .map_err(|_| ProgramError::InvalidInstructionData)?; + let ata_packed_idx = ata_index.saturating_sub(SYSTEM_ACCOUNTS_OFFSET); + + if use_delegate { + let first_input = parsed + .in_token_data + .first() + .ok_or(ProgramError::InvalidInstructionData)?; + if !first_input.has_delegate() { + msg!("delegate mode but input has no delegate"); + return Err(ProgramError::InvalidAccountData); + } + let delegate_packed_idx = first_input.delegate as usize; + let delegate = accounts + .get(delegate_packed_idx + SYSTEM_ACCOUNTS_OFFSET) + .ok_or(ProgramError::NotEnoughAccountKeys)?; + + // CHECK: signer + check_signer(delegate).map_err(|_| { + msg!("delegate not signer"); + ProgramError::MissingRequiredSignature + })?; + + // CHECK: same delegate and compressed owner = ATA + for input in parsed.in_token_data.iter() { + if input.owner as usize != ata_packed_idx + || !input.has_delegate() + || input.delegate as usize != delegate_packed_idx + { + msg!("input owner/delegate mismatch"); + return Err(ProgramError::InvalidAccountData); + } + } + } else { + // CHECK: signer + check_signer(owner_wallet).map_err(|_| { + msg!("owner not signer"); + ProgramError::MissingRequiredSignature + })?; + + // CHECK: compressed owner = ATA + for input in parsed.in_token_data.iter() { + if input.owner as usize != ata_packed_idx { + msg!("input owner mismatch"); + return Err(ProgramError::InvalidAccountData); + } + } + } + + // self-CPI with ATA as signer + let mut ix_data = Vec::with_capacity(1 + transfer2_data.len()); + ix_data.push(101u8); + ix_data.extend_from_slice(transfer2_data); + + let account_metas: Vec<_> = accounts + .iter() + .enumerate() + .map(|(i, acc)| { + AccountMeta::new( + acc.key(), + acc.is_writable(), + acc.is_signer() || i == ata_index, + ) + }) + .collect(); + + let bump_seed = [ata_bump]; + let ata_seeds = [ + Seed::from(owner_wallet.key().as_ref()), + Seed::from(LIGHT_CPI_SIGNER.program_id.as_ref()), + Seed::from(mint.key().as_ref()), + Seed::from(bump_seed.as_ref()), + ]; + + slice_invoke_signed( + &Instruction { + program_id: &LIGHT_CPI_SIGNER.program_id, + accounts: &account_metas, + data: &ix_data, + }, + accounts, + &[Signer::from(&ata_seeds)], + ) + .map_err(|_| { + msg!("self-CPI failed"); + ProgramError::InvalidArgument + }) +} diff --git a/sdk-libs/compressed-token-sdk/src/ctoken/create_ata.rs b/sdk-libs/compressed-token-sdk/src/ctoken/create_ata.rs index 892b311846..4c6324c556 100644 --- a/sdk-libs/compressed-token-sdk/src/ctoken/create_ata.rs +++ b/sdk-libs/compressed-token-sdk/src/ctoken/create_ata.rs @@ -99,7 +99,7 @@ impl CreateAssociatedTokenAccount { 0 }, write_top_up: config.lamports_per_write.unwrap_or(0), - compress_to_account_pubkey: None, + compress_to_account_pubkey: config.compress_to_account_pubkey.clone(), }); let instruction_data = CreateAssociatedTokenAccountInstructionData { @@ -241,7 +241,7 @@ impl<'info> From<&CreateAssociatedTokenAccountInfos<'info>> for CreateAssociated rent_sponsor: *config.rent_sponsor.key, pre_pay_num_epochs: config.pre_pay_num_epochs, lamports_per_write: config.lamports_per_write, - compress_to_account_pubkey: None, + compress_to_account_pubkey: config.compress_to_account_pubkey.clone(), token_account_version: config.token_account_version, }), idempotent: account_infos.idempotent, @@ -320,7 +320,7 @@ impl CreateAssociatedTokenAccount2 { 0 }, write_top_up: config.lamports_per_write.unwrap_or(0), - compress_to_account_pubkey: None, + compress_to_account_pubkey: config.compress_to_account_pubkey.clone(), }); let instruction_data = CreateAssociatedTokenAccount2InstructionData { @@ -444,7 +444,7 @@ impl<'info> From<&CreateAssociatedTokenAccount2Infos<'info>> for CreateAssociate rent_sponsor: *config.rent_sponsor.key, pre_pay_num_epochs: config.pre_pay_num_epochs, lamports_per_write: config.lamports_per_write, - compress_to_account_pubkey: None, + compress_to_account_pubkey: config.compress_to_account_pubkey.clone(), token_account_version: config.token_account_version, }), idempotent: account_infos.idempotent, diff --git a/sdk-libs/compressed-token-sdk/src/ctoken/decompress_ata.rs b/sdk-libs/compressed-token-sdk/src/ctoken/decompress_ata.rs new file mode 100644 index 0000000000..d7341c12a8 --- /dev/null +++ b/sdk-libs/compressed-token-sdk/src/ctoken/decompress_ata.rs @@ -0,0 +1,117 @@ +//! Instruction builder for Transfer2WithAta (decompress ATA-owned compressed tokens). +//! +//! This transforms a base Transfer2 instruction into a Transfer2WithAta instruction, +//! which enables decompression of tokens where owner = ATA pubkey. + +use light_compressed_token_types::constants::TRANSFER2_WITH_ATA; +use light_ctoken_types::COMPRESSED_TOKEN_PROGRAM_ID; +use solana_instruction::{AccountMeta, Instruction}; +use solana_pubkey::Pubkey; + +use super::derive_ctoken_ata; +use crate::error::{Result, TokenSdkError}; + +/// Input for transforming a Transfer2 instruction into Transfer2WithAta +#[derive(Debug, Clone)] +pub struct DecompressAtaParams { + /// The wallet that owns the ATA (used for ATA derivation) + pub owner_wallet: Pubkey, + /// The mint of the tokens + pub mint: Pubkey, + /// If true, use delegate mode (delegate must sign). + /// If false, use owner mode (owner_wallet must sign). + pub use_delegate: bool, + /// The delegate pubkey (only required if use_delegate=true) + pub delegate: Option, +} + +/// Transforms a Transfer2 instruction into a Transfer2WithAta instruction. +/// +/// This wraps an existing Transfer2 decompress instruction to work with +/// ATA-owned compressed tokens (created with compress_to_pubkey=true). +/// +/// The on-chain program derives the ATA from [owner_wallet, program_id, mint], +/// validates all inputs have that ATA as owner, and performs a self-CPI +/// to Transfer2 with the ATA signed. +/// +/// # Arguments +/// * `transfer2_ix` - A base Transfer2 instruction for decompression +/// * `params` - Parameters for the transformation +/// +/// # Returns +/// A Transfer2WithAta instruction ready to be sent +pub fn create_decompress_ata_instruction( + mut transfer2_ix: Instruction, + params: DecompressAtaParams, +) -> Result { + let (derived_ata, ata_bump) = derive_ctoken_ata(¶ms.owner_wallet, ¶ms.mint); + + // Find or add wallet to accounts + let wallet_index = match transfer2_ix + .accounts + .iter() + .position(|m| m.pubkey == params.owner_wallet) + { + Some(idx) => { + if !params.use_delegate { + transfer2_ix.accounts[idx].is_signer = true; + } + idx as u8 + } + None => { + let idx = transfer2_ix.accounts.len() as u8; + transfer2_ix.accounts.push(AccountMeta::new_readonly( + params.owner_wallet, + !params.use_delegate, // is_signer only if owner mode + )); + idx + } + }; + + // In delegate mode, mark delegate as signer + if params.use_delegate { + let delegate = params.delegate.ok_or(TokenSdkError::InvalidAccountData)?; + if let Some(idx) = transfer2_ix + .accounts + .iter() + .position(|m| m.pubkey == delegate) + { + transfer2_ix.accounts[idx].is_signer = true; + } + } + + // Find mint index (must exist in accounts) + let mint_index = transfer2_ix + .accounts + .iter() + .position(|m| m.pubkey == params.mint) + .ok_or(TokenSdkError::InvalidAccountData)? as u8; + + // Find ATA index (must exist in accounts as decompress destination) + let ata_index = transfer2_ix + .accounts + .iter() + .position(|m| m.pubkey == derived_ata) + .ok_or(TokenSdkError::InvalidAccountData)? as u8; + + // Un-mark ATA as signer - it's a PDA, on-chain will sign via CPI + if let Some(ata_meta) = transfer2_ix.accounts.get_mut(ata_index as usize) { + ata_meta.is_signer = false; + } + + // Modify instruction data: + // - Change discriminator to Transfer2WithAta + // - Append: [wallet_idx, mint_idx, ata_idx, bump, use_delegate] + transfer2_ix.data[0] = TRANSFER2_WITH_ATA; + transfer2_ix.data.push(wallet_index); + transfer2_ix.data.push(mint_index); + transfer2_ix.data.push(ata_index); + transfer2_ix.data.push(ata_bump); + transfer2_ix.data.push(params.use_delegate as u8); + + Ok(Instruction { + program_id: Pubkey::from(COMPRESSED_TOKEN_PROGRAM_ID), + accounts: transfer2_ix.accounts, + data: transfer2_ix.data, + }) +} diff --git a/sdk-libs/compressed-token-sdk/src/ctoken/mod.rs b/sdk-libs/compressed-token-sdk/src/ctoken/mod.rs index 130e4326ac..61730077e5 100644 --- a/sdk-libs/compressed-token-sdk/src/ctoken/mod.rs +++ b/sdk-libs/compressed-token-sdk/src/ctoken/mod.rs @@ -3,6 +3,7 @@ mod compressible; mod create; mod create_ata; mod create_cmint; +mod decompress_ata; mod mint_to; mod transfer_ctoken; mod transfer_ctoken_spl; @@ -14,6 +15,7 @@ pub use compressible::{CompressibleParams, CompressibleParamsInfos}; pub use create::*; pub use create_ata::*; pub use create_cmint::*; +pub use decompress_ata::{create_decompress_ata_instruction, DecompressAtaParams}; use light_compressed_token_types::POOL_SEED; use light_compressible::config::CompressibleConfig; pub use light_ctoken_types::{ diff --git a/sdk-libs/compressed-token-types/src/constants.rs b/sdk-libs/compressed-token-types/src/constants.rs index b4eaafc8c1..f72de18c43 100644 --- a/sdk-libs/compressed-token-types/src/constants.rs +++ b/sdk-libs/compressed-token-types/src/constants.rs @@ -50,3 +50,4 @@ pub const THAW: [u8; 8] = [226, 249, 34, 57, 189, 21, 177, 101]; pub const CREATE_TOKEN_POOL: [u8; 8] = [23, 169, 27, 122, 147, 169, 209, 152]; pub const CREATE_ADDITIONAL_TOKEN_POOL: [u8; 8] = [114, 143, 210, 73, 96, 115, 1, 228]; pub const TRANSFER2: u8 = 101; +pub const TRANSFER2_WITH_ATA: u8 = 108; diff --git a/sdk-libs/program-test/src/utils/mod.rs b/sdk-libs/program-test/src/utils/mod.rs index 2309bd1055..a03e49924e 100644 --- a/sdk-libs/program-test/src/utils/mod.rs +++ b/sdk-libs/program-test/src/utils/mod.rs @@ -8,4 +8,4 @@ pub mod setup_light_programs; pub mod tree_accounts; pub mod simulation; -pub use simulation::simulate_cu; +pub use simulation::{simulate_cu, simulate_cu_multi}; diff --git a/sdk-libs/program-test/src/utils/simulation.rs b/sdk-libs/program-test/src/utils/simulation.rs index a5af0b331e..f2d2b6df9e 100644 --- a/sdk-libs/program-test/src/utils/simulation.rs +++ b/sdk-libs/program-test/src/utils/simulation.rs @@ -13,16 +13,33 @@ pub async fn simulate_cu( rpc: &mut LightProgramTest, payer: &Keypair, instruction: &Instruction, +) -> u64 { + simulate_cu_multi(rpc, payer, instruction, &[]).await +} + +/// Simulate a transaction with multiple signers and return the compute units consumed. +/// +/// This is a test utility function for measuring transaction costs. +/// The payer is always included as a signer. Additional signers can be passed via `additional_signers`. +pub async fn simulate_cu_multi( + rpc: &mut LightProgramTest, + payer: &Keypair, + instruction: &Instruction, + additional_signers: &[&Keypair], ) -> u64 { let blockhash = rpc .get_latest_blockhash() .await .expect("Failed to get latest blockhash") .0; + + let mut signers: Vec<&Keypair> = vec![payer]; + signers.extend(additional_signers); + let tx = Transaction::new_signed_with_payer( std::slice::from_ref(instruction), Some(&payer.pubkey()), - &[payer], + &signers, blockhash, ); let simulate_tx = VersionedTransaction::from(tx); diff --git a/sdk-libs/token-client/src/instructions/mod.rs b/sdk-libs/token-client/src/instructions/mod.rs index 0fce0b5387..5c1b3c73d3 100644 --- a/sdk-libs/token-client/src/instructions/mod.rs +++ b/sdk-libs/token-client/src/instructions/mod.rs @@ -2,4 +2,5 @@ pub mod create_mint; pub mod mint_action; pub mod mint_to_compressed; pub mod transfer2; +pub mod transfer2_with_ata; pub mod update_compressed_mint; diff --git a/sdk-libs/token-client/src/instructions/transfer2_with_ata.rs b/sdk-libs/token-client/src/instructions/transfer2_with_ata.rs new file mode 100644 index 0000000000..47b3a6e2e8 --- /dev/null +++ b/sdk-libs/token-client/src/instructions/transfer2_with_ata.rs @@ -0,0 +1,115 @@ +//! SDK helper for building Transfer2WithAta instructions via RPC. +//! +//! This module provides the high-level interface for decompressing ATA-owned +//! compressed tokens. It fetches required data via RPC and delegates the +//! instruction building to the compressed-token-sdk. + +use light_client::{ + indexer::{CompressedTokenAccount, Indexer}, + rpc::Rpc, +}; +use light_compressed_token_sdk::{ + ctoken::{create_decompress_ata_instruction, derive_ctoken_ata, DecompressAtaParams}, + error::TokenSdkError, +}; +use solana_instruction::Instruction; +use solana_pubkey::Pubkey; + +use super::transfer2::{ + create_generic_transfer2_instruction, DecompressInput, Transfer2InstructionType, +}; + +/// Input for decompressing ATA-owned compressed tokens +#[derive(Debug, Clone)] +pub struct DecompressAtaInput { + /// Compressed token accounts to decompress (ALL must have owner = ATA) + pub compressed_token_accounts: Vec, + /// The wallet that owns the ATA (used for ATA derivation, must sign if no delegate) + pub owner_wallet: Pubkey, + /// The mint of the tokens + pub mint: Pubkey, + /// The destination CToken ATA to decompress into + pub destination_ata: Pubkey, + /// Amount to decompress (if None, decompress full balance) + pub decompress_amount: Option, + /// If true, use delegate mode (delegate from inputs must sign). + /// If false, use owner mode (owner_wallet must sign). + pub use_delegate: bool, +} + +/// Creates a Transfer2WithAta instruction for decompressing ATA-owned compressed tokens. +/// +/// This fetches required account data via RPC, builds a base Transfer2 decompress +/// instruction, and then transforms it into a Transfer2WithAta instruction. +pub async fn create_decompress_ata_instruction_rpc( + rpc: &mut R, + input: DecompressAtaInput, + payer: Pubkey, +) -> Result { + if input.compressed_token_accounts.is_empty() { + return Err(TokenSdkError::InvalidAccountData); + } + + // Derive ATA and validate + let (derived_ata, _) = derive_ctoken_ata(&input.owner_wallet, &input.mint); + if input.destination_ata != derived_ata { + return Err(TokenSdkError::InvalidAccountData); + } + + // Validate all inputs have owner = ATA + for account in &input.compressed_token_accounts { + if account.token.owner != derived_ata { + return Err(TokenSdkError::InvalidAccountData); + } + } + + // If delegate mode, validate all inputs have same delegate set + let delegate = if input.use_delegate { + let first_delegate = input.compressed_token_accounts[0] + .token + .delegate + .ok_or(TokenSdkError::InvalidAccountData)?; + for account in &input.compressed_token_accounts { + if account.token.delegate != Some(first_delegate) { + return Err(TokenSdkError::InvalidAccountData); + } + } + Some(first_delegate) + } else { + None + }; + + // Calculate decompress amount + let total_balance: u64 = input + .compressed_token_accounts + .iter() + .map(|acc| acc.token.amount) + .sum(); + let decompress_amount = input.decompress_amount.unwrap_or(total_balance); + + // Build base Transfer2 decompress instruction via RPC + let transfer2_ix = create_generic_transfer2_instruction( + rpc, + vec![Transfer2InstructionType::Decompress(DecompressInput { + compressed_token_account: input.compressed_token_accounts, + decompress_amount, + solana_token_account: derived_ata, + amount: decompress_amount, + pool_index: None, + })], + payer, + true, // filter zero outputs + ) + .await?; + + // Transform to Transfer2WithAta using SDK + create_decompress_ata_instruction( + transfer2_ix, + DecompressAtaParams { + owner_wallet: input.owner_wallet, + mint: input.mint, + use_delegate: input.use_delegate, + delegate, + }, + ) +}