[Zellic] 4.2 Remediation strategy for missing input validation

g16ckt/examples/pairing_gate_counts.rs

@@ -211,7 +211,8 @@ fn main() {
             move || {
                 run_and_print("test_ell_montgomery", inputs, move |ctx, w| {
                     let f0 = fq12_one_const();
-                    let coeffs = pairing::ell_coeffs_montgomery(ctx, &w.g2);
+                    // ignoring _is_valid because this function is used only for benchmarking over valid inputs
+                    let (coeffs, _is_valid) = pairing::ell_coeffs_montgomery(ctx, &w.g2);
                     // Take first coeff triple and evaluate once
                     let c = coeffs.into_iter().next().unwrap();
                     let _f1 = pairing::ell_montgomery(ctx, &f0, &c, &w.g1);

g16ckt/src/gadgets/bn254/final_exponentiation.rs

@@ -5,10 +5,14 @@
 //! the structure used in the main branch.
 
 use ark_ec::bn::BnConfig;
-use ark_ff::{BitIteratorBE, Field};
+use ark_ff::{AdditiveGroup, BitIteratorBE, Field};
 use circuit_component_macro::component;
 
-use crate::{CircuitContext, gadgets::bn254::fq12::Fq12};
+use crate::{
+    CircuitContext,
+    circuit::TRUE_WIRE,
+    gadgets::bn254::fq12::{Fq12, ValidFq12},
+};
 
 pub fn conjugate_native(f: ark_bn254::Fq12) -> ark_bn254::Fq12 {
     ark_bn254::Fq12::new(f.c0, -f.c1)
@@ -98,7 +102,16 @@ pub fn exp_by_neg_x_montgomery<C: CircuitContext>(circuit: &mut C, f: &Fq12) ->
 }
 
 #[component]
-pub fn final_exponentiation_montgomery<C: CircuitContext>(circuit: &mut C, f: &Fq12) -> Fq12 {
+pub fn final_exponentiation_montgomery<C: CircuitContext>(circuit: &mut C, f: &Fq12) -> ValidFq12 {
+    let is_zero = Fq12::equal_constant(circuit, f, &ark_bn254::Fq12::ZERO);
+    let is_valid = circuit.issue_wire();
+    circuit.add_gate(crate::Gate {
+        wire_a: is_zero,
+        wire_b: TRUE_WIRE,
+        wire_c: is_valid,
+        gate_type: crate::GateType::Xor,
+    });
+
     let f_inv = Fq12::inverse_montgomery(circuit, f);
     let f_conjugate = Fq12::conjugate(circuit, f);
     let u = Fq12::mul_montgomery(circuit, &f_inv, &f_conjugate);
@@ -127,7 +140,9 @@ pub fn final_exponentiation_montgomery<C: CircuitContext>(circuit: &mut C, f: &F
     let y18 = Fq12::mul_montgomery(circuit, &r2, &y11);
     let y19 = Fq12::frobenius_montgomery(circuit, &y18, 3);
 
-    Fq12::mul_montgomery(circuit, &y19, &y17)
+    let f = Fq12::mul_montgomery(circuit, &y19, &y17);
+
+    ValidFq12 { f, is_valid }
 }
 
 #[cfg(test)]
@@ -171,6 +186,7 @@ mod tests {
         }
         struct Out {
             value: ark_bn254::Fq12,
+            valid: bool,
         }
         fn encode_fq6_to_wires<M: CircuitMode<WireValue = bool>>(
             val: &ark_bn254::Fq6,
@@ -260,7 +276,7 @@ mod tests {
             }
         }
         impl CircuitOutput<ExecuteMode> for Out {
-            type WireRepr = Fq12Wires;
+            type WireRepr = ValidFq12;
             fn decode(wires: Self::WireRepr, cache: &mut ExecuteMode) -> Self {
                 fn decode_fq6_from_wires(
                     wires: &Fq6Wires,
@@ -298,10 +314,12 @@ mod tests {
                         ark_bn254::Fq2::new(ark_bn254::Fq::from(c2_c0), ark_bn254::Fq::from(c2_c1));
                     ark_bn254::Fq6::new(c0, c1, c2)
                 }
-                let c0 = decode_fq6_from_wires(&wires.0[0], cache);
-                let c1 = decode_fq6_from_wires(&wires.0[1], cache);
+                let c0 = decode_fq6_from_wires(&wires.f.0[0], cache);
+                let c1 = decode_fq6_from_wires(&wires.f.0[1], cache);
+                let valid = cache.lookup_wire(wires.is_valid).expect("missing wire");
                 Self {
                     value: ark_bn254::Fq12::new(c0, c1),
+                    valid,
                 }
             }
         }
@@ -312,6 +330,27 @@ mod tests {
                 final_exponentiation_montgomery(ctx, &input.f)
             });
 
-        assert_eq!(result.output_value.value, expected_m);
+        assert!(
+            result.output_value.valid,
+            "final_exponentiation_montgomery input should be valid"
+        );
+        assert_eq!(
+            result.output_value.value, expected_m,
+            "final_exponentiation_montgomery output should be valid"
+        );
+
+        // Test for non-invertible element
+        let input = In {
+            f: ark_bn254::Fq12::ZERO,
+        };
+        let result: StreamingResult<_, _, Out> =
+            CircuitBuilder::streaming_execute(input, 10_000, |ctx, input| {
+                final_exponentiation_montgomery(ctx, &input.f)
+            });
+
+        assert!(
+            !result.output_value.valid,
+            "final_exponentiation_montgomery input should be invalid"
+        );
     }
 }

g16ckt/src/gadgets/bn254/fq.rs

@@ -297,14 +297,24 @@ impl Fq {
             &BigUint::from_str(Self::MODULUS_ADD_1_DIV_4).unwrap(),
         )
     }
+
+    /// Return a>b in standard form given inputs in montgomery form
+    pub fn greater_than<C: CircuitContext>(circuit: &mut C, a: &Fq, b: &Fq) -> WireId {
+        // First convert the inputs 'a' and 'b' back to standard form
+        let a = Fq::mul_by_constant_montgomery(circuit, a, &ark_bn254::Fq::ONE);
+        let b = Fq::mul_by_constant_montgomery(circuit, b, &ark_bn254::Fq::ONE);
+        // only now perform comparison
+        bigint::greater_than(circuit, &a, &b)
+    }
 }
 
 #[cfg(test)]
 pub(super) mod tests {
     use std::{array, iter};
 
     use ark_ff::AdditiveGroup;
-    use rand::Rng;
+    use rand::{Rng, SeedableRng};
+    use rand_chacha::ChaCha20Rng;
     use test_log::test;
     use tracing::trace;
 
@@ -680,6 +690,29 @@ pub(super) mod tests {
         assert_eq!(result.output_value.value, expected_c);
     }
 
+    #[test]
+    fn test_fq_sqrt_montgomery_roundtrip() {
+        let mut rng = ChaCha20Rng::seed_from_u64(42);
+        for _ in 0..5 {
+            let aa_v = Fq::random(&mut rng);
+            let sqrt_exists = aa_v.sqrt().is_some();
+
+            let aa_montgomery = Fq::as_montgomery(aa_v);
+            let input = FqInput::new([aa_montgomery]);
+
+            let result =
+                CircuitBuilder::streaming_execute::<_, _, FqOutput>(input, 10_000, |ctx, input| {
+                    let [aa_wire] = input;
+                    let sqrt = Fq::sqrt_montgomery(ctx, aa_wire);
+                    Fq::square_montgomery(ctx, &sqrt)
+                });
+
+            let calc_aa_montgomery = result.output_value.value;
+
+            assert_eq!(sqrt_exists, calc_aa_montgomery == aa_montgomery);
+        }
+    }
+
     #[test]
     fn test_fq_multiplexer() {
         let w = 1;

g16ckt/src/gadgets/bn254/fq12.rs

@@ -5,7 +5,7 @@ use rand::Rng;
 use super::fq6::Fq6Components;
 use crate::{
     CircuitContext, Gate, WireId,
-    circuit::{FromWires, WiresObject},
+    circuit::{FromWires, WiresArity, WiresObject},
     gadgets::bn254::{fq::Fq, fq2::Fq2, fq6::Fq6},
 };
 
@@ -447,6 +447,58 @@ impl Fq12 {
     }
 }
 
+/// Analogous to Option<Fq12> where `is_valid` carries `false` if variable is None
+#[derive(Clone, Debug)]
+pub struct ValidFq12 {
+    pub f: Fq12,
+    pub is_valid: WireId,
+}
+
+impl WiresObject for ValidFq12 {
+    fn to_wires_vec(&self) -> Vec<WireId> {
+        let mut wires: Vec<WireId> = self.f.0[0]
+            .to_wires_vec()
+            .into_iter()
+            .chain(self.f.0[1].to_wires_vec())
+            .collect();
+        wires.push(self.is_valid);
+        wires
+    }
+
+    fn clone_from(&self, mut wire_gen: &mut impl FnMut() -> WireId) -> Self {
+        ValidFq12 {
+            f: Fq12([
+                self.f.0[0].clone_from(&mut wire_gen),
+                self.f.0[1].clone_from(&mut wire_gen),
+            ]),
+            is_valid: wire_gen(),
+        }
+    }
+}
+
+impl FromWires for ValidFq12 {
+    fn from_wires(wires: &[WireId]) -> Option<Self> {
+        if wires.len() == ValidFq12::ARITY {
+            let mid = Fq6::N_BITS;
+            let fq6_1 = Fq6::from_wires(&wires[..mid])?;
+            let fq6_2 = Fq6::from_wires(&wires[mid..2 * mid])?;
+            let is_valid_wires = &wires[2 * mid..];
+            assert_eq!(is_valid_wires.len(), 1); // single is valid wire
+            let res = ValidFq12 {
+                f: Fq12([fq6_1, fq6_2]),
+                is_valid: is_valid_wires[0],
+            };
+            Some(res)
+        } else {
+            None
+        }
+    }
+}
+
+impl WiresArity for ValidFq12 {
+    const ARITY: usize = Fq12::N_BITS + 1;
+}
+
 #[cfg(test)]
 mod tests {
     use std::{array, str::FromStr};

g16ckt/src/gadgets/bn254/fq2.rs

@@ -15,7 +15,8 @@ use crate::{
     CircuitContext, Gate, WireId,
     circuit::WiresObject,
     gadgets::{
-        bigint::{BigIntWires, select},
+        basic,
+        bigint::{self, BigIntWires, select},
         bn254::{fp254impl::Fp254Impl, fq::Fq},
     },
 };
@@ -444,6 +445,14 @@ impl Fq2 {
 
         Fq2::from_components(c0_final, c1_final)
     }
+
+    /// Return a>b in standard form given inputs in montgomery form
+    pub fn greater_than<C: CircuitContext>(circuit: &mut C, a: &Fq2, b: &Fq2) -> WireId {
+        let c1_equal = bigint::equal(circuit, a.c1(), b.c1());
+        let c1_greater = Fq::greater_than(circuit, a.c1(), b.c1());
+        let c0_greater = Fq::greater_than(circuit, a.c0(), b.c0());
+        basic::selector(circuit, c0_greater, c1_greater, c1_equal)
+    }
 }
 
 #[cfg(test)]