newton-chainio 0.5.2

//! Diagnostic utilities for debugging BLS aggregation issues.
//!
//! This module provides functions to verify cached G2 keys against on-chain values,
//! helping diagnose BN254 pairing check failures.

use alloy::primitives::Address;
use ark_ec::AffineRepr;
use eigensdk::{
    common::get_provider,
    crypto_bls::{alloy_registry_g2_point_to_g2_affine, BlsG2Point},
    types::operator::OperatorPubKeys,
    utils::slashing::middleware::bls_apk_registry::BLSApkRegistry,
};
use newton_core::operator_registry::OperatorRegistry;
use tracing::{debug, error, warn};

/// Result of G2 key verification for a single operator.
#[derive(Debug, Clone)]
pub struct G2VerificationResult {
    /// The operator's address.
    pub operator_addr: Address,
    /// Whether the cached G2 key matches the on-chain value.
    pub matches: bool,
    /// The cached G2 public key (if available).
    pub cached_g2: Option<BlsG2Point>,
    /// The on-chain G2 public key (if successfully fetched).
    pub onchain_g2: Option<BlsG2Point>,
}

/// Verifies cached G2 keys against fresh on-chain reads.
///
/// This diagnostic function compares G2 keys from a cached operator state
/// against what's currently stored on-chain in the BLSApkRegistry.
///
/// # Arguments
///
/// * `operator_registry_address` - Address of the OperatorRegistry contract
/// * `rpc_url` - HTTP RPC URL for chain queries
/// * `cached_operators` - Map of operator addresses to their cached public keys
///
/// # Returns
///
/// A tuple of (all_match: bool, results: Vec<G2VerificationResult>)
pub async fn verify_g2_keys_against_chain(
    operator_registry_address: Address,
    rpc_url: &str,
    cached_operators: &[(Address, OperatorPubKeys)],
) -> eyre::Result<(bool, Vec<G2VerificationResult>)> {
    debug!(
        "[G2_VERIFY] Starting G2 key verification for {} operators",
        cached_operators.len()
    );

    let provider = get_provider(rpc_url);
    let operator_registry = OperatorRegistry::new(operator_registry_address, provider.clone());

    let bls_apk_registry_address = operator_registry
        .blsApkRegistry()
        .call()
        .await
        .map_err(|e| eyre::eyre!("Failed to get BLS APK registry address: {}", e))?;

    let bls_apk_registry = BLSApkRegistry::new(bls_apk_registry_address, provider);

    let mut results = Vec::new();
    let mut all_match = true;

    for (operator_addr, cached_keys) in cached_operators {
        let onchain_g2_result = bls_apk_registry.getOperatorPubkeyG2(*operator_addr).call().await;

        let result = match onchain_g2_result {
            Ok(onchain_g2_point) => {
                let onchain_g2 = BlsG2Point::new(alloy_registry_g2_point_to_g2_affine(onchain_g2_point));
                let cached_g2 = &cached_keys.g2_pub_key;

                let cached_g2_affine = cached_g2.g2();
                let onchain_g2_affine = onchain_g2.g2();

                let matches = cached_g2_affine == onchain_g2_affine;

                if matches {
                    debug!("[G2_VERIFY] Operator {} G2 key MATCH", operator_addr);
                } else {
                    error!(
                        "[G2_VERIFY] Operator {} G2 key MISMATCH - cached != on-chain",
                        operator_addr
                    );
                    // Log detailed coordinates for debugging
                    if let (Some(cx), Some(cy)) = (cached_g2_affine.x(), cached_g2_affine.y()) {
                        error!(
                            "[G2_VERIFY] Cached G2: X_c0={} X_c1={} Y_c0={} Y_c1={}",
                            cx.c0, cx.c1, cy.c0, cy.c1
                        );
                    }
                    if let (Some(ox), Some(oy)) = (onchain_g2_affine.x(), onchain_g2_affine.y()) {
                        error!(
                            "[G2_VERIFY] OnChain G2: X_c0={} X_c1={} Y_c0={} Y_c1={}",
                            ox.c0, ox.c1, oy.c0, oy.c1
                        );
                    }
                    all_match = false;
                }

                G2VerificationResult {
                    operator_addr: *operator_addr,
                    matches,
                    cached_g2: Some(cached_g2.clone()),
                    onchain_g2: Some(onchain_g2),
                }
            }
            Err(e) => {
                error!(
                    "[G2_VERIFY] Failed to fetch on-chain G2 for operator {}: {}",
                    operator_addr, e
                );
                all_match = false;
                G2VerificationResult {
                    operator_addr: *operator_addr,
                    matches: false,
                    cached_g2: Some(cached_keys.g2_pub_key.clone()),
                    onchain_g2: None,
                }
            }
        };

        results.push(result);
    }

    debug!(
        "[G2_VERIFY] Verification complete: all_match={}, verified={} operators",
        all_match,
        results.len()
    );

    Ok((all_match, results))
}

/// Verifies G2 keys for operators given their addresses and cached keys.
///
/// This is a convenience wrapper that takes operator addresses directly.
///
/// # Arguments
///
/// * `operator_registry_address` - Address of the OperatorRegistry contract
/// * `rpc_url` - HTTP RPC URL for chain queries
/// * `operator_addresses` - List of operator addresses to verify
/// * `cached_keys_by_addr` - Function to look up cached keys by address
///
/// # Returns
///
/// A tuple of (all_match: bool, mismatch_count: usize)
pub async fn verify_signer_g2_keys_by_address(
    operator_registry_address: Address,
    rpc_url: &str,
    operators_with_keys: Vec<(Address, OperatorPubKeys)>,
) -> eyre::Result<(bool, usize)> {
    if operators_with_keys.is_empty() {
        warn!("[G2_VERIFY] No operators provided for verification");
        return Ok((true, 0));
    }

    let (all_match, results) =
        verify_g2_keys_against_chain(operator_registry_address, rpc_url, &operators_with_keys).await?;

    let mismatch_count = results.iter().filter(|r| !r.matches).count();
    Ok((all_match, mismatch_count))
}