tally_sdk/
ata.rs

1//! Associated Token Account (ATA) computation and token program detection utilities
2
3use crate::{error::Result, TallyError};
4use anchor_client::solana_client::rpc_client::RpcClient;
5use anchor_client::solana_sdk::commitment_config::CommitmentConfig;
6use anchor_client::solana_sdk::pubkey::Pubkey;
7use anchor_client::solana_sdk::{account::Account, program_pack::Pack};
8use spl_associated_token_account::get_associated_token_address;
9use spl_token::state::{Account as TokenAccount, Mint};
10
11/// Token program variants supported by the SDK
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum TokenProgram {
14    /// Classic SPL Token program
15    Token,
16    /// Token Extensions (Token-2022) program
17    Token2022,
18}
19
20impl TokenProgram {
21    /// Get the program ID for this token program variant
22    #[must_use]
23    pub const fn program_id(&self) -> Pubkey {
24        match self {
25            Self::Token => spl_token::id(),
26            Self::Token2022 => spl_token_2022::id(),
27        }
28    }
29}
30
31/// Get the associated token address for a wallet and mint
32///
33/// # Arguments
34/// * `wallet` - The wallet pubkey
35/// * `mint` - The token mint pubkey
36///
37/// # Returns
38/// * `Ok(Pubkey)` - The associated token address
39/// * `Err(TallyError)` - If computation fails
40pub fn get_associated_token_address_for_mint(wallet: &Pubkey, mint: &Pubkey) -> Result<Pubkey> {
41    Ok(get_associated_token_address(wallet, mint))
42}
43
44/// Get the associated token address with explicit token program
45///
46/// # Arguments
47/// * `wallet` - The wallet pubkey
48/// * `mint` - The token mint pubkey
49/// * `token_program` - The token program to use
50///
51/// # Returns
52/// * `Ok(Pubkey)` - The associated token address
53/// * `Err(TallyError)` - If computation fails
54pub fn get_associated_token_address_with_program(
55    wallet: &Pubkey,
56    mint: &Pubkey,
57    token_program: TokenProgram,
58) -> Result<Pubkey> {
59    Ok(
60        spl_associated_token_account::get_associated_token_address_with_program_id(
61            wallet,
62            mint,
63            &token_program.program_id(),
64        ),
65    )
66}
67
68/// Detect the token program used by a mint
69///
70/// # Arguments
71/// * `rpc_client` - RPC client for account queries
72/// * `mint` - The mint pubkey to check
73///
74/// # Returns
75/// * `Ok(TokenProgram)` - The detected token program
76/// * `Err(TallyError)` - If detection fails
77pub fn detect_token_program(rpc_client: &RpcClient, mint: &Pubkey) -> Result<TokenProgram> {
78    let account = rpc_client
79        .get_account_with_commitment(mint, CommitmentConfig::confirmed())
80        .map_err(|e| TallyError::Generic(format!("Failed to fetch mint account: {e}")))?
81        .value
82        .ok_or_else(|| TallyError::AccountNotFound(mint.to_string()))?;
83
84    // Check the owner to determine token program
85    if account.owner == spl_token::id() {
86        Ok(TokenProgram::Token)
87    } else if account.owner == spl_token_2022::id() {
88        Ok(TokenProgram::Token2022)
89    } else {
90        Err(TallyError::TokenProgramDetectionFailed {
91            mint: mint.to_string(),
92        })
93    }
94}
95
96/// Get mint information from SPL Token program
97///
98/// # Arguments
99/// * `account` - The mint account data
100///
101/// # Returns
102/// * `Ok(Mint)` - The parsed mint information
103/// * `Err(TallyError)` - If parsing fails
104pub fn parse_mint_account(account: &Account) -> Result<Mint> {
105    Mint::unpack(&account.data)
106        .map_err(|e| TallyError::Generic(format!("Failed to parse SPL Token mint: {e}")))
107}
108
109/// Get token account information from SPL Token program
110///
111/// # Arguments
112/// * `account` - The token account data
113///
114/// # Returns
115/// * `Ok(TokenAccount)` - The parsed token account information
116/// * `Err(TallyError)` - If parsing fails
117pub fn parse_token_account(account: &Account) -> Result<TokenAccount> {
118    TokenAccount::unpack(&account.data)
119        .map_err(|e| TallyError::Generic(format!("Failed to parse SPL Token account: {e}")))
120}
121
122/// Get the correct ATA for a wallet and mint, detecting token program automatically
123///
124/// # Arguments
125/// * `rpc_client` - RPC client for token program detection
126/// * `wallet` - The wallet pubkey
127/// * `mint` - The token mint pubkey
128///
129/// # Returns
130/// * `Ok((Pubkey, TokenProgram))` - The ATA address and detected token program
131/// * `Err(TallyError)` - If detection or computation fails
132pub fn get_ata_with_program_detection(
133    rpc_client: &RpcClient,
134    wallet: &Pubkey,
135    mint: &Pubkey,
136) -> Result<(Pubkey, TokenProgram)> {
137    let token_program = detect_token_program(rpc_client, mint)?;
138    let ata = get_associated_token_address_with_program(wallet, mint, token_program)?;
139    Ok((ata, token_program))
140}
141
142/// Check if a token account exists and get its information
143///
144/// # Arguments
145/// * `rpc_client` - RPC client for account queries
146/// * `token_account` - The token account pubkey to check
147///
148/// # Returns
149/// * `Ok(Some((TokenAccount, TokenProgram)))` - If account exists and can be parsed
150/// * `Ok(None)` - If account doesn't exist
151/// * `Err(TallyError)` - If fetching or parsing fails
152pub fn get_token_account_info(
153    rpc_client: &RpcClient,
154    token_account: &Pubkey,
155) -> Result<Option<(TokenAccount, TokenProgram)>> {
156    let Some(account) = rpc_client
157        .get_account_with_commitment(token_account, CommitmentConfig::confirmed())
158        .map_err(|e| TallyError::Generic(format!("Failed to fetch token account: {e}")))?
159        .value
160    else {
161        return Ok(None);
162    };
163
164    // Detect token program from account owner
165    let token_program = if account.owner == spl_token::id() {
166        TokenProgram::Token
167    } else if account.owner == spl_token_2022::id() {
168        TokenProgram::Token2022
169    } else {
170        return Err(TallyError::InvalidTokenProgram {
171            expected: "SPL Token or Token-2022".to_string(),
172            found: account.owner.to_string(),
173        });
174    };
175
176    // For now, only parse SPL Token accounts
177    if token_program == TokenProgram::Token {
178        let token_account_data = parse_token_account(&account)?;
179        Ok(Some((token_account_data, token_program)))
180    } else {
181        Err(TallyError::Generic(
182            "Token-2022 account parsing not yet implemented".to_string(),
183        ))
184    }
185}
186
187/// Create instruction data for creating an associated token account
188///
189/// # Arguments
190/// * `payer` - The account that will pay for account creation
191/// * `wallet` - The wallet that will own the token account
192/// * `mint` - The token mint
193/// * `token_program` - The token program to use
194///
195/// # Returns
196/// * `Ok(anchor_client::solana_sdk::instruction::Instruction)` - The create ATA instruction
197/// * `Err(TallyError)` - If instruction creation fails
198pub fn create_associated_token_account_instruction(
199    payer: &Pubkey,
200    wallet: &Pubkey,
201    mint: &Pubkey,
202    token_program: TokenProgram,
203) -> Result<anchor_client::solana_sdk::instruction::Instruction> {
204    Ok(
205        spl_associated_token_account::instruction::create_associated_token_account(
206            payer,
207            wallet,
208            mint,
209            &token_program.program_id(),
210        ),
211    )
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217    use anchor_client::solana_sdk::signature::{Keypair, Signer};
218    use std::str::FromStr;
219
220    #[test]
221    fn test_token_program_program_id() {
222        assert_eq!(TokenProgram::Token.program_id(), spl_token::id());
223        assert_eq!(TokenProgram::Token2022.program_id(), spl_token_2022::id());
224    }
225
226    #[test]
227    fn test_get_associated_token_address() {
228        let wallet = Pubkey::from(Keypair::new().pubkey().to_bytes());
229        let mint = Pubkey::from(Keypair::new().pubkey().to_bytes());
230
231        let ata = get_associated_token_address_for_mint(&wallet, &mint).unwrap();
232
233        // Should be deterministic
234        let ata2 = get_associated_token_address_for_mint(&wallet, &mint).unwrap();
235        assert_eq!(ata, ata2);
236
237        // Different wallets should have different ATAs
238        let wallet2 = Pubkey::from(Keypair::new().pubkey().to_bytes());
239        let ata3 = get_associated_token_address_for_mint(&wallet2, &mint).unwrap();
240        assert_ne!(ata, ata3);
241    }
242
243    #[test]
244    fn test_get_associated_token_address_with_program() {
245        let wallet = Pubkey::from(Keypair::new().pubkey().to_bytes());
246        let mint = Pubkey::from(Keypair::new().pubkey().to_bytes());
247
248        let ata_token =
249            get_associated_token_address_with_program(&wallet, &mint, TokenProgram::Token).unwrap();
250        let ata_token2022 =
251            get_associated_token_address_with_program(&wallet, &mint, TokenProgram::Token2022)
252                .unwrap();
253
254        // Different token programs should produce different ATAs
255        assert_ne!(ata_token, ata_token2022);
256    }
257
258    #[test]
259    fn test_usdc_mainnet_ata() {
260        // Test with real USDC mint on mainnet
261        let usdc_mint = Pubkey::from_str("EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v").unwrap();
262        let wallet = Pubkey::from(Keypair::new().pubkey().to_bytes());
263
264        let ata = get_associated_token_address_for_mint(&wallet, &usdc_mint).unwrap();
265
266        // Should be a valid pubkey
267        assert!(ata != Pubkey::default());
268
269        // Test with specific token program
270        let ata_token =
271            get_associated_token_address_with_program(&wallet, &usdc_mint, TokenProgram::Token)
272                .unwrap();
273
274        // USDC is on classic token program, so should match
275        assert_eq!(ata, ata_token);
276    }
277
278    #[test]
279    fn test_create_ata_instruction() {
280        let payer = Pubkey::from(Keypair::new().pubkey().to_bytes());
281        let wallet = Pubkey::from(Keypair::new().pubkey().to_bytes());
282        let mint = Pubkey::from(Keypair::new().pubkey().to_bytes());
283
284        let ix_token = create_associated_token_account_instruction(
285            &payer,
286            &wallet,
287            &mint,
288            TokenProgram::Token,
289        )
290        .unwrap();
291
292        let ix_token2022 = create_associated_token_account_instruction(
293            &payer,
294            &wallet,
295            &mint,
296            TokenProgram::Token2022,
297        )
298        .unwrap();
299
300        // Instructions should have same program ID (ATA program)
301        assert_eq!(ix_token.program_id, spl_associated_token_account::id());
302        assert_eq!(ix_token2022.program_id, spl_associated_token_account::id());
303
304        // The data might be the same since the ATA program handles both token programs
305        // but the accounts will include different token program IDs
306        assert!(!ix_token.accounts.is_empty());
307        assert!(!ix_token2022.accounts.is_empty());
308    }
309}
tally_sdk/ata.rs

tally_sdk/
ata.rs
