diff --git a/zk-token-sdk/src/errors.rs b/zk-token-sdk/src/errors.rs index e9e36e43f73c18..a5cdaeb39d09d4 100644 --- a/zk-token-sdk/src/errors.rs +++ b/zk-token-sdk/src/errors.rs @@ -20,6 +20,8 @@ pub enum ProofError { PubkeyDeserialization, #[error("ciphertext does not exist in instruction data")] MissingCiphertext, + #[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 8867823cfee393..c099401b51c8c6 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::ProofError, + instruction::batched_range_proof::MAX_COMMITMENTS, range_proof::RangeProof, }, std::convert::TryInto, @@ -74,6 +75,12 @@ impl ZkProofData for BatchedRangeProofU128Data { #[cfg(not(target_os = "solana"))] fn verify_proof(&self) -> Result<(), ProofError> { 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(ProofError::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 da25267cded33a..a9b1c3c13fae66 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::ProofError, + instruction::batched_range_proof::MAX_COMMITMENTS, range_proof::RangeProof, }, std::convert::TryInto, @@ -72,6 +73,12 @@ impl ZkProofData for BatchedRangeProofU256Data { #[cfg(not(target_os = "solana"))] fn verify_proof(&self) -> Result<(), ProofError> { 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(ProofError::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 ae34771aa5697a..6f7c773bfd8da5 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::ProofError, + instruction::batched_range_proof::MAX_COMMITMENTS, range_proof::RangeProof, }, std::convert::TryInto, @@ -74,6 +75,12 @@ impl ZkProofData for BatchedRangeProofU64Data { #[cfg(not(target_os = "solana"))] fn verify_proof(&self) -> Result<(), ProofError> { 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(ProofError::IllegalCommitmentLength); + } + let mut transcript = self.context_data().new_transcript(); let proof: RangeProof = self.proof.try_into()?;