aspens 0.8.0

//! Curve-agnostic wallet abstraction for signing.
//!
//! Wraps EVM (secp256k1, via Alloy) and Solana (Ed25519, via solana-sdk)
//! keys behind a single interface so call sites don't need to branch on
//! curve type. Solana support is gated behind the `solana` feature.

use alloy_signer::Signer;
use alloy_signer_local::PrivateKeySigner;
use eyre::{Result, eyre};

#[cfg(feature = "solana")]
use solana_sdk::signature::{Keypair, Signer as SolanaSigner};

/// Cryptographic curve used by a wallet.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CurveType {
    /// secp256k1 ECDSA — EVM-compatible chains
    Secp256k1,
    /// Ed25519 EdDSA — Solana-compatible chains. Constructing/using a wallet
    /// of this curve requires the `solana` feature.
    Ed25519,
}

/// A wallet that can sign messages on EVM (always) or Solana (with `solana`
/// feature) chains.
pub enum Wallet {
    /// EVM wallet backed by an Alloy `PrivateKeySigner` (secp256k1).
    Evm(PrivateKeySigner),
    /// Solana wallet backed by a `solana-sdk` `Keypair` (Ed25519).
    #[cfg(feature = "solana")]
    Solana(Box<Keypair>),
}

impl Wallet {
    /// Load an EVM wallet from a hex private key (with or without `0x` prefix).
    pub fn from_evm_hex(hex_key: &str) -> Result<Self> {
        let signer: PrivateKeySigner = hex_key
            .parse()
            .map_err(|e| eyre!("invalid EVM private key: {}", e))?;
        Ok(Wallet::Evm(signer))
    }

    /// Load a Solana wallet from a base58-encoded keypair string
    /// (the standard `solana-keygen` output format).
    #[cfg(feature = "solana")]
    pub fn from_solana_base58(b58: &str) -> Result<Self> {
        let bytes = bs58::decode(b58.trim())
            .into_vec()
            .map_err(|e| eyre!("invalid base58 keypair: {}", e))?;
        if bytes.len() != 64 {
            return Err(eyre!(
                "Solana keypair must be 64 bytes, got {}",
                bytes.len()
            ));
        }
        let keypair = Keypair::try_from(bytes.as_slice())
            .map_err(|e| eyre!("invalid Solana keypair bytes: {}", e))?;
        Ok(Wallet::Solana(Box::new(keypair)))
    }

    /// Load a Solana wallet from a JSON byte array (alternate `solana-keygen` format,
    /// e.g. `[12,34,56,...]` — 64 bytes).
    #[cfg(feature = "solana")]
    pub fn from_solana_json(json: &str) -> Result<Self> {
        let bytes: Vec<u8> =
            serde_json::from_str(json).map_err(|e| eyre!("invalid Solana keypair JSON: {}", e))?;
        if bytes.len() != 64 {
            return Err(eyre!(
                "Solana keypair must be 64 bytes, got {}",
                bytes.len()
            ));
        }
        let keypair = Keypair::try_from(bytes.as_slice())
            .map_err(|e| eyre!("invalid Solana keypair bytes: {}", e))?;
        Ok(Wallet::Solana(Box::new(keypair)))
    }

    /// Return the wallet's curve type.
    pub fn curve(&self) -> CurveType {
        match self {
            Wallet::Evm(_) => CurveType::Secp256k1,
            #[cfg(feature = "solana")]
            Wallet::Solana(_) => CurveType::Ed25519,
        }
    }

    /// Return the wallet's address as a string.
    /// - EVM: checksummed `0x`-prefixed hex
    /// - Solana: base58-encoded 32-byte public key
    pub fn address(&self) -> String {
        match self {
            Wallet::Evm(s) => s.address().to_checksum(None),
            #[cfg(feature = "solana")]
            Wallet::Solana(kp) => kp.pubkey().to_string(),
        }
    }

    /// Sign an arbitrary message and return the raw signature bytes.
    /// - EVM: 65-byte ECDSA signature (r || s || v)
    /// - Solana: 64-byte Ed25519 signature
    pub async fn sign_message(&self, msg: &[u8]) -> Result<Vec<u8>> {
        match self {
            Wallet::Evm(s) => {
                let sig = s.sign_message(msg).await?;
                Ok(sig.as_bytes().to_vec())
            }
            #[cfg(feature = "solana")]
            Wallet::Solana(kp) => {
                let sig = kp.sign_message(msg);
                Ok(sig.as_ref().to_vec())
            }
        }
    }

    /// Borrow as an EVM signer, if this is an EVM wallet.
    pub fn as_evm(&self) -> Option<&PrivateKeySigner> {
        match self {
            Wallet::Evm(s) => Some(s),
            #[cfg(feature = "solana")]
            Wallet::Solana(_) => None,
        }
    }

    /// Borrow as a Solana keypair, if this is a Solana wallet.
    #[cfg(feature = "solana")]
    pub fn as_solana(&self) -> Option<&Keypair> {
        match self {
            Wallet::Evm(_) => None,
            Wallet::Solana(kp) => Some(kp),
        }
    }
}

/// Load a trader wallet from environment variables based on the requested curve.
///
/// - `Secp256k1`: reads `TRADER_PRIVKEY` (hex)
/// - `Ed25519`: reads `TRADER_PRIVKEY_SOLANA` (base58 keypair) — requires the
///   `solana` feature
pub fn load_trader_wallet(curve: CurveType) -> Result<Wallet> {
    match curve {
        CurveType::Secp256k1 => {
            let key = std::env::var("TRADER_PRIVKEY")
                .map_err(|_| eyre!("TRADER_PRIVKEY not set in environment"))?;
            Wallet::from_evm_hex(&key)
        }
        CurveType::Ed25519 => {
            #[cfg(feature = "solana")]
            {
                let key = std::env::var("TRADER_PRIVKEY_SOLANA")
                    .map_err(|_| eyre!("TRADER_PRIVKEY_SOLANA not set in environment"))?;
                Wallet::from_solana_base58(&key).or_else(|_| Wallet::from_solana_json(&key))
            }
            #[cfg(not(feature = "solana"))]
            {
                Err(eyre!(
                    "Ed25519/Solana wallets require the `solana` feature to be enabled"
                ))
            }
        }
    }
}

// The three helpers below take generated proto types from
// `crate::commands::config`, which only exists under the `client`
// feature. They are gated accordingly so lean-signing consumers
// (`--no-default-features --features evm,solana`) can still build.

/// Map a chain's architecture to the curve its keys live on.
///
/// Solana chains use Ed25519; everything else (EVM today, others in the
/// future) uses secp256k1. Mirrors the dispatch in
/// [`crate::chain_client::ChainClient::from_chain_config`] and the privileged
/// signing paths in `commands::trading`.
#[cfg(feature = "client")]
pub fn chain_curve(chain: &crate::commands::config::config_pb::Chain) -> CurveType {
    if chain
        .architecture
        .eq_ignore_ascii_case(crate::chain_client::ARCH_SOLANA)
    {
        CurveType::Ed25519
    } else {
        CurveType::Secp256k1
    }
}

/// Load the trader wallet whose curve matches `chain`'s architecture.
///
/// Equivalent to `load_trader_wallet(chain_curve(chain))`. Provided as a
/// helper so binaries don't have to inline the architecture→curve dispatch.
#[cfg(feature = "client")]
pub fn load_trader_wallet_for_chain(
    chain: &crate::commands::config::config_pb::Chain,
) -> Result<Wallet> {
    load_trader_wallet(chain_curve(chain))
}

/// Look up `network` in `config` and load the matching trader wallet for
/// its architecture. Common shape for binaries that take a network as a
/// CLI / REPL arg.
#[cfg(feature = "client")]
pub fn load_trader_wallet_for_network(
    config: &crate::commands::config::config_pb::GetConfigResponse,
    network: &str,
) -> Result<Wallet> {
    let chain = config
        .get_chain(network)
        .ok_or_else(|| eyre!("Chain '{}' not found in server configuration", network))?;
    load_trader_wallet_for_chain(chain)
}

/// Load an admin wallet from environment variables based on the requested curve.
///
/// - `Secp256k1`: reads `ADMIN_PRIVKEY` (hex)
/// - `Ed25519`: reads `ADMIN_PRIVKEY_SOLANA` (base58 keypair) — requires the
///   `solana` feature
pub fn load_admin_wallet(curve: CurveType) -> Result<Wallet> {
    match curve {
        CurveType::Secp256k1 => {
            let key = std::env::var("ADMIN_PRIVKEY")
                .map_err(|_| eyre!("ADMIN_PRIVKEY not set in environment"))?;
            Wallet::from_evm_hex(&key)
        }
        CurveType::Ed25519 => {
            #[cfg(feature = "solana")]
            {
                let key = std::env::var("ADMIN_PRIVKEY_SOLANA")
                    .map_err(|_| eyre!("ADMIN_PRIVKEY_SOLANA not set in environment"))?;
                Wallet::from_solana_base58(&key).or_else(|_| Wallet::from_solana_json(&key))
            }
            #[cfg(not(feature = "solana"))]
            {
                Err(eyre!(
                    "Ed25519/Solana wallets require the `solana` feature to be enabled"
                ))
            }
        }
    }
}

/// Env var holding the Solana operator-admin keypair — the authority on
/// `instance.operator_admin`, which gates `set_withdraw_epoch_cap` and
/// `set_settle_epoch_cap`.
///
/// Deliberately its OWN variable, not `TRADER_PRIVKEY_SOLANA` or
/// `ADMIN_PRIVKEY_SOLANA`: the caps only contain a compromised or
/// buggy TEE while its authority is a key nothing else in the stack holds, so
/// conflating it with the trading wallet would quietly dissolve the guarantee
/// the cap exists to provide. Keep this key offline.
#[cfg(feature = "solana")]
pub const OPERATOR_ADMIN_PRIVKEY_SOLANA_ENV: &str = "OPERATOR_ADMIN_PRIVKEY_SOLANA";

/// Load the Solana operator-admin wallet from
/// [`OPERATOR_ADMIN_PRIVKEY_SOLANA_ENV`].
///
/// Accepts either `solana-keygen` format — a base58 64-byte keypair string or
/// a JSON byte array (`[12,34,…]`, e.g. the contents of an `id.json`) — so an
/// operator can paste whichever form their key custody produced.
#[cfg(feature = "solana")]
pub fn load_operator_admin_wallet_solana() -> Result<Wallet> {
    let key = std::env::var(OPERATOR_ADMIN_PRIVKEY_SOLANA_ENV).map_err(|_| {
        eyre!(
            "{} not set in environment — it must hold the Solana keypair for the \
             instance's `operator_admin` (base58 or JSON byte array). This is \
             intentionally NOT the trader/admin key.",
            OPERATOR_ADMIN_PRIVKEY_SOLANA_ENV
        )
    })?;
    Wallet::from_solana_base58(&key)
        .or_else(|_| Wallet::from_solana_json(&key))
        .map_err(|e| {
            eyre!(
                "{} is not a valid Solana keypair (expected base58 or a JSON \
                 byte array of 64 bytes): {}",
                OPERATOR_ADMIN_PRIVKEY_SOLANA_ENV,
                e
            )
        })
}

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

    // Anvil test key #0
    const TEST_EVM_KEY: &str = "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80";

    #[cfg(feature = "solana")]
    fn fresh_solana_keypair_b58() -> String {
        let kp = Keypair::new();
        bs58::encode(kp.to_bytes()).into_string()
    }

    #[test]
    fn evm_wallet_address_is_hex() {
        let w = Wallet::from_evm_hex(TEST_EVM_KEY).unwrap();
        let addr = w.address();
        assert!(addr.starts_with("0x"));
        assert_eq!(addr.len(), 42);
        assert_eq!(w.curve(), CurveType::Secp256k1);
    }

    #[cfg(feature = "solana")]
    #[test]
    fn solana_wallet_address_is_base58() {
        let b58 = fresh_solana_keypair_b58();
        let w = Wallet::from_solana_base58(&b58).unwrap();
        let addr = w.address();
        assert!(!addr.is_empty());
        assert!(!addr.starts_with("0x"));
        assert!(addr.len() >= 32 && addr.len() <= 44);
        for c in addr.chars() {
            assert!(c.is_ascii_alphanumeric());
            assert!(c != '0' && c != 'O' && c != 'I' && c != 'l');
        }
        assert_eq!(w.curve(), CurveType::Ed25519);
    }

    #[tokio::test]
    async fn evm_sign_message_is_65_bytes() {
        let w = Wallet::from_evm_hex(TEST_EVM_KEY).unwrap();
        let sig = w.sign_message(b"hello").await.unwrap();
        assert_eq!(sig.len(), 65, "EVM signature should be 65 bytes");
    }

    #[cfg(feature = "solana")]
    #[tokio::test]
    async fn solana_sign_message_is_64_bytes() {
        let w = Wallet::from_solana_base58(&fresh_solana_keypair_b58()).unwrap();
        let sig = w.sign_message(b"hello").await.unwrap();
        assert_eq!(sig.len(), 64, "Ed25519 signature should be 64 bytes");
    }

    /// The operator-admin key is loaded format-agnostically (base58 OR a JSON
    /// byte array), so an operator whose custody produced an `id.json` doesn't
    /// have to re-encode it. Exercises both branches of the `or_else` in
    /// [`load_operator_admin_wallet_solana`] without touching the process env.
    #[cfg(feature = "solana")]
    #[test]
    fn solana_keypair_loads_from_either_format() {
        let kp = Keypair::new();
        let expected = kp.pubkey().to_string();

        let b58 = bs58::encode(kp.to_bytes()).into_string();
        let json = serde_json::to_string(&kp.to_bytes().to_vec()).unwrap();

        for key in [b58, json] {
            let w = Wallet::from_solana_base58(&key)
                .or_else(|_| Wallet::from_solana_json(&key))
                .expect("both solana-keygen formats must load");
            assert_eq!(w.address(), expected);
        }
    }

    #[cfg(feature = "solana")]
    #[test]
    fn solana_wallet_rejects_short_key() {
        let short = bs58::encode(vec![0u8; 32]).into_string();
        assert!(Wallet::from_solana_base58(&short).is_err());
    }

    #[test]
    fn evm_wallet_rejects_invalid_hex() {
        assert!(Wallet::from_evm_hex("not-hex").is_err());
    }

    #[cfg(not(feature = "solana"))]
    #[test]
    fn ed25519_load_errors_without_feature() {
        assert!(load_trader_wallet(CurveType::Ed25519).is_err());
    }
}