aspens 0.7.0

//! Chain-aware RPC client that dispatches between Alloy (EVM) and
//! solana-client (Solana) based on the chain's `architecture` field.

use alloy::primitives::{Address, Uint};
use alloy::providers::{Provider, ProviderBuilder};
use alloy_chains::NamedChain;
use eyre::Result;
use url::Url;

#[cfg(feature = "solana")]
use eyre::eyre;
#[cfg(feature = "solana")]
use solana_client::nonblocking::rpc_client::RpcClient as SolanaRpcClient;
#[cfg(feature = "solana")]
use solana_sdk::pubkey::Pubkey;
#[cfg(feature = "solana")]
use std::str::FromStr;

use crate::commands::config::config_pb::{Chain, Token};

/// Architecture string used in chain config for Solana chains.
pub const ARCH_SOLANA: &str = "Solana";
/// Architecture string used in chain config for EVM chains.
pub const ARCH_EVM: &str = "EVM";

/// A curve-aware RPC client.
pub enum ChainClient {
    /// EVM provider (Alloy).
    Evm {
        /// HTTP(S) JSON-RPC endpoint for the EVM chain.
        rpc_url: String,
        /// EIP-155 chain id used to construct the Alloy provider.
        chain_id: u32,
    },
    /// Solana RPC client.
    #[cfg(feature = "solana")]
    Solana {
        /// Async Solana JSON-RPC client.
        client: SolanaRpcClient,
    },
}

impl ChainClient {
    /// Build a `ChainClient` from a chain config entry.
    ///
    /// Dispatches on `chain.architecture`:
    /// - `"EVM"` (or empty/anything else for backward compat) → Alloy provider
    /// - `"Solana"` → Solana RPC client (requires the `solana` feature)
    ///
    /// The endpoint goes through [`resolve_rpc_url`], NOT `chain.rpc_url`
    /// directly: the server masks `rpc_url` (it can embed an API key), so the
    /// raw field is usually the literal `"********"`. Resolving here rather
    /// than at each call site means no caller can reintroduce that — reaching
    /// for `chain.rpc_url` was what silently broke every wallet and native
    /// balance read while the deposited column beside them stayed correct.
    pub fn from_chain_config(chain: &Chain) -> Result<Self> {
        let rpc_url = resolve_rpc_url(&chain.network, &chain.rpc_url)?;
        if chain.architecture.eq_ignore_ascii_case(ARCH_SOLANA) {
            #[cfg(feature = "solana")]
            {
                return Ok(ChainClient::Solana {
                    client: SolanaRpcClient::new(rpc_url),
                });
            }
            #[cfg(not(feature = "solana"))]
            {
                return Err(eyre::eyre!(
                    "chain '{}' is Solana but the `solana` feature is disabled",
                    chain.network
                ));
            }
        }
        Ok(ChainClient::Evm {
            rpc_url,
            chain_id: chain.chain_id,
        })
    }

    /// Query the native gas balance for an address. Returns a `u128` in
    /// the smallest unit (wei for EVM, lamports for Solana).
    pub async fn native_balance(&self, address: &str) -> Result<u128> {
        match self {
            ChainClient::Evm { rpc_url, .. } => {
                let url = Url::parse(rpc_url)?;
                let provider = ProviderBuilder::new().connect_http(url);
                let addr: Address = address.parse()?;
                let balance: Uint<256, 4> = provider.get_balance(addr).await?;
                Ok(balance.try_into().unwrap_or(u128::MAX))
            }
            #[cfg(feature = "solana")]
            ChainClient::Solana { client } => {
                let pubkey = Pubkey::from_str(address)
                    .map_err(|e| eyre!("invalid Solana address: {}", e))?;
                let lamports = client.get_balance(&pubkey).await?;
                Ok(lamports as u128)
            }
        }
    }

    /// Query the token balance for an owner on this chain.
    ///
    /// - EVM: ERC-20 `balanceOf(owner)`
    /// - Solana: SPL token account balance for the associated token account
    ///   derived from the owner pubkey and the token mint.
    pub async fn token_balance(&self, token: &Token, owner: &str) -> Result<u128> {
        match self {
            ChainClient::Evm { rpc_url, chain_id } => {
                use crate::evm::rpc::IERC20;
                let url = Url::parse(rpc_url)?;
                // Apply alloy's chain metadata only when the id is one it
                // knows. This used to `unwrap_or(BaseSepolia)`, which meant an
                // unknown id (HyperEVM testnet 998, anvil, any new rollup) read
                // balances through a DIFFERENT chain's provider config.
                // `balanceOf` is a plain eth_call, so no metadata is needed.
                let provider = match named_chain_for(*chain_id as u64) {
                    Some(named) => ProviderBuilder::new()
                        .with_chain(named)
                        .connect_http(url)
                        .erased(),
                    None => ProviderBuilder::new().connect_http(url).erased(),
                };
                let token_addr: Address = token.address.parse()?;
                let owner_addr: Address = owner.parse()?;
                let contract = IERC20::new(token_addr, &provider);
                let result: Uint<256, 4> = contract.balanceOf(owner_addr).call().await?;
                Ok(result.try_into().unwrap_or(u128::MAX))
            }
            #[cfg(feature = "solana")]
            ChainClient::Solana { client } => {
                let owner_pubkey = Pubkey::from_str(owner)
                    .map_err(|e| eyre!("invalid Solana owner address: {}", e))?;
                let mint_pubkey = Pubkey::from_str(&token.address)
                    .map_err(|e| eyre!("invalid Solana mint address: {}", e))?;
                let ata =
                    crate::solana::derive_associated_token_account(&owner_pubkey, &mint_pubkey);

                match client.get_token_account_balance(&ata).await {
                    Ok(b) => b
                        .amount
                        .parse::<u128>()
                        .map_err(|e| eyre!("invalid Solana token balance: {}", e)),
                    Err(_) => Ok(0),
                }
            }
        }
    }
}

/// The usable RPC endpoint for `chain`, honouring the client-side override.
///
/// Prefer this over reading `chain.rpc_url` directly. The arborter MASKS that
/// field in `GetConfig` (it can embed an API key), so the raw value is normally
/// the literal `********` — parsing it yields the memorable and thoroughly
/// unhelpful "relative URL without a base", several layers from the cause.
pub fn chain_rpc_url(chain: &crate::commands::config::config_pb::Chain) -> Result<String> {
    resolve_rpc_url(&chain.network, &chain.rpc_url)
}

/// alloy's [`NamedChain`] for `chain_id`, or `None` when alloy does not know it.
///
/// `None` is a normal, working outcome — not a failure. A chain outside alloy's
/// registry (HyperEVM testnet 998, local anvil, a fresh rollup) is still a
/// standard EVM chain: the RPC reports its own id, transactions are signed
/// externally, and the only thing lost is alloy's chain-specific metadata.
/// Callers should build the provider WITHOUT `.with_chain(..)` in that case,
/// which is what the arborter does (`chain-evm/src/market.rs`).
///
/// What callers must NOT do is substitute some other chain. An earlier version
/// of `ChainClient::token_balance` fell back to `NamedChain::BaseSepolia`, so
/// every provider built for an unknown id silently carried a different chain's
/// metadata. Wrong-chain is strictly worse than no-chain.
pub fn named_chain_for(chain_id: u64) -> Option<NamedChain> {
    NamedChain::try_from(chain_id).ok()
}

/// The env-var key a client sets to supply its own RPC endpoint for `network`:
/// `ASPENS_RPC_URL_<NETWORK>`, with `network` upper-cased and every
/// non-alphanumeric byte replaced by `_` (e.g. `base-sepolia` →
/// `ASPENS_RPC_URL_BASE_SEPOLIA`, `anvil-1` → `ASPENS_RPC_URL_ANVIL_1`).
pub fn rpc_override_env_key(network: &str) -> String {
    let suffix: String = network
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() {
                c.to_ascii_uppercase()
            } else {
                '_'
            }
        })
        .collect();
    format!("ASPENS_RPC_URL_{suffix}")
}

/// Resolve the RPC endpoint to use for `network`.
///
/// The arborter masks `rpc_url` in its `GetConfig` response (it can embed an API
/// key), so the server value is usually unusable. Resolution order:
/// 1. the per-network override env var ([`rpc_override_env_key`]), if set; else
/// 2. `server_rpc_url`, when it is a real URL (an unmasked server / local
///    fixture); else
/// 3. an actionable error naming the env var to set.
pub fn resolve_rpc_url(network: &str, server_rpc_url: &str) -> Result<String> {
    let override_val = std::env::var(rpc_override_env_key(network)).ok();
    resolve_rpc_url_with(network, override_val.as_deref(), server_rpc_url)
}

/// Core of [`resolve_rpc_url`] with the override value injected, so the
/// precedence logic is testable without mutating the process environment.
fn resolve_rpc_url_with(
    network: &str,
    override_val: Option<&str>,
    server_rpc_url: &str,
) -> Result<String> {
    if let Some(v) = override_val {
        let v = v.trim();
        if !v.is_empty() {
            return Ok(v.to_string());
        }
    }
    let server = server_rpc_url.trim();
    if !server.is_empty() && Url::parse(server).is_ok() {
        return Ok(server.to_string());
    }
    Err(eyre::eyre!(
        "no usable RPC endpoint for chain '{network}': the server masks rpc_url in its config \
         (it can embed an API key). Set {} to your own RPC URL for '{network}'.",
        rpc_override_env_key(network)
    ))
}

#[cfg(test)]
mod named_chain_tests {
    use super::*;

    /// alloy's registry is the source of truth for ids it KNOWS. These two are
    /// pinned because the stack trades on both and they behave differently.
    #[test]
    fn known_ids_resolve_to_their_named_chain() {
        assert_eq!(named_chain_for(114), Some(NamedChain::FlareCoston2));
        // Hyperliquid MAINNET is in the registry...
        assert_eq!(named_chain_for(999), Some(NamedChain::Hyperliquid));
    }

    /// ...its TESTNET (998) is not, and that must be `None` rather than an
    /// error or a substitute. A generic EVM chain is perfectly usable — the RPC
    /// supplies the chain id — it just gets no alloy metadata.
    #[test]
    fn hyperevm_testnet_is_unnamed_not_an_error() {
        assert_eq!(named_chain_for(998), None);
    }

    /// The bug this replaces: an unknown id silently became BaseSepolia, so
    /// every provider built for HyperEVM testnet carried another chain's
    /// metadata. Wrong-chain is worse than no-chain — assert we never return a
    /// chain that is not the one asked for.
    #[test]
    fn an_unknown_id_never_substitutes_a_different_chain() {
        for id in [998u64, 31337, 1337, 424242] {
            match named_chain_for(id) {
                None => {}
                Some(c) => assert_eq!(
                    u64::from(c),
                    id,
                    "id {id} resolved to a DIFFERENT chain ({c:?}) — a silent substitution"
                ),
            }
        }
    }
}

#[cfg(test)]
mod rpc_resolve_tests {
    use super::*;

    #[test]
    fn env_key_uppercases_and_sanitizes() {
        assert_eq!(
            rpc_override_env_key("base-sepolia"),
            "ASPENS_RPC_URL_BASE_SEPOLIA"
        );
        assert_eq!(rpc_override_env_key("anvil-1"), "ASPENS_RPC_URL_ANVIL_1");
        assert_eq!(
            rpc_override_env_key("solana-devnet"),
            "ASPENS_RPC_URL_SOLANA_DEVNET"
        );
    }

    #[test]
    fn override_wins_over_masked_server_value() {
        let got = resolve_rpc_url_with("net", Some("https://my.rpc/v2/key"), "********").unwrap();
        assert_eq!(got, "https://my.rpc/v2/key");
    }

    #[test]
    fn blank_override_falls_through_to_usable_server_value() {
        let got = resolve_rpc_url_with("net", Some("   "), "https://server.example").unwrap();
        assert_eq!(got, "https://server.example");
    }

    #[test]
    fn no_override_uses_unmasked_server_value() {
        let got = resolve_rpc_url_with("net", None, "http://localhost:8545").unwrap();
        assert_eq!(got, "http://localhost:8545");
    }

    #[test]
    fn masked_value_without_override_errors_with_env_key() {
        let err = resolve_rpc_url_with("base-sepolia", None, "********")
            .unwrap_err()
            .to_string();
        assert!(
            err.contains("ASPENS_RPC_URL_BASE_SEPOLIA"),
            "actionable: {err}"
        );
    }

    #[test]
    fn empty_server_without_override_errors() {
        assert!(resolve_rpc_url_with("net", None, "").is_err());
    }

    /// RPC-MASK-2 regression. `from_chain_config` used to take
    /// `chain.rpc_url` verbatim, so against a masking server it happily
    /// returned a client holding `"********"`. Every read through it then died
    /// in `Url::parse` with `relative URL without a base` — and because the
    /// deposited/locked reads resolve the endpoint separately, `balance`
    /// printed `error` in the wallet column next to a correct deposited
    /// figure. Fail at construction instead, naming the env key to set.
    ///
    /// The network name is unique to this test so the override key cannot be
    /// set by another test running in parallel (`resolve_rpc_url` reads the
    /// process env, which is global).
    #[test]
    fn from_chain_config_rejects_a_masked_url_rather_than_building_a_broken_client() {
        let chain = Chain {
            network: "rpc-mask-2-regression".to_string(),
            rpc_url: "********".to_string(),
            chain_id: 114,
            ..Default::default()
        };
        // Matched rather than `expect_err`: that needs `T: Debug`, and adding
        // `Debug` to a public client type just to satisfy a test is the wrong
        // trade — the Solana variant would have to derive it too.
        let err = match ChainClient::from_chain_config(&chain) {
            Ok(_) => panic!("a masked rpc_url must not yield a client"),
            Err(e) => e.to_string(),
        };
        assert!(
            err.contains("ASPENS_RPC_URL_RPC_MASK_2_REGRESSION"),
            "error should name the env key to set, got: {err}"
        );
    }

    /// The unmasked case is unchanged: a stack that serves a real `rpc_url`
    /// (local anvil, any non-masking server) still works with no override set.
    #[test]
    fn from_chain_config_keeps_using_a_usable_server_url() {
        let chain = Chain {
            network: "rpc-mask-2-unmasked".to_string(),
            rpc_url: "http://localhost:8545".to_string(),
            chain_id: 31337,
            ..Default::default()
        };
        match ChainClient::from_chain_config(&chain).expect("usable url builds a client") {
            ChainClient::Evm { rpc_url, chain_id } => {
                assert_eq!(rpc_url, "http://localhost:8545");
                assert_eq!(chain_id, 31337);
            }
            #[cfg(feature = "solana")]
            ChainClient::Solana { .. } => panic!("expected an EVM client"),
        }
    }
}

#[cfg(all(test, feature = "solana"))]
mod tests {
    use super::*;
    use crate::solana::derive_associated_token_account;

    #[test]
    fn ata_derivation_matches_known_pair() {
        let owner = Pubkey::from_str("11111111111111111111111111111112").unwrap();
        let mint = Pubkey::from_str("So11111111111111111111111111111111111111112").unwrap();
        let ata = derive_associated_token_account(&owner, &mint);
        assert_ne!(ata, Pubkey::default());
    }
}