aspens 0.8.0

//! Authentication module for admin operations
//!
//! This module provides EIP-712 signature-based authentication to obtain
//! JWT tokens for admin operations on the Aspens platform.

/// Generated protobuf bindings for the `arborter_auth.v1` service.
#[allow(missing_docs)]
pub mod auth_pb {
    include!("../../../proto/generated/xyz.aspens.arborter_auth.v1.rs");
}

use alloy::primitives::{Address, B256, U256, keccak256};
use alloy::signers::{Signer, local::PrivateKeySigner};
use auth_pb::auth_service_client::AuthServiceClient;
use auth_pb::{AuthRequest, AuthResponse, InitializeAdminRequest, InitializeAdminResponse};
use eyre::{Result, eyre};
use std::time::{SystemTime, UNIX_EPOCH};

use crate::grpc::create_channel;
use crate::wallet::{CurveType, Wallet};

/// EIP-712 domain separator for Arborter authentication
const EIP712_DOMAIN_NAME: &str = "Arborter";
const EIP712_DOMAIN_VERSION: &str = "1";

/// Authentication response containing JWT token and metadata
#[derive(Debug, Clone)]
pub struct AuthToken {
    /// JWT token for authenticated requests
    pub jwt_token: String,
    /// Unix timestamp when the token expires (in seconds)
    pub expires_at: u64,
    /// The address that was authenticated
    pub address: String,
}

impl From<AuthResponse> for AuthToken {
    fn from(response: AuthResponse) -> Self {
        Self {
            jwt_token: response.jwt_token,
            expires_at: response.expires_at,
            address: response.address,
        }
    }
}

impl From<InitializeAdminResponse> for AuthToken {
    fn from(response: InitializeAdminResponse) -> Self {
        Self {
            jwt_token: response.jwt_token,
            expires_at: response.expires_at,
            address: response.address,
        }
    }
}

/// Initialize the first admin on a fresh Aspens stack
///
/// This can only be called once when no admin exists. It sets up
/// the initial admin and returns a JWT token for that admin.
///
/// # Arguments
/// * `url` - The Aspens stack gRPC URL
/// * `address` - The Ethereum address to set as admin
pub async fn initialize_admin(url: String, address: String) -> Result<AuthToken> {
    let channel = create_channel(&url).await?;

    let mut client = AuthServiceClient::new(channel);

    let request = tonic::Request::new(InitializeAdminRequest { address });

    let response = client.initialize_admin(request).await?;

    Ok(response.into_inner().into())
}

/// Authenticate with a curve-agnostic wallet to obtain a JWT token.
///
/// - **EVM (Secp256k1)**: signs an EIP-712 typed data digest (existing behavior).
/// - **Solana (Ed25519)**: signs the canonical message bytes
///   `address||timestamp||nonce` with raw Ed25519. The arborter auth service
///   must accept Ed25519 signatures for Solana wallets.
///
/// # Arguments
/// * `url` - The Aspens stack gRPC URL
/// * `wallet` - The wallet to authenticate with
/// * `chain_id` - The chain ID for EIP-712 domain (EVM only, defaults to 1)
pub async fn authenticate_with_wallet(
    url: String,
    wallet: &Wallet,
    chain_id: Option<u64>,
) -> Result<AuthToken> {
    let address_str = wallet.address();

    // Generate timestamp and nonce
    let timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
    let nonce = generate_nonce();

    let signature = match wallet.curve() {
        CurveType::Secp256k1 => {
            // EIP-712 path
            let evm_signer = wallet
                .as_evm()
                .ok_or_else(|| eyre!("expected EVM wallet"))?;
            let address: Address = address_str.parse()?;
            sign_auth_message(evm_signer, address, timestamp, &nonce, chain_id).await?
        }
        CurveType::Ed25519 => {
            // Solana path: sign the server's literal auth message
            let msg = build_solana_auth_message(&address_str, timestamp, &nonce);
            let sig_bytes = wallet.sign_message(msg.as_bytes()).await?;
            format!("0x{}", hex::encode(sig_bytes))
        }
    };

    // Connect to gRPC service
    let channel = create_channel(&url).await?;

    let mut client = AuthServiceClient::new(channel);

    let request = tonic::Request::new(AuthRequest {
        address: address_str,
        timestamp,
        nonce,
        signature,
    });

    let response = client.authenticate_with_signature(request).await?;

    Ok(response.into_inner().into())
}

/// The Solana admin auth message, byte-for-byte what the arborter verifies:
/// a literal text block, signed raw (no prefix, no pre-hash) with Ed25519.
pub fn build_solana_auth_message(address: &str, timestamp: u64, nonce: &str) -> String {
    format!(
        "Arborter v1 Authentication\n\nAddress: {address}\nTimestamp: {timestamp}\nNonce: {nonce}"
    )
}

/// Generate a random nonce for authentication
fn generate_nonce() -> String {
    use std::time::Instant;
    // Use a combination of timestamp and random-ish data
    let instant = Instant::now();
    let nanos = instant.elapsed().as_nanos();
    let timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos();
    format!("{:x}{:x}", timestamp, nanos)
}

/// Sign an authentication message using EIP-712 typed data
///
/// The typed data structure matches the server's expected format:
/// - Domain: { name: "Arborter", version: "1", chainId: <chain_id> }
/// - Message: { address, timestamp, nonce }
async fn sign_auth_message(
    signer: &PrivateKeySigner,
    address: Address,
    timestamp: u64,
    nonce: &str,
    chain_id: Option<u64>,
) -> Result<String> {
    let chain_id = chain_id.unwrap_or(1);

    // Compute domain separator
    let domain_separator = compute_domain_separator(chain_id);

    // Compute struct hash for the authentication message
    let struct_hash = compute_auth_struct_hash(address, timestamp, nonce);

    // Compute final EIP-712 hash: keccak256("\x19\x01" || domainSeparator || structHash)
    let mut digest_input = Vec::with_capacity(66);
    digest_input.extend_from_slice(&[0x19, 0x01]);
    digest_input.extend_from_slice(domain_separator.as_slice());
    digest_input.extend_from_slice(struct_hash.as_slice());

    let digest = keccak256(&digest_input);

    // Sign the digest
    let signature = signer.sign_hash(&digest).await?;

    // Return as hex string with 0x prefix
    Ok(format!("0x{}", hex::encode(signature.as_bytes())))
}

/// Compute EIP-712 domain separator
///
/// domainSeparator = keccak256(
///     keccak256("EIP712Domain(string name,string version,uint256 chainId)") ||
///     keccak256(name) ||
///     keccak256(version) ||
///     chainId
/// )
fn compute_domain_separator(chain_id: u64) -> B256 {
    let type_hash = keccak256(b"EIP712Domain(string name,string version,uint256 chainId)");
    let name_hash = keccak256(EIP712_DOMAIN_NAME.as_bytes());
    let version_hash = keccak256(EIP712_DOMAIN_VERSION.as_bytes());

    let mut encoded = Vec::with_capacity(128);
    encoded.extend_from_slice(type_hash.as_slice());
    encoded.extend_from_slice(name_hash.as_slice());
    encoded.extend_from_slice(version_hash.as_slice());
    encoded.extend_from_slice(&U256::from(chain_id).to_be_bytes::<32>());

    keccak256(&encoded)
}

/// Compute struct hash for authentication message
///
/// structHash = keccak256(
///     keccak256("AuthRequest(address address,uint64 timestamp,string nonce)") ||
///     address ||
///     timestamp ||
///     keccak256(nonce)
/// )
fn compute_auth_struct_hash(address: Address, timestamp: u64, nonce: &str) -> B256 {
    let type_hash = keccak256(b"AuthRequest(address address,uint64 timestamp,string nonce)");
    let nonce_hash = keccak256(nonce.as_bytes());

    let mut encoded = Vec::with_capacity(128);
    encoded.extend_from_slice(type_hash.as_slice());
    // Address is left-padded to 32 bytes
    encoded.extend_from_slice(&[0u8; 12]);
    encoded.extend_from_slice(address.as_slice());
    // timestamp is uint64, but still encoded as 32 bytes (left-padded)
    encoded.extend_from_slice(&[0u8; 24]);
    encoded.extend_from_slice(&timestamp.to_be_bytes());
    encoded.extend_from_slice(nonce_hash.as_slice());

    keccak256(&encoded)
}

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

    #[test]
    fn test_nonce_generation() {
        let nonce1 = generate_nonce();
        let nonce2 = generate_nonce();
        // Nonces should be non-empty
        assert!(!nonce1.is_empty());
        assert!(!nonce2.is_empty());
    }

    #[test]
    fn test_domain_separator() {
        // Just verify it computes without panicking
        let separator = compute_domain_separator(1);
        assert!(!separator.is_zero());
    }

    #[test]
    fn solana_auth_message_matches_the_server_literal() {
        let msg = build_solana_auth_message(
            "7xKXtg2CW87d97TXJSDpbD5jBkheTqA83TZRuJosgAsU",
            1_756_800_000,
            "n-1",
        );
        assert_eq!(
            msg,
            "Arborter v1 Authentication\n\nAddress: 7xKXtg2CW87d97TXJSDpbD5jBkheTqA83TZRuJosgAsU\nTimestamp: 1756800000\nNonce: n-1"
        );
    }
}