[zk-sdk] Add input validation to sigma proof generation functions

zk-sdk-wasm-js/src/proof_data/zero_ciphertext.rs

@@ -128,8 +128,8 @@ mod tests {
 
         // Proof for an invalid encryption of 1
         let one_ciphertext = keypair.pubkey().encrypt_u64(1);
-        let proof_invalid = ZeroCiphertextProofData::new(&keypair, &one_ciphertext).unwrap();
-        assert!(proof_invalid.verify().is_err());
+        let result = ZeroCiphertextProofData::new(&keypair, &one_ciphertext);
+        assert!(result.is_err());
     }
 
     #[wasm_bindgen_test]

zk-sdk/src/range_proof/pod.rs

@@ -16,7 +16,7 @@ use {
 };
 
 /// The `RangeProof` type as a `Pod` restricted to proofs on 64-bit numbers.
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 #[repr(transparent)]
 pub struct PodRangeProofU64(pub(crate) [u8; RANGE_PROOF_U64_LEN]);
 
@@ -63,7 +63,7 @@ impl_from_str!(
 impl_from_bytes!(TYPE = PodRangeProofU64, BYTES_LEN = RANGE_PROOF_U64_LEN);
 
 /// The `RangeProof` type as a `Pod` restricted to proofs on 128-bit numbers.
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 #[repr(transparent)]
 pub struct PodRangeProofU128(pub(crate) [u8; RANGE_PROOF_U128_LEN]);
 
@@ -110,7 +110,7 @@ impl_from_str!(
 impl_from_bytes!(TYPE = PodRangeProofU128, BYTES_LEN = RANGE_PROOF_U128_LEN);
 
 /// The `RangeProof` type as a `Pod` restricted to proofs on 256-bit numbers.
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 #[repr(transparent)]
 pub struct PodRangeProofU256(pub(crate) [u8; RANGE_PROOF_U256_LEN]);
 

zk-sdk/src/sigma_proofs/batched_grouped_ciphertext_validity/handles_2.rs

@@ -29,12 +29,12 @@ use {
 use {
     crate::{
         sigma_proofs::{
-            errors::ValidityProofVerificationError,
+            errors::{SigmaProofVerificationError, ValidityProofVerificationError},
             grouped_ciphertext_validity::GroupedCiphertext2HandlesValidityProof,
         },
         transcript::TranscriptProtocol,
     },
-    curve25519_dalek::scalar::Scalar,
+    curve25519_dalek::{scalar::Scalar, traits::IsIdentity},
     merlin::Transcript,
 };
 
@@ -112,6 +112,16 @@ impl BatchedGroupedCiphertext2HandlesValidityProof {
         grouped_ciphertext_hi: &GroupedElGamalCiphertext<2>,
         transcript: &mut Transcript,
     ) -> Result<(), ValidityProofVerificationError> {
+        // We reject if the first public key or the commitments are the identity point.
+        // We allow the second public key to be an identity point as it is often the auditor's
+        // public key in the token-2022 program that can be the identity.
+        if first_pubkey.get_point().is_identity()
+            || grouped_ciphertext_lo.commitment.get_point().is_identity()
+            || grouped_ciphertext_hi.commitment.get_point().is_identity()
+        {
+            return Err(SigmaProofVerificationError::IdentityPoint.into());
+        }
+
         Self::hash_context_into_transcript(
             first_pubkey,
             second_pubkey,

zk-sdk/src/sigma_proofs/batched_grouped_ciphertext_validity/handles_3.rs

@@ -33,13 +33,13 @@ use {
 use {
     crate::{
         sigma_proofs::{
-            errors::ValidityProofVerificationError,
+            errors::{SigmaProofVerificationError, ValidityProofVerificationError},
             grouped_ciphertext_validity::GroupedCiphertext3HandlesValidityProof,
         },
         transcript::TranscriptProtocol,
         UNIT_LEN,
     },
-    curve25519_dalek::scalar::Scalar,
+    curve25519_dalek::{scalar::Scalar, traits::IsIdentity},
     merlin::Transcript,
 };
 
@@ -118,6 +118,17 @@ impl BatchedGroupedCiphertext3HandlesValidityProof {
         grouped_ciphertext_hi: &GroupedElGamalCiphertext<3>,
         transcript: &mut Transcript,
     ) -> Result<(), ValidityProofVerificationError> {
+        // We reject if the public keys or the commitments are the identity point.
+        // The exception is the third public key, which is often the auditor's
+        // public key in the tokne-2022 program that can be the identity.
+        if first_pubkey.get_point().is_identity()
+            || second_pubkey.get_point().is_identity()
+            || grouped_ciphertext_lo.commitment.get_point().is_identity()
+            || grouped_ciphertext_hi.commitment.get_point().is_identity()
+        {
+            return Err(SigmaProofVerificationError::IdentityPoint.into());
+        }
+
         Self::hash_context_into_transcript(
             first_pubkey,
             second_pubkey,

zk-sdk/src/sigma_proofs/ciphertext_ciphertext_equality.rs

@@ -152,6 +152,14 @@ impl CiphertextCiphertextEqualityProof {
         second_ciphertext: &ElGamalCiphertext,
         transcript: &mut Transcript,
     ) -> Result<(), EqualityProofVerificationError> {
+        if first_pubkey.get_point().is_identity()
+            || second_pubkey.get_point().is_identity()
+            || first_ciphertext.commitment.get_point().is_identity()
+            || second_ciphertext.commitment.get_point().is_identity()
+        {
+            return Err(SigmaProofVerificationError::IdentityPoint.into());
+        }
+
         Self::hash_context_into_transcript(
             first_pubkey,
             second_pubkey,

zk-sdk/src/sigma_proofs/ciphertext_commitment_equality.rs

@@ -143,6 +143,13 @@ impl CiphertextCommitmentEqualityProof {
         commitment: &PedersenCommitment,
         transcript: &mut Transcript,
     ) -> Result<(), EqualityProofVerificationError> {
+        if pubkey.get_point().is_identity()
+            || ciphertext.commitment.get_point().is_identity()
+            || commitment.get_point().is_identity()
+        {
+            return Err(SigmaProofVerificationError::IdentityPoint.into());
+        }
+
         Self::hash_context_into_transcript(pubkey, ciphertext, commitment, transcript);
         transcript.ciphertext_commitment_equality_proof_domain_separator();
 
@@ -345,11 +352,10 @@ mod test {
     }
 
     #[test]
-    fn test_ciphertext_commitment_equality_proof_edge_cases() {
-        // if ElGamal public key zero (public key is invalid), then the proof should always reject
+    fn test_ciphertext_commitment_equality_proof_identity_inputs() {
+        // ElGamal public key zero (already invalid, but now checks IdentityPoint)
         let public = ElGamalPubkey::try_from([0u8; 32].as_slice()).unwrap();
         let secret = ElGamalSecretKey::new_rand();
-
         let elgamal_keypair = ElGamalKeypair::new_for_tests(public, secret);
 
         let message: u64 = 55;
@@ -368,19 +374,19 @@ mod test {
             &mut prover_transcript,
         );
 
-        assert!(proof
-            .verify(
-                elgamal_keypair.pubkey(),
-                &ciphertext,
-                &commitment,
-                &mut verifier_transcript
-            )
-            .is_err());
+        let result = proof.verify(
+            elgamal_keypair.pubkey(),
+            &ciphertext,
+            &commitment,
+            &mut verifier_transcript,
+        );
+        assert_eq!(
+            result.unwrap_err(),
+            EqualityProofVerificationError::from(SigmaProofVerificationError::IdentityPoint)
+        );
 
-        // if ciphertext is all-zero (valid commitment of 0) and commitment is also all-zero, then
-        // the proof should still accept
+        // Ciphertext and Commitment are all-zero
         let elgamal_keypair = ElGamalKeypair::new_rand();
-
         let message: u64 = 0;
         let ciphertext = ElGamalCiphertext::from_bytes(&[0u8; 64]).unwrap();
         let commitment = PedersenCommitment::from_bytes(&[0u8; 32]).unwrap();
@@ -398,19 +404,19 @@ mod test {
             &mut prover_transcript,
         );
 
-        proof
-            .verify(
-                elgamal_keypair.pubkey(),
-                &ciphertext,
-                &commitment,
-                &mut verifier_transcript,
-            )
-            .unwrap();
+        let result = proof.verify(
+            elgamal_keypair.pubkey(),
+            &ciphertext,
+            &commitment,
+            &mut verifier_transcript,
+        );
+        assert_eq!(
+            result.unwrap_err(),
+            EqualityProofVerificationError::from(SigmaProofVerificationError::IdentityPoint)
+        );
 
-        // if commitment is all-zero and the ciphertext is a correct encryption of 0, then the
-        // proof should still accept
+        // Only Commitment is zero
         let elgamal_keypair = ElGamalKeypair::new_rand();
-
         let message: u64 = 0;
         let ciphertext = elgamal_keypair.pubkey().encrypt(message);
         let commitment = PedersenCommitment::from_bytes(&[0u8; 32]).unwrap();
@@ -428,43 +434,16 @@ mod test {
             &mut prover_transcript,
         );
 
-        proof
-            .verify(
-                elgamal_keypair.pubkey(),
-                &ciphertext,
-                &commitment,
-                &mut verifier_transcript,
-            )
-            .unwrap();
-
-        // if ciphertext is all zero and commitment correctly encodes 0, then the proof should
-        // still accept
-        let elgamal_keypair = ElGamalKeypair::new_rand();
-
-        let message: u64 = 0;
-        let ciphertext = ElGamalCiphertext::from_bytes(&[0u8; 64]).unwrap();
-        let (commitment, opening) = Pedersen::new(message);
-
-        let mut prover_transcript = Transcript::new_zk_elgamal_transcript(b"Test");
-        let mut verifier_transcript = Transcript::new_zk_elgamal_transcript(b"Test");
-
-        let proof = CiphertextCommitmentEqualityProof::new(
-            &elgamal_keypair,
+        let result = proof.verify(
+            elgamal_keypair.pubkey(),
             &ciphertext,
             &commitment,
-            &opening,
-            message,
-            &mut prover_transcript,
+            &mut verifier_transcript,
+        );
+        assert_eq!(
+            result.unwrap_err(),
+            EqualityProofVerificationError::from(SigmaProofVerificationError::IdentityPoint)
         );
-
-        proof
-            .verify(
-                elgamal_keypair.pubkey(),
-                &ciphertext,
-                &commitment,
-                &mut verifier_transcript,
-            )
-            .unwrap();
     }
 
     #[test]

zk-sdk/src/sigma_proofs/errors.rs

@@ -11,8 +11,8 @@ pub enum SigmaProofVerificationError {
     MultiscalarMul,
     #[error("transcript failed to produce a challenge")]
     Transcript(#[from] TranscriptError),
-    #[error("public key is the identity")]
-    PubkeyIsIdentity,
+    #[error("input point is the identity")]
+    IdentityPoint,
 }
 
 macro_rules! impl_from_transcript_error {