diff --git a/magicblock-chainlink/src/chainlink/fetch_cloner/ata_projection.rs b/magicblock-chainlink/src/chainlink/fetch_cloner/ata_projection.rs index bec28bef5..250801688 100644 --- a/magicblock-chainlink/src/chainlink/fetch_cloner/ata_projection.rs +++ b/magicblock-chainlink/src/chainlink/fetch_cloner/ata_projection.rs @@ -4,8 +4,8 @@ use dlp_api::state::DelegationRecord; use futures_util::future::join_all; use magicblock_accounts_db::traits::AccountsBank; use magicblock_core::token_programs::{ - is_ata, try_derive_ata_address_and_bump, try_derive_eata_address_and_bump, - AtaInfo, + is_ata, try_derive_eata_address_and_bump, try_derive_supported_ata_pubkeys, + AtaInfo, EphemeralAta, MaybeIntoAta, EATA_PROGRAM_ID, }; use magicblock_metrics::metrics; use solana_account::{AccountSharedData, ReadableAccount}; @@ -22,8 +22,8 @@ use super::{ use crate::{ cloner::{AccountCloneRequest, Cloner, DelegationActions}, remote_account_provider::{ - ChainPubsubClient, ChainRpcClient, ResolvedAccountSharedData, - SubscriptionReason, + ChainPubsubClient, ChainRpcClient, MatchSlotsConfig, RemoteAccount, + ResolvedAccountSharedData, SubscriptionReason, }, }; @@ -58,18 +58,257 @@ fn ata_info_from_layout( let mint = Pubkey::new_from_array(data[0..32].try_into().ok()?); let wallet_owner = Pubkey::new_from_array(data[32..64].try_into().ok()?); - let (derived_ata, _) = - try_derive_ata_address_and_bump(&wallet_owner, &mint)?; - if derived_ata != *ata_pubkey { + let ata_pubkeys = try_derive_supported_ata_pubkeys(&wallet_owner, &mint); + if ata_pubkeys.contains(ata_pubkey) { + return Some(AtaInfo { + mint, + owner: wallet_owner, + }); + } + + None +} + +pub(crate) fn is_known_empty_eata( + this: &FetchCloner, + eata_pubkey: &Pubkey, +) -> bool +where + T: ChainRpcClient, + U: ChainPubsubClient, + V: AccountsBank, + C: Cloner, +{ + this.known_empty_eatas.lock().get(eata_pubkey).is_some() +} + +pub(crate) fn mark_eata_empty( + this: &FetchCloner, + eata_pubkey: Pubkey, +) where + T: ChainRpcClient, + U: ChainPubsubClient, + V: AccountsBank, + C: Cloner, +{ + this.known_empty_eatas.lock().put(eata_pubkey, ()); +} + +pub(crate) fn maybe_build_projected_ata_clone_request_from_eata_sub_update< + T, + U, + V, + C, +>( + this: &FetchCloner, + eata_pubkey: Pubkey, + eata_account: &AccountSharedData, + deleg_record: Option<&DelegationRecord>, + delegation_actions: &DelegationActions, +) -> Option +where + T: ChainRpcClient, + U: ChainPubsubClient, + V: AccountsBank, + C: Cloner, +{ + let deleg_record = deleg_record?; + + if deleg_record.authority != this.validator_pubkey { return None; } + let (wallet_owner, mint) = delegation::parse_raw_eata_pda( + &eata_pubkey, + eata_account.data(), + deleg_record.owner, + )?; + let ata_pubkeys = try_derive_supported_ata_pubkeys(&wallet_owner, &mint); - Some(AtaInfo { - mint, - owner: wallet_owner, + // eATA updates only carry the projected balance fields. The in-bank ATA is + // required as the base so the clone preserves the actual token program + // owner and any Token-2022 account layout extensions. + let mut ata_pubkey = None; + let mut in_bank_ata = None; + for candidate_pubkey in ata_pubkeys.token_2022_first().into_iter().flatten() + { + if let Some(candidate_account) = + this.accounts_bank.get_account(&candidate_pubkey) + { + ata_pubkey = Some(candidate_pubkey); + in_bank_ata = Some(candidate_account); + break; + } + } + let in_bank_ata = in_bank_ata.as_ref()?; + let ata_pubkey = ata_pubkey?; + if in_bank_ata.delegated() || in_bank_ata.undelegating() { + return None; + } + let projected_ata = maybe_project_delegated_ata_from_eata( + this, + in_bank_ata, + eata_account, + deleg_record, + )?; + Some(AccountCloneRequest { + pubkey: ata_pubkey, + account: projected_ata, + commit_frequency_ms: None, + delegation_actions: delegation_actions.clone(), + delegated_to_other: None, }) } +pub(crate) async fn maybe_project_ata_from_subscription_update( + this: &FetchCloner, + ata_pubkey: Pubkey, + ata_account: AccountSharedData, +) -> ( + AccountSharedData, + Option<(DelegationRecord, Option)>, +) +where + T: ChainRpcClient, + U: ChainPubsubClient, + V: AccountsBank, + C: Cloner, +{ + let Some(ata_info) = is_ata(&ata_pubkey, &ata_account) else { + return (ata_account, None); + }; + + let Some((eata_pubkey, _)) = + try_derive_eata_address_and_bump(&ata_info.owner, &ata_info.mint) + else { + return (ata_account, None); + }; + + let was_watching = this.remote_account_provider.is_watching(&eata_pubkey); + + // Ensure before cache checks; this keeps the subscription LRU warm + // without refcounting the projection reason on every ATA update. + let subscribed = match this + .ensure_subscription(&eata_pubkey, SubscriptionReason::AtaProjection) + .await + { + Ok(()) => true, + Err(err) => { + warn!( + pubkey = %eata_pubkey, + error = ?err, + "Failed to subscribe to derived eATA" + ); + false + } + }; + + // Known-empty eATAs skip the fetch only if the subscription was already live. + if was_watching && subscribed && is_known_empty_eata(this, &eata_pubkey) { + return (ata_account, None); + } + + let (eata_account, definitively_not_found) = match this + .remote_account_provider + .try_get_multi_until_slots_match( + &[eata_pubkey], + Some(MatchSlotsConfig { + min_context_slot: Some(ata_account.remote_slot()), + ..Default::default() + }), + metrics::AccountFetchOrigin::ProjectAta, + ) + .await + { + Ok(mut accounts) => { + let popped = accounts.pop(); + // Only `NotFound` proves absence; stale, missing, or failed fetches retry later. + let nf = matches!(popped, Some(RemoteAccount::NotFound(_))); + let fresh = popped.and_then(|a| a.fresh_account()); + (fresh, nf) + } + Err(err) => { + debug!( + pubkey = %eata_pubkey, + error = ?err, + "Failed to fetch eATA for projection" + ); + (None, false) + } + }; + + let Some(eata_account) = eata_account else { + // Cache absence only after a confirmed NotFound and live subscription. + if definitively_not_found && subscribed { + mark_eata_empty(this, eata_pubkey); + } + return (ata_account, None); + }; + + let deleg_record = delegation::fetch_and_parse_delegation_record( + this, + eata_pubkey, + ata_account.remote_slot().max(eata_account.remote_slot()), + metrics::AccountFetchOrigin::ProjectAta, + ) + .await; + + let Some(deleg_record) = deleg_record else { + return (ata_account, None); + }; + let (deleg_record, delegation_actions) = deleg_record; + + if let Some(projected_ata) = maybe_project_delegated_ata_from_eata( + this, + &ata_account, + &eata_account, + &deleg_record, + ) { + return (projected_ata, Some((deleg_record, delegation_actions))); + } + (ata_account, Some((deleg_record, delegation_actions))) +} + +pub(crate) fn maybe_project_delegated_ata_from_eata( + this: &FetchCloner, + ata_account: &AccountSharedData, + eata_account: &AccountSharedData, + deleg_record: &DelegationRecord, +) -> Option +where + T: ChainRpcClient, + U: ChainPubsubClient, + V: AccountsBank, + C: Cloner, +{ + if deleg_record.authority != this.validator_pubkey { + return None; + } + + // Projecting from eATA must preserve the base ATA's owner and data length. + // That is what keeps Token-2022 accounts from being rebuilt as legacy SPL + // Token accounts when the eATA itself only stores owner, mint, and amount. + let projected_from_base_ata = if deleg_record.owner == EATA_PROGRAM_ID { + EphemeralAta::try_from_account_data(eata_account.data()) + .and_then(|eata| eata.project_into_ata_account(ata_account)) + } else { + None + }; + + let mut projected_ata = match projected_from_base_ata + .or_else(|| eata_account.maybe_into_ata(deleg_record.owner)) + { + Some(projected_ata) => projected_ata, + None => { + return None; + } + }; + let projected_slot = + ata_account.remote_slot().max(eata_account.remote_slot()); + projected_ata.set_remote_slot(projected_slot); + projected_ata.set_delegated(true); + Some(projected_ata) +} + /// Resolves ATAs with eATA projection. /// For each detected ATA, we derive the eATA PDA, subscribe to both, /// and, if the ATA is delegated to us and the eATA exists, we clone the eATA data @@ -230,8 +469,9 @@ where delegation::get_delegated_to_other(this, &deleg_record); commit_frequency_ms = Some(deleg_record.commit_frequency_ms); - if let Some(projected_ata) = this - .maybe_project_delegated_ata_from_eata( + if let Some(projected_ata) = + maybe_project_delegated_ata_from_eata( + this, input.ata_account.account_shared_data(), eata_shared, &deleg_record, diff --git a/magicblock-chainlink/src/chainlink/fetch_cloner/delegation.rs b/magicblock-chainlink/src/chainlink/fetch_cloner/delegation.rs index a167f80f7..5067ddb88 100644 --- a/magicblock-chainlink/src/chainlink/fetch_cloner/delegation.rs +++ b/magicblock-chainlink/src/chainlink/fetch_cloner/delegation.rs @@ -3,7 +3,9 @@ use dlp_api::{ pda::delegation_record_pda_from_delegated_account, state::DelegationRecord, }; use magicblock_accounts_db::traits::AccountsBank; -use magicblock_core::token_programs::{derive_eata, EATA_PROGRAM_ID}; +use magicblock_core::token_programs::{ + try_derive_eata_address_and_bump, EphemeralAta, EATA_PROGRAM_ID, +}; use magicblock_metrics::metrics; use solana_account::ReadableAccount; use solana_keypair::Keypair; @@ -136,14 +138,15 @@ pub(crate) fn parse_raw_eata_pda( data: &[u8], owner_program: Pubkey, ) -> Option<(Pubkey, Pubkey)> { - if owner_program != EATA_PROGRAM_ID || data.len() < 72 { + if owner_program != EATA_PROGRAM_ID { return None; } - let wallet_owner = Pubkey::new_from_array(data[0..32].try_into().ok()?); - let mint = Pubkey::new_from_array(data[32..64].try_into().ok()?); - (derive_eata(&wallet_owner, &mint) == *account_pubkey) - .then_some((wallet_owner, mint)) + let eata = EphemeralAta::try_from_account_data(data)?; + let (derived_eata, bump) = + try_derive_eata_address_and_bump(&eata.owner, &eata.mint)?; + (derived_eata == *account_pubkey && bump == eata.bump) + .then_some((eata.owner, eata.mint)) } pub(crate) fn get_delegated_to_other( diff --git a/magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs b/magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs index 157929eff..d5908b47c 100644 --- a/magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs +++ b/magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs @@ -16,8 +16,7 @@ use magicblock_accounts_db::traits::AccountsBank; use magicblock_aml::RiskService; use magicblock_config::config::AllowedProgram; use magicblock_core::token_programs::{ - is_ata, try_derive_ata_address_and_bump, try_derive_eata_address_and_bump, - MaybeIntoAta, EATA_PROGRAM_ID, + try_derive_supported_ata_pubkeys, EATA_PROGRAM_ID, }; use magicblock_metrics::metrics::{self, AccountFetchOrigin}; use parking_lot::Mutex as PlMutex; @@ -912,19 +911,23 @@ where } let delegation_actions = delegation_actions.unwrap_or_default(); - let greedy_ata_pubkey = delegation::parse_raw_eata_pda( + let greedy_ata_pubkeys = delegation::parse_raw_eata_pda( &pubkey, account.data(), deleg_record.owner, ) - .and_then(|(wallet_owner, mint)| { - try_derive_ata_address_and_bump(&wallet_owner, &mint) - .map(|(ata_pubkey, _)| ata_pubkey) - }); - let mut pubkeys_to_clone = vec![pubkey]; - if let Some(ata_pubkey) = greedy_ata_pubkey { - pubkeys_to_clone.push(ata_pubkey); - } + .map(|(wallet_owner, mint)| { + try_derive_supported_ata_pubkeys(&wallet_owner, &mint) + .token_2022_first() + .into_iter() + .flatten() + .collect::>() + }) + .unwrap_or_default(); + let mut pubkeys_to_clone = + Vec::with_capacity(1 + greedy_ata_pubkeys.len()); + pubkeys_to_clone.push(pubkey); + pubkeys_to_clone.extend(greedy_ata_pubkeys.iter().copied()); match self .fetch_and_clone_accounts_with_dedup( @@ -990,8 +993,10 @@ where true } } else { - let cloned_ata_pubkey = - greedy_ata_pubkey.filter(|ata_pubkey| { + let cloned_ata_pubkey = greedy_ata_pubkeys + .iter() + .copied() + .find(|ata_pubkey| { self.accounts_bank .get_account(ata_pubkey) .is_some_and(|account_in_bank| { @@ -1334,51 +1339,18 @@ where deleg_record: Option<&DelegationRecord>, delegation_actions: &DelegationActions, ) -> Option { - let deleg_record = deleg_record?; - - if deleg_record.authority != self.validator_pubkey { - return None; - } - let (wallet_owner, mint) = delegation::parse_raw_eata_pda( - &eata_pubkey, - eata_account.data(), - deleg_record.owner, - )?; - let (ata_pubkey, _) = - try_derive_ata_address_and_bump(&wallet_owner, &mint)?; - - let projected_ata = self.maybe_project_delegated_ata_from_eata( - // Intentional: in this subscription-update path there is no separate ATA, so - // maybe_project_delegated_ata_from_eata uses eata_account as ata_account; with deleg_record, - // ata_account only affects projected slot via max(), so passing eata_account twice is correct. - eata_account, + ata_projection::maybe_build_projected_ata_clone_request_from_eata_sub_update( + self, + eata_pubkey, eata_account, deleg_record, - )?; - - if let Some(in_bank_ata) = self.accounts_bank.get_account(&ata_pubkey) { - if in_bank_ata.delegated() || in_bank_ata.undelegating() { - return None; - } - if in_bank_ata.remote_slot() >= projected_ata.remote_slot() { - return None; - } - } - Some(AccountCloneRequest { - pubkey: ata_pubkey, - account: projected_ata, - commit_frequency_ms: None, - delegation_actions: delegation_actions.clone(), - delegated_to_other: None, - }) + delegation_actions, + ) } + #[cfg(test)] fn is_known_empty_eata(&self, eata_pubkey: &Pubkey) -> bool { - self.known_empty_eatas.lock().get(eata_pubkey).is_some() - } - - fn mark_eata_empty(&self, eata_pubkey: Pubkey) { - self.known_empty_eatas.lock().put(eata_pubkey, ()); + ata_projection::is_known_empty_eata(self, eata_pubkey) } async fn maybe_project_ata_from_subscription_update( @@ -1389,127 +1361,12 @@ where AccountSharedData, Option<(DelegationRecord, Option)>, ) { - let Some(ata_info) = is_ata(&ata_pubkey, &ata_account) else { - return (ata_account, None); - }; - - let Some((eata_pubkey, _)) = - try_derive_eata_address_and_bump(&ata_info.owner, &ata_info.mint) - else { - return (ata_account, None); - }; - - let was_watching = - self.remote_account_provider.is_watching(&eata_pubkey); - - // Ensure before cache checks; this keeps the subscription LRU warm - // without refcounting the projection reason on every ATA update. - let subscribed = match self - .ensure_subscription( - &eata_pubkey, - SubscriptionReason::AtaProjection, - ) - .await - { - Ok(()) => true, - Err(err) => { - warn!( - pubkey = %eata_pubkey, - error = ?err, - "Failed to subscribe to derived eATA" - ); - false - } - }; - - // Known-empty eATAs skip the fetch only if the subscription was already live. - if was_watching && subscribed && self.is_known_empty_eata(&eata_pubkey) - { - return (ata_account, None); - } - - let (eata_account, definitively_not_found) = match self - .remote_account_provider - .try_get_multi_until_slots_match( - &[eata_pubkey], - Some(MatchSlotsConfig { - min_context_slot: Some(ata_account.remote_slot()), - ..Default::default() - }), - AccountFetchOrigin::ProjectAta, - ) - .await - { - Ok(mut accounts) => { - let popped = accounts.pop(); - // Only `NotFound` proves absence; stale, missing, or failed fetches retry later. - let nf = matches!(popped, Some(RemoteAccount::NotFound(_))); - let fresh = popped.and_then(|a| a.fresh_account()); - (fresh, nf) - } - Err(err) => { - debug!( - pubkey = %eata_pubkey, - error = ?err, - "Failed to fetch eATA for projection" - ); - (None, false) - } - }; - - let Some(eata_account) = eata_account else { - // Cache absence only after a confirmed NotFound and live subscription. - if definitively_not_found && subscribed { - self.mark_eata_empty(eata_pubkey); - } - return (ata_account, None); - }; - - let deleg_record = delegation::fetch_and_parse_delegation_record( + ata_projection::maybe_project_ata_from_subscription_update( self, - eata_pubkey, - ata_account.remote_slot().max(eata_account.remote_slot()), - AccountFetchOrigin::ProjectAta, + ata_pubkey, + ata_account, ) - .await; - - let Some(deleg_record) = deleg_record else { - return (ata_account, None); - }; - let (deleg_record, delegation_actions) = deleg_record; - - if let Some(projected_ata) = self.maybe_project_delegated_ata_from_eata( - &ata_account, - &eata_account, - &deleg_record, - ) { - return (projected_ata, Some((deleg_record, delegation_actions))); - } - (ata_account, Some((deleg_record, delegation_actions))) - } - - fn maybe_project_delegated_ata_from_eata( - &self, - ata_account: &AccountSharedData, - eata_account: &AccountSharedData, - deleg_record: &DelegationRecord, - ) -> Option { - if deleg_record.authority != self.validator_pubkey { - return None; - } - - let mut projected_ata = - match eata_account.maybe_into_ata(deleg_record.owner) { - Some(projected_ata) => projected_ata, - None => { - return None; - } - }; - let projected_slot = - ata_account.remote_slot().max(eata_account.remote_slot()); - projected_ata.set_remote_slot(projected_slot); - projected_ata.set_delegated(true); - Some(projected_ata) + .await } /// Parses a delegation record from account data bytes. diff --git a/magicblock-chainlink/src/chainlink/fetch_cloner/tests.rs b/magicblock-chainlink/src/chainlink/fetch_cloner/tests.rs index 95a1462a5..e4f0be583 100644 --- a/magicblock-chainlink/src/chainlink/fetch_cloner/tests.rs +++ b/magicblock-chainlink/src/chainlink/fetch_cloner/tests.rs @@ -5,6 +5,13 @@ use solana_account::{Account, AccountSharedData, WritableAccount}; use solana_keypair::Keypair; use solana_sdk_ids::system_program; use solana_signer::Signer; +use spl_token_2022::{ + extension::{ + immutable_owner::ImmutableOwner, BaseStateWithExtensions, + StateWithExtensions, + }, + state::Account as Token2022Account, +}; use tokio::sync::mpsc; use super::*; @@ -28,8 +35,10 @@ use crate::{ add_invalid_delegation_record_for, delegation_record_to_vec, }, eatas::{ - create_ata_account, create_eata_account, derive_ata, derive_eata, - EATA_PROGRAM_ID, + create_ata_account, create_eata_account, + create_token_2022_ata_account, derive_ata, + derive_ata_with_token_program, derive_eata, EATA_PROGRAM_ID, + TOKEN_2022_PROGRAM_ID, }, init_logger, rpc_client_mock::{ChainRpcClientMock, ChainRpcClientMockBuilder}, @@ -187,6 +196,19 @@ where } } +fn insert_plain_ata_in_bank( + accounts_bank: &Arc, + ata_pubkey: Pubkey, + wallet_owner: &Pubkey, + mint: &Pubkey, + remote_slot: u64, +) { + let mut ata_account = + AccountSharedData::from(create_ata_account(wallet_owner, mint)); + ata_account.set_remote_slot(remote_slot); + accounts_bank.insert(ata_pubkey, ata_account); +} + fn create_non_raw_eata_owned_account( pubkey: Pubkey, data_len: usize, @@ -3136,6 +3158,13 @@ async fn test_out_of_order_delegated_eata_subscription_update_still_projects_ata in_bank_eata.set_owner(EATA_PROGRAM_ID); in_bank_eata.set_remote_slot(CURRENT_SLOT); accounts_bank.insert(eata_pubkey, in_bank_eata); + insert_plain_ata_in_bank( + &accounts_bank, + ata_pubkey, + &wallet_owner, + &mint, + CURRENT_SLOT, + ); use crate::remote_account_provider::{ RemoteAccount, RemoteAccountUpdateSource, @@ -3156,7 +3185,11 @@ async fn test_out_of_order_delegated_eata_subscription_update_still_projects_ata const POLL_INTERVAL: std::time::Duration = Duration::from_millis(10); const TIMEOUT: std::time::Duration = Duration::from_millis(500); tokio::time::timeout(TIMEOUT, async { - while accounts_bank.get_account(&ata_pubkey).is_none() { + while !accounts_bank + .get_account(&ata_pubkey) + .map(|account| account.delegated()) + .unwrap_or(false) + { tokio::time::sleep(POLL_INTERVAL).await; } }) @@ -3219,6 +3252,13 @@ async fn test_out_of_order_delegated_eata_update_clones_action_dependencies() { in_bank_eata.set_owner(EATA_PROGRAM_ID); in_bank_eata.set_remote_slot(CURRENT_SLOT); accounts_bank.insert(eata_pubkey, in_bank_eata); + insert_plain_ata_in_bank( + &accounts_bank, + ata_pubkey, + &wallet_owner, + &mint, + CURRENT_SLOT, + ); use crate::remote_account_provider::{ RemoteAccount, RemoteAccountUpdateSource, @@ -3240,10 +3280,13 @@ async fn test_out_of_order_delegated_eata_update_clones_action_dependencies() { const TIMEOUT: std::time::Duration = Duration::from_millis(500); tokio::time::timeout(TIMEOUT, async { loop { - let has_ata = accounts_bank.get_account(&ata_pubkey).is_some(); + let has_projected_ata = accounts_bank + .get_account(&ata_pubkey) + .map(|account| account.delegated()) + .unwrap_or(false); let has_action_program = accounts_bank.get_account(&action_program_pubkey).is_some(); - if has_ata && has_action_program { + if has_projected_ata && has_action_program { break; } tokio::time::sleep(POLL_INTERVAL).await; @@ -3392,6 +3435,13 @@ async fn test_delegated_eata_subscription_update_clones_raw_eata_and_projects_at validator_pubkey, EATA_PROGRAM_ID, ); + insert_plain_ata_in_bank( + &accounts_bank, + ata_pubkey, + &wallet_owner, + &mint, + CURRENT_SLOT, + ); use crate::remote_account_provider::{ RemoteAccount, RemoteAccountUpdateSource, @@ -3414,8 +3464,11 @@ async fn test_delegated_eata_subscription_update_clones_raw_eata_and_projects_at tokio::time::timeout(TIMEOUT, async { loop { let has_eata = accounts_bank.get_account(&eata_pubkey).is_some(); - let has_ata = accounts_bank.get_account(&ata_pubkey).is_some(); - if has_eata && has_ata { + let has_projected_ata = accounts_bank + .get_account(&ata_pubkey) + .map(|account| account.delegated()) + .unwrap_or(false); + if has_eata && has_projected_ata { break; } tokio::time::sleep(POLL_INTERVAL).await; @@ -3603,6 +3656,13 @@ async fn test_delegated_eata_subscription_update_clones_action_dependencies() { EATA_PROGRAM_ID, action_program_pubkey, ); + insert_plain_ata_in_bank( + &accounts_bank, + ata_pubkey, + &wallet_owner, + &mint, + CURRENT_SLOT, + ); use crate::remote_account_provider::{ RemoteAccount, RemoteAccountUpdateSource, @@ -3625,10 +3685,13 @@ async fn test_delegated_eata_subscription_update_clones_action_dependencies() { tokio::time::timeout(TIMEOUT, async { loop { let has_eata = accounts_bank.get_account(&eata_pubkey).is_some(); - let has_ata = accounts_bank.get_account(&ata_pubkey).is_some(); + let has_projected_ata = accounts_bank + .get_account(&ata_pubkey) + .map(|account| account.delegated()) + .unwrap_or(false); let has_action_program = accounts_bank.get_account(&action_program_pubkey).is_some(); - if has_eata && has_ata && has_action_program { + if has_eata && has_projected_ata && has_action_program { break; } tokio::time::sleep(POLL_INTERVAL).await; @@ -3660,6 +3723,7 @@ async fn test_projected_ata_clone_request_from_eata_update_keeps_actions() { let eata_account = create_eata_account(&wallet_owner, &mint, 777, true); let FetcherTestCtx { + accounts_bank, fetch_cloner, rpc_client, .. @@ -3677,6 +3741,13 @@ async fn test_projected_ata_clone_request_from_eata_update_keeps_actions() { EATA_PROGRAM_ID, action_program_pubkey, ); + insert_plain_ata_in_bank( + &accounts_bank, + ata_pubkey, + &wallet_owner, + &mint, + CURRENT_SLOT, + ); let (deleg_record, delegation_actions) = fetch_cloner .fetch_and_parse_delegation_record( @@ -3710,6 +3781,63 @@ async fn test_projected_ata_clone_request_from_eata_update_keeps_actions() { ); } +#[tokio::test] +async fn test_projected_ata_clone_request_from_eata_update_requires_ata_in_bank( +) { + init_logger(); + let validator_keypair = Keypair::new(); + let validator_pubkey = validator_keypair.pubkey(); + let wallet_owner = random_pubkey(); + let mint = random_pubkey(); + const CURRENT_SLOT: u64 = 100; + + let eata_pubkey = derive_eata(&wallet_owner, &mint); + let eata_account = create_eata_account(&wallet_owner, &mint, 777, true); + + let FetcherTestCtx { + fetch_cloner, + rpc_client, + .. + } = setup( + [(eata_pubkey, eata_account.clone())], + CURRENT_SLOT, + validator_keypair.insecure_clone(), + ) + .await; + + add_delegation_record_for( + &rpc_client, + eata_pubkey, + validator_pubkey, + EATA_PROGRAM_ID, + ); + + let (deleg_record, _) = fetch_cloner + .fetch_and_parse_delegation_record( + eata_pubkey, + CURRENT_SLOT, + AccountFetchOrigin::GetAccount, + ) + .await + .expect("delegation record should resolve"); + + let mut eata_shared = AccountSharedData::from(eata_account); + eata_shared.set_remote_slot(CURRENT_SLOT); + + let projected_ata_request = fetch_cloner + .maybe_build_projected_ata_clone_request_from_eata_sub_update( + eata_pubkey, + &eata_shared, + Some(&deleg_record), + &DelegationActions::default(), + ); + + assert!( + projected_ata_request.is_none(), + "delegated eATA updates should not synthesize a projected ATA without an ATA already in the bank", + ); +} + #[tokio::test] async fn test_fetch_and_parse_delegation_record_releases_direct_ref_when_already_watched( ) { @@ -3864,6 +3992,327 @@ async fn test_delegated_eata_update_does_not_override_delegated_ata_in_bank() { ); } +#[tokio::test] +async fn test_delegated_eata_update_projects_existing_plain_ata_in_bank() { + init_logger(); + let validator_keypair = Keypair::new(); + let validator_pubkey = validator_keypair.pubkey(); + let wallet_owner = random_pubkey(); + let mint = random_pubkey(); + const EATA_SLOT: u64 = 100; + const PLAIN_ATA_SLOT: u64 = EATA_SLOT + 5; + const EATA_AMOUNT: u64 = 777; + const PLAIN_ATA_AMOUNT: u64 = 999; + + let eata_pubkey = derive_eata(&wallet_owner, &mint); + let ata_pubkey = derive_ata(&wallet_owner, &mint); + let eata_account = + create_eata_account(&wallet_owner, &mint, EATA_AMOUNT, true); + + let FetcherTestCtx { + accounts_bank, + rpc_client, + subscription_tx, + .. + } = setup( + [(eata_pubkey, eata_account.clone())], + EATA_SLOT, + validator_keypair.insecure_clone(), + ) + .await; + + add_delegation_record_for( + &rpc_client, + eata_pubkey, + validator_pubkey, + EATA_PROGRAM_ID, + ); + + let mut plain_ata = create_ata_account(&wallet_owner, &mint); + plain_ata.data[64..72].copy_from_slice(&PLAIN_ATA_AMOUNT.to_le_bytes()); + let mut plain_ata_shared = AccountSharedData::from(plain_ata); + plain_ata_shared.set_remote_slot(PLAIN_ATA_SLOT); + accounts_bank.insert(ata_pubkey, plain_ata_shared); + + use crate::remote_account_provider::{ + RemoteAccount, RemoteAccountUpdateSource, + }; + + subscription_tx + .send(ForwardedSubscriptionUpdate { + pubkey: eata_pubkey, + account: RemoteAccount::from_fresh_account( + eata_account, + EATA_SLOT, + RemoteAccountUpdateSource::Subscription, + ), + }) + .await + .unwrap(); + + const POLL_INTERVAL: std::time::Duration = Duration::from_millis(10); + const TIMEOUT: std::time::Duration = Duration::from_millis(500); + tokio::time::timeout(TIMEOUT, async { + loop { + if accounts_bank + .get_account(&ata_pubkey) + .is_some_and(|account| account.delegated()) + { + break; + } + tokio::time::sleep(POLL_INTERVAL).await; + } + }) + .await + .expect("timed out waiting for existing ATA projection"); + + let projected_ata = accounts_bank + .get_account(&ata_pubkey) + .expect("ATA should exist in bank"); + assert!(projected_ata.delegated()); + assert_eq!( + projected_ata.remote_slot(), + PLAIN_ATA_SLOT, + "Projected ATA should preserve the freshest source slot", + ); + + let ata_data = projected_ata.data(); + let projected_mint = + Pubkey::new_from_array(ata_data[0..32].try_into().unwrap()); + let projected_owner = + Pubkey::new_from_array(ata_data[32..64].try_into().unwrap()); + let projected_amount = + u64::from_le_bytes(ata_data[64..72].try_into().unwrap()); + assert_eq!(projected_mint, mint); + assert_eq!(projected_owner, wallet_owner); + assert_eq!(projected_amount, EATA_AMOUNT); +} + +#[tokio::test] +async fn test_delegated_eata_update_projects_existing_token_2022_ata_in_bank() { + init_logger(); + let validator_keypair = Keypair::new(); + let validator_pubkey = validator_keypair.pubkey(); + let wallet_owner = random_pubkey(); + let mint = random_pubkey(); + const EATA_SLOT: u64 = 100; + const PLAIN_ATA_SLOT: u64 = EATA_SLOT + 5; + const EATA_AMOUNT: u64 = 777; + const PLAIN_ATA_AMOUNT: u64 = 999; + const LEGACY_ATA_AMOUNT: u64 = 555; + + let eata_pubkey = derive_eata(&wallet_owner, &mint); + let legacy_ata_pubkey = derive_ata(&wallet_owner, &mint); + let token_2022_ata_pubkey = derive_ata_with_token_program( + &wallet_owner, + &mint, + &TOKEN_2022_PROGRAM_ID, + ); + let eata_account = + create_eata_account(&wallet_owner, &mint, EATA_AMOUNT, true); + + let FetcherTestCtx { + accounts_bank, + rpc_client, + subscription_tx, + .. + } = setup( + [(eata_pubkey, eata_account.clone())], + EATA_SLOT, + validator_keypair.insecure_clone(), + ) + .await; + + add_delegation_record_for( + &rpc_client, + eata_pubkey, + validator_pubkey, + EATA_PROGRAM_ID, + ); + + let mut plain_ata = create_token_2022_ata_account(&wallet_owner, &mint); + plain_ata.data[64..72].copy_from_slice(&PLAIN_ATA_AMOUNT.to_le_bytes()); + let expected_len = plain_ata.data.len(); + let mut plain_ata_shared = AccountSharedData::from(plain_ata); + plain_ata_shared.set_remote_slot(PLAIN_ATA_SLOT); + accounts_bank.insert(token_2022_ata_pubkey, plain_ata_shared); + + let mut legacy_ata = create_ata_account(&wallet_owner, &mint); + legacy_ata.data[64..72].copy_from_slice(&LEGACY_ATA_AMOUNT.to_le_bytes()); + let mut legacy_ata_shared = AccountSharedData::from(legacy_ata); + legacy_ata_shared.set_remote_slot(PLAIN_ATA_SLOT + 1); + accounts_bank.insert(legacy_ata_pubkey, legacy_ata_shared); + + use crate::remote_account_provider::{ + RemoteAccount, RemoteAccountUpdateSource, + }; + + subscription_tx + .send(ForwardedSubscriptionUpdate { + pubkey: eata_pubkey, + account: RemoteAccount::from_fresh_account( + eata_account, + EATA_SLOT, + RemoteAccountUpdateSource::Subscription, + ), + }) + .await + .unwrap(); + + const POLL_INTERVAL: std::time::Duration = Duration::from_millis(10); + const TIMEOUT: std::time::Duration = Duration::from_millis(500); + tokio::time::timeout(TIMEOUT, async { + loop { + if accounts_bank + .get_account(&token_2022_ata_pubkey) + .is_some_and(|account| account.delegated()) + { + break; + } + tokio::time::sleep(POLL_INTERVAL).await; + } + }) + .await + .expect("timed out waiting for existing Token-2022 ATA projection"); + + let projected_ata = accounts_bank + .get_account(&token_2022_ata_pubkey) + .expect("Token-2022 ATA should exist in bank"); + assert!(projected_ata.delegated()); + assert_eq!(*projected_ata.owner(), TOKEN_2022_PROGRAM_ID); + assert_eq!(projected_ata.data().len(), expected_len); + assert_eq!( + projected_ata.remote_slot(), + PLAIN_ATA_SLOT, + "Projected ATA should preserve the freshest source slot", + ); + + let ata_data = projected_ata.data(); + let projected_mint = + Pubkey::new_from_array(ata_data[0..32].try_into().unwrap()); + let projected_owner = + Pubkey::new_from_array(ata_data[32..64].try_into().unwrap()); + let projected_amount = + u64::from_le_bytes(ata_data[64..72].try_into().unwrap()); + let projected_token_account = + StateWithExtensions::::unpack(ata_data) + .expect("unpack projected Token-2022 ATA"); + assert_eq!(projected_mint, mint); + assert_eq!(projected_owner, wallet_owner); + assert_eq!(projected_amount, EATA_AMOUNT); + projected_token_account + .get_extension::() + .expect("projected Token-2022 ATA preserves ImmutableOwner"); + let legacy_ata = accounts_bank + .get_account(&legacy_ata_pubkey) + .expect("legacy ATA should remain in bank"); + let legacy_amount = + u64::from_le_bytes(legacy_ata.data()[64..72].try_into().unwrap()); + assert!(!legacy_ata.delegated()); + assert_eq!(legacy_amount, LEGACY_ATA_AMOUNT); +} + +#[tokio::test] +async fn test_greedy_delegated_eata_update_projects_remote_token_2022_ata() { + init_logger(); + let validator_keypair = Keypair::new(); + let validator_pubkey = validator_keypair.pubkey(); + let wallet_owner = random_pubkey(); + let mint = random_pubkey(); + const EATA_SLOT: u64 = 100; + const EATA_AMOUNT: u64 = 777; + + let eata_pubkey = derive_eata(&wallet_owner, &mint); + let legacy_ata_pubkey = derive_ata(&wallet_owner, &mint); + let token_2022_ata_pubkey = derive_ata_with_token_program( + &wallet_owner, + &mint, + &TOKEN_2022_PROGRAM_ID, + ); + let eata_account = + create_eata_account(&wallet_owner, &mint, EATA_AMOUNT, true); + let token_2022_ata_account = + create_token_2022_ata_account(&wallet_owner, &mint); + let expected_len = token_2022_ata_account.data.len(); + + let FetcherTestCtx { + accounts_bank, + rpc_client, + subscription_tx, + .. + } = setup( + [ + (eata_pubkey, eata_account.clone()), + (token_2022_ata_pubkey, token_2022_ata_account), + ], + EATA_SLOT, + validator_keypair.insecure_clone(), + ) + .await; + + add_delegation_record_for( + &rpc_client, + eata_pubkey, + validator_pubkey, + EATA_PROGRAM_ID, + ); + + use crate::remote_account_provider::{ + RemoteAccount, RemoteAccountUpdateSource, + }; + + subscription_tx + .send(ForwardedSubscriptionUpdate { + pubkey: eata_pubkey, + account: RemoteAccount::from_fresh_account( + eata_account, + EATA_SLOT, + RemoteAccountUpdateSource::Subscription, + ), + }) + .await + .unwrap(); + + const POLL_INTERVAL: std::time::Duration = Duration::from_millis(10); + const TIMEOUT: std::time::Duration = Duration::from_millis(500); + tokio::time::timeout(TIMEOUT, async { + loop { + if accounts_bank + .get_account(&token_2022_ata_pubkey) + .is_some_and(|account| account.delegated()) + { + break; + } + tokio::time::sleep(POLL_INTERVAL).await; + } + }) + .await + .expect("timed out waiting for greedy Token-2022 ATA projection"); + + let projected_ata = accounts_bank + .get_account(&token_2022_ata_pubkey) + .expect("Token-2022 ATA should be projected"); + assert!(projected_ata.delegated()); + assert_eq!(*projected_ata.owner(), TOKEN_2022_PROGRAM_ID); + assert_eq!(projected_ata.data().len(), expected_len); + + let ata_data = projected_ata.data(); + let projected_mint = + Pubkey::new_from_array(ata_data[0..32].try_into().unwrap()); + let projected_owner = + Pubkey::new_from_array(ata_data[32..64].try_into().unwrap()); + let projected_amount = + u64::from_le_bytes(ata_data[64..72].try_into().unwrap()); + assert_eq!(projected_mint, mint); + assert_eq!(projected_owner, wallet_owner); + assert_eq!(projected_amount, EATA_AMOUNT); + + assert!( + accounts_bank.get_account(&legacy_ata_pubkey).is_none(), + "Token-2022 eATA projection must not synthesize a legacy ATA" + ); +} + #[tokio::test] async fn test_fetch_subscription_race_duplicate_clone() { // This test validates that pending clone ownership prevents duplicate @@ -4283,6 +4732,78 @@ async fn test_ata_projection_releases_ata_direct_ref_after_fetch() { ); } +#[tokio::test] +async fn test_token_2022_ata_projection_preserves_token_program_and_layout() { + init_logger(); + let validator_keypair = Keypair::new(); + let validator_pubkey = validator_keypair.pubkey(); + let wallet_owner = random_pubkey(); + let mint = random_pubkey(); + const CURRENT_SLOT: u64 = 100; + const AMOUNT: u64 = 777; + + let ata_pubkey = derive_ata_with_token_program( + &wallet_owner, + &mint, + &TOKEN_2022_PROGRAM_ID, + ); + let eata_pubkey = derive_eata(&wallet_owner, &mint); + let ata_account = create_token_2022_ata_account(&wallet_owner, &mint); + let eata_account = create_eata_account(&wallet_owner, &mint, AMOUNT, true); + let expected_len = ata_account.data.len(); + + let FetcherTestCtx { + accounts_bank, + fetch_cloner, + rpc_client, + .. + } = setup( + [(ata_pubkey, ata_account), (eata_pubkey, eata_account)], + CURRENT_SLOT, + validator_keypair.insecure_clone(), + ) + .await; + + add_delegation_record_with_slot_for( + &rpc_client, + eata_pubkey, + validator_pubkey, + EATA_PROGRAM_ID, + CURRENT_SLOT + 1, + ); + + let result = fetch_cloner + .fetch_and_clone_accounts_with_dedup( + &[ata_pubkey], + None, + None, + AccountFetchOrigin::GetAccount, + None, + ) + .await + .expect("Token-2022 ATA projection fetch should not fail"); + assert!(result.is_ok(), "Token-2022 ATA projection should succeed"); + + let projected_ata = accounts_bank + .get_account(&ata_pubkey) + .expect("Token-2022 ATA should be projected"); + assert!(projected_ata.delegated()); + assert_eq!(*projected_ata.owner(), TOKEN_2022_PROGRAM_ID); + assert_eq!(projected_ata.data().len(), expected_len); + assert_eq!(projected_ata.remote_slot(), CURRENT_SLOT); + + let ata_data = projected_ata.data(); + let projected_mint = + Pubkey::new_from_array(ata_data[0..32].try_into().unwrap()); + let projected_owner = + Pubkey::new_from_array(ata_data[32..64].try_into().unwrap()); + let projected_amount = + u64::from_le_bytes(ata_data[64..72].try_into().unwrap()); + assert_eq!(projected_mint, mint); + assert_eq!(projected_owner, wallet_owner); + assert_eq!(projected_amount, AMOUNT); +} + #[tokio::test] async fn test_fetch_keeps_undelegating_projected_ata_in_bank() { init_logger(); diff --git a/magicblock-chainlink/src/testing/eatas.rs b/magicblock-chainlink/src/testing/eatas.rs index ac3f64e93..f3cc1c3eb 100644 --- a/magicblock-chainlink/src/testing/eatas.rs +++ b/magicblock-chainlink/src/testing/eatas.rs @@ -1,11 +1,19 @@ pub use magicblock_core::token_programs::{ - derive_ata, derive_eata, EphemeralAta, EATA_PROGRAM_ID, TOKEN_PROGRAM_ID, + derive_ata, derive_ata_with_token_program, derive_eata, EphemeralAta, + EATA_PROGRAM_ID, TOKEN_2022_PROGRAM_ID, TOKEN_PROGRAM_ID, }; use solana_account::Account; use solana_program::{program_option::COption, program_pack::Pack}; use solana_pubkey::Pubkey; use solana_rent::Rent; use spl_token::state::{Account as SplAccount, AccountState}; +use spl_token_2022::{ + extension::{ + set_account_type, BaseStateWithExtensionsMut, ExtensionType, + StateWithExtensionsMut, + }, + state::Account as Token2022Account, +}; /// Creates a test ATA (Associated Token Account) with initialized state and zero balance. /// @@ -13,6 +21,71 @@ use spl_token::state::{Account as SplAccount, AccountState}; /// * `owner` - The public key of the account owner /// * `mint` - The public key of the token mint pub fn create_ata_account(owner: &Pubkey, mint: &Pubkey) -> Account { + create_ata_account_with_token_program( + owner, + mint, + TOKEN_PROGRAM_ID, + SplAccount::LEN, + ) +} + +pub fn create_token_2022_ata_account(owner: &Pubkey, mint: &Pubkey) -> Account { + // Default Token-2022 ATAs carry ImmutableOwner; mint-specific extensions can require more. + create_token_2022_ata_account_with_extensions( + owner, + mint, + &[ExtensionType::ImmutableOwner], + ) +} + +pub fn create_token_2022_ata_account_with_extensions( + owner: &Pubkey, + mint: &Pubkey, + account_extensions: &[ExtensionType], +) -> Account { + // Token-2022 account length depends on account extensions required by the mint. + let data_len = + ExtensionType::try_calculate_account_len::( + account_extensions, + ) + .expect("calculate Token-2022 account length"); + let mut account = create_ata_account_with_token_program( + owner, + mint, + TOKEN_2022_PROGRAM_ID, + data_len, + ); + initialize_token_2022_account_extensions( + &mut account.data, + account_extensions, + ); + account +} + +fn initialize_token_2022_account_extensions( + data: &mut [u8], + account_extensions: &[ExtensionType], +) { + if account_extensions.is_empty() { + return; + } + set_account_type::(data) + .expect("set Token-2022 account type"); + let mut state = StateWithExtensionsMut::::unpack(data) + .expect("unpack Token-2022 account"); + for extension_type in account_extensions { + state + .init_account_extension_from_type(*extension_type) + .expect("initialize Token-2022 account extension"); + } +} + +fn create_ata_account_with_token_program( + owner: &Pubkey, + mint: &Pubkey, + token_program: Pubkey, + data_len: usize, +) -> Account { let token_account = SplAccount { mint: *mint, owner: *owner, @@ -24,12 +97,15 @@ pub fn create_ata_account(owner: &Pubkey, mint: &Pubkey) -> Account { close_authority: COption::None, }; - let mut data = vec![0u8; SplAccount::LEN]; - SplAccount::pack(token_account, &mut data).expect("pack spl token account"); + let mut packed = vec![0u8; SplAccount::LEN]; + SplAccount::pack(token_account, &mut packed) + .expect("pack spl token account"); + let mut data = vec![0u8; data_len.max(SplAccount::LEN)]; + data[..SplAccount::LEN].copy_from_slice(&packed); let lamports = Rent::default().minimum_balance(data.len()); Account { - owner: TOKEN_PROGRAM_ID, + owner: token_program, data, lamports, executable: false, @@ -43,11 +119,19 @@ pub fn create_eata_account( amount: u64, delegate: bool, ) -> Account { - let mut data = Vec::with_capacity(64 + 8); - data.extend_from_slice(owner.as_ref()); - data.extend_from_slice(mint.as_ref()); - data.extend_from_slice(&amount.to_le_bytes()); - let lamports = Rent::default().minimum_balance(data.len()); + let bump = + magicblock_core::token_programs::try_derive_eata_address_and_bump( + owner, mint, + ) + .expect("derive eATA") + .1; + let eata_account: Account = EphemeralAta { + owner: *owner, + mint: *mint, + amount, + bump, + } + .into(); let account_owner = if delegate { dlp_api::id() @@ -57,8 +141,8 @@ pub fn create_eata_account( Account { owner: account_owner, - data, - lamports, + data: eata_account.data, + lamports: eata_account.lamports, ..Default::default() } } diff --git a/magicblock-core/src/token_programs.rs b/magicblock-core/src/token_programs.rs index 6218a6dfa..6e2184534 100644 --- a/magicblock-core/src/token_programs.rs +++ b/magicblock-core/src/token_programs.rs @@ -1,4 +1,6 @@ -use solana_account::{Account, AccountSharedData, ReadableAccount}; +use solana_account::{ + Account, AccountSharedData, ReadableAccount, WritableAccount, +}; use solana_program::{ program_error::ProgramError, program_option::COption, program_pack::Pack, rent::Rent, @@ -12,6 +14,10 @@ use spl_token::state::{Account as SplAccount, AccountState}; pub const TOKEN_PROGRAM_ID: Pubkey = pubkey!("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"); +// Token-2022 Program ID (Tokenz...) +pub const TOKEN_2022_PROGRAM_ID: Pubkey = + pubkey!("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb"); + // Associated Token Account Program ID (ATokenG...) pub const ASSOCIATED_TOKEN_PROGRAM_ID: Pubkey = pubkey!("ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL"); @@ -20,6 +26,9 @@ pub const ASSOCIATED_TOKEN_PROGRAM_ID: Pubkey = pub const EATA_PROGRAM_ID: Pubkey = pubkey!("SPLxh1LVZzEkX99H6rqYizhytLWPZVV296zyYDPagv2"); +pub const EPHEMERAL_ATA_LEN: usize = 80; +const LEGACY_EPHEMERAL_ATA_LEN: usize = 72; + /// Derives the standard Associated Token Account (ATA) address for the given wallet owner and token mint. /// /// # Arguments @@ -29,8 +38,16 @@ pub const EATA_PROGRAM_ID: Pubkey = /// # Returns /// The derived ATA address as `Pubkey`. pub fn derive_ata(owner: &Pubkey, mint: &Pubkey) -> Pubkey { + derive_ata_with_token_program(owner, mint, &TOKEN_PROGRAM_ID) +} + +pub fn derive_ata_with_token_program( + owner: &Pubkey, + mint: &Pubkey, + token_program: &Pubkey, +) -> Pubkey { Pubkey::find_program_address( - &[owner.as_ref(), TOKEN_PROGRAM_ID.as_ref(), mint.as_ref()], + &[owner.as_ref(), token_program.as_ref(), mint.as_ref()], &ASSOCIATED_TOKEN_PROGRAM_ID, ) .0 @@ -47,13 +64,62 @@ pub fn derive_ata(owner: &Pubkey, mint: &Pubkey) -> Pubkey { pub fn try_derive_ata_address_and_bump( owner: &Pubkey, mint: &Pubkey, +) -> Option<(Pubkey, u8)> { + try_derive_ata_address_and_bump_with_token_program( + owner, + mint, + &TOKEN_PROGRAM_ID, + ) +} + +pub fn try_derive_ata_address_and_bump_with_token_program( + owner: &Pubkey, + mint: &Pubkey, + token_program: &Pubkey, ) -> Option<(Pubkey, u8)> { Pubkey::try_find_program_address( - &[owner.as_ref(), TOKEN_PROGRAM_ID.as_ref(), mint.as_ref()], + &[owner.as_ref(), token_program.as_ref(), mint.as_ref()], &ASSOCIATED_TOKEN_PROGRAM_ID, ) } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SupportedAtaPubkeys { + pub legacy: Option, + pub token_2022: Option, +} + +impl SupportedAtaPubkeys { + pub fn token_2022_first(&self) -> [Option; 2] { + [self.token_2022, self.legacy] + } + + pub fn contains(&self, pubkey: &Pubkey) -> bool { + self.legacy.as_ref() == Some(pubkey) + || self.token_2022.as_ref() == Some(pubkey) + } +} + +pub fn try_derive_supported_ata_pubkeys( + owner: &Pubkey, + mint: &Pubkey, +) -> SupportedAtaPubkeys { + SupportedAtaPubkeys { + legacy: try_derive_ata_address_and_bump_with_token_program( + owner, + mint, + &TOKEN_PROGRAM_ID, + ) + .map(|(pubkey, _)| pubkey), + token_2022: try_derive_ata_address_and_bump_with_token_program( + owner, + mint, + &TOKEN_2022_PROGRAM_ID, + ) + .map(|(pubkey, _)| pubkey), + } +} + /// Derives the Enhanced Associated Token Account (eATA) Program Derived Address (PDA) for the given wallet owner and token mint. /// /// # Arguments @@ -107,7 +173,7 @@ pub fn is_ata( // The account must be owned by the SPL Token program (legacy) or Token-2022 let token_program_owner = account.owner(); let is_spl_token = *token_program_owner == spl_token::id(); - let is_token_2022 = *token_program_owner == spl_token_2022::id(); + let is_token_2022 = *token_program_owner == TOKEN_2022_PROGRAM_ID; if !(is_spl_token || is_token_2022) { return None; } @@ -155,7 +221,10 @@ pub fn try_remap_ata_to_eata( pubkey: &Pubkey, account: &AccountSharedData, ) -> Option<(Pubkey, EphemeralAta)> { - if account.owner() != &TOKEN_PROGRAM_ID || !account.delegated() { + let token_program_owner = account.owner(); + let is_spl_token = *token_program_owner == TOKEN_PROGRAM_ID; + let is_token_2022 = *token_program_owner == TOKEN_2022_PROGRAM_ID; + if !(is_spl_token || is_token_2022) || !account.delegated() { return None; } @@ -168,7 +237,8 @@ pub fn try_remap_ata_to_eata( let owner = Pubkey::new_from_array(data[32..64].try_into().ok()?); let amount = u64::from_le_bytes(data[64..72].try_into().ok()?); - let ata = derive_ata(&owner, &mint); + let (eata_pubkey, bump) = try_derive_eata_address_and_bump(&owner, &mint)?; + let ata = derive_ata_with_token_program(&owner, &mint, token_program_owner); if ata != *pubkey { return None; } @@ -177,9 +247,10 @@ pub fn try_remap_ata_to_eata( owner, mint, amount, + bump, }; - Some((derive_eata(&owner, &mint), eata)) + Some((eata_pubkey, eata)) } // ---------------- eATA -> ATA projection helpers ---------------- @@ -204,6 +275,60 @@ pub struct EphemeralAta { pub mint: Pubkey, /// The amount of tokens this account holds. pub amount: u64, + /// The bump of the eATA PDA. + pub bump: u8, +} + +impl EphemeralAta { + pub fn try_from_account_data(data: &[u8]) -> Option { + let owner = Pubkey::new_from_array(data.get(0..32)?.try_into().ok()?); + let mint = Pubkey::new_from_array(data.get(32..64)?.try_into().ok()?); + if mint == Pubkey::default() { + return None; + } + let amount = u64::from_le_bytes(data.get(64..72)?.try_into().ok()?); + let bump = match data.len() { + EPHEMERAL_ATA_LEN => data[72], + LEGACY_EPHEMERAL_ATA_LEN => { + try_derive_eata_address_and_bump(&owner, &mint)?.1 + } + _ => return None, + }; + + Some(Self { + owner, + mint, + amount, + bump, + }) + } + + pub fn project_into_ata_account( + &self, + ata_account: &AccountSharedData, + ) -> Option { + let token_program_owner = ata_account.owner(); + let is_spl_token = *token_program_owner == TOKEN_PROGRAM_ID; + let is_token_2022 = *token_program_owner == TOKEN_2022_PROGRAM_ID; + if !(is_spl_token || is_token_2022) { + return None; + } + + let data = ata_account.data(); + if data.len() < 72 { + return None; + } + if &data[0..32] != self.mint.as_ref() + || &data[32..64] != self.owner.as_ref() + { + return None; + } + + let mut projected = ata_account.clone(); + projected.data_as_mut_slice()[64..72] + .copy_from_slice(&self.amount.to_le_bytes()); + Some(projected) + } } impl TryFrom for AccountSharedData { @@ -240,11 +365,12 @@ impl TryFrom for AccountSharedData { impl From for Account { fn from(val: EphemeralAta) -> Self { - // Encode as: owner(32) | mint(32) | amount(8) - let mut data = Vec::with_capacity(72); + let mut data = Vec::with_capacity(EPHEMERAL_ATA_LEN); data.extend_from_slice(val.owner.as_ref()); data.extend_from_slice(val.mint.as_ref()); data.extend_from_slice(&val.amount.to_le_bytes()); + data.push(val.bump); + data.extend_from_slice(&[0; 7]); Account { lamports: Rent::default().minimum_balance(data.len()), @@ -265,19 +391,7 @@ impl MaybeIntoAta for AccountSharedData { return None; } - let data = self.data(); - // Expect at least owner(32) + mint(32) + amount(8) - if data.len() < 72 { - return None; - } - let owner = Pubkey::new_from_array(data[0..32].try_into().ok()?); - let mint = Pubkey::new_from_array(data[32..64].try_into().ok()?); - let amount = u64::from_le_bytes(data[64..72].try_into().ok()?); - let eata = EphemeralAta { - owner, - mint, - amount, - }; + let eata = EphemeralAta::try_from_account_data(self.data())?; eata.try_into().ok() } } diff --git a/programs/magicblock/src/clone_account/process_clone.rs b/programs/magicblock/src/clone_account/process_clone.rs index 578dbd490..8c2fafd93 100644 --- a/programs/magicblock/src/clone_account/process_clone.rs +++ b/programs/magicblock/src/clone_account/process_clone.rs @@ -66,11 +66,6 @@ pub(crate) fn process_clone_account( "CloneAccount: actions_tx_sig={}", actions_tx_sig ); - } else { - ic_msg!( - invoke_context, - "CloneAccount did not receive actions_tx_sig" - ); } let current_lamports = account.lamports(); diff --git a/programs/magicblock/src/magic_scheduled_base_intent.rs b/programs/magicblock/src/magic_scheduled_base_intent.rs index e399c52a3..88050bf62 100644 --- a/programs/magicblock/src/magic_scheduled_base_intent.rs +++ b/programs/magicblock/src/magic_scheduled_base_intent.rs @@ -2,7 +2,9 @@ use std::collections::{HashMap, HashSet}; use magicblock_core::{ intent::{BaseActionCallback, CommittedAccount}, - token_programs::{EATA_PROGRAM_ID, TOKEN_PROGRAM_ID}, + token_programs::{ + EATA_PROGRAM_ID, TOKEN_2022_PROGRAM_ID, TOKEN_PROGRAM_ID, + }, Slot, }; use magicblock_magic_program_api::args::{ @@ -1193,9 +1195,10 @@ pub(crate) fn validate_commit_schedule_permissions( signers: &HashSet, ) -> Result<(), InstructionError> { let validator_id = effective_validator_authority_id(); - let is_eata_token_program_call = parent_program_id - == Some(&EATA_PROGRAM_ID) - && committee_owner == &TOKEN_PROGRAM_ID; + let is_token_account_owner = committee_owner == &TOKEN_PROGRAM_ID + || committee_owner == &TOKEN_2022_PROGRAM_ID; + let is_eata_token_program_call = + parent_program_id == Some(&EATA_PROGRAM_ID) && is_token_account_owner; if parent_program_id != Some(committee_owner) && !signers.contains(committee_pubkey) && !signers.contains(&validator_id) diff --git a/programs/magicblock/src/schedule_transactions/process_schedule_commit_tests.rs b/programs/magicblock/src/schedule_transactions/process_schedule_commit_tests.rs index 3aa357ea1..eeb724135 100644 --- a/programs/magicblock/src/schedule_transactions/process_schedule_commit_tests.rs +++ b/programs/magicblock/src/schedule_transactions/process_schedule_commit_tests.rs @@ -276,8 +276,13 @@ mod tests { // ---------- Helpers for ATA/eATA remapping tests ---------- // Use shared SPL/ATA/eATA constants and helpers // Reuse test helper to create proper SPL ATA account data - use magicblock_chainlink::testing::eatas::create_ata_account; - use magicblock_core::token_programs::{derive_ata, derive_eata}; + use magicblock_chainlink::testing::eatas::{ + create_ata_account, create_token_2022_ata_account, + }; + use magicblock_core::token_programs::{ + derive_ata, derive_ata_with_token_program, derive_eata, + EATA_PROGRAM_ID, TOKEN_2022_PROGRAM_ID, + }; use serial_test::serial; use solana_seed_derivable::SeedDerivable; use test_kit::init_logger; @@ -295,6 +300,16 @@ mod tests { acc } + fn make_delegated_token_2022_ata_account( + owner: &Pubkey, + mint: &Pubkey, + ) -> AccountSharedData { + let ata_account = create_token_2022_ata_account(owner, mint); + let mut acc = AccountSharedData::from(ata_account); + acc.set_delegated(true); + acc + } + #[test] #[serial] fn test_schedule_commit_single_account_success() { @@ -651,6 +666,69 @@ mod tests { ); } + #[test] + #[serial] + fn test_schedule_commit_allows_token_2022_ata_from_eata_parent() { + init_logger!(); + + let payer = + Keypair::from_seed(b"schedule_commit_token_2022_ata_eata_parent") + .unwrap(); + let eata_parent_owned_committee = Pubkey::new_unique(); + let wallet_owner = Pubkey::new_unique(); + let mint = Pubkey::new_unique(); + let ata_pubkey = derive_ata_with_token_program( + &wallet_owner, + &mint, + &TOKEN_2022_PROGRAM_ID, + ); + let eata_pubkey = derive_eata(&wallet_owner, &mint); + + let (mut account_data, mut transaction_accounts) = + prepare_transaction_with_single_committee( + &payer, + EATA_PROGRAM_ID, + eata_parent_owned_committee, + ); + account_data.insert( + ata_pubkey, + make_delegated_token_2022_ata_account(&wallet_owner, &mint), + ); + + let ix = instruction_from_account_metas(vec![ + AccountMeta::new(payer.pubkey(), true), + AccountMeta::new(MAGIC_CONTEXT_PUBKEY, false), + AccountMeta::new_readonly(eata_parent_owned_committee, false), + AccountMeta::new_readonly(ata_pubkey, false), + ]); + extend_transaction_accounts_from_ix( + &ix, + &mut account_data, + &mut transaction_accounts, + ); + + let processed_scheduled = process_instruction( + ix.data.as_slice(), + transaction_accounts, + ix.accounts, + Ok(()), + ); + let magic_context_acc = assert_non_accepted_actions( + &processed_scheduled, + &payer.pubkey(), + 1, + ); + let magic_context = + bincode::deserialize::(magic_context_acc.data()) + .unwrap(); + let scheduled = &magic_context.scheduled_base_intents[0]; + + assert_eq!( + scheduled.get_all_committed_pubkeys(), + vec![eata_parent_owned_committee, eata_pubkey] + ); + } + #[test] #[serial] fn test_schedule_commit_and_undelegate_remaps_delegated_ata_to_eata() {