[zk-token-sdk] add pubkey proof (#28392)

* add pubkey proof

* add pubkey sigma proof

* add docs for the sigma proof functions

* add pod public key sigma proof

* add public-key validity proof instruction

* add public-key validity proof instruction

* add VerifyPubkeyValidity instruction

* cargo fmt

programs/zk-token-proof/src/lib.rs

@@ -70,5 +70,9 @@ pub fn process_instruction(
             ic_msg!(invoke_context, "VerifyTransferWithFee");
             verify::<TransferWithFeeData>(invoke_context)
         }
+        ProofInstruction::VerifyPubkeyValidity => {
+            ic_msg!(invoke_context, "VerifyPubkeyValidity");
+            verify::<PubkeyValidityData>(invoke_context)
+        }
     }
 }

zk-token-sdk/src/errors.rs

@@ -32,6 +32,8 @@ pub enum ProofError {
     DiscreteLogThreads,
     #[error("discrete log batch size too large")]
     DiscreteLogBatchSize,
+    #[error("public-key sigma proof failed to verify")]
+    PubkeySigmaProof,
 }
 
 #[derive(Error, Clone, Debug, Eq, PartialEq)]
@@ -68,3 +70,9 @@ impl From<ValidityProofError> for ProofError {
         Self::ValidityProof
     }
 }
+
+impl From<PubkeySigmaProofError> for ProofError {
+    fn from(_err: PubkeySigmaProofError) -> Self {
+        Self::PubkeySigmaProof
+    }
+}

zk-token-sdk/src/instruction/close_account.rs

@@ -32,7 +32,7 @@ pub struct CloseAccountData {
     pub ciphertext: pod::ElGamalCiphertext, // 64 bytes
 
     /// Proof that the source account available balance is zero
-    pub proof: CloseAccountProof, // 64 bytes
+    pub proof: CloseAccountProof, // 96 bytes
 }
 
 #[cfg(not(target_os = "solana"))]

zk-token-sdk/src/instruction/mod.rs

@@ -1,4 +1,5 @@
 pub mod close_account;
+pub mod pubkey_validity;
 pub mod transfer;
 pub mod transfer_with_fee;
 pub mod withdraw;
@@ -17,7 +18,7 @@ use {
     subtle::ConstantTimeEq,
 };
 pub use {
-    close_account::CloseAccountData, transfer::TransferData,
+    close_account::CloseAccountData, pubkey_validity::PubkeyValidityData, transfer::TransferData,
     transfer_with_fee::TransferWithFeeData, withdraw::WithdrawData,
     withdraw_withheld::WithdrawWithheldTokensData,
 };

zk-token-sdk/src/instruction/pubkey_validity.rs

@@ -0,0 +1,105 @@
+use {
+    crate::zk_token_elgamal::pod,
+    bytemuck::{Pod, Zeroable},
+};
+#[cfg(not(target_os = "solana"))]
+use {
+    crate::{
+        encryption::elgamal::{ElGamalKeypair, ElGamalPubkey},
+        errors::ProofError,
+        instruction::Verifiable,
+        sigma_proofs::pubkey_proof::PubkeySigmaProof,
+        transcript::TranscriptProtocol,
+    },
+    merlin::Transcript,
+    std::convert::TryInto,
+};
+
+/// This struct includes the cryptographic proof *and* the account data information needed to
+/// verify the proof
+///
+/// - The pre-instruction should call PubkeyValidityData::verify_proof(&self)
+/// - The actual program should check that the public key in this struct is consistent with what is
+/// stored in the confidential token account
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct PubkeyValidityData {
+    /// The public key to be proved
+    pub pubkey: pod::ElGamalPubkey,
+
+    /// Proof that the public key is well-formed
+    pub proof: PubkeyValidityProof, // 64 bytes
+}
+
+#[cfg(not(target_os = "solana"))]
+impl PubkeyValidityData {
+    pub fn new(keypair: &ElGamalKeypair) -> Result<Self, ProofError> {
+        let pod_pubkey = pod::ElGamalPubkey(keypair.public.to_bytes());
+
+        let mut transcript = PubkeyValidityProof::transcript_new(&pod_pubkey);
+
+        let proof = PubkeyValidityProof::new(keypair, &mut transcript);
+
+        Ok(PubkeyValidityData {
+            pubkey: pod_pubkey,
+            proof,
+        })
+    }
+}
+
+#[cfg(not(target_os = "solana"))]
+impl Verifiable for PubkeyValidityData {
+    fn verify(&self) -> Result<(), ProofError> {
+        let mut transcript = PubkeyValidityProof::transcript_new(&self.pubkey);
+        let pubkey = self.pubkey.try_into()?;
+        self.proof.verify(&pubkey, &mut transcript)
+    }
+}
+
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+#[allow(non_snake_case)]
+pub struct PubkeyValidityProof {
+    /// Associated public-key sigma proof
+    pub proof: pod::PubkeySigmaProof,
+}
+
+#[allow(non_snake_case)]
+#[cfg(not(target_os = "solana"))]
+impl PubkeyValidityProof {
+    fn transcript_new(pubkey: &pod::ElGamalPubkey) -> Transcript {
+        let mut transcript = Transcript::new(b"PubkeyProof");
+        transcript.append_pubkey(b"pubkey", pubkey);
+        transcript
+    }
+
+    pub fn new(keypair: &ElGamalKeypair, transcript: &mut Transcript) -> Self {
+        let proof = PubkeySigmaProof::new(keypair, transcript);
+        Self {
+            proof: proof.into(),
+        }
+    }
+
+    pub fn verify(
+        &self,
+        pubkey: &ElGamalPubkey,
+        transcript: &mut Transcript,
+    ) -> Result<(), ProofError> {
+        let proof: PubkeySigmaProof = self.proof.try_into()?;
+        proof.verify(pubkey, transcript)?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_pubkey_validity_correctness() {
+        let keypair = ElGamalKeypair::new_rand();
+
+        let pubkey_validity_data = PubkeyValidityData::new(&keypair).unwrap();
+        assert!(pubkey_validity_data.verify().is_ok());
+    }
+}

zk-token-sdk/src/instruction/withdraw.rs

@@ -43,8 +43,8 @@ pub struct WithdrawData {
     pub proof: WithdrawProof, // 736 bytes
 }
 
+#[cfg(not(target_os = "solana"))]
 impl WithdrawData {
-    #[cfg(not(target_os = "solana"))]
     pub fn new(
         amount: u64,
         keypair: &ElGamalKeypair,

zk-token-sdk/src/instruction/withdraw_withheld.rs

@@ -39,8 +39,8 @@ pub struct WithdrawWithheldTokensData {
     pub proof: WithdrawWithheldTokensProof,
 }
 
+#[cfg(not(target_os = "solana"))]
 impl WithdrawWithheldTokensData {
-    #[cfg(not(target_os = "solana"))]
     pub fn new(
         withdraw_withheld_authority_keypair: &ElGamalKeypair,
         destination_pubkey: &ElGamalPubkey,

zk-token-sdk/src/sigma_proofs/errors.rs

@@ -72,3 +72,21 @@ impl From<TranscriptError> for FeeSigmaProofError {
         Self::Transcript
     }
 }
+
+#[derive(Error, Clone, Debug, Eq, PartialEq)]
+pub enum PubkeySigmaProofError {
+    #[error("the required algebraic relation does not hold")]
+    AlgebraicRelation,
+    #[error("malformed proof")]
+    Format,
+    #[error("multiscalar multiplication failed")]
+    MultiscalarMul,
+    #[error("transcript failed to produce a challenge")]
+    Transcript,
+}
+
+impl From<TranscriptError> for PubkeySigmaProofError {
+    fn from(_err: TranscriptError) -> Self {
+        Self::Transcript
+    }
+}

zk-token-sdk/src/sigma_proofs/mod.rs

@@ -18,5 +18,6 @@
 pub mod equality_proof;
 pub mod errors;
 pub mod fee_proof;
+pub mod pubkey_proof;
 pub mod validity_proof;
 pub mod zero_balance_proof;

zk-token-sdk/src/sigma_proofs/pubkey_proof.rs

@@ -0,0 +1,153 @@
+//! The public-key (validity) proof system.
+//!
+//! A public-key proof is defined with respect to an ElGamal public key. The proof certifies that a
+//! given public key is a valid ElGamal public key (i.e. the prover knows a corresponding secret
+//! key). To generate the proof, a prover must prove the secret key for the public key.
+//!
+//! The protocol guarantees computational soundness (by the hardness of discrete log) and perfect
+//! zero-knowledge in the random oracle model.
+
+#[cfg(not(target_os = "solana"))]
+use {
+    crate::encryption::{
+        elgamal::{ElGamalKeypair, ElGamalPubkey},
+        pedersen::H,
+    },
+    rand::rngs::OsRng,
+    zeroize::Zeroize,
+};
+use {
+    crate::{sigma_proofs::errors::PubkeySigmaProofError, transcript::TranscriptProtocol},
+    arrayref::{array_ref, array_refs},
+    curve25519_dalek::{
+        ristretto::{CompressedRistretto, RistrettoPoint},
+        scalar::Scalar,
+        traits::{IsIdentity, VartimeMultiscalarMul},
+    },
+    merlin::Transcript,
+};
+
+/// Public-key proof.
+///
+/// Contains all the elliptic curve and scalar components that make up the sigma protocol.
+#[allow(non_snake_case)]
+#[derive(Clone)]
+pub struct PubkeySigmaProof {
+    Y: CompressedRistretto,
+    z: Scalar,
+}
+
+#[allow(non_snake_case)]
+#[cfg(not(target_os = "solana"))]
+impl PubkeySigmaProof {
+    /// Public-key proof constructor.
+    ///
+    /// The function does *not* hash the public key and ciphertext 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 `OsRng` internally to generate random scalars.
+    ///
+    /// This function panics if the provided keypair is not valid (i.e. secret key is not
+    /// invertible).
+    ///
+    /// * `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) -> Self {
+        transcript.pubkey_proof_domain_sep();
+
+        // extract the relevant scalar and Ristretto points from the input
+        let s = elgamal_keypair.secret.get_scalar();
+
+        assert!(s != &Scalar::zero());
+        let s_inv = s.invert();
+
+        // generate a random masking factor that also serves as a nonce
+        let mut y = Scalar::random(&mut OsRng);
+        let Y = (&y * &(*H)).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;
+
+        y.zeroize();
+
+        Self { Y, z }
+    }
+
+    /// Public-key proof verifier.
+    ///
+    /// * `elgamal_pubkey` - The ElGamal public key to be proved
+    /// * `transcript` - The transcript that does the bookkeeping for the Fiat-Shamir heuristic
+    pub fn verify(
+        self,
+        elgamal_pubkey: &ElGamalPubkey,
+        transcript: &mut Transcript,
+    ) -> Result<(), PubkeySigmaProofError> {
+        transcript.pubkey_proof_domain_sep();
+
+        // extract the relvant scalar and Ristretto points from the input
+        let P = elgamal_pubkey.get_point();
+
+        // include Y to transcript and extract challenge
+        transcript.validate_and_append_point(b"Y", &self.Y)?;
+        let c = transcript.challenge_scalar(b"c");
+
+        // check that the required algebraic condition holds
+        let Y = self.Y.decompress().ok_or(PubkeySigmaProofError::Format)?;
+
+        let check = RistrettoPoint::vartime_multiscalar_mul(
+            vec![&self.z, &(-&c), &(-&Scalar::one())],
+            vec![&(*H), P, &Y],
+        );
+
+        if check.is_identity() {
+            Ok(())
+        } else {
+            Err(PubkeySigmaProofError::AlgebraicRelation)
+        }
+    }
+
+    pub fn to_bytes(&self) -> [u8; 64] {
+        let mut buf = [0_u8; 64];
+        buf[..32].copy_from_slice(self.Y.as_bytes());
+        buf[32..64].copy_from_slice(self.z.as_bytes());
+        buf
+    }
+
+    pub fn from_bytes(bytes: &[u8]) -> Result<Self, PubkeySigmaProofError> {
+        if bytes.len() != 64 {
+            return Err(PubkeySigmaProofError::Format);
+        }
+
+        let bytes = array_ref![bytes, 0, 64];
+        let (Y, z) = array_refs![bytes, 32, 32];
+
+        let Y = CompressedRistretto::from_slice(Y);
+        let z = Scalar::from_canonical_bytes(*z).ok_or(PubkeySigmaProofError::Format)?;
+
+        Ok(PubkeySigmaProof { Y, z })
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_pubkey_proof_correctness() {
+        let keypair = ElGamalKeypair::new_rand();
+
+        let mut prover_transcript = Transcript::new(b"test");
+        let mut verifier_transcript = Transcript::new(b"test");
+
+        let proof = PubkeySigmaProof::new(&keypair, &mut prover_transcript);
+        assert!(proof
+            .verify(&keypair.public, &mut verifier_transcript)
+            .is_ok());
+    }
+}

zk-token-sdk/src/sigma_proofs/zero_balance_proof.rs

@@ -61,6 +61,8 @@ impl ZeroBalanceProof {
         ciphertext: &ElGamalCiphertext,
         transcript: &mut Transcript,
     ) -> Self {
+        transcript.zero_balance_proof_domain_sep();
+
         // extract the relevant scalar and Ristretto points from the input
         let P = elgamal_keypair.public.get_point();
         let s = elgamal_keypair.secret.get_scalar();
@@ -98,6 +100,8 @@ impl ZeroBalanceProof {
         ciphertext: &ElGamalCiphertext,
         transcript: &mut Transcript,
     ) -> Result<(), ZeroBalanceProofError> {
+        transcript.zero_balance_proof_domain_sep();
+
         // extract the relevant scalar and Ristretto points from the input
         let P = elgamal_pubkey.get_point();
         let C = ciphertext.commitment.get_point();
@@ -112,7 +116,7 @@ impl ZeroBalanceProof {
 
         let w_negated = -&w;
 
-        // decompress R or return verification error
+        // decompress Y or return verification error
         let Y_P = self.Y_P.decompress().ok_or(ZeroBalanceProofError::Format)?;
         let Y_D = self.Y_D.decompress().ok_or(ZeroBalanceProofError::Format)?;
 

zk-token-sdk/src/transcript.rs

@@ -58,6 +58,9 @@ pub trait TranscriptProtocol {
     /// Append a domain separator for fee sigma proof.
     fn fee_sigma_proof_domain_sep(&mut self);
 
+    /// Append a domain separator for public-key proof.
+    fn pubkey_proof_domain_sep(&mut self);
+
     /// Check that a point is not the identity, then append it to the
     /// transcript.  Otherwise, return an error.
     fn validate_and_append_point(
@@ -161,4 +164,8 @@ impl TranscriptProtocol for Transcript {
     fn fee_sigma_proof_domain_sep(&mut self) {
         self.append_message(b"dom-sep", b"fee-sigma-proof")
     }
+
+    fn pubkey_proof_domain_sep(&mut self) {
+        self.append_message(b"dom-sep", b"pubkey-proof")
+    }
 }

zk-token-sdk/src/zk_token_elgamal/convert.rs

@@ -67,6 +67,7 @@ mod target_arch {
                 equality_proof::{CtxtCommEqualityProof, CtxtCtxtEqualityProof},
                 errors::*,
                 fee_proof::FeeSigmaProof,
+                pubkey_proof::PubkeySigmaProof,
                 validity_proof::{AggregatedValidityProof, ValidityProof},
                 zero_balance_proof::ZeroBalanceProof,
             },
@@ -272,6 +273,20 @@ mod target_arch {
         }
     }
 
+    impl From<PubkeySigmaProof> for pod::PubkeySigmaProof {
+        fn from(proof: PubkeySigmaProof) -> Self {
+            Self(proof.to_bytes())
+        }
+    }
+
+    impl TryFrom<pod::PubkeySigmaProof> for PubkeySigmaProof {
+        type Error = PubkeySigmaProofError;
+
+        fn try_from(pod: pod::PubkeySigmaProof) -> Result<Self, Self::Error> {
+            Self::from_bytes(&pod.0)
+        }
+    }
+
     impl TryFrom<RangeProof> for pod::RangeProof64 {
         type Error = RangeProofError;
 

zk-token-sdk/src/zk_token_elgamal/pod.rs

@@ -146,6 +146,11 @@ unsafe impl Pod for ZeroBalanceProof {}
 #[repr(transparent)]
 pub struct FeeSigmaProof(pub [u8; 256]);
 
+/// Serialization of public-key sigma proof
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(transparent)]
+pub struct PubkeySigmaProof(pub [u8; 64]);
+
 /// Serialization of range proofs for 64-bit numbers (for `Withdraw` instruction)
 #[derive(Clone, Copy)]
 #[repr(transparent)]

zk-token-sdk/src/zk_token_proof_instruction.rs

@@ -59,6 +59,16 @@ pub enum ProofInstruction {
     ///   `TransferWithFeeData`
     ///
     VerifyTransferWithFee,
+
+    /// Verify a `PubkeyValidityData` struct
+    ///
+    /// Accounts expected by this instruction:
+    ///   None
+    ///
+    /// Data expected by this instruction:
+    ///   `PubkeyValidityData`
+    ///
+    VerifyPubkeyValidity,
 }
 
 impl ProofInstruction {
@@ -104,3 +114,7 @@ pub fn verify_transfer(proof_data: &TransferData) -> Instruction {
 pub fn verify_transfer_with_fee(proof_data: &TransferWithFeeData) -> Instruction {
     ProofInstruction::VerifyTransferWithFee.encode(proof_data)
 }
+
+pub fn verify_pubkey_validity(proof_data: &PubkeyValidityData) -> Instruction {
+    ProofInstruction::VerifyPubkeyValidity.encode(proof_data)
+}