From e251b8607c8fdbde1000fa8fd11677446f5beb22 Mon Sep 17 00:00:00 2001
From: samkim-crypto <skim13@cs.stanford.edu>
Date: Tue, 21 Nov 2023 08:48:15 +0900
Subject: [PATCH] [zk-token-sdk] Add a length check on range proof commitment
 length (#34165)

add a length check on range proof commitment length
---
 zk-token-sdk/src/errors.rs                                 | 2 ++
 .../batched_range_proof/batched_range_proof_u128.rs        | 7 +++++++
 .../batched_range_proof/batched_range_proof_u256.rs        | 7 +++++++
 .../batched_range_proof/batched_range_proof_u64.rs         | 7 +++++++
 4 files changed, 23 insertions(+)

diff --git a/zk-token-sdk/src/errors.rs b/zk-token-sdk/src/errors.rs
index 328a68ca7..ad43b680d 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 96fbc18a6..916245c2f 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<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()?;
 
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 1a6dad666..1bdba644f 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<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()?;
 
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 2a9fae57e..94b76b5be 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<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()?;