Skip to main content

coldstar_network/
token.rs

1//! SPL Token balance fetcher for Coldstar.
2//!
3//! Ported from the Python `src/token_fetcher.py` in `coldstar-devsyrem`.
4//! Supports both the original Token Program and Token-2022 (Token Extensions).
5
6use serde::{Deserialize, Serialize};
7use serde_json::{json, Value};
8
9use crate::{RpcError, Result, SolanaRpcClient, validate_address};
10
11// ---------------------------------------------------------------------------
12// Well-known token registry
13// ---------------------------------------------------------------------------
14
15/// A known SPL token with its mint address and metadata.
16#[derive(Debug, Clone)]
17pub struct KnownToken {
18    /// Human-readable symbol (e.g. "USDC").
19    pub symbol: &'static str,
20    /// Mint address (base-58).
21    pub mint: &'static str,
22    /// Number of decimal places.
23    pub decimals: u8,
24}
25
26/// Well-known SPL token mints (mainnet).
27pub const KNOWN_TOKENS: &[KnownToken] = &[
28    KnownToken {
29        symbol: "USDC",
30        mint: "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
31        decimals: 6,
32    },
33    KnownToken {
34        symbol: "USDT",
35        mint: "Es9vMFrzaCERmJfrF4H2FYD4KCoNkY11McCe8BenwNYB",
36        decimals: 6,
37    },
38    KnownToken {
39        symbol: "RAY",
40        mint: "4k3Dyjzvzp8eMZWUXbBCjEvwSkkk59S5iCNLY3QrkX6R",
41        decimals: 6,
42    },
43    KnownToken {
44        symbol: "BONK",
45        mint: "DezXAZ8z7PnrnRJjz3wXBoRgixCa6xjnB7YaB1pPB263",
46        decimals: 5,
47    },
48    KnownToken {
49        symbol: "JUP",
50        mint: "JUPyiwrYJFskUPiHa7hkeR8VUtAeFoSYbKedZNsDvCN",
51        decimals: 6,
52    },
53];
54
55/// Devnet USDC mint (separate from mainnet).
56pub const USDC_DEVNET_MINT: &str = "Gh9ZwEmdLJ8DscKNTkTqPbNwLNNBjuSzaG9Vp2KGtKJr";
57
58/// Token-2022 (Token Extensions) program ID.
59pub const TOKEN_2022_PROGRAM_ID: &str = "TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb";
60
61/// Original Token Program ID.
62pub const TOKEN_PROGRAM_ID: &str = "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA";
63
64// ---------------------------------------------------------------------------
65// Data types
66// ---------------------------------------------------------------------------
67
68/// A single token account as returned by `getTokenAccountsByOwner`.
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct TokenAccount {
71    /// The token account address (base-58).
72    pub address: String,
73    /// The mint address of the token (base-58).
74    pub mint: String,
75    /// Owner of the token account (the wallet).
76    pub owner: String,
77    /// Raw token amount as a string (to avoid precision loss).
78    pub amount: String,
79    /// Human-readable balance with decimals applied.
80    pub ui_amount: Option<f64>,
81    /// Number of decimals for this token.
82    pub decimals: u8,
83    /// Whether this is a Token-2022 account.
84    pub is_token_2022: bool,
85    /// Whether confidential transfers are configured (Token-2022 only).
86    pub confidential_transfer: bool,
87}
88
89/// Aggregated token balance for display.
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct TokenBalance {
92    /// Token mint address.
93    pub mint: String,
94    /// Human-readable symbol (e.g. "USDC") or "Unknown".
95    pub symbol: String,
96    /// Human-readable balance.
97    pub balance: f64,
98    /// Number of decimal places.
99    pub decimals: u8,
100    /// Whether this token is in the known registry.
101    pub is_known: bool,
102    /// Whether this is a Token-2022 account.
103    pub is_token_2022: bool,
104    /// Whether confidential transfers are enabled (Token-2022 only).
105    pub confidential_transfer: bool,
106}
107
108// ---------------------------------------------------------------------------
109// TokenFetcher
110// ---------------------------------------------------------------------------
111
112/// Fetches SPL token accounts and balances for a wallet.
113///
114/// Supports both the original Token Program and Token-2022.
115pub struct TokenFetcher<'a> {
116    client: &'a SolanaRpcClient,
117}
118
119impl<'a> TokenFetcher<'a> {
120    /// Create a new fetcher backed by `client`.
121    pub fn new(client: &'a SolanaRpcClient) -> Self {
122        Self { client }
123    }
124
125    /// Fetch all **original Token Program** accounts for `owner`.
126    pub fn get_token_accounts(&self, owner: &str) -> Result<Vec<TokenAccount>> {
127        validate_address(owner)?;
128        self.fetch_token_accounts_for_program(owner, TOKEN_PROGRAM_ID, false)
129    }
130
131    /// Fetch all **Token-2022** accounts for `owner`.
132    pub fn get_token_accounts_2022(&self, owner: &str) -> Result<Vec<TokenAccount>> {
133        validate_address(owner)?;
134        self.fetch_token_accounts_for_program(owner, TOKEN_2022_PROGRAM_ID, true)
135    }
136
137    /// Fetch and merge token balances from both Token and Token-2022 programs.
138    ///
139    /// Results are sorted: known tokens first, then by descending balance.
140    pub fn get_all_token_balances(&self, owner: &str) -> Result<Vec<TokenBalance>> {
141        validate_address(owner)?;
142
143        let mut balances = Vec::new();
144
145        // Original Token Program accounts.
146        let accounts = self.fetch_token_accounts_for_program(owner, TOKEN_PROGRAM_ID, false)?;
147        for acct in &accounts {
148            balances.push(token_account_to_balance(acct));
149        }
150
151        // Token-2022 accounts.
152        let accounts_2022 =
153            self.fetch_token_accounts_for_program(owner, TOKEN_2022_PROGRAM_ID, true)?;
154        for acct in &accounts_2022 {
155            balances.push(token_account_to_balance(acct));
156        }
157
158        // Sort: known tokens first, then by descending balance.
159        balances.sort_by(|a, b| {
160            match (a.is_known, b.is_known) {
161                (true, false) => std::cmp::Ordering::Less,
162                (false, true) => std::cmp::Ordering::Greater,
163                _ => b
164                    .balance
165                    .partial_cmp(&a.balance)
166                    .unwrap_or(std::cmp::Ordering::Equal),
167            }
168        });
169
170        Ok(balances)
171    }
172
173    // ── Internal ─────────────────────────────────────────────────────
174
175    fn fetch_token_accounts_for_program(
176        &self,
177        owner: &str,
178        program_id: &str,
179        is_2022: bool,
180    ) -> Result<Vec<TokenAccount>> {
181        let payload = json!({
182            "jsonrpc": "2.0",
183            "id": 1,
184            "method": "getTokenAccountsByOwner",
185            "params": [
186                owner,
187                {"programId": program_id},
188                {"encoding": "jsonParsed"}
189            ]
190        });
191
192        let resp = self
193            .client
194            .http_client()
195            .post(self.client.rpc_url())
196            .header("Content-Type", "application/json")
197            .json(&payload)
198            .send()
199            .map_err(RpcError::Http)?;
200
201        let body: Value = resp.json().map_err(RpcError::Http)?;
202
203        if let Some(err) = body.get("error") {
204            let code = err.get("code").and_then(Value::as_i64).unwrap_or(-1);
205            let message = err
206                .get("message")
207                .and_then(Value::as_str)
208                .unwrap_or("Unknown RPC error")
209                .to_string();
210            return Err(RpcError::Rpc { code, message });
211        }
212
213        let entries = body
214            .get("result")
215            .and_then(|r| r.get("value"))
216            .and_then(Value::as_array)
217            .cloned()
218            .unwrap_or_default();
219
220        let mut accounts = Vec::with_capacity(entries.len());
221        for entry in &entries {
222            if let Some(acct) = parse_token_account(entry, is_2022) {
223                accounts.push(acct);
224            }
225        }
226
227        Ok(accounts)
228    }
229}
230
231// ---------------------------------------------------------------------------
232// Parsing helpers
233// ---------------------------------------------------------------------------
234
235/// Parse a single token account entry from the RPC response.
236fn parse_token_account(entry: &Value, is_2022: bool) -> Option<TokenAccount> {
237    let address = entry.get("pubkey")?.as_str()?.to_string();
238
239    let parsed = entry
240        .get("account")?
241        .get("data")?
242        .get("parsed")?;
243
244    let info = parsed.get("info")?;
245
246    let mint = info.get("mint")?.as_str()?.to_string();
247    let owner = info.get("owner")?.as_str()?.to_string();
248
249    let token_amount = info.get("tokenAmount")?;
250    let amount = token_amount
251        .get("amount")
252        .and_then(Value::as_str)
253        .unwrap_or("0")
254        .to_string();
255    let ui_amount = token_amount.get("uiAmount").and_then(Value::as_f64);
256    let decimals = token_amount
257        .get("decimals")
258        .and_then(Value::as_u64)
259        .unwrap_or(0) as u8;
260
261    // Check for confidential transfer extension (Token-2022 only).
262    let confidential_transfer = if is_2022 {
263        info.get("extensions")
264            .and_then(Value::as_array)
265            .map(|exts| {
266                exts.iter().any(|ext| {
267                    ext.get("extension")
268                        .and_then(Value::as_str)
269                        .map(|e| e == "confidentialTransferAccount")
270                        .unwrap_or(false)
271                })
272            })
273            .unwrap_or(false)
274    } else {
275        false
276    };
277
278    Some(TokenAccount {
279        address,
280        mint,
281        owner,
282        amount,
283        ui_amount,
284        decimals,
285        is_token_2022: is_2022,
286        confidential_transfer,
287    })
288}
289
290/// Convert a [`TokenAccount`] into a [`TokenBalance`] for display,
291/// looking up the known token registry for symbol resolution.
292fn token_account_to_balance(acct: &TokenAccount) -> TokenBalance {
293    let known = KNOWN_TOKENS
294        .iter()
295        .find(|t| t.mint == acct.mint);
296
297    TokenBalance {
298        mint: acct.mint.clone(),
299        symbol: known.map(|t| t.symbol.to_string()).unwrap_or_else(|| "Unknown".into()),
300        balance: acct.ui_amount.unwrap_or(0.0),
301        decimals: acct.decimals,
302        is_known: known.is_some(),
303        is_token_2022: acct.is_token_2022,
304        confidential_transfer: acct.confidential_transfer,
305    }
306}
307
308// ---------------------------------------------------------------------------
309// Lookup helper
310// ---------------------------------------------------------------------------
311
312/// Look up a known token by its symbol (case-insensitive).
313pub fn lookup_token(symbol: &str) -> Option<&'static KnownToken> {
314    let upper = symbol.to_uppercase();
315    KNOWN_TOKENS.iter().find(|t| t.symbol == upper)
316}
317
318/// Look up a known token by its mint address.
319pub fn lookup_token_by_mint(mint: &str) -> Option<&'static KnownToken> {
320    KNOWN_TOKENS.iter().find(|t| t.mint == mint)
321}
322
323// ---------------------------------------------------------------------------
324// Tests
325// ---------------------------------------------------------------------------
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330
331    #[test]
332    fn known_tokens_count() {
333        assert_eq!(KNOWN_TOKENS.len(), 5);
334    }
335
336    #[test]
337    fn lookup_usdc() {
338        let token = lookup_token("USDC").unwrap();
339        assert_eq!(token.decimals, 6);
340        assert_eq!(
341            token.mint,
342            "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v"
343        );
344    }
345
346    #[test]
347    fn lookup_by_mint_bonk() {
348        let token =
349            lookup_token_by_mint("DezXAZ8z7PnrnRJjz3wXBoRgixCa6xjnB7YaB1pPB263").unwrap();
350        assert_eq!(token.symbol, "BONK");
351        assert_eq!(token.decimals, 5);
352    }
353
354    #[test]
355    fn lookup_unknown_returns_none() {
356        assert!(lookup_token("FAKE").is_none());
357        assert!(lookup_token_by_mint("notamint").is_none());
358    }
359
360    #[test]
361    fn token_balance_serializes() {
362        let bal = TokenBalance {
363            mint: "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v".into(),
364            symbol: "USDC".into(),
365            balance: 100.5,
366            decimals: 6,
367            is_known: true,
368            is_token_2022: false,
369            confidential_transfer: false,
370        };
371        let json = serde_json::to_string(&bal).unwrap();
372        assert!(json.contains("USDC"));
373        assert!(json.contains("100.5"));
374    }
375
376    #[test]
377    fn token_program_ids_are_valid_base58() {
378        // Solana base-58 addresses are 32-44 characters.
379        assert!(TOKEN_PROGRAM_ID.len() >= 32 && TOKEN_PROGRAM_ID.len() <= 44);
380        assert!(TOKEN_2022_PROGRAM_ID.len() >= 32 && TOKEN_2022_PROGRAM_ID.len() <= 44);
381    }
382}