add a length check on range proof commitment length

Author: samkim-crypto

zk-token-sdk/src/errors.rs

@@ -36,6 +36,8 @@ pub enum ProofVerificationError {
     ElGamal(#[from] ElGamalError),
     #[error("Invalid proof context")]
     ProofContext,
+    #[error("illegal commitment length")]
+    IllegalCommitmentLength,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]

zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u128.rs

@@ -5,6 +5,7 @@ use {
     crate::{
         encryption::pedersen::{PedersenCommitment, PedersenOpening},
         errors::{ProofGenerationError, ProofVerificationError},
+        instruction::batched_range_proof::MAX_COMMITMENTS,
         range_proof::RangeProof,
     },
     std::convert::TryInto,
@@ -77,6 +78,12 @@ impl ZkProofData<BatchedRangeProofContext> for BatchedRangeProofU128Data {
     #[cfg(not(target_os = "solana"))]
     fn verify_proof(&self) -> Result<(), ProofVerificationError> {
         let (commitments, bit_lengths) = self.context.try_into()?;
+        let num_commitments = commitments.len();
+
+        if num_commitments > MAX_COMMITMENTS || num_commitments != bit_lengths.len() {
+            return Err(ProofVerificationError::IllegalCommitmentLength);
+        }
+
         let mut transcript = self.context_data().new_transcript();
         let proof: RangeProof = self.proof.try_into()?;
 

zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u256.rs

@@ -5,6 +5,7 @@ use {
     crate::{
         encryption::pedersen::{PedersenCommitment, PedersenOpening},
         errors::{ProofGenerationError, ProofVerificationError},
+        instruction::batched_range_proof::MAX_COMMITMENTS,
         range_proof::RangeProof,
     },
     std::convert::TryInto,
@@ -74,6 +75,12 @@ impl ZkProofData<BatchedRangeProofContext> for BatchedRangeProofU256Data {
     #[cfg(not(target_os = "solana"))]
     fn verify_proof(&self) -> Result<(), ProofVerificationError> {
         let (commitments, bit_lengths) = self.context.try_into()?;
+        let num_commitments = commitments.len();
+
+        if num_commitments > MAX_COMMITMENTS || num_commitments != bit_lengths.len() {
+            return Err(ProofVerificationError::IllegalCommitmentLength);
+        }
+
         let mut transcript = self.context_data().new_transcript();
         let proof: RangeProof = self.proof.try_into()?;
 

zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u64.rs

@@ -5,6 +5,7 @@ use {
     crate::{
         encryption::pedersen::{PedersenCommitment, PedersenOpening},
         errors::{ProofGenerationError, ProofVerificationError},
+        instruction::batched_range_proof::MAX_COMMITMENTS,
         range_proof::RangeProof,
     },
     std::convert::TryInto,
@@ -76,6 +77,12 @@ impl ZkProofData<BatchedRangeProofContext> for BatchedRangeProofU64Data {
     #[cfg(not(target_os = "solana"))]
     fn verify_proof(&self) -> Result<(), ProofVerificationError> {
         let (commitments, bit_lengths) = self.context.try_into()?;
+        let num_commitments = commitments.len();
+
+        if num_commitments > MAX_COMMITMENTS || num_commitments != bit_lengths.len() {
+            return Err(ProofVerificationError::IllegalCommitmentLength);
+        }
+
         let mut transcript = self.context_data().new_transcript();
         let proof: RangeProof = self.proof.try_into()?;