arcis-compiler 0.9.7

//! Arcis implementation of <https://github.com/solana-program/zk-elgamal-proof/blob/main/zk-sdk/src/sigma_proofs/pubkey_validity.rs>

use crate::{
    core::{
        circuits::boolean::{boolean_value::BooleanValue, byte::Byte},
        global_value::{
            curve_value::{CompressedCurveValue, CurveValue},
            value::FieldValue,
        },
    },
    traits::{Invert, Reveal, ToLeBytes},
    utils::{
        field::ScalarField,
        zkp::{
            elgamal::{ElGamalKeypair, ElGamalPubkey},
            transcript::Transcript,
            util::UNIT_LEN,
        },
    },
};
use std::sync::LazyLock;
use zk_elgamal_proof::encryption::{pedersen::H, ELGAMAL_PUBKEY_LEN};

pub const PUBKEY_VALIDITY_PROOF_LEN: usize = 64;

pub const PUBKEY_VALIDITY_PROOF_CONTEXT_LEN: usize = ELGAMAL_PUBKEY_LEN;

pub const PUBKEY_VALIDITY_PROOF_DATA_LEN: usize =
    PUBKEY_VALIDITY_PROOF_CONTEXT_LEN + PUBKEY_VALIDITY_PROOF_LEN;

/// The instruction data that is needed for the `ProofInstruction::VerifyPubkeyValidity`
/// instruction.
///
/// It includes the cryptographic proof as well as the context data information needed to verify
/// the proof.
#[derive(Clone, Copy)]
pub struct PubkeyValidityProofData {
    pub context: PubkeyValidityProofContext,
    pub proof: PubkeyValidityProof,
}

/// The context data needed to verify a pubkey validity proof.
#[derive(Clone, Copy)]
pub struct PubkeyValidityProofContext {
    /// The public key to be proved
    pub pubkey: ElGamalPubkey,
}

impl PubkeyValidityProofContext {
    fn new_transcript(&self) -> Transcript<BooleanValue> {
        let mut transcript = Transcript::new(b"pubkey-validity-instruction");
        transcript.append_point(b"pubkey", &self.pubkey.get_point().compress());
        transcript
    }

    pub fn to_bytes(&self) -> [Byte<BooleanValue>; PUBKEY_VALIDITY_PROOF_CONTEXT_LEN] {
        self.pubkey.get_point().compress().to_bytes()
    }
}

impl PubkeyValidityProofData {
    pub fn new(keypair: &ElGamalKeypair) -> Self {
        let context = PubkeyValidityProofContext {
            pubkey: *keypair.pubkey(),
        };
        let mut transcript = context.new_transcript();
        let proof = PubkeyValidityProof::new(keypair, &mut transcript);
        PubkeyValidityProofData { context, proof }
    }

    pub fn to_bytes(&self) -> [Byte<BooleanValue>; PUBKEY_VALIDITY_PROOF_DATA_LEN] {
        let mut bytes = [Byte::<BooleanValue>::from(0); PUBKEY_VALIDITY_PROOF_DATA_LEN];
        bytes[..PUBKEY_VALIDITY_PROOF_CONTEXT_LEN].copy_from_slice(&self.context.to_bytes());
        bytes[PUBKEY_VALIDITY_PROOF_CONTEXT_LEN..].copy_from_slice(&self.proof.to_bytes());
        bytes
    }
}

/// Public-key proof.
///
/// Contains all the elliptic curve and scalar components that make up the sigma protocol.
#[allow(non_snake_case, dead_code)]
#[derive(Clone, Copy)]
pub struct PubkeyValidityProof {
    pub(crate) Y: CompressedCurveValue,
    pub(crate) z: FieldValue<ScalarField>,
}

#[allow(non_snake_case)]
impl PubkeyValidityProof {
    /// Creates a public key validity proof.
    ///
    /// The function does *not* hash the public key into the transcript. For
    /// security, the caller (the main protocol) should hash these public key components prior to
    /// invoking this constructor.
    ///
    /// This function is randomized. It uses random singlets internally to generate randomness.
    ///
    /// * `elgamal_keypair` = The ElGamal keypair that pertains to the ElGamal public key to be
    ///   proved
    /// * `transcript` - The transcript that does the bookkeeping for the Fiat-Shamir heuristic
    pub fn new(
        elgamal_keypair: &ElGamalKeypair,
        transcript: &mut Transcript<BooleanValue>,
    ) -> Self {
        transcript.pubkey_proof_domain_separator();

        // extract the relevant scalar and Ristretto points from the input
        let s = elgamal_keypair.secret().get_scalar();

        let s_inv = s.invert(true);

        // generate a random masking factor that also serves as a nonce
        let y = FieldValue::<ScalarField>::random();
        let Y = (y * CurveValue::from(*LazyLock::force(&H)))
            .reveal()
            .compress();

        // record masking factors in transcript and get challenges
        transcript.append_point(b"Y", &Y);
        let c = transcript.challenge_scalar(b"c");

        // compute masked secret key
        let z = (c * s_inv + y).reveal();

        Self { Y, z }
    }

    pub fn to_bytes(&self) -> [Byte<BooleanValue>; PUBKEY_VALIDITY_PROOF_LEN] {
        let mut buf = [Byte::<BooleanValue>::from(0); PUBKEY_VALIDITY_PROOF_LEN];
        let mut chunks = buf.chunks_mut(UNIT_LEN);
        chunks.next().unwrap().copy_from_slice(&self.Y.to_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.z.to_le_bytes());
        buf
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        core::{
            bounds::FieldBounds,
            expressions::{
                curve_expr::{CurveExpr, InputInfo},
                domain::Domain,
                expr::EvalValue,
                field_expr::FieldExpr,
                InputKind,
            },
            global_value::{
                curve_value::CurveValue,
                global_expr_store::with_local_expr_store_as_global,
                value::FieldValue,
            },
            ir_builder::{ExprStore, IRBuilder},
        },
        utils::{
            curve_point::CurvePoint,
            field::ScalarField,
            zkp::{
                elgamal::{ElGamalKeypair, ElGamalPubkey, ElGamalSecretKey},
                pubkey_validity::PubkeyValidityProofData,
            },
        },
    };
    use group::GroupEncoding;
    use primitives::algebra::elliptic_curve::{Curve as AsyncMPCCurve, Curve25519Ristretto};
    use rand::Rng;
    use rustc_hash::FxHashMap;
    use std::rc::Rc;
    use zk_elgamal_proof::{
        encryption::elgamal::{
            ElGamalKeypair as SolanaElGamalKeypair,
            ElGamalSecretKey as SolanaElGamalSecretKey,
        },
        zk_elgamal_proof_program::proof_data::{
            PubkeyValidityProofData as SolanaPubkeyValidityProofData,
            ZkProofData,
        },
    };

    #[test]
    #[allow(non_snake_case)]
    fn test_pubkey_validity_proof() {
        let rng = &mut crate::utils::test_rng::get();

        // random ElGamal keypair
        let mut keypair = SolanaElGamalKeypair::new_rand();
        let pubkey = *keypair.pubkey();

        // we flip a coin and generate an invalid proof if the bit is false
        let is_valid_proof = rng.gen_bool(0.5);
        if !is_valid_proof {
            keypair =
                SolanaElGamalKeypair::new_for_tests(pubkey, SolanaElGamalSecretKey::new_rand());
        }
        let secret = keypair.secret();

        let solana_proof_data = SolanaPubkeyValidityProofData::new(&keypair).unwrap();
        assert_eq!(solana_proof_data.verify_proof().is_ok(), is_valid_proof);

        let mut expr_store = IRBuilder::new(true);

        // add inputs
        let mut input_vals = FxHashMap::<usize, EvalValue>::default();
        let _ = <IRBuilder as ExprStore<ScalarField>>::push_curve(
            &mut expr_store,
            CurveExpr::Input(0, Rc::new(InputInfo::from(InputKind::Plaintext))),
        );
        input_vals.insert(
            0,
            EvalValue::Curve(CurvePoint::new(
                <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(
                    &pubkey.get_point().to_bytes(),
                )
                .unwrap(),
            )),
        );
        let _ = expr_store.push_field(FieldExpr::Input(
            1,
            FieldBounds::<ScalarField>::All.as_input_info(InputKind::Secret),
        ));
        input_vals.insert(
            1,
            EvalValue::Scalar(ScalarField::from(*secret.get_scalar())),
        );

        let outputs = with_local_expr_store_as_global(
            || {
                let pubkey = CurveValue::new(0);
                let secret = FieldValue::<ScalarField>::from_id(1);

                let arcis_proof_data =
                    PubkeyValidityProofData::new(&ElGamalKeypair::new_from_inner(
                        ElGamalPubkey::new_from_inner(pubkey),
                        ElGamalSecretKey::new(secret),
                    ));

                arcis_proof_data
                    .to_bytes()
                    .into_iter()
                    .map(|byte| FieldValue::<ScalarField>::from(byte).get_id())
                    .collect::<Vec<usize>>()
            },
            &mut expr_store,
        );

        let ir = expr_store.into_ir(outputs);
        let result = ir
            .eval(rng, &mut input_vals)
            .map(|x| {
                x.into_iter()
                    .map(ScalarField::unwrap)
                    .collect::<Vec<ScalarField>>()
            })
            .unwrap();

        let arcis_proof_data_bytes = result
            .iter()
            .map(|byte| byte.to_le_bytes()[0])
            .collect::<Vec<u8>>();

        let arcis_proof_data =
            SolanaPubkeyValidityProofData::from_bytes(&arcis_proof_data_bytes).unwrap();

        let arcis_verification = arcis_proof_data.verify_proof();

        assert_eq!(arcis_verification.is_ok(), is_valid_proof);
    }
}