From b6625fa0b08bccd47f8b9cbaaa8fa0245bf30892 Mon Sep 17 00:00:00 2001 From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com> Date: Wed, 29 Nov 2023 16:50:05 -0800 Subject: [PATCH] [fix] soundness bug in `BasicDynLookupConfig::assign_virtual_table_to_raw` (#224) * fix: `BasicDynLookupConfig::assign_virtual_table_to_raw` * feat: add safety check on all assigned_advice HashMap insertions --- .../gates/flex_gate/threads/single_phase.rs | 10 +++- halo2-base/src/utils/halo2.rs | 49 ++++++++----------- .../src/virtual_region/copy_constraints.rs | 13 +++-- halo2-base/src/virtual_region/lookups.rs | 5 +- .../src/virtual_region/lookups/basic.rs | 25 ++++++---- .../virtual_region/tests/lookups/memory.rs | 40 ++++++++------- 6 files changed, 79 insertions(+), 63 deletions(-) diff --git a/halo2-base/src/gates/flex_gate/threads/single_phase.rs b/halo2-base/src/gates/flex_gate/threads/single_phase.rs index f9359814..a554d727 100644 --- a/halo2-base/src/gates/flex_gate/threads/single_phase.rs +++ b/halo2-base/src/gates/flex_gate/threads/single_phase.rs @@ -219,9 +219,15 @@ pub fn assign_with_constraints( .assign_advice(|| "", column, row_offset, || value.map(|v| *v)) .unwrap() .cell(); - copy_manager + if let Some(old_cell) = copy_manager .assigned_advices - .insert(ContextCell::new(ctx.type_id, ctx.context_id, i), cell); + .insert(ContextCell::new(ctx.type_id, ctx.context_id, i), cell) + { + assert!( + old_cell.row_offset == cell.row_offset && old_cell.column == cell.column, + "Trying to overwrite virtual cell with a different raw cell" + ); + } // If selector enabled and row_offset is valid add break point, account for break point overlap, and enforce equality constraint for gate outputs. // ⚠️ This assumes overlap is of form: gate enabled at `i - delta` and `i`, where `delta = ROTATIONS - 1`. We currently do not support `delta < ROTATIONS - 1`. diff --git a/halo2-base/src/utils/halo2.rs b/halo2-base/src/utils/halo2.rs index 463b128f..a3781342 100644 --- a/halo2-base/src/utils/halo2.rs +++ b/halo2-base/src/utils/halo2.rs @@ -1,9 +1,11 @@ +use std::collections::hash_map::Entry; + use crate::ff::Field; use crate::halo2_proofs::{ circuit::{AssignedCell, Cell, Region, Value}, plonk::{Advice, Assigned, Column, Fixed}, }; -use crate::virtual_region::copy_constraints::{CopyConstraintManager, SharedCopyConstraintManager}; +use crate::virtual_region::copy_constraints::{CopyConstraintManager, EXTERNAL_CELL_TYPE_ID}; use crate::AssignedValue; /// Raw (physical) assigned cell in Plonkish arithmetization. @@ -74,30 +76,11 @@ pub fn raw_constrain_equal(region: &mut Region, left: Cell, right: region.constrain_equal(left, right).unwrap(); } -/// Assign virtual cell to raw halo2 cell in column `column` at row offset `offset` within the `region`. -/// Stores the mapping between `virtual_cell` and the raw assigned cell in `copy_manager`, if provided. -/// -/// `copy_manager` **must** be provided unless you are only doing witness generation -/// without constraints. -pub fn assign_virtual_to_raw<'v, F: Field + Ord>( - region: &mut Region, - column: Column, - offset: usize, - virtual_cell: AssignedValue, - copy_manager: Option<&SharedCopyConstraintManager>, -) -> Halo2AssignedCell<'v, F> { - let raw = raw_assign_advice(region, column, offset, Value::known(virtual_cell.value)); - if let Some(copy_manager) = copy_manager { - let mut copy_manager = copy_manager.lock().unwrap(); - let cell = virtual_cell.cell.unwrap(); - copy_manager.assigned_advices.insert(cell, raw.cell()); - drop(copy_manager); - } - raw -} - -/// Constrains that `virtual` is equal to `external`. The `virtual` cell must have -/// **already** been raw assigned, with the raw assigned cell stored in `copy_manager`. +/// Constrains that `virtual_cell` is equal to `external_cell`. The `virtual_cell` must have +/// already been raw assigned with the raw assigned cell stored in `copy_manager` +/// **unless** it is marked an external-only cell with type id [EXTERNAL_CELL_TYPE_ID]. +/// * When the virtual cell has already been assigned, the assigned cell is constrained to be equal to the external cell. +/// * When the virtual cell has not been assigned **and** it is marked as an external cell, it is assigned to `external_cell` and the mapping is stored in `copy_manager`. /// /// This should only be called when `witness_gen_only` is false, otherwise it will panic. /// @@ -107,9 +90,19 @@ pub fn constrain_virtual_equals_external( region: &mut Region, virtual_cell: AssignedValue, external_cell: Cell, - copy_manager: &CopyConstraintManager, + copy_manager: &mut CopyConstraintManager, ) { let ctx_cell = virtual_cell.cell.unwrap(); - let acell = copy_manager.assigned_advices.get(&ctx_cell).expect("cell not assigned"); - region.constrain_equal(*acell, external_cell); + match copy_manager.assigned_advices.entry(ctx_cell) { + Entry::Occupied(acell) => { + // The virtual cell has already been assigned, so we can constrain it to equal the external cell. + region.constrain_equal(*acell.get(), external_cell); + } + Entry::Vacant(assigned) => { + // The virtual cell **must** be an external cell + assert_eq!(ctx_cell.type_id, EXTERNAL_CELL_TYPE_ID); + // We map the virtual cell to point to the raw external cell in `copy_manager` + assigned.insert(external_cell); + } + } } diff --git a/halo2-base/src/virtual_region/copy_constraints.rs b/halo2-base/src/virtual_region/copy_constraints.rs index 5ab80d3b..11a77944 100644 --- a/halo2-base/src/virtual_region/copy_constraints.rs +++ b/halo2-base/src/virtual_region/copy_constraints.rs @@ -15,6 +15,9 @@ use crate::{ff::Field, ContextCell}; use super::manager::VirtualRegionManager; +/// Type ID to distinguish external raw Halo2 cells. **This Type ID must be unique.** +pub const EXTERNAL_CELL_TYPE_ID: &str = "halo2-base:External Raw Halo2 Cell"; + /// Thread-safe shared global manager for all copy constraints. pub type SharedCopyConstraintManager = Arc>>; @@ -86,11 +89,15 @@ impl CopyConstraintManager { } fn load_external_cell_impl(&mut self, cell: Option) -> ContextCell { - let context_cell = - ContextCell::new("halo2-base:External Raw Halo2 Cell", 0, self.external_cell_count); + let context_cell = ContextCell::new(EXTERNAL_CELL_TYPE_ID, 0, self.external_cell_count); self.external_cell_count += 1; if let Some(cell) = cell { - self.assigned_advices.insert(context_cell, cell); + if let Some(old_cell) = self.assigned_advices.insert(context_cell, cell) { + assert!( + old_cell.row_offset == cell.row_offset && old_cell.column == cell.column, + "External cell already assigned" + ) + } } context_cell } diff --git a/halo2-base/src/virtual_region/lookups.rs b/halo2-base/src/virtual_region/lookups.rs index 3e301921..7823a573 100644 --- a/halo2-base/src/virtual_region/lookups.rs +++ b/halo2-base/src/virtual_region/lookups.rs @@ -125,7 +125,8 @@ impl VirtualRegionManager type Config = Vec<[Column; ADVICE_COLS]>; fn assign_raw(&self, config: &Self::Config, region: &mut Region) { - let copy_manager = (!self.witness_gen_only).then(|| self.copy_manager().lock().unwrap()); + let mut copy_manager = + (!self.witness_gen_only).then(|| self.copy_manager().lock().unwrap()); let cells_to_lookup = self.cells_to_lookup.lock().unwrap(); // Copy the cells to the config columns, going left to right, then top to bottom. // Will panic if out of rows @@ -139,7 +140,7 @@ impl VirtualRegionManager for (advice, &column) in advices.iter().zip(config[lookup_col].iter()) { let bcell = raw_assign_advice(region, column, lookup_offset, Value::known(advice.value)); - if let Some(copy_manager) = copy_manager.as_ref() { + if let Some(copy_manager) = copy_manager.as_mut() { constrain_virtual_equals_external(region, *advice, bcell.cell(), copy_manager); } } diff --git a/halo2-base/src/virtual_region/lookups/basic.rs b/halo2-base/src/virtual_region/lookups/basic.rs index 3b214545..f5299c38 100644 --- a/halo2-base/src/virtual_region/lookups/basic.rs +++ b/halo2-base/src/virtual_region/lookups/basic.rs @@ -8,10 +8,7 @@ use crate::{ poly::Rotation, }, utils::{ - halo2::{ - assign_virtual_to_raw, constrain_virtual_equals_external, raw_assign_advice, - raw_assign_fixed, - }, + halo2::{constrain_virtual_equals_external, raw_assign_advice, raw_assign_fixed}, ScalarField, }, virtual_region::copy_constraints::SharedCopyConstraintManager, @@ -83,7 +80,7 @@ impl BasicDynLookupConfig { Self { table_is_enabled, table, to_lookup } } - /// Assign managed lookups + /// Assign managed lookups. The `keys` must have already been raw assigned beforehand. /// /// `copy_manager` **must** be provided unless you are only doing witness generation /// without constraints. @@ -114,6 +111,8 @@ impl BasicDynLookupConfig { .unwrap(); } + /// Assign managed lookups. The `keys` must have already been raw assigned beforehand. + /// /// `copy_manager` **must** be provided unless you are only doing witness generation /// without constraints. pub fn assign_virtual_to_lookup_to_raw_from_offset( @@ -123,7 +122,7 @@ impl BasicDynLookupConfig { mut offset: usize, copy_manager: Option<&SharedCopyConstraintManager>, ) { - let copy_manager = copy_manager.map(|c| c.lock().unwrap()); + let mut copy_manager = copy_manager.map(|c| c.lock().unwrap()); // Copied from `LookupAnyManager::assign_raw` but modified to set `key_is_enabled` to 1. // Copy the cells to the config columns, going left to right, then top to bottom. // Will panic if out of rows @@ -138,7 +137,7 @@ impl BasicDynLookupConfig { raw_assign_fixed(region, key_is_enabled_col, offset, F::ONE); for (advice, column) in zip(key, key_col) { let bcell = raw_assign_advice(region, column, offset, Value::known(advice.value)); - if let Some(copy_manager) = copy_manager.as_ref() { + if let Some(copy_manager) = copy_manager.as_mut() { constrain_virtual_equals_external(region, advice, bcell.cell(), copy_manager); } } @@ -147,7 +146,7 @@ impl BasicDynLookupConfig { } } - /// Assign virtual table to raw. + /// Assign virtual table to raw. The `rows` must have already been raw assigned beforehand. /// /// `copy_manager` **must** be provided unless you are only doing witness generation /// without constraints. @@ -178,6 +177,8 @@ impl BasicDynLookupConfig { .unwrap(); } + /// Assign virtual table to raw. The `rows` must have already been raw assigned beforehand. + /// /// `copy_manager` **must** be provided unless you are only doing witness generation /// without constraints. pub fn assign_virtual_table_to_raw_from_offset( @@ -187,11 +188,15 @@ impl BasicDynLookupConfig { mut offset: usize, copy_manager: Option<&SharedCopyConstraintManager>, ) { + let mut copy_manager = copy_manager.map(|c| c.lock().unwrap()); for row in rows { // Enable this row in the table raw_assign_fixed(region, self.table_is_enabled, offset, F::ONE); - for (col, virtual_cell) in self.table.into_iter().zip(row) { - assign_virtual_to_raw(region, col, offset, virtual_cell, copy_manager); + for (advice, column) in zip(row, self.table) { + let bcell = raw_assign_advice(region, column, offset, Value::known(advice.value)); + if let Some(copy_manager) = copy_manager.as_mut() { + constrain_virtual_equals_external(region, advice, bcell.cell(), copy_manager); + } } offset += 1; } diff --git a/halo2-base/src/virtual_region/tests/lookups/memory.rs b/halo2-base/src/virtual_region/tests/lookups/memory.rs index d1adc3b9..8b94a6e5 100644 --- a/halo2-base/src/virtual_region/tests/lookups/memory.rs +++ b/halo2-base/src/virtual_region/tests/lookups/memory.rs @@ -1,5 +1,3 @@ -use std::any::TypeId; - use crate::{ halo2_proofs::{ arithmetic::Field, @@ -8,7 +6,9 @@ use crate::{ halo2curves::bn256::Fr, plonk::{keygen_pk, keygen_vk, Assigned, Circuit, ConstraintSystem, Error}, }, - virtual_region::lookups::basic::BasicDynLookupConfig, + virtual_region::{ + copy_constraints::EXTERNAL_CELL_TYPE_ID, lookups::basic::BasicDynLookupConfig, + }, AssignedValue, ContextCell, }; use halo2_proofs_axiom::plonk::FirstPhase; @@ -109,34 +109,38 @@ impl Circuit for RAMCircuit { config: Self::Config, mut layouter: impl Layouter, ) -> Result<(), Error> { + layouter.assign_region( + || "cpu", + |mut region| { + self.cpu.assign_raw( + &(config.cpu.basic_gates[0].clone(), config.cpu.max_rows), + &mut region, + ); + Ok(()) + }, + )?; + + let copy_manager = (!self.cpu.witness_gen_only()).then_some(&self.cpu.copy_manager); + // Make purely virtual cells so we can raw assign them let memory = self.memory.iter().enumerate().map(|(i, value)| { let idx = Assigned::Trivial(F::from(i as u64)); - let idx = - AssignedValue { value: idx, cell: Some(ContextCell::new("RAM Config", 0, i)) }; + let idx = AssignedValue { + value: idx, + cell: Some(ContextCell::new(EXTERNAL_CELL_TYPE_ID, 0, i)), + }; let value = Assigned::Trivial(*value); - let value = AssignedValue { value, cell: Some(ContextCell::new("RAM Config", 1, i)) }; + let value = + AssignedValue { value, cell: Some(ContextCell::new(EXTERNAL_CELL_TYPE_ID, 1, i)) }; [idx, value] }); - let copy_manager = (!self.cpu.witness_gen_only()).then_some(&self.cpu.copy_manager); - config.memory.assign_virtual_table_to_raw( layouter.namespace(|| "memory"), memory, copy_manager, ); - layouter.assign_region( - || "cpu", - |mut region| { - self.cpu.assign_raw( - &(config.cpu.basic_gates[0].clone(), config.cpu.max_rows), - &mut region, - ); - Ok(()) - }, - )?; config.memory.assign_virtual_to_lookup_to_raw( layouter.namespace(|| "memory accesses"), self.mem_access.clone(),