[zk-token-sdk] Restrict range proof generator length and prevent 0-bit range proof

zk-token-sdk/src/range_proof/errors.rs

@@ -2,7 +2,10 @@
 use {crate::errors::TranscriptError, thiserror::Error};
 
 #[derive(Error, Clone, Debug, Eq, PartialEq)]
-pub enum RangeProofGenerationError {}
+pub enum RangeProofGenerationError {
+    #[error("maximum generator length exceeded")]
+    MaximumGeneratorLengthExceeded,
+}
 
 #[derive(Error, Clone, Debug, Eq, PartialEq)]
 pub enum RangeProofVerificationError {
@@ -20,4 +23,12 @@ pub enum RangeProofVerificationError {
     InvalidBitSize,
     #[error("insufficient generators for the proof")]
     InvalidGeneratorsLength,
+    #[error("maximum generator length exceeded")]
+    MaximumGeneratorLengthExceeded,
+}
+
+#[derive(Error, Clone, Debug, Eq, PartialEq)]
+pub enum RangeProofGeneratorError {
+    #[error("maximum generator length exceeded")]
+    MaximumGeneratorLengthExceeded,
 }

zk-token-sdk/src/range_proof/generators.rs

@@ -1,11 +1,15 @@
 use {
+    crate::range_proof::errors::RangeProofGeneratorError,
     curve25519_dalek::{
         digest::{ExtendableOutput, Update, XofReader},
         ristretto::RistrettoPoint,
     },
     sha3::{Sha3XofReader, Shake256},
 };
 
+#[cfg(not(target_os = "solana"))]
+const MAX_GENERATOR_LENGTH: usize = u32::MAX as usize;
+
 /// Generators for Pedersen vector commitments that are used for inner-product proofs.
 struct GeneratorsChain {
     reader: Sha3XofReader,
@@ -67,37 +71,44 @@ pub struct BulletproofGens {
 }
 
 impl BulletproofGens {
-    pub fn new(gens_capacity: usize) -> Self {
+    pub fn new(gens_capacity: usize) -> Result<Self, RangeProofGeneratorError> {
         let mut gens = BulletproofGens {
             gens_capacity: 0,
             G_vec: Vec::new(),
             H_vec: Vec::new(),
         };
-        gens.increase_capacity(gens_capacity);
-        gens
+        gens.increase_capacity(gens_capacity)?;
+        Ok(gens)
     }
 
     /// Increases the generators' capacity to the amount specified.
     /// If less than or equal to the current capacity, does nothing.
-    pub fn increase_capacity(&mut self, new_capacity: usize) {
+    pub fn increase_capacity(
+        &mut self,
+        new_capacity: usize,
+    ) -> Result<(), RangeProofGeneratorError> {
         if self.gens_capacity >= new_capacity {
-            return;
+            return Ok(());
+        }
+
+        if new_capacity > MAX_GENERATOR_LENGTH {
+            return Err(RangeProofGeneratorError::MaximumGeneratorLengthExceeded);
         }
 
-        let label = [b'G'];
         self.G_vec.extend(
-            &mut GeneratorsChain::new(&[label, [b'G']].concat())
+            &mut GeneratorsChain::new(&[b'G'])
                 .fast_forward(self.gens_capacity)
                 .take(new_capacity - self.gens_capacity),
         );
 
         self.H_vec.extend(
-            &mut GeneratorsChain::new(&[label, [b'H']].concat())
+            &mut GeneratorsChain::new(&[b'H'])
                 .fast_forward(self.gens_capacity)
                 .take(new_capacity - self.gens_capacity),
         );
 
         self.gens_capacity = new_capacity;
+        Ok(())
     }
 
     #[allow(non_snake_case)]

zk-token-sdk/src/range_proof/inner_product.rs

@@ -205,7 +205,7 @@ impl InnerProductProof {
         transcript: &mut Transcript,
     ) -> Result<(Vec<Scalar>, Vec<Scalar>, Vec<Scalar>), RangeProofVerificationError> {
         let lg_n = self.L_vec.len();
-        if lg_n >= 32 {
+        if lg_n == 0 || lg_n >= 32 {
             // 4 billion multiplications should be enough for anyone
             // and this check prevents overflow in 1<<lg_n below.
             return Err(RangeProofVerificationError::InvalidBitSize);
@@ -410,7 +410,7 @@ mod tests {
     fn test_basic_correctness() {
         let n = 32;
 
-        let bp_gens = BulletproofGens::new(n);
+        let bp_gens = BulletproofGens::new(n).unwrap();
         let G: Vec<RistrettoPoint> = bp_gens.G(n).cloned().collect();
         let H: Vec<RistrettoPoint> = bp_gens.H(n).cloned().collect();
 

zk-token-sdk/src/range_proof/mod.rs

@@ -82,7 +82,8 @@ impl RangeProof {
         let nm: usize = bit_lengths.iter().sum();
         assert!(nm.is_power_of_two());
 
-        let bp_gens = BulletproofGens::new(nm);
+        let bp_gens = BulletproofGens::new(nm)
+            .map_err(|_| RangeProofGenerationError::MaximumGeneratorLengthExceeded)?;
 
         // bit-decompose values and generate their Pedersen vector commitment
         let a_blinding = Scalar::random(&mut OsRng);
@@ -241,7 +242,8 @@ impl RangeProof {
 
         let m = bit_lengths.len();
         let nm: usize = bit_lengths.iter().sum();
-        let bp_gens = BulletproofGens::new(nm);
+        let bp_gens = BulletproofGens::new(nm)
+            .map_err(|_| RangeProofVerificationError::MaximumGeneratorLengthExceeded)?;
 
         if !nm.is_power_of_two() {
             return Err(RangeProofVerificationError::InvalidBitSize);