newton-chainio 0.5.2

//! Identity data management: EIP-712 signing, on-chain fetch, and domain dispatch.
//!
//! # Domain-Flexible Architecture
//!
//! Identity data is stored on-chain as an opaque encrypted blob via
//! `EncryptedIdentityData { string data }`. The identity domain (bytes32) stored
//! on-chain determines how the decrypted blob is interpreted. Operators call
//! `getLinkedDomains(policyClient, clientUser)` to enumerate all linked domains.
//!
//! Domain dispatch uses compile-time enum variants. Each domain defines:
//! - A Rust struct matching the expected JSON schema
//! - A conversion to `KycIdentityData` (or other domain-specific Rego type)
//!
//! ## Adding a New Identity Domain
//!
//! 1. Add a variant to [`IdentityDomain`] enum
//! 2. Add a domain struct (e.g., `SocialIdentityData`) implementing `Deserialize`
//! 3. Add a conversion to the corresponding regorus `PolicyDomainData` impl
//! 4. Add the `from_json` arm in [`deserialize_identity_data`]
//! 5. Register Rego extensions in `libs/regorus/src/extensions/identity.rs`
//!
//! ## EIP-712 Typing
//!
//! Currently all domains use a single `EncryptedIdentityData { string data }`
//! EIP-712 struct for on-chain signing. This keeps the contract interface simple
//! since data is encrypted anyway. If per-domain type hashes become necessary
//! (e.g., for selective disclosure or domain-specific signature semantics),
//! each domain could define its own sol! struct. The trade-off: per-domain
//! structs enable tighter on-chain validation but increase contract complexity
//! and require ABI updates per new domain.

use alloy::{
    primitives::{Address, Bytes, FixedBytes, Signature, B256},
    sol,
    sol_types::{eip712_domain, SolStruct},
};
use newton_core::{
    identity_registry::IdentityRegistry,
    newton_prover_task_manager::NewtonMessage::Intent,
    rego::{KycIdentityData as KycIdentityDataRego, PolicyDomainData},
};
use serde::{Deserialize, Serialize};
use thiserror::Error;

/// object for passing around data for parse_and_evaluate_task
#[derive(Debug, Clone, Copy)]
pub struct IdentityContext<'a> {
    /// rpc url for reading onchain data
    pub rpc_url: &'a str,
    /// the policy client from the task, for looking up in the identity registry
    pub policy_client: &'a Address,
    /// the intent, for signature recovery
    pub intent: &'a Intent,
    /// the signature of the intent to recover the signer
    pub intent_signature: &'a Bytes,
    /// the address of the identity registry
    pub identity_registry: &'a Address,
    /// the timestamp to use for stamping on the rego data
    pub timestamp: u64,
}

// ── Known Identity Domains ─────────────────────────────────────────────
//
// Compile-time enum for domain dispatch. Each variant maps a bytes32 domain
// key to a specific data shape. External contributors add new variants here
// and implement the corresponding deserialization + Rego extension.

/// Well-known identity domain identifiers.
///
/// Each variant corresponds to a `bytes32` domain identifier used by `IdentityRegistry`.
/// The byte representation is the keccak256 of the domain name string, left-padded to 32 bytes,
/// matching the Solidity convention: `bytes32(keccak256("kyc"))`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IdentityDomain {
    /// Know-Your-Customer data (status, country, birthdate, expiration, etc.)
    Kyc,
}

impl IdentityDomain {
    /// Parse a bytes32 domain identifier into a known domain variant.
    ///
    /// Returns `None` for unrecognized domains. Callers should treat unknown
    /// domains as an error or skip identity processing, depending on context.
    pub fn from_bytes32(domain: &FixedBytes<32>) -> Option<Self> {
        let kyc_hash = alloy::primitives::keccak256(b"kyc");
        if *domain == kyc_hash {
            return Some(Self::Kyc);
        }

        None
    }

    /// Returns the human-readable domain name.
    pub fn name(&self) -> &'static str {
        match self {
            Self::Kyc => "kyc",
        }
    }
}

// ── KYC Domain Sol Struct ──────────────────────────────────────────────
//
// This struct is used for deserializing decrypted KYC identity blobs and
// converting them to the Rego-compatible type. It is NOT used for on-chain
// EIP-712 signing (that uses EncryptedIdentityData).
sol! {
    #[derive(Debug, Default, Serialize, Deserialize)]
    struct KycIdentityDataSol {
        string status;
        string selected_country_code;
        string address_subdivision;
        string address_country_code;
        string birthdate;
        string expiration_date;
        string issue_date;
        string issuing_authority;
    }
}

// ── EncryptedIdentityData (domain-agnostic) ────────────────────────────
//
// Single EIP-712 struct for all domains. Data is an opaque encrypted blob;
// the domain bytes32 determines interpretation after decryption.
sol! {
    #[derive(Debug, Default, Serialize, Deserialize)]
    struct EncryptedIdentityData {
        string data;
    }
}

// ── RegisterIdentityData ───────────────────────────────────────────────
//
// EIP-712 struct matching the on-chain REGISTER_IDENTITY_TYPEHASH:
// registerIdentityData(address identityOwner,bytes32 identityDomain,string dataRefId,uint256 deadline)
sol! {
    #[derive(Debug, Default)]
    struct registerIdentityData {
        address identityOwner;
        bytes32 identityDomain;
        string dataRefId;
        uint256 deadline;
    }
}

// ── LinkIdentity Typed Data ─────────────────────────────────────────────
//
// EIP-712 structs matching the on-chain LINK_SIGNER_TYPEHASH and
// LINK_USER_TYPEHASH constants in IdentityRegistry.
sol! {
    #[derive(Debug, Default)]
    struct linkIdentitySigner {
        address identityOwner;
        address policyClient;
        address clientUser;
        bytes32[] identityDomains;
        uint256 identityOwnerNonce;
        uint256 deadline;
    }

    #[derive(Debug, Default)]
    struct linkIdentityUser {
        address identityOwner;
        address policyClient;
        address clientUser;
        bytes32[] identityDomains;
        uint256 clientUserNonce;
        uint256 deadline;
    }
}

impl KycIdentityDataSol {
    /// Convert to the Rego-compatible type by adding a YYYY-MM-DD reference timestamp.
    pub fn to_rego_data(self, timestamp: String) -> KycIdentityDataRego {
        KycIdentityDataRego {
            reference_date: timestamp,
            status: self.status,
            selected_country_code: self.selected_country_code,
            address_subdivision: self.address_subdivision,
            address_country_code: self.address_country_code,
            birthdate: self.birthdate,
            expiration_date: self.expiration_date,
            issue_date: self.issue_date,
            issuing_authority: self.issuing_authority,
        }
    }
}

/// Deserialize a raw JSON identity blob into a domain-specific Rego data type.
///
/// This is the domain dispatch point: the `domain` bytes32 determines which
/// struct the JSON is parsed into, and the result is converted to the
/// corresponding Rego extension type.
///
/// # Arguments
/// * `domain` - The identity domain identifier (bytes32 from on-chain IdentityRegistry)
/// * `json_str` - The decrypted identity data as a JSON string
/// * `timestamp` - Reference date in YYYY-MM-DD format for age/expiry calculations
///
/// # Errors
/// Returns `IdentityDataError::DeserializationError` if the domain is unrecognized
/// or the JSON doesn't match the expected schema for that domain.
pub fn deserialize_identity_data(
    domain: &FixedBytes<32>,
    json_str: &str,
    timestamp: String,
) -> Result<Box<dyn PolicyDomainData>, IdentityDataError> {
    let identity_domain = IdentityDomain::from_bytes32(domain)
        .ok_or_else(|| IdentityDataError::DeserializationError(format!("unrecognized identity domain: {}", domain)))?;

    match identity_domain {
        IdentityDomain::Kyc => {
            let kyc: KycIdentityDataSol = serde_json::from_str(json_str).map_err(|e| {
                IdentityDataError::DeserializationError(format!("failed to parse KYC identity data: {e}"))
            })?;
            Ok(Box::new(kyc.to_rego_data(timestamp)))
        } // Future domains: add arms here, each returning Box<dyn PolicyDomainData>.
    }
}

/// Error types for identity EIP712 operations
#[derive(Debug, Error)]
pub enum IdentityDataError {
    /// Invalid signature format
    #[error("Invalid signature format: {0}")]
    InvalidSignature(String),
    /// Failed to recover signer
    #[error("Failed to recover signer: {0}")]
    SignerRecoveryFailed(String),
    /// Failed to fetch domain
    #[error("Failed to fetch signing domain: {0}")]
    FetchDomainFailed(String),
    /// Failed to get the identity linked eoa
    #[error("Failed to get the identity linked eoa for policy client: {0} client user: {1} and domain: {2} error {3}")]
    FailedToGetIdentityLinkedEOA(String, String, String, String),
    /// Failed to get intent signing domain
    #[error("Failed to get intent signing domain for client: {0} reason: {1}")]
    FailedToGetSigningDomain(String, String),
    /// Failed to encode intent signature
    #[error("Failed to encode intent signature: {0} reason: {1}")]
    FailedToEncodeIntentSig(String, String),
    /// Signer recovery failed
    #[error("Failed to recover signer for signature: {0} reason: {1}")]
    FailedToRecoverSigner(String, String),
    /// Failed to find a linked identity
    #[error("Failed to get linked identity for client user: {0} and domain: {1} error: {2}")]
    FailedToGetLinkedIdentity(String, String, String),
    /// Failed to get identity data for the owner
    #[error("Failed to get identity data for owner: {0} and domain: {1} error: {2}")]
    FailedToGetIdentityData(String, String, String),
    /// Error for converting identity_domain json into bytes32
    #[error("Error deserializing identity_domain: {0}")]
    DeserializationError(String),
    /// Refuse to sign over a partial identity-domain set.
    #[error("too many linked identity domains for policy_client={policy_client} signer={intent_signer}: got {count}, limit {limit}")]
    TooManyDomains {
        /// Policy client whose linked-domain fanout exceeded the cap.
        policy_client: String,
        /// Recovered intent signer used for the registry lookup.
        intent_signer: String,
        /// Reported domain count from the registry.
        count: usize,
        /// Hard per-request domain cap.
        limit: usize,
    },
}

/// EIP712 domain name for the IdentityRegistry contract.
pub const REGISTRY_DOMAIN_NAME: &str = "IdentityRegistry";
/// EIP712 domain version for the IdentityRegistry contract.
pub const REGISTRY_DOMAIN_VERSION: &str = "1";

/// EIP712 domain configuration for EncryptedIdentityData signing
#[derive(Debug)]
pub struct RegistryEip712Domain {
    /// Name of the domain
    pub name: String,
    /// Version of the domain
    pub version: String,
    /// Chain ID of the domain
    pub chain_id: u64,
    /// Verifying contract of the domain
    pub verifying_contract: Address,
}

impl Default for RegistryEip712Domain {
    fn default() -> Self {
        Self {
            name: REGISTRY_DOMAIN_NAME.to_string(),
            version: REGISTRY_DOMAIN_VERSION.to_string(),
            chain_id: 0,                       // Must be set explicitly
            verifying_contract: Address::ZERO, // Must be set explicitly
        }
    }
}

impl RegistryEip712Domain {
    /// Create a new [`registryEip712Domain`] with the standard IdentityRegistry name and version.
    pub fn new(chain_id: u64, verifying_contract: Address) -> Self {
        Self {
            name: REGISTRY_DOMAIN_NAME.to_string(),
            version: REGISTRY_DOMAIN_VERSION.to_string(),
            chain_id,
            verifying_contract,
        }
    }
}

/// Compute EIP712 hash for an EncryptedIdentityData struct
///
/// This computes the EIP712 message hash according to EIP-712 standard:
/// hash = keccak256("\x19\x01" || domain_separator || struct_hash)
pub fn compute_identity_eip712_hash(data: &EncryptedIdentityData, domain: &RegistryEip712Domain) -> B256 {
    let domain_alloy = eip712_domain! {
        name: domain.name.clone(),
        version: domain.version.clone(),
        chain_id: domain.chain_id,
        verifying_contract: domain.verifying_contract,
    };

    data.eip712_signing_hash(&domain_alloy)
}

/// Compute the EIP-712 digest for a `registerIdentityData` gateway authorization.
///
/// Used by the gateway to produce a signed authorization that the on-chain
/// `IdentityRegistry.registerIdentityData()` call verifies against the gateway's key.
pub fn compute_register_identity_eip712_digest(data: &registerIdentityData, domain: &RegistryEip712Domain) -> B256 {
    let domain_alloy = eip712_domain! {
        name: domain.name.clone(),
        version: domain.version.clone(),
        chain_id: domain.chain_id,
        verifying_contract: domain.verifying_contract,
    };

    data.eip712_signing_hash(&domain_alloy)
}

/// Recover signer address from EIP712 signature on the encrypted data
pub fn recover_identity_signer(
    identity: &EncryptedIdentityData,
    domain: &RegistryEip712Domain,
    signature_bytes: &Bytes,
) -> Result<Address, IdentityDataError> {
    // Compute EIP712 hash
    let eip712_hash = compute_identity_eip712_hash(identity, domain);

    // Recover signer from signature
    let signature = Signature::try_from(signature_bytes.as_ref())
        .map_err(|e| IdentityDataError::InvalidSignature(e.to_string()))?;

    signature
        .recover_address_from_prehash(&eip712_hash)
        .map_err(|e| IdentityDataError::SignerRecoveryFailed(e.to_string()))
}

/// Fetch ALL identity data for a policy client from the on-chain IdentityRegistry.
///
/// Recovers the intent signer from the EIP-712 signature, calls `getLinkedDomains(policy_client, intent_signer)`
/// to enumerate all linked domains, then for each domain fetches the identity data reference.
///
/// Returns a vec of (domain, data_ref_id) pairs. Returns an empty vec if no domains are linked.
pub async fn fetch_all_identity_data(
    ctx: IdentityContext<'_>,
) -> Result<Vec<(FixedBytes<32>, String)>, IdentityDataError> {
    use alloy::{
        primitives::Signature,
        sol_types::{eip712_domain, SolStruct},
    };
    use eigensdk::common::get_provider;
    use tracing::{debug, info, warn};

    use crate::newton_core::eip712_upgradeable::EIP712Upgradeable;

    if ctx.identity_registry == &Address::ZERO {
        return Ok(Vec::new());
    }

    let provider = get_provider(ctx.rpc_url);

    // Identity resolution requires recovering the intent signer via EIP-712.
    // Policy clients that don't support identity data won't implement EIP-712,
    // so probe for support before attempting full resolution.
    let policy_client_instance = EIP712Upgradeable::new(*ctx.policy_client, provider.clone());
    let supports_eip712 = policy_client_instance.eip712Domain().call().await.ok();
    let domain_return = match supports_eip712 {
        Some(d) => d,
        None => {
            debug!(
                policy_client = %ctx.policy_client,
                "policy client does not support EIP-712, skipping identity resolution"
            );
            return Ok(Vec::new());
        }
    };

    let domain = eip712_domain! {
        name: domain_return.name,
        version: domain_return.version,
        chain_id: domain_return.chainId.to::<u64>(),
        verifying_contract: domain_return.verifyingContract,
    };
    debug!("fetch_all_identity_data: fetched domain: {:?}", domain);

    let intent_hash = ctx.intent.eip712_signing_hash(&domain);
    debug!("fetch_all_identity_data: calculated intent hash: {:?}", intent_hash);

    let signature = Signature::try_from(ctx.intent_signature.as_ref())
        .map_err(|e| IdentityDataError::FailedToEncodeIntentSig(ctx.intent_signature.to_string(), e.to_string()))?;

    let intent_signer = signature
        .recover_address_from_prehash(&intent_hash)
        .map_err(|e| IdentityDataError::FailedToRecoverSigner(ctx.intent_signature.to_string(), e.to_string()))?;

    let identity_registry = IdentityRegistry::new(*ctx.identity_registry, provider.clone());

    // Get all linked domains for this policy client + intent signer
    let linked_domains = identity_registry
        .getLinkedDomains(*ctx.policy_client, intent_signer)
        .call()
        .await
        .map_err(|e| {
            IdentityDataError::FailedToGetLinkedIdentity(
                intent_signer.to_string(),
                "multiple domains".to_string(),
                e.to_string(),
            )
        })?;

    // Octane #8 sub-finding: cap per-request linked-domain enumeration. A
    // malicious policy client can link unbounded domains; iterating all of them
    // amplifies RPC + DB work and risks operator DoS.
    const MAX_DOMAINS_PER_REQUEST: usize = 256;
    if linked_domains.len() > MAX_DOMAINS_PER_REQUEST {
        return Err(IdentityDataError::TooManyDomains {
            policy_client: ctx.policy_client.to_string(),
            intent_signer: intent_signer.to_string(),
            count: linked_domains.len(),
            limit: MAX_DOMAINS_PER_REQUEST,
        });
    }

    let mut results = Vec::new();
    for identity_domain in linked_domains {
        // Get owner EOA for this domain
        let owner_eoa = identity_registry
            .policyClientLinks(*ctx.policy_client, intent_signer, identity_domain)
            .call()
            .await
            .map_err(|e| {
                IdentityDataError::FailedToGetIdentityLinkedEOA(
                    ctx.policy_client.to_string(),
                    intent_signer.to_string(),
                    identity_domain.to_string(),
                    e.to_string(),
                )
            })?;

        // Skip if no owner linked
        if owner_eoa == Address::ZERO {
            warn!(
                domain = %identity_domain,
                intent_signer = %intent_signer,
                "no linked identity owner for domain, skipping"
            );
            continue;
        }

        // Get identity data for this owner + domain
        let identity_data = identity_registry
            .identityData(owner_eoa, identity_domain)
            .call()
            .await
            .map_err(|e| {
                IdentityDataError::FailedToGetIdentityData(
                    owner_eoa.to_string(),
                    identity_domain.to_string(),
                    e.to_string(),
                )
            })?;

        // Skip if empty data
        if identity_data.is_empty() {
            warn!(
                domain = %identity_domain,
                owner = %owner_eoa,
                "empty identity data for domain, skipping"
            );
            continue;
        }

        info!(
            domain = %identity_domain,
            owner = %owner_eoa,
            "fetch_all_identity_data: successfully sourced identity data"
        );
        results.push((identity_domain, identity_data));
    }

    Ok(results)
}

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

    #[test]
    fn kyc_data_conversion_to_rego() {
        let id = KycIdentityDataSol {
            status: "a".to_string(),
            selected_country_code: "b".to_string(),
            address_subdivision: "c".to_string(),
            address_country_code: "d".to_string(),
            birthdate: "e".to_string(),
            expiration_date: "f".to_string(),
            issue_date: "g".to_string(),
            issuing_authority: "h".to_string(),
        };

        let timestamp = "2025-01-15".to_string();
        let id_rego = id.clone().to_rego_data(timestamp.clone());

        assert_eq!(id.status, id_rego.status);
        assert_eq!(id.selected_country_code, id_rego.selected_country_code);
        assert_eq!(id.address_subdivision, id_rego.address_subdivision);
        assert_eq!(id.address_country_code, id_rego.address_country_code);
        assert_eq!(id.birthdate, id_rego.birthdate);
        assert_eq!(id.expiration_date, id_rego.expiration_date);
        assert_eq!(id.issue_date, id_rego.issue_date);
        assert_eq!(id.issuing_authority, id_rego.issuing_authority);
        assert_eq!(timestamp, id_rego.reference_date);
    }

    #[test]
    fn identity_domain_kyc_roundtrip() {
        let kyc_hash = keccak256(b"kyc");
        let domain = IdentityDomain::from_bytes32(&kyc_hash);
        assert_eq!(domain, Some(IdentityDomain::Kyc));
        assert_eq!(domain.unwrap().name(), "kyc");
    }

    #[test]
    fn identity_domain_unknown_returns_none() {
        let unknown = keccak256(b"unknown_domain_xyz");
        assert_eq!(IdentityDomain::from_bytes32(&unknown), None);
    }

    #[test]
    fn deserialize_kyc_identity_data_valid() {
        let kyc_hash = keccak256(b"kyc");
        let json = r#"{
            "status": "approved",
            "selected_country_code": "US",
            "address_subdivision": "CA",
            "address_country_code": "US",
            "birthdate": "1990-01-15",
            "expiration_date": "2030-12-31",
            "issue_date": "2020-06-01",
            "issuing_authority": "DMV"
        }"#;

        let result = deserialize_identity_data(&kyc_hash, json, "2025-03-14".to_string());
        assert!(result.is_ok());

        let rego = result.unwrap();
        assert_eq!(rego.domain_name(), "kyc");
        let fields = rego.to_field_map();
        assert_eq!(
            fields.get("status").unwrap(),
            &newton_core::rego::Value::from("approved")
        );
        assert_eq!(
            fields.get("birthdate").unwrap(),
            &newton_core::rego::Value::from("1990-01-15")
        );
    }

    #[test]
    fn deserialize_identity_data_unknown_domain_errors() {
        let unknown = keccak256(b"nope");
        let result = deserialize_identity_data(&unknown, "{}", "2025-01-01".to_string());
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("unrecognized identity domain"));
    }

    #[test]
    fn deserialize_kyc_identity_data_malformed_json_errors() {
        let kyc_hash = keccak256(b"kyc");
        let result = deserialize_identity_data(&kyc_hash, "not json", "2025-01-01".to_string());
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("failed to parse KYC"));
    }

    #[test]
    fn encrypted_identity_data_typehash_matches_canonical_string() {
        // This catches drift in the Rust `sol!` struct when no on-chain
        // ENCRYPTED_IDENTITY_TYPEHASH constant is exposed for integration checks.
        let rust_typehash = EncryptedIdentityData::default().eip712_type_hash();
        let canonical_typehash = keccak256("EncryptedIdentityData(string data)".as_bytes());
        assert_eq!(rust_typehash, canonical_typehash);
    }
}