Fix: token2022 support

magicblock-chainlink/src/chainlink/fetch_cloner/ata_projection.rs

@@ -4,8 +4,9 @@ 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_ata_address_and_bump_with_token_program,
+    try_derive_eata_address_and_bump, AtaInfo, TOKEN_2022_PROGRAM_ID,
+    TOKEN_PROGRAM_ID,
 };
 use magicblock_metrics::metrics;
 use solana_account::{AccountSharedData, ReadableAccount};
@@ -58,16 +59,22 @@ 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 {
-        return None;
+    for token_program in [TOKEN_PROGRAM_ID, TOKEN_2022_PROGRAM_ID] {
+        if try_derive_ata_address_and_bump_with_token_program(
+            &wallet_owner,
+            &mint,
+            &token_program,
+        )
+        .is_some_and(|(derived_ata, _)| derived_ata == *ata_pubkey)
+        {
+            return Some(AtaInfo {
+                mint,
+                owner: wallet_owner,
+            });
+        }
     }
 
-    Some(AtaInfo {
-        mint,
-        owner: wallet_owner,
-    })
+    None
 }
 
 /// Resolves ATAs with eATA projection.

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<T, U, V, C>(

magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs

@@ -16,8 +16,10 @@ 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,
+    is_ata, try_derive_ata_address_and_bump,
+    try_derive_ata_address_and_bump_with_token_program,
+    try_derive_eata_address_and_bump, EphemeralAta, MaybeIntoAta,
+    EATA_PROGRAM_ID, TOKEN_2022_PROGRAM_ID, TOKEN_PROGRAM_ID,
 };
 use magicblock_metrics::metrics::{self, AccountFetchOrigin};
 use parking_lot::Mutex as PlMutex;
@@ -1344,26 +1346,43 @@ where
             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,
-            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;
+        let legacy_ata_pubkey =
+            try_derive_ata_address_and_bump_with_token_program(
+                &wallet_owner,
+                &mint,
+                &TOKEN_PROGRAM_ID,
+            )?
+            .0;
+        let token_2022_ata_pubkey =
+            try_derive_ata_address_and_bump_with_token_program(
+                &wallet_owner,
+                &mint,
+                &TOKEN_2022_PROGRAM_ID,
+            )?
+            .0;
+
+        let mut ata_pubkey = legacy_ata_pubkey;
+        let mut in_bank_ata = None;
+        for candidate_pubkey in [token_2022_ata_pubkey, legacy_ata_pubkey] {
+            if let Some(candidate_account) =
+                self.accounts_bank.get_account(&candidate_pubkey)
+            {
+                ata_pubkey = candidate_pubkey;
+                in_bank_ata = Some(candidate_account);
+                break;
             }
-            if in_bank_ata.remote_slot() >= projected_ata.remote_slot() {
+        }
+        if let Some(in_bank_ata) = &in_bank_ata {
+            if in_bank_ata.delegated() || in_bank_ata.undelegating() {
                 return None;
             }
         }
+        let ata_account = in_bank_ata.as_ref().unwrap_or(eata_account);
+        let projected_ata = self.maybe_project_delegated_ata_from_eata(
+            ata_account,
+            eata_account,
+            deleg_record,
+        )?;
         Some(AccountCloneRequest {
             pubkey: ata_pubkey,
             account: projected_ata,
@@ -1498,13 +1517,21 @@ where
             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_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);