fix: protect against weak Fiat-Shamir (#54)

* fix: protect against weak Fiat-Shamir

* add log

* feat: switch to row hash to avoid DFT overhead

* feat: switch to bincode for stability

* use inner pattern for code quality

* Revert "feat: switch to bincode for stability"

vkey size increases a lot

This reverts commit 1e1f7ad.

Author: jonathanpwang

Cargo.toml

@@ -108,6 +108,7 @@ async-trait = "0.1.83"
 getset = "0.1.3"
 rand = { version = "0.8.5", default-features = false }
 hex = { version = "0.4.3", default-features = false }
+bitcode = "0.6.5"
 
 # default-features = false for no_std
 itertools = { version = "0.14.0", default-features = false }

crates/stark-backend/Cargo.toml

@@ -28,8 +28,8 @@ derive-new.workspace = true
 metrics = { workspace = true, optional = true }
 cfg-if.workspace = true
 thiserror.workspace = true
-async-trait.workspace = true
 rustc-hash.workspace = true
+bitcode = { workspace = true, features = ["serde"] }
 
 [target.'cfg(unix)'.dependencies]
 tikv-jemallocator = { version = "0.6", optional = true }

crates/stark-backend/src/keygen/mod.rs

@@ -2,9 +2,10 @@ use std::{collections::HashMap, iter::zip, sync::Arc};
 
 use itertools::Itertools;
 use p3_commit::Pcs;
-use p3_field::{Field, FieldExtensionAlgebra};
-use p3_matrix::Matrix;
+use p3_field::{Field, FieldAlgebra, FieldExtensionAlgebra};
+use p3_matrix::{dense::RowMajorMatrix, Matrix};
 use tracing::instrument;
+use types::MultiStarkVerifyingKey0;
 
 use crate::{
     air_builders::symbolic::{get_symbolic_builder, SymbolicRapBuilder},
@@ -195,11 +196,33 @@ impl<'a, SC: StarkGenericConfig> MultiStarkKeygenBuilder<'a, SC> {
             threshold: log_up_security_params.max_interaction_count,
         });
 
+        let pre_vk: MultiStarkVerifyingKey0<SC> = MultiStarkVerifyingKey0 {
+            per_air: pk_per_air.iter().map(|pk| pk.vk.clone()).collect(),
+            trace_height_constraints: trace_height_constraints.clone(),
+            log_up_pow_bits: log_up_security_params.log_up_pow_bits,
+        };
+        // To protect against weak Fiat-Shamir, we hash the "pre"-verifying key and include it in the
+        // final verifying key. This just needs to commit to the verifying key and does not need to be
+        // verified by the verifier, so we just use bincode to serialize it.
+        let vk_bytes = bitcode::serialize(&pre_vk).unwrap();
+        tracing::info!("pre-vkey: {} bytes", vk_bytes.len());
+        // Purely to get type compatibility and convenience, we hash using pcs.commit as a single row
+        let vk_as_row = RowMajorMatrix::new_row(
+            vk_bytes
+                .into_iter()
+                .map(Val::<SC>::from_canonical_u8)
+                .collect(),
+        );
+        let pcs = self.config.pcs();
+        let deg_1_domain = pcs.natural_domain_for_degree(1);
+        let (vk_pre_hash, _) = pcs.commit(vec![(deg_1_domain, vk_as_row)]);
+
         MultiStarkProvingKey {
             per_air: pk_per_air,
             trace_height_constraints,
             max_constraint_degree: self.max_constraint_degree,
             log_up_pow_bits: log_up_security_params.log_up_pow_bits,
+            vk_pre_hash,
         }
     }
 }

crates/stark-backend/src/keygen/types.rs

@@ -84,6 +84,22 @@ pub struct StarkVerifyingKey<Val, Com> {
     deserialize = "Com<SC>: Deserialize<'de>"
 ))]
 pub struct MultiStarkVerifyingKey<SC: StarkGenericConfig> {
+    /// All parts of the verifying key needed by the verifier, except
+    /// the `pre_hash` used to initialize the Fiat-Shamir transcript.
+    pub inner: MultiStarkVerifyingKey0<SC>,
+    /// The hash of all other parts of the verifying key. The Fiat-Shamir hasher will
+    /// initialize by observing this hash.
+    pub pre_hash: Com<SC>,
+}
+
+/// Everything in [MultiStarkVerifyingKey] except the `pre_hash` used to initialize the Fiat-Shamir transcript.
+#[derive(Derivative, Serialize, Deserialize)]
+#[derivative(Clone(bound = "Com<SC>: Clone"))]
+#[serde(bound(
+    serialize = "Com<SC>: Serialize",
+    deserialize = "Com<SC>: Deserialize<'de>"
+))]
+pub struct MultiStarkVerifyingKey0<SC: StarkGenericConfig> {
     pub per_air: Vec<StarkVerifyingKey<Val<SC>, Com<SC>>>,
     pub trace_height_constraints: Vec<LinearConstraint>,
     pub log_up_pow_bits: usize,
@@ -129,6 +145,8 @@ pub struct MultiStarkProvingKey<SC: StarkGenericConfig> {
     /// Maximum degree of constraints across all AIRs
     pub max_constraint_degree: usize,
     pub log_up_pow_bits: usize,
+    /// See [MultiStarkVerifyingKey]
+    pub vk_pre_hash: Com<SC>,
 }
 
 impl<Val, Com> StarkVerifyingKey<Val, Com> {
@@ -148,6 +166,13 @@ impl<Val, Com> StarkVerifyingKey<Val, Com> {
 impl<SC: StarkGenericConfig> MultiStarkProvingKey<SC> {
     pub fn get_vk(&self) -> MultiStarkVerifyingKey<SC> {
         MultiStarkVerifyingKey {
+            inner: self.get_vk0(),
+            pre_hash: self.vk_pre_hash.clone(),
+        }
+    }
+
+    fn get_vk0(&self) -> MultiStarkVerifyingKey0<SC> {
+        MultiStarkVerifyingKey0 {
             per_air: self.per_air.iter().map(|pk| pk.vk.clone()).collect(),
             trace_height_constraints: self.trace_height_constraints.clone(),
             log_up_pow_bits: self.log_up_pow_bits,

crates/stark-backend/src/keygen/view.rs

@@ -11,17 +11,19 @@ pub(crate) struct MultiStarkVerifyingKeyView<'a, Val, Com> {
     /// Trace height constraints are *not* filtered by AIR. When computing the dot product, this
     /// will be indexed into by air_id.
     pub trace_height_constraints: &'a [LinearConstraint],
+    pub pre_hash: Com,
 }
 
 impl<SC: StarkGenericConfig> MultiStarkVerifyingKey<SC> {
     /// Returns a view with all airs.
     pub(crate) fn full_view(&self) -> MultiStarkVerifyingKeyView<Val<SC>, Com<SC>> {
-        self.view(&(0..self.per_air.len()).collect_vec())
+        self.view(&(0..self.inner.per_air.len()).collect_vec())
     }
     pub(crate) fn view(&self, air_ids: &[usize]) -> MultiStarkVerifyingKeyView<Val<SC>, Com<SC>> {
         MultiStarkVerifyingKeyView {
-            per_air: air_ids.iter().map(|&id| &self.per_air[id]).collect(),
-            trace_height_constraints: &self.trace_height_constraints,
+            per_air: air_ids.iter().map(|&id| &self.inner.per_air[id]).collect(),
+            trace_height_constraints: &self.inner.trace_height_constraints,
+            pre_hash: self.pre_hash.clone(),
         }
     }
 }

crates/stark-backend/src/prover/coordinator.rs

@@ -84,6 +84,7 @@ where
         #[cfg(feature = "bench-metrics")]
         let start = std::time::Instant::now();
         assert!(mpk.validate(&ctx), "Invalid proof input");
+        self.challenger.observe(mpk.vk_pre_hash.clone());
 
         let num_air = ctx.per_air.len();
         info!(num_air);
@@ -307,6 +308,7 @@ impl<'a, PB: ProverBackend> DeviceMultiStarkProvingKey<'a, PB> {
         MultiStarkVerifyingKeyView::new(
             self.per_air.iter().map(|pk| pk.vk).collect(),
             &self.trace_height_constraints,
+            self.vk_pre_hash.clone(),
         )
     }
 }

crates/stark-backend/src/prover/cpu/mod.rs

@@ -443,7 +443,12 @@ where
                 }
             })
             .collect();
-        DeviceMultiStarkProvingKey::new(air_ids, per_air, mpk.trace_height_constraints.clone())
+        DeviceMultiStarkProvingKey::new(
+            air_ids,
+            per_air,
+            mpk.trace_height_constraints.clone(),
+            mpk.vk_pre_hash.clone(),
+        )
     }
     fn transport_matrix_to_device(
         &self,

crates/stark-backend/src/prover/types.rs

@@ -19,19 +19,22 @@ pub struct DeviceMultiStarkProvingKey<'a, PB: ProverBackend> {
     /// Each [LinearConstraint] is indexed by AIR ID.
     /// **Caution**: the linear constraints are **not** filtered for only the AIRs appearing in `per_air`.
     pub trace_height_constraints: Vec<LinearConstraint>,
+    pub vk_pre_hash: PB::Commitment,
 }
 
 impl<'a, PB: ProverBackend> DeviceMultiStarkProvingKey<'a, PB> {
     pub fn new(
         air_ids: Vec<usize>,
         per_air: Vec<DeviceStarkProvingKey<'a, PB>>,
         trace_height_constraints: Vec<LinearConstraint>,
+        vk_pre_hash: PB::Commitment,
     ) -> Self {
         assert_eq!(air_ids.len(), per_air.len());
         Self {
             air_ids,
             per_air,
             trace_height_constraints,
+            vk_pre_hash,
         }
     }
 }

crates/stark-backend/src/verifier/mod.rs

@@ -57,6 +57,7 @@ impl<'c, SC: StarkGenericConfig> MultiTraceStarkVerifier<'c, SC> {
         mvk: &MultiStarkVerifyingKeyView<Val<SC>, Com<SC>>,
         proof: &Proof<SC>,
     ) -> Result<(), VerificationError> {
+        challenger.observe(mvk.pre_hash.clone());
         // Enforce trace height linear inequalities
         for constraint in mvk.trace_height_constraints {
             let sum = proof

crates/stark-backend/tests/interaction/mod.rs

@@ -46,7 +46,7 @@ fn test_interaction_trace_height_constraints() {
     keygen_builder.add_air(Arc::new(sender_air_2));
     keygen_builder.add_air(Arc::new(sender_air_3));
     let pk = keygen_builder.generate_pk();
-    let vk = pk.get_vk();
+    let vk = pk.get_vk().inner;
 
     assert_eq!(vk.trace_height_constraints.len(), 3);