diff --git a/zk-token-sdk/src/errors.rs b/zk-token-sdk/src/errors.rs index 328a68ca755437..ad43b680dc9b0d 100644 --- a/zk-token-sdk/src/errors.rs +++ b/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)] diff --git a/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u128.rs b/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u128.rs index 96fbc18a6b612e..916245c2f31611 100644 --- a/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u128.rs +++ b/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 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()?; diff --git a/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u256.rs b/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u256.rs index 1a6dad66656487..1bdba644f3c296 100644 --- a/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u256.rs +++ b/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 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()?; diff --git a/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u64.rs b/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u64.rs index 2a9fae57e46d6f..94b76b5beff89d 100644 --- a/zk-token-sdk/src/instruction/batched_range_proof/batched_range_proof_u64.rs +++ b/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 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()?;