From 616a6b39190ea9d9410deea7bf0f2a86c49299de Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Mon, 19 Sep 2022 20:54:45 -0700
Subject: [PATCH] Validate EVM proof shape

---
 evm/src/stark.rs    |  4 ++++
 evm/src/verifier.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/evm/src/stark.rs b/evm/src/stark.rs
index 1af8a5e2..00818bbe 100644
--- a/evm/src/stark.rs
+++ b/evm/src/stark.rs
@@ -76,6 +76,10 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
         1.max(self.constraint_degree() - 1)
     }
 
+    fn num_quotient_polys(&self, config: &StarkConfig) -> usize {
+        self.quotient_degree_factor() * config.num_challenges
+    }
+
     /// Computes the FRI instance used to prove this Stark.
     fn fri_instance(
         &self,
diff --git a/evm/src/verifier.rs b/evm/src/verifier.rs
index 3f5a5a88..a9e4d30b 100644
--- a/evm/src/verifier.rs
+++ b/evm/src/verifier.rs
@@ -119,6 +119,7 @@ where
     [(); S::COLUMNS]:,
     [(); C::Hasher::HASH_SIZE]:,
 {
+    validate_proof_shape(&stark, proof, config, ctl_vars.len())?;
     let StarkOpeningSet {
         local_values,
         next_values,
@@ -204,6 +205,49 @@ where
     Ok(())
 }
 
+fn validate_proof_shape<F, C, S, const D: usize>(
+    stark: &S,
+    proof: &StarkProof<F, C, D>,
+    config: &StarkConfig,
+    num_ctls: usize,
+) -> anyhow::Result<()>
+where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+    [(); S::COLUMNS]:,
+    [(); C::Hasher::HASH_SIZE]:,
+{
+    let StarkProof {
+        trace_cap,
+        permutation_ctl_zs_cap,
+        quotient_polys_cap,
+        openings,
+        // The shape of the opening proof will be checked in the FRI verifier (see
+        // validate_fri_proof_shape), so we ignore it here.
+        opening_proof: _,
+    } = proof;
+
+    let degree_bits = proof.recover_degree_bits(config);
+    let fri_params = config.fri_params(degree_bits);
+    let cap_height = fri_params.config.cap_height;
+    ensure!(trace_cap.height() == cap_height);
+    ensure!(permutation_ctl_zs_cap.height() == cap_height);
+    ensure!(quotient_polys_cap.height() == cap_height);
+
+    ensure!(openings.local_values.len() == S::COLUMNS);
+    ensure!(openings.next_values.len() == S::COLUMNS);
+    let num_ctl_zs = num_ctls * config.num_challenges;
+    let num_zs = num_ctl_zs + stark.num_permutation_batches(config);
+    ensure!(openings.permutation_ctl_zs.len() == num_zs);
+    ensure!(openings.permutation_ctl_zs_next.len() == num_zs);
+    ensure!(openings.ctl_zs_last.len() == num_ctl_zs);
+    let num_quotient_polys = stark.num_quotient_polys(config);
+    ensure!(openings.quotient_polys.len() == num_quotient_polys);
+
+    Ok(())
+}
+
 /// Evaluate the Lagrange polynomials `L_0` and `L_(n-1)` at a point `x`.
 /// `L_0(x) = (x^n - 1)/(n * (x - 1))`
 /// `L_(n-1)(x) = (x^n - 1)/(n * (g * x - 1))`, with `g` the first element of the subgroup.