poseidon: Fix padding to follow sponge construction

Sponge constructions pad the entire input message and then split it into
rate-sized chunks. The previous implementation was using an incorrect
duplex-like hybrid where padding was applied to each chunked input. We
now use an enum to distinguish message and padding words being absorbed
into the sponge.

This also fixes two previous bugs:
- If a `ConstantLength` hash had a length greater than the permutation's
  rate but not a multiple of it, no padding would be generated and the
  circuit would fail to create proofs.
- If a sponge usage required more output than the permutation's rate,
  the squeeze-side permutations would in some cases incorrectly apply
  padding, when it should instead use the prior state as-is. We now add
  zeroes instead.

This change doesn't alter the Orchard circuit, because it doesn't need
any padding cells, only takes a single field element as output, and
padding is still assigned in the same region as before.
This commit is contained in:
Jack Grigg 2021-12-10 02:07:34 +00:00
parent fdeb2fb817
commit ae72813f77
3 changed files with 79 additions and 92 deletions

View File

@ -5,6 +5,7 @@ use std::convert::TryInto;
use std::fmt;
use std::marker::PhantomData;
use group::ff::Field;
use halo2::{
arithmetic::FieldExt,
circuit::{AssignedCell, Chip, Layouter},
@ -18,6 +19,15 @@ use crate::primitives::poseidon::{
Absorbing, ConstantLength, Domain, Spec, SpongeMode, SpongeRate, Squeezing, State,
};
/// A word from the padded input to a Poseidon sponge.
#[derive(Clone, Debug)]
pub enum PaddedWord<F: Field> {
/// A message word provided by the prover.
Message(AssignedCell<F, F>),
/// A padding word, that will be fixed in the circuit parameters.
Padding(F),
}
/// The set of circuit instructions required to use the Poseidon permutation.
pub trait PoseidonInstructions<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>:
Chip<F>
@ -39,7 +49,7 @@ pub trait PoseidonInstructions<F: FieldExt, S: Spec<F, T, RATE>, const T: usize,
pub trait PoseidonSpongeInstructions<
F: FieldExt,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
>: PoseidonInstructions<F, S, T, RATE>
@ -48,12 +58,12 @@ pub trait PoseidonSpongeInstructions<
fn initial_state(&self, layouter: &mut impl Layouter<F>)
-> Result<State<Self::Word, T>, Error>;
/// Pads the given input (according to the specified domain) and adds it to the state.
fn pad_and_add(
/// Adds the given input to the state.
fn add_input(
&self,
layouter: &mut impl Layouter<F>,
initial_state: &State<Self::Word, T>,
input: &SpongeRate<Self::Word, RATE>,
input: &SpongeRate<PaddedWord<F>, RATE>,
) -> Result<State<Self::Word, T>, Error>;
/// Extracts sponge output from the given state.
@ -95,16 +105,16 @@ fn poseidon_sponge<
F: FieldExt,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
>(
chip: &PoseidonChip,
mut layouter: impl Layouter<F>,
state: &mut State<PoseidonChip::Word, T>,
input: &SpongeRate<PoseidonChip::Word, RATE>,
input: &SpongeRate<PaddedWord<F>, RATE>,
) -> Result<SpongeRate<PoseidonChip::Word, RATE>, Error> {
*state = chip.pad_and_add(&mut layouter, state, input)?;
*state = chip.add_input(&mut layouter, state, input)?;
*state = chip.permute(&mut layouter, state)?;
Ok(PoseidonChip::get_output(state))
}
@ -116,7 +126,7 @@ pub struct Sponge<
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
M: SpongeMode,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> {
@ -130,10 +140,10 @@ impl<
F: FieldExt,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> Sponge<F, PoseidonChip, S, Absorbing<PoseidonChip::Word, RATE>, D, T, RATE>
> Sponge<F, PoseidonChip, S, Absorbing<PaddedWord<F>, RATE>, D, T, RATE>
{
/// Constructs a new duplex sponge for the given Poseidon specification.
pub fn new(chip: PoseidonChip, mut layouter: impl Layouter<F>) -> Result<Self, Error> {
@ -155,11 +165,11 @@ impl<
pub fn absorb(
&mut self,
mut layouter: impl Layouter<F>,
value: AssignedCell<F, F>,
value: PaddedWord<F>,
) -> Result<(), Error> {
for entry in self.mode.0.iter_mut() {
if entry.is_none() {
*entry = Some(value.into());
*entry = Some(value);
return Ok(());
}
}
@ -171,7 +181,7 @@ impl<
&mut self.state,
&self.mode.0,
)?;
self.mode = Absorbing::init_with(value.into());
self.mode = Absorbing::init_with(value);
Ok(())
}
@ -203,7 +213,7 @@ impl<
F: FieldExt,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> Sponge<F, PoseidonChip, S, Squeezing<PoseidonChip::Word, RATE>, D, T, RATE>
@ -222,7 +232,11 @@ impl<
&self.chip,
layouter.namespace(|| "PoseidonSponge"),
&mut self.state,
&self.mode.0,
&(0..RATE)
.map(|_| Some(PaddedWord::Padding(F::zero())))
.collect::<Vec<_>>()
.try_into()
.unwrap(),
)?);
}
}
@ -234,18 +248,18 @@ pub struct Hash<
F: FieldExt,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> {
sponge: Sponge<F, PoseidonChip, S, Absorbing<PoseidonChip::Word, RATE>, D, T, RATE>,
sponge: Sponge<F, PoseidonChip, S, Absorbing<PaddedWord<F>, RATE>, D, T, RATE>,
}
impl<
F: FieldExt,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> Hash<F, PoseidonChip, S, D, T, RATE>
@ -271,7 +285,11 @@ impl<
mut layouter: impl Layouter<F>,
message: [AssignedCell<F, F>; L],
) -> Result<AssignedCell<F, F>, Error> {
for (i, value) in array::IntoIter::new(message).enumerate() {
for (i, value) in array::IntoIter::new(message)
.map(PaddedWord::Message)
.chain(<ConstantLength<L> as Domain<F, RATE>>::padding(L).map(PaddedWord::Padding))
.enumerate()
{
self.sponge
.absorb(layouter.namespace(|| format!("absorb_{}", i)), value)?;
}

View File

@ -8,7 +8,7 @@ use halo2::{
poly::Rotation,
};
use super::{PoseidonInstructions, PoseidonSpongeInstructions};
use super::{PaddedWord, PoseidonInstructions, PoseidonSpongeInstructions};
use crate::circuit::gadget::utilities::Var;
use crate::primitives::poseidon::{Domain, Mds, Spec, SpongeRate, State};
@ -271,7 +271,7 @@ impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize
impl<
F: FieldExt,
S: Spec<F, WIDTH, RATE>,
D: Domain<F, WIDTH, RATE>,
D: Domain<F, RATE>,
const WIDTH: usize,
const RATE: usize,
> PoseidonSpongeInstructions<F, S, D, WIDTH, RATE> for Pow5Chip<F, WIDTH, RATE>
@ -309,15 +309,15 @@ impl<
Ok(state.try_into().unwrap())
}
fn pad_and_add(
fn add_input(
&self,
layouter: &mut impl Layouter<F>,
initial_state: &State<Self::Word, WIDTH>,
input: &SpongeRate<Self::Word, RATE>,
input: &SpongeRate<PaddedWord<F>, RATE>,
) -> Result<State<Self::Word, WIDTH>, Error> {
let config = self.config();
layouter.assign_region(
|| format!("pad-and-add for domain {}", D::name()),
|| format!("add input for domain {}", D::name()),
|mut region| {
config.s_pad_and_add.enable(&mut region, 1)?;
@ -337,19 +337,17 @@ impl<
(0..WIDTH).map(load_state_word).collect();
let initial_state = initial_state?;
let padding_values = D::padding();
// Load the input and padding into this region.
// Load the input into this region.
let load_input_word = |i: usize| {
let constraint_var = match (input[i].clone(), padding_values[i]) {
(Some(word), None) => word.0,
(None, Some(padding_value)) => region.assign_fixed(
let constraint_var = match input[i].clone() {
Some(PaddedWord::Message(word)) => word,
Some(PaddedWord::Padding(padding_value)) => region.assign_fixed(
|| format!("load pad_{}", i),
config.rc_b[i],
1,
|| Ok(padding_value),
)?,
_ => panic!("Input and padding don't match"),
_ => panic!("Input is not padded"),
};
constraint_var
.copy_advice(

View File

@ -127,11 +127,14 @@ pub(crate) fn permute<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RA
fn poseidon_sponge<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>(
state: &mut State<F, T>,
input: &SpongeRate<F, RATE>,
pad_and_add: &dyn Fn(&mut State<F, T>, &SpongeRate<F, RATE>),
mds_matrix: &Mds<F, T>,
round_constants: &[[F; T]],
) -> SpongeRate<F, RATE> {
pad_and_add(state, input);
// `Iterator::zip` short-circuits when one iterator completes, so this will only
// mutate the rate portion of the state.
for (word, value) in state.iter_mut().zip(input.iter()) {
*word += value.expect("poseidon_sponge is called with a padded input");
}
permute::<F, S, T, RATE>(state, mds_matrix, round_constants);
@ -179,7 +182,6 @@ pub(crate) struct Sponge<
> {
mode: M,
state: State<F, T>,
pad_and_add: Box<dyn Fn(&mut State<F, T>, &SpongeRate<F, RATE>)>,
mds_matrix: Mds<F, T>,
round_constants: Vec<[F; T]>,
_marker: PhantomData<(S, M)>,
@ -189,10 +191,7 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
Sponge<F, S, Absorbing<F, RATE>, T, RATE>
{
/// Constructs a new sponge for the given Poseidon specification.
pub(crate) fn new(
initial_capacity_element: F,
pad_and_add: Box<dyn Fn(&mut State<F, T>, &SpongeRate<F, RATE>)>,
) -> Self {
pub(crate) fn new(initial_capacity_element: F) -> Self {
let (round_constants, mds_matrix, _) = S::constants();
let mode = Absorbing([None; RATE]);
@ -202,7 +201,6 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
Sponge {
mode,
state,
pad_and_add,
mds_matrix,
round_constants,
_marker: PhantomData::default(),
@ -222,7 +220,6 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
let _ = poseidon_sponge::<F, S, T, RATE>(
&mut self.state,
&self.mode.0,
&self.pad_and_add,
&self.mds_matrix,
&self.round_constants,
);
@ -234,7 +231,6 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
let mode = Squeezing(poseidon_sponge::<F, S, T, RATE>(
&mut self.state,
&self.mode.0,
&self.pad_and_add,
&self.mds_matrix,
&self.round_constants,
));
@ -242,7 +238,6 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
Sponge {
mode,
state: self.state,
pad_and_add: self.pad_and_add,
mds_matrix: self.mds_matrix,
round_constants: self.round_constants,
_marker: PhantomData::default(),
@ -266,7 +261,6 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
self.mode = Squeezing(poseidon_sponge::<F, S, T, RATE>(
&mut self.state,
&self.mode.0,
&self.pad_and_add,
&self.mds_matrix,
&self.round_constants,
));
@ -275,19 +269,18 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
}
/// A domain in which a Poseidon hash function is being used.
pub trait Domain<F: FieldExt, const T: usize, const RATE: usize> {
pub trait Domain<F: FieldExt, const RATE: usize> {
/// Iterator that outputs padding field elements.
type Padding: IntoIterator<Item = F>;
/// The name of this domain, for debug formatting purposes.
fn name() -> String;
/// The initial capacity element, encoding this domain.
fn initial_capacity_element() -> F;
/// The padding that will be added to each state word by [`Domain::pad_and_add`].
fn padding() -> SpongeRate<F, RATE>;
/// Returns a function that will update the given state with the given input to a
/// duplex permutation round, applying padding according to this domain specification.
fn pad_and_add() -> Box<dyn Fn(&mut State<F, T>, &SpongeRate<F, RATE>)>;
/// Returns the padding to be appended to the input.
fn padding(input_len: usize) -> Self::Padding;
}
/// A Poseidon hash function used with constant input length.
@ -296,9 +289,9 @@ pub trait Domain<F: FieldExt, const T: usize, const RATE: usize> {
#[derive(Clone, Copy, Debug)]
pub struct ConstantLength<const L: usize>;
impl<F: FieldExt, const T: usize, const RATE: usize, const L: usize> Domain<F, T, RATE>
for ConstantLength<L>
{
impl<F: FieldExt, const RATE: usize, const L: usize> Domain<F, RATE> for ConstantLength<L> {
type Padding = iter::Take<iter::Repeat<F>>;
fn name() -> String {
format!("ConstantLength<{}>", L)
}
@ -309,28 +302,14 @@ impl<F: FieldExt, const T: usize, const RATE: usize, const L: usize> Domain<F, T
F::from_u128((L as u128) << 64)
}
fn padding() -> SpongeRate<F, RATE> {
// For constant-input-length hashing, padding consists of the field elements being
// zero.
let mut padding = [None; RATE];
for word in padding.iter_mut().skip(L) {
*word = Some(F::zero());
}
padding
}
fn pad_and_add() -> Box<dyn Fn(&mut State<F, T>, &SpongeRate<F, RATE>)> {
Box::new(|state, input| {
// `Iterator::zip` short-circuits when one iterator completes, so this will only
// mutate the rate portion of the state.
for (word, value) in state.iter_mut().zip(input.iter()) {
// For constant-input-length hashing, padding consists of the field
// elements being zero, so we don't add anything to the state.
if let Some(value) = value {
*word += value;
}
}
})
fn padding(input_len: usize) -> Self::Padding {
assert_eq!(input_len, L);
// For constant-input-length hashing, we pad the input with zeroes to a multiple
// of RATE. On its own this is not sponge-compliant padding, but the Poseidon
// authors instead encode the constant length into the capacity element, ensuring
// that inputs of different lengths do not share the same permutation.
let k = (L + RATE - 1) / RATE;
iter::repeat(F::zero()).take(k * RATE - L)
}
}
@ -338,7 +317,7 @@ impl<F: FieldExt, const T: usize, const RATE: usize, const L: usize> Domain<F, T
pub struct Hash<
F: FieldExt,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> {
@ -346,13 +325,8 @@ pub struct Hash<
_domain: PhantomData<D>,
}
impl<
F: FieldExt,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
const T: usize,
const RATE: usize,
> fmt::Debug for Hash<F, S, D, T, RATE>
impl<F: FieldExt, S: Spec<F, T, RATE>, D: Domain<F, RATE>, const T: usize, const RATE: usize>
fmt::Debug for Hash<F, S, D, T, RATE>
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Hash")
@ -365,18 +339,13 @@ impl<
}
}
impl<
F: FieldExt,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
const T: usize,
const RATE: usize,
> Hash<F, S, D, T, RATE>
impl<F: FieldExt, S: Spec<F, T, RATE>, D: Domain<F, RATE>, const T: usize, const RATE: usize>
Hash<F, S, D, T, RATE>
{
/// Initializes a new hasher.
pub fn init() -> Self {
Hash {
sponge: Sponge::new(D::initial_capacity_element(), D::pad_and_add()),
sponge: Sponge::new(D::initial_capacity_element()),
_domain: PhantomData::default(),
}
}
@ -387,7 +356,9 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize, const
{
/// Hashes the given input.
pub fn hash(mut self, message: [F; L]) -> F {
for value in array::IntoIter::new(message) {
for value in
array::IntoIter::new(message).chain(<ConstantLength<L> as Domain<F, RATE>>::padding(L))
{
self.sponge.absorb(value);
}
self.sponge.finish_absorbing().squeeze()