From 13e90acecb084ce4c34cfb644731cc3fafc4d3f5 Mon Sep 17 00:00:00 2001
From: moana <moana@dusk.network>
Date: Mon, 14 Aug 2023 15:33:14 +0200
Subject: [PATCH] Add test and fix for panic on prove

This fixes a panic on proof creation when the circuit to prove is
larger than in the provers circuit description

Resolves: #760
---
 CHANGELOG.md    |  5 ++++
 src/composer.rs |  5 ++++
 src/error.rs    |  8 ++++++-
 tests/size.rs   | 64 +++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 81 insertions(+), 1 deletion(-)
 create mode 100644 tests/size.rs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d47a621a..f615f589 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Fix panic when creating proof for circuit with different circuit size [#760]
+
 ## [0.14.1] - 2022-06-28
 
 ### Added
@@ -478,6 +482,7 @@ is necessary since `rkyv/validation` was required as a bound.
 - Proof system module.
 
 <!-- ISSUES -->
+[#760]: https://github.com/dusk-network/plonk/issues/760
 [#752]: https://github.com/dusk-network/plonk/pull/752
 [#738]: https://github.com/dusk-network/plonk/issues/738
 [#746]: https://github.com/dusk-network/plonk/issues/746
diff --git a/src/composer.rs b/src/composer.rs
index 461cf894..ba8df417 100644
--- a/src/composer.rs
+++ b/src/composer.rs
@@ -1066,6 +1066,11 @@ pub trait Composer: Sized + Index<Witness, Output = BlsScalar> {
 
         circuit.circuit(&mut builder)?;
 
+        // assert that the circuit has the expected amount of constraints
+        if builder.constraints() != constraints {
+            return Err(Error::InvalidCircuitSize);
+        }
+
         builder.runtime().event(RuntimeEvent::ProofFinished);
 
         Ok(builder)
diff --git a/src/error.rs b/src/error.rs
index bec2832f..cddd26ef 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -9,7 +9,7 @@
 use dusk_bytes::Error as DuskBytesError;
 
 /// Defines all possible errors that can be encountered in PLONK.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub enum Error {
     // FFT errors
     /// This error occurs when an error triggers on any of the fft module
@@ -35,6 +35,9 @@ pub enum Error {
     /// This error occurs when the Prover structure already contains a
     /// preprocessed circuit inside, but you call preprocess again.
     CircuitAlreadyPreprocessed,
+    /// This error occurs when the circuit for the proof has a different size
+    /// than the prover circuit description
+    InvalidCircuitSize,
 
     // Preprocessing errors
     /// This error occurs when an error triggers during the preprocessing
@@ -121,6 +124,9 @@ impl std::fmt::Display for Error {
             Self::CircuitAlreadyPreprocessed => {
                 write!(f, "circuit has already been preprocessed")
             }
+            Self::InvalidCircuitSize => {
+                write!(f, "circuit size doesn't match with circuit description")
+            }
             Self::DegreeIsZero => {
                 write!(f, "cannot create PublicParameters with max degree 0")
             }
diff --git a/tests/size.rs b/tests/size.rs
new file mode 100644
index 00000000..e8c9c3f7
--- /dev/null
+++ b/tests/size.rs
@@ -0,0 +1,64 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+//
+// Copyright (c) DUSK NETWORK. All rights reserved.
+
+use dusk_plonk::prelude::*;
+use rand::rngs::StdRng;
+use rand::SeedableRng;
+
+const CAPACITY: usize = 1 << 4;
+const LABEL: &[u8] = b"check_public_inputs";
+
+#[derive(Default)]
+pub struct TestSize {
+    witnesses: Vec<BlsScalar>,
+    sum: BlsScalar,
+}
+
+impl TestSize {
+    pub fn new(witnesses: Vec<BlsScalar>, sum: BlsScalar) -> Self {
+        Self { witnesses, sum }
+    }
+}
+
+impl Circuit for TestSize {
+    fn circuit<C>(&self, composer: &mut C) -> Result<(), Error>
+    where
+        C: Composer,
+    {
+        let sum = self.witnesses.iter().fold(C::ZERO, |acc, scalar| {
+            let w = composer.append_witness(*scalar);
+            let constraint = Constraint::new().left(1).a(acc).right(1).b(w);
+            composer.gate_add(constraint)
+        });
+
+        let expected_sum = composer.append_witness(self.sum);
+        composer.assert_equal(sum, expected_sum);
+
+        Ok(())
+    }
+}
+
+#[test]
+fn size() {
+    let rng = &mut StdRng::seed_from_u64(0x10b);
+    let pp = PublicParameters::setup(CAPACITY, rng)
+        .expect("Creation of public parameter shouldn't fail");
+
+    // compiling the default version of TestSize, which only one gate: sum = 0
+    let (prover, _verifier) = Compiler::compile::<TestSize>(&pp, LABEL)
+        .expect("It should be possible to compile the prover and verifier");
+
+    // Create circuit with more gates
+    let pi: Vec<BlsScalar> = [BlsScalar::one(); 5].into();
+    let sum = pi.iter().sum();
+    let circuit = TestSize::new(pi, sum);
+    let result = prover.prove(rng, &circuit);
+    assert_eq!(
+        result,
+        Err(Error::InvalidCircuitSize),
+        "proof creation for different sized circuit shouldn't be possible"
+    );
+}