khadhroony-bot3/kb-pipeline/src/solana_token_2022_crypto_preflight.rs

// file: kb-pipeline/src/solana_token_2022_crypto_preflight.rs
// version: 3

//! Bounded cryptographic preflight for Token-2022 proof context-state accounts.

/// Maximum number of distinct proof context-state accounts accepted by one request.
pub const MAX_TOKEN_2022_PROOF_CONTEXTS: usize = 8;
/// Generic metadata bytes retained by every ZK ElGamal proof context-state account.
pub const ZK_PROOF_CONTEXT_META_BYTES: usize = 33;

/// One exact pre-verified proof context-state requirement.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Token2022ProofContextRequirement {
    /// Stable semantic role used in diagnostics.
    pub role: std::string::String,
    /// Canonical proof context-state account.
    pub account: kb_lib::MdPubkey,
    /// Exact proof kind expected by the Token-2022 builder.
    pub proof_type: kb_lib::ExSolanaCoreZkElGamalProofType,
    /// Optional authority retained by the context-state account.
    pub expected_authority: std::option::Option<kb_lib::MdPubkey>,
}

/// One bounded cryptographic preflight request.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Token2022CryptographicPreflightRequest {
    /// Endpoint role used by every account read.
    pub query_role: std::string::String,
    /// Optional minimum RPC context slot.
    pub min_context_slot: std::option::Option<u64>,
    /// Ordered proof context-state requirements.
    pub proof_contexts: std::vec::Vec<Token2022ProofContextRequirement>,
}

/// One validated proof context-state result.
#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
pub struct Token2022ProofContextReport {
    /// Stable semantic role.
    pub role: std::string::String,
    /// Canonical proof context-state account.
    pub account: kb_lib::MdPubkey,
    /// Official one-byte proof discriminator.
    pub proof_discriminator: u8,
    /// Exact context-state size for the proof type.
    pub context_state_bytes: usize,
    /// Context slot returned by the endpoint.
    pub context_slot: u64,
    /// Ordered successful checks.
    pub checks: std::vec::Vec<std::string::String>,
}

/// Complete cryptographic preflight report.
#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
pub struct Token2022CryptographicPreflightReport {
    /// Commitment used for all reads.
    pub commitment: std::string::String,
    /// Highest context slot observed across all reads.
    pub context_slot: u64,
    /// Ordered validated proof contexts.
    pub proof_contexts: std::vec::Vec<Token2022ProofContextReport>,
}

/// Reads and validates all pre-verified proof context-state accounts required by one operation.
pub async fn inspect_token_2022_cryptographic_preflight(
    pool: &kb_onchain_transport::HttpEndpointPool,
    request: &crate::Token2022CryptographicPreflightRequest,
) -> kb_core::Result<crate::Token2022CryptographicPreflightReport> {
    let requirements = match validate_proof_context_requirements(request) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let mut context_slot = request.min_context_slot.unwrap_or(0);
    let mut reports = std::vec::Vec::with_capacity(requirements.len());
    for requirement in requirements {
        let expected_size = requirement.proof_type.context_state_size();
        let config = match kb_onchain_transport::GetAccountInfoConfig::new_with_data(
            kb_onchain_transport::RpcCommitmentLevel::Confirmed,
            request.min_context_slot,
            expected_size,
        ) {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        let result = match pool
            .get_account_info_for_role(request.query_role.as_str(), &requirement.account, &config)
            .await
        {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        let report = match validate_proof_context_account(&requirement, &result) {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        context_slot = context_slot.max(report.context_slot);
        reports.push(report);
    }
    return std::result::Result::Ok(crate::Token2022CryptographicPreflightReport {
        commitment: "confirmed".to_string(),
        context_slot,
        proof_contexts: reports,
    });
}

fn validate_proof_context_requirements(
    request: &crate::Token2022CryptographicPreflightRequest,
) -> kb_core::Result<std::vec::Vec<crate::Token2022ProofContextRequirement>> {
    if request.query_role.trim().is_empty() {
        return std::result::Result::Err(kb_core::Error::config(
            "Token-2022 cryptographic preflight query_role must not be empty",
        ));
    }
    if request.proof_contexts.len() > crate::MAX_TOKEN_2022_PROOF_CONTEXTS {
        return std::result::Result::Err(kb_core::Error::new(
            "token_2022_proof_context_limit_exceeded",
            format!(
                "Token-2022 cryptographic preflight accepts at most {} proof contexts",
                crate::MAX_TOKEN_2022_PROOF_CONTEXTS
            ),
        ));
    }
    let mut indexes = std::collections::BTreeMap::<std::string::String, usize>::new();
    let mut unique = std::vec::Vec::<crate::Token2022ProofContextRequirement>::new();
    for requirement in &request.proof_contexts {
        if requirement.role.trim().is_empty() {
            return std::result::Result::Err(kb_core::Error::config(
                "Token-2022 proof context role must not be empty",
            ));
        }
        if let std::option::Option::Some(index) = indexes.get(requirement.account.0.as_str()) {
            if &unique[*index] != requirement {
                return std::result::Result::Err(kb_core::Error::new(
                    "token_2022_proof_context_conflicting_duplicate",
                    format!(
                        "Token-2022 proof context {} has conflicting requirements",
                        requirement.account.0
                    ),
                ));
            }
            continue;
        }
        indexes.insert(requirement.account.0.clone(), unique.len());
        unique.push(requirement.clone());
    }
    return std::result::Result::Ok(unique);
}

fn validate_proof_context_account(
    requirement: &crate::Token2022ProofContextRequirement,
    result: &kb_onchain_transport::AccountInfoResult,
) -> kb_core::Result<crate::Token2022ProofContextReport> {
    let account = match result.account.as_ref() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(kb_core::Error::new(
                "token_2022_proof_context_missing",
                format!(
                    "Token-2022 proof context account {} does not exist",
                    requirement.account.0
                ),
            ));
        },
    };
    if account.owner.0 != kb_program_ids::ZK_ELGAMAL_PROOF_PROGRAM_ID {
        return std::result::Result::Err(kb_core::Error::new(
            "token_2022_proof_context_owner_mismatch",
            format!(
                "Token-2022 proof context owner must be {}, got {}",
                kb_program_ids::ZK_ELGAMAL_PROOF_PROGRAM_ID,
                account.owner.0
            ),
        ));
    }
    if account.executable {
        return std::result::Result::Err(kb_core::Error::new(
            "token_2022_proof_context_executable",
            "Token-2022 proof context account must not be executable",
        ));
    }
    let expected_size = requirement.proof_type.context_state_size();
    if account.space != expected_size as u64 || account.data.len() != expected_size {
        return std::result::Result::Err(kb_core::Error::new(
            "token_2022_proof_context_size_mismatch",
            format!(
                "Token-2022 proof context {} must contain exactly {expected_size} bytes, got space {} and data {}",
                requirement.account.0,
                account.space,
                account.data.len()
            ),
        ));
    }
    if account.data.len() < crate::ZK_PROOF_CONTEXT_META_BYTES {
        return std::result::Result::Err(kb_core::Error::invalid_state(
            "Token-2022 proof context is shorter than its generic metadata",
        ));
    }
    let discriminator = account.data[32];
    if discriminator != requirement.proof_type.discriminator() {
        return std::result::Result::Err(kb_core::Error::new(
            "token_2022_proof_context_type_mismatch",
            format!(
                "Token-2022 proof context {} discriminator {} does not match expected {}",
                requirement.account.0,
                discriminator,
                requirement.proof_type.discriminator()
            ),
        ));
    }
    let mut checks = std::vec![
        "zk_program_owner".to_string(),
        "not_executable".to_string(),
        "exact_context_state_size".to_string(),
        "proof_type_discriminator".to_string(),
    ];
    if let std::option::Option::Some(expected_authority) = requirement.expected_authority.as_ref() {
        let retained_authority = bs58::encode(&account.data[..32]).into_string();
        if retained_authority != expected_authority.0 {
            return std::result::Result::Err(kb_core::Error::new(
                "token_2022_proof_context_authority_mismatch",
                format!(
                    "Token-2022 proof context {} retains authority {}, expected {}",
                    requirement.account.0, retained_authority, expected_authority.0
                ),
            ));
        }
        checks.push("retained_authority".to_string());
    }
    return std::result::Result::Ok(crate::Token2022ProofContextReport {
        role: requirement.role.clone(),
        account: requirement.account.clone(),
        proof_discriminator: discriminator,
        context_state_bytes: expected_size,
        context_slot: result.context.slot,
        checks,
    });
}

#[cfg(test)]
mod tests {
    fn pubkey(byte: u8) -> kb_lib::MdPubkey {
        return kb_lib::MdPubkey(bs58::encode([byte; 32]).into_string());
    }

    fn requirement(
        account: kb_lib::MdPubkey,
        proof_type: kb_lib::ExSolanaCoreZkElGamalProofType,
    ) -> crate::Token2022ProofContextRequirement {
        return crate::Token2022ProofContextRequirement {
            role: "equality_proof".to_string(),
            account,
            proof_type,
            expected_authority: std::option::Option::Some(pubkey(9)),
        };
    }

    fn account_result(
        proof_type: kb_lib::ExSolanaCoreZkElGamalProofType,
    ) -> kb_onchain_transport::AccountInfoResult {
        let mut data = std::vec![0u8; proof_type.context_state_size()];
        data[..32].copy_from_slice(&[9u8; 32]);
        data[32] = proof_type.discriminator();
        return kb_onchain_transport::AccountInfoResult {
            context: kb_onchain_transport::RpcResponseContext {
                slot: 77,
                api_version: std::option::Option::None,
            },
            account: std::option::Option::Some(kb_onchain_transport::AccountInfoValue {
                lamports: 1,
                owner: kb_lib::MdProgramId(kb_program_ids::ZK_ELGAMAL_PROOF_PROGRAM_ID.to_string()),
                executable: false,
                rent_epoch: 0,
                space: data.len() as u64,
                data,
            }),
        };
    }

    #[test]
    fn exact_duplicates_are_deduplicated_and_conflicts_fail_closed() {
        let item = requirement(
            pubkey(1),
            kb_lib::ExSolanaCoreZkElGamalProofType::CiphertextCiphertextEquality,
        );
        let request = crate::Token2022CryptographicPreflightRequest {
            query_role: "query".to_string(),
            min_context_slot: std::option::Option::Some(7),
            proof_contexts: std::vec![item.clone(), item.clone()],
        };
        let unique = super::validate_proof_context_requirements(&request);
        assert_eq!(unique.as_ref().map(std::vec::Vec::len), std::result::Result::Ok(1));
        let conflict = crate::Token2022CryptographicPreflightRequest {
            query_role: "query".to_string(),
            min_context_slot: std::option::Option::None,
            proof_contexts: std::vec![
                item,
                requirement(
                    pubkey(1),
                    kb_lib::ExSolanaCoreZkElGamalProofType::BatchedRangeProofU128,
                ),
            ],
        };
        assert!(super::validate_proof_context_requirements(&conflict).is_err());
    }

    #[test]
    fn proof_context_validation_checks_owner_size_type_and_authority() {
        let proof_type = kb_lib::ExSolanaCoreZkElGamalProofType::CiphertextCiphertextEquality;
        let context_requirement = requirement(pubkey(1), proof_type);
        let result = account_result(proof_type);
        let report = super::validate_proof_context_account(&context_requirement, &result);
        assert_eq!(
            report.as_ref().map(|value| return value.checks.len()),
            std::result::Result::Ok(5)
        );
        let wrong_type =
            requirement(pubkey(1), kb_lib::ExSolanaCoreZkElGamalProofType::BatchedRangeProofU128);
        assert!(super::validate_proof_context_account(&wrong_type, &result,).is_err());
    }
}