polynode 0.13.11

//! Wallet signer abstraction for EIP-712 typed data signing.

use alloy_primitives::Address;
use alloy_signer::Signer;
use alloy_signer_local::PrivateKeySigner as AlloySigner;
use async_trait::async_trait;
use std::collections::BTreeSet;

use super::types::Eip712Payload;
use crate::error::{Error, Result};

/// Trait for signing EIP-712 typed data and raw messages.
///
/// Implement this trait for custom signer backends (e.g. HSM, remote KMS, Privy).
/// The SDK provides [`PrivateKeySigner`] for raw private keys.
#[async_trait]
pub trait TradingSigner: Send + Sync {
    /// The EOA address (signer address, not the Safe/Proxy).
    fn address(&self) -> Address;

    /// Sign EIP-712 typed data, returning the 65-byte signature.
    async fn sign_typed_data(&self, payload: &Eip712Payload) -> Result<Vec<u8>>;

    /// Sign a raw message (personal_sign), returning the 65-byte signature.
    async fn sign_message(&self, message: &[u8]) -> Result<Vec<u8>>;

    /// Sign a raw 32-byte hash, returning the 65-byte signature (r, s, v).
    /// Used for on-chain transaction signing (wrap/unwrap, approvals).
    /// Default implementation returns an error — override in signers that support raw hash signing.
    async fn sign_hash(&self, _hash: &[u8; 32]) -> Result<Vec<u8>> {
        Err(Error::Signing(
            "sign_hash not supported by this signer".into(),
        ))
    }
}

/// Signer backed by a raw secp256k1 private key.
pub struct PrivateKeySigner {
    inner: AlloySigner,
}

impl PrivateKeySigner {
    /// Create from a hex-encoded private key (with or without 0x prefix).
    pub fn from_hex(key: &str) -> Result<Self> {
        let key = key.strip_prefix("0x").unwrap_or(key);
        let signer: AlloySigner = key
            .parse()
            .map_err(|e| Error::Signing(format!("Invalid private key: {}", e)))?;
        Ok(Self { inner: signer })
    }

    /// Generate a new random private key. Returns (signer, hex_private_key).
    pub fn generate() -> (Self, String) {
        let signer = AlloySigner::random();
        let key_hex = hex::encode(signer.credential().to_bytes());
        (Self { inner: signer }, format!("0x{}", key_hex))
    }

    /// Get the underlying alloy signer for direct use.
    pub fn inner(&self) -> &AlloySigner {
        &self.inner
    }
}

#[async_trait]
impl TradingSigner for PrivateKeySigner {
    fn address(&self) -> Address {
        self.inner.address()
    }

    async fn sign_typed_data(&self, payload: &Eip712Payload) -> Result<Vec<u8>> {
        // Build the EIP-712 hash manually:
        // 1. Hash the domain separator
        // 2. Hash the struct data
        // 3. Combine: keccak256(0x1901 || domainSeparator || structHash)
        // 4. Sign the result
        let signing_hash = compute_eip712_hash(payload)?;
        let sig = self
            .inner
            .sign_hash(&signing_hash)
            .await
            .map_err(|e| Error::Signing(format!("sign_typed_data failed: {}", e)))?;
        let mut bytes = sig.as_bytes().to_vec();
        // Normalize v to raw recovery id (0/1): CLOB expects v=0/1, not Ethereum's 27/28
        if bytes.len() == 65 && bytes[64] >= 27 {
            bytes[64] -= 27;
        }
        Ok(bytes)
    }

    async fn sign_message(&self, message: &[u8]) -> Result<Vec<u8>> {
        let sig = self
            .inner
            .sign_message(message)
            .await
            .map_err(|e| Error::Signing(format!("sign_message failed: {}", e)))?;
        Ok(sig.as_bytes().to_vec())
    }

    async fn sign_hash(&self, hash: &[u8; 32]) -> Result<Vec<u8>> {
        let hash = alloy_primitives::B256::from_slice(hash);
        let sig = self
            .inner
            .sign_hash(&hash)
            .await
            .map_err(|e| Error::Signing(format!("sign_hash failed: {}", e)))?;
        Ok(sig.as_bytes().to_vec())
    }
}

/// Compute the EIP-712 signing hash from a payload.
///
/// This implements the EIP-712 hashing algorithm:
/// `keccak256(0x1901 || domainSeparator || structHash)`
pub fn compute_eip712_hash(payload: &Eip712Payload) -> Result<alloy_primitives::B256> {
    use alloy_primitives::keccak256;

    // Build domain separator from the JSON domain object
    let domain = &payload.domain;
    let domain_sep = hash_eip712_domain(domain)?;

    // Filter EIP712Domain from types before hashing (matches TS SDK behavior)
    let mut filtered_types = payload.types.clone();
    if let Some(obj) = filtered_types.as_object_mut() {
        obj.remove("EIP712Domain");
    }

    // Hash the struct
    let struct_hash = hash_eip712_struct(&payload.primary_type, &filtered_types, &payload.message)?;

    // Final hash: keccak256(0x1901 || domainSeparator || structHash)
    let mut buf = Vec::with_capacity(2 + 32 + 32);
    buf.extend_from_slice(&[0x19, 0x01]);
    buf.extend_from_slice(domain_sep.as_slice());
    buf.extend_from_slice(struct_hash.as_slice());
    Ok(keccak256(&buf))
}

fn hash_eip712_domain(domain: &serde_json::Value) -> Result<alloy_primitives::B256> {
    use alloy_primitives::keccak256;

    // Build the domain type string dynamically based on which fields are present.
    // EIP-712 spec: the type hash must match exactly the fields in the domain object.
    let mut type_parts = Vec::new();
    if domain.get("name").is_some() {
        type_parts.push("string name");
    }
    if domain.get("version").is_some() {
        type_parts.push("string version");
    }
    if domain.get("chainId").is_some() {
        type_parts.push("uint256 chainId");
    }
    if domain.get("verifyingContract").is_some() {
        type_parts.push("address verifyingContract");
    }
    if domain.get("salt").is_some() {
        type_parts.push("bytes32 salt");
    }
    let type_string = format!("EIP712Domain({})", type_parts.join(","));
    let type_hash = keccak256(type_string.as_bytes());

    let mut encoded = Vec::new();
    encoded.extend_from_slice(type_hash.as_slice());

    // name
    if let Some(name) = domain.get("name").and_then(|v| v.as_str()) {
        encoded.extend_from_slice(keccak256(name.as_bytes()).as_slice());
    }
    // version
    if let Some(version) = domain.get("version").and_then(|v| v.as_str()) {
        encoded.extend_from_slice(keccak256(version.as_bytes()).as_slice());
    }
    // chainId
    if let Some(chain_id) = domain.get("chainId") {
        let id = chain_id.as_u64().unwrap_or(0);
        let mut word = [0u8; 32];
        word[24..32].copy_from_slice(&id.to_be_bytes());
        encoded.extend_from_slice(&word);
    }
    // verifyingContract
    if let Some(addr) = domain.get("verifyingContract").and_then(|v| v.as_str()) {
        let addr_bytes = parse_address(addr)?;
        let mut word = [0u8; 32];
        word[12..32].copy_from_slice(&addr_bytes);
        encoded.extend_from_slice(&word);
    }

    Ok(keccak256(&encoded))
}

pub fn hash_eip712_struct(
    primary_type: &str,
    types: &serde_json::Value,
    message: &serde_json::Value,
) -> Result<alloy_primitives::B256> {
    use alloy_primitives::keccak256;

    // Build the type string for the primary type
    let type_string = build_type_string(primary_type, types)?;
    let type_hash = keccak256(type_string.as_bytes());

    let mut encoded = Vec::new();
    encoded.extend_from_slice(type_hash.as_slice());

    // Encode each field
    if let Some(fields) = types.get(primary_type).and_then(|v| v.as_array()) {
        for field in fields {
            let name = field.get("name").and_then(|v| v.as_str()).unwrap_or("");
            let typ = field.get("type").and_then(|v| v.as_str()).unwrap_or("");
            let value = &message[name];
            encode_field(typ, value, types, &mut encoded)?;
        }
    }

    Ok(keccak256(&encoded))
}

fn build_type_string(primary_type: &str, types: &serde_json::Value) -> Result<String> {
    let mut ordered_types = vec![primary_type.to_string()];
    let mut deps = BTreeSet::new();
    collect_type_dependencies(primary_type, types, &mut deps)?;
    deps.remove(primary_type);
    ordered_types.extend(deps);

    let mut result = String::new();
    for type_name in ordered_types {
        result.push_str(&encode_type_definition(&type_name, types)?);
    }
    Ok(result)
}

fn encode_type_definition(type_name: &str, types: &serde_json::Value) -> Result<String> {
    let fields = types
        .get(type_name)
        .and_then(|v| v.as_array())
        .ok_or_else(|| Error::Signing(format!("Missing type definition for {}", type_name)))?;

    let mut result = format!("{}(", type_name);
    for (i, field) in fields.iter().enumerate() {
        if i > 0 {
            result.push(',');
        }
        let typ = field.get("type").and_then(|v| v.as_str()).unwrap_or("");
        let name = field.get("name").and_then(|v| v.as_str()).unwrap_or("");
        result.push_str(typ);
        result.push(' ');
        result.push_str(name);
    }
    result.push(')');
    Ok(result)
}

fn collect_type_dependencies(
    type_name: &str,
    types: &serde_json::Value,
    deps: &mut BTreeSet<String>,
) -> Result<()> {
    let Some(fields) = types.get(type_name).and_then(|v| v.as_array()) else {
        return Ok(());
    };

    for field in fields {
        let typ = field.get("type").and_then(|v| v.as_str()).unwrap_or("");
        let base = base_type(typ);
        if base != type_name && base != "EIP712Domain" && types.get(base).is_some() {
            if deps.insert(base.to_string()) {
                collect_type_dependencies(base, types, deps)?;
            }
        }
    }

    Ok(())
}

fn base_type(typ: &str) -> &str {
    typ.split('[').next().unwrap_or(typ)
}

fn encode_field(
    typ: &str,
    value: &serde_json::Value,
    types: &serde_json::Value,
    out: &mut Vec<u8>,
) -> Result<()> {
    use alloy_primitives::keccak256;

    if typ.contains('[') {
        let base = base_type(typ);
        let values = value
            .as_array()
            .ok_or_else(|| Error::Signing(format!("Expected array for {}", typ)))?;
        let mut encoded_items = Vec::with_capacity(values.len() * 32);
        for item in values {
            let word = encode_array_item(base, item, types)?;
            encoded_items.extend_from_slice(&word);
        }
        out.extend_from_slice(keccak256(&encoded_items).as_slice());
        return Ok(());
    }

    match typ {
        "string" => {
            let s = value.as_str().unwrap_or("");
            out.extend_from_slice(keccak256(s.as_bytes()).as_slice());
        }
        "bytes" => {
            let s = value.as_str().unwrap_or("0x");
            let bytes = hex::decode(s.strip_prefix("0x").unwrap_or(s))
                .map_err(|e| Error::Signing(format!("Invalid bytes: {}", e)))?;
            out.extend_from_slice(keccak256(&bytes).as_slice());
        }
        "address" => {
            let addr = value
                .as_str()
                .unwrap_or("0x0000000000000000000000000000000000000000");
            let addr_bytes = parse_address(addr)?;
            let mut word = [0u8; 32];
            word[12..32].copy_from_slice(&addr_bytes);
            out.extend_from_slice(&word);
        }
        "bool" => {
            let b = value.as_bool().unwrap_or(false);
            let mut word = [0u8; 32];
            if b {
                word[31] = 1;
            }
            out.extend_from_slice(&word);
        }
        t if t.starts_with("uint") || t.starts_with("int") => {
            // Handle uint256, int256, uint8, etc. — encode as 32-byte big-endian word.
            // Values can be JSON numbers or decimal/hex strings.
            let word = encode_uint_value(value)?;
            out.extend_from_slice(&word);
        }
        t if t.starts_with("bytes") && t.len() > 5 => {
            // Fixed-size bytes (bytes1..bytes32)
            let s = value.as_str().unwrap_or("0x");
            let bytes = hex::decode(s.strip_prefix("0x").unwrap_or(s))
                .map_err(|e| Error::Signing(format!("Invalid {}: {}", t, e)))?;
            let mut word = [0u8; 32];
            let len = bytes.len().min(32);
            word[..len].copy_from_slice(&bytes[..len]);
            out.extend_from_slice(&word);
        }
        _ => {
            // Struct type — hash recursively
            if types.get(typ).is_some() {
                let hash = hash_eip712_struct(typ, types, value)?;
                out.extend_from_slice(hash.as_slice());
            } else {
                // Unknown type, treat as bytes32(0)
                out.extend_from_slice(&[0u8; 32]);
            }
        }
    }
    Ok(())
}

fn encode_array_item(
    typ: &str,
    value: &serde_json::Value,
    types: &serde_json::Value,
) -> Result<[u8; 32]> {
    use alloy_primitives::keccak256;

    let mut word = [0u8; 32];
    match typ {
        "string" => {
            let s = value.as_str().unwrap_or("");
            word.copy_from_slice(keccak256(s.as_bytes()).as_slice());
        }
        "bytes" => {
            let s = value.as_str().unwrap_or("0x");
            let bytes = hex::decode(s.strip_prefix("0x").unwrap_or(s))
                .map_err(|e| Error::Signing(format!("Invalid bytes: {}", e)))?;
            word.copy_from_slice(keccak256(&bytes).as_slice());
        }
        "address" => {
            let addr = value
                .as_str()
                .unwrap_or("0x0000000000000000000000000000000000000000");
            let addr_bytes = parse_address(addr)?;
            word[12..32].copy_from_slice(&addr_bytes);
        }
        "bool" => {
            if value.as_bool().unwrap_or(false) {
                word[31] = 1;
            }
        }
        t if t.starts_with("uint") || t.starts_with("int") => {
            word = encode_uint_value(value)?;
        }
        t if t.starts_with("bytes") && t.len() > 5 => {
            let s = value.as_str().unwrap_or("0x");
            let bytes = hex::decode(s.strip_prefix("0x").unwrap_or(s))
                .map_err(|e| Error::Signing(format!("Invalid {}: {}", t, e)))?;
            let len = bytes.len().min(32);
            word[..len].copy_from_slice(&bytes[..len]);
        }
        _ => {
            if types.get(typ).is_some() {
                let hash = hash_eip712_struct(typ, types, value)?;
                word.copy_from_slice(hash.as_slice());
            }
        }
    }
    Ok(word)
}

fn parse_address(addr: &str) -> Result<[u8; 20]> {
    let hex_str = addr.strip_prefix("0x").unwrap_or(addr);
    let bytes = hex::decode(hex_str)
        .map_err(|e| Error::Signing(format!("Invalid address {}: {}", addr, e)))?;
    if bytes.len() != 20 {
        return Err(Error::Signing(format!(
            "Address wrong length: {} bytes",
            bytes.len()
        )));
    }
    let mut arr = [0u8; 20];
    arr.copy_from_slice(&bytes);
    Ok(arr)
}

/// Encode a uint/int value to a 32-byte big-endian word.
/// Handles JSON numbers, decimal strings (any size up to uint256), and hex strings.
fn encode_uint_value(value: &serde_json::Value) -> Result<[u8; 32]> {
    use alloy_primitives::U256;

    let mut word = [0u8; 32];

    if let Some(n) = value.as_u64() {
        word[24..32].copy_from_slice(&n.to_be_bytes());
    } else if let Some(n) = value.as_i64() {
        if n >= 0 {
            word[24..32].copy_from_slice(&n.to_be_bytes());
        } else {
            // Signed negative: two's complement
            let bytes = n.to_be_bytes();
            word[..24].fill(0xff);
            word[24..32].copy_from_slice(&bytes);
        }
    } else if let Some(s) = value.as_str() {
        if s.starts_with("0x") || s.starts_with("0X") {
            // Hex string
            let v = U256::from_str_radix(&s[2..], 16)
                .map_err(|e| Error::Signing(format!("Invalid hex uint: {} ({})", s, e)))?;
            word = v.to_be_bytes::<32>();
        } else {
            // Decimal string — use U256 for arbitrary precision
            let v = U256::from_str_radix(s, 10)
                .map_err(|e| Error::Signing(format!("Invalid decimal uint: {} ({})", s, e)))?;
            word = v.to_be_bytes::<32>();
        }
    }

    Ok(word)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::trading::types::Eip712Payload;

    #[test]
    fn deposit_wallet_batch_hash_matches_viem() {
        let payload = Eip712Payload {
            domain: serde_json::json!({
                "name": "DepositWallet",
                "version": "1",
                "chainId": 137,
                "verifyingContract": "0x575d2ea47fd5b97988736529ad8319feab69369b",
            }),
            types: serde_json::json!({
                "Call": [
                    {"name": "target", "type": "address"},
                    {"name": "value", "type": "uint256"},
                    {"name": "data", "type": "bytes"}
                ],
                "Batch": [
                    {"name": "wallet", "type": "address"},
                    {"name": "nonce", "type": "uint256"},
                    {"name": "deadline", "type": "uint256"},
                    {"name": "calls", "type": "Call[]"}
                ]
            }),
            primary_type: "Batch".into(),
            message: serde_json::json!({
                "wallet": "0x575d2ea47fd5b97988736529ad8319feab69369b",
                "nonce": "9",
                "deadline": "1770000000",
                "calls": [
                    {
                        "target": "0xc011a7e12a19f7b1f670d46f03b03f3342e82dfb",
                        "value": "0",
                        "data": "0x095ea7b3000000000000000000000000e111180000d2663c0091e4f400237545b87b996bffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
                    },
                    {
                        "target": "0x4D97DCd97eC945f40cF65F87097ACe5EA0476045",
                        "value": "0",
                        "data": "0xa22cb465000000000000000000000000e111180000d2663c0091e4f400237545b87b996b0000000000000000000000000000000000000000000000000000000000000001"
                    }
                ]
            }),
        };

        let hash = compute_eip712_hash(&payload).unwrap();
        assert_eq!(
            format!("{hash:#x}"),
            "0x58786694f0e97d1bc51769d55d389180c254a882625ede1b75c5844f4aafba7b"
        );
    }
}