fix: issue with under-constraint byte-address

Basically, the bug allows applying sponge permutation on a given chunk of packed field elements in arbitrary order. That is, prover, instead of supplying given chunk c: [F; 8] in the intended order [c_1, c_2 ...c_8]` to Poseidon2Sponge, can instead supply it in any way he wants!

Here is how:
Say prover packs two field elements out of address region (0...14) with content [1, 2 ... 14]
So corresponding packed values would be
let packed_1 = Poseidon2PreimagePack<F> {
    pub clk: some_clk,
    pub byte_addr: 0,
    pub fe_addr: 0,
    pub bytes: [1, 2 ..7], // 2261738553347342
    pub is_executed: 1,
}
let packed_2 = Poseidon2PreimagePack<F> {
    pub clk: some_clk,
    pub byte_addr: 8,
    pub fe_addr: 1,
    pub bytes: [8, 9, .. 14], // 283686952306183
    pub is_executed: 1,
}
The idea is that there is no constraint tying byte_addr with fe_addr.
That is, the malicious prover can flip the above fe_addrs by (packed_1.fe_addr, packed_2.fe_addr) = (packed_2.fe_addr, packed_1.fe_addr).
It won't be detected anywhere!
Hence when supplying the above to Poseidon2Sponge, he can emulate that he is supplying 283686952306183 before 2261738553347342.
Simply because the former pack has fe_addr=0, and the latter pack has fe_addr=1 after modification.
So instead of sponge permutation acting on [2261738553347342, 283686952306183...],
it acts on [283686952306183, 2261738553347342...], hence leading to different output.
In the former codebase, this can't happen because of CTL against Memory stark, which ensures that input to Poseidon2Sponge is coming in correct order (according to their addresses). But in the current implementation, there is no restriction, because fe_addr is not a real address as you pointed out!
I think there is an easy way out though. We change the interpretation of input_addr in Poseidon2Sponge
pub struct Poseidon2Sponge<T> {
    pub clk: T,
    pub ops: Ops<T>,
    pub input_addr: T,
    pub output_addr: T,
    pub input_len: T,
    pub preimage: [T; WIDTH],
    pub output: [T; WIDTH],
    pub gen_output: T,
    pub input_addr_padded: T,
}
Earlier this said "Hey, I am taking 8 bytes in region input_addr...input_addr+8, feeding it to current sponge preimage , apply sponge permutation, and I get output .  I have also ordered the trace in increasing order of input_addr which you can check. Hence, you can be sure I am feeding it in correct order."

Author: reshmem

Diff for: circuits/src/memory/columns.rs

@@ -119,7 +119,7 @@ impl<F: RichField> From<&Poseidon2Sponge<F>> for Vec<Memory<F>> {
             // each Field element in preimage represents packed data (packed bytes)
             (0..Poseidon2Permutation::<F>::RATE)
                 .flat_map(|fe_index_inside_preimage| {
-                    let base_address = value.input_addr_padded
+                    let base_address = value.input_addr
                         + MozakPoseidon2::data_capacity_fe::<F>()
                             * F::from_canonical_usize(fe_index_inside_preimage);
                     let unpacked = MozakPoseidon2::unpack_to_field_elements(

Diff for: circuits/src/poseidon2_preimage_pack/columns.rs

@@ -15,7 +15,6 @@ use crate::stark::mozak_stark::{Poseidon2PreimagePackTable, TableWithTypedOutput
 pub struct Poseidon2PreimagePack<F> {
     pub clk: F,
     pub byte_addr: F,
-    pub fe_addr: F,
     pub bytes: [F; MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT],
     pub is_executed: F,
 }
@@ -40,8 +39,7 @@ impl<F: RichField> From<&Poseidon2Sponge<F>> for Vec<Poseidon2PreimagePack<F>> {
                 [..MozakPoseidon2::FIELD_ELEMENTS_RATE]
                 .try_into()
                 .expect("Should succeed since preimage can't be empty");
-            let mut byte_base_address = value.input_addr_padded;
-            let mut fe_base_addr = value.input_addr;
+            let mut byte_base_address = value.input_addr;
             // For each FE of preimage we have BYTES_COUNT bytes
             preimage
                 .iter()
@@ -50,15 +48,9 @@ impl<F: RichField> From<&Poseidon2Sponge<F>> for Vec<Poseidon2PreimagePack<F>> {
                     let byte_addr = byte_base_address;
                     // increase by DATA_CAP the byte base address after each iteration
                     byte_base_address += MozakPoseidon2::data_capacity_fe();
-                    // specific field-el address
-                    let fe_addr = fe_base_addr;
-                    // increase by 1 after each iteration
-                    fe_base_addr += F::ONE;
-
                     Poseidon2PreimagePack {
                         clk: value.clk,
                         byte_addr,
-                        fe_addr,
                         bytes: MozakPoseidon2::unpack_to_field_elements(fe),
                         is_executed: F::ONE,
                     }
@@ -74,7 +66,6 @@ columns_view_impl!(Poseidon2SpongePreimagePackCtl);
 pub struct Poseidon2SpongePreimagePackCtl<T> {
     pub clk: T,
     pub value: T,
-    pub fe_addr: T,
     pub byte_addr: T,
 }
 #[must_use]
@@ -84,7 +75,6 @@ pub fn lookup_for_poseidon2_sponge() -> TableWithTypedOutput<Poseidon2SpongePrei
         Poseidon2SpongePreimagePackCtl {
             clk: PACK.clk,
             value: ColumnWithTypedInput::reduce_with_powers(PACK.bytes, 1 << 8),
-            fe_addr: PACK.fe_addr,
             byte_addr: PACK.byte_addr,
         },
         PACK.is_executed,

Diff for: circuits/src/poseidon2_sponge/columns.rs

@@ -31,7 +31,6 @@ pub struct Poseidon2Sponge<T> {
     pub preimage: [T; WIDTH],
     pub output: [T; WIDTH],
     pub gen_output: T,
-    pub input_addr_padded: T,
 }
 
 columns_view_impl!(Poseidon2Sponge);
@@ -116,8 +115,7 @@ pub fn lookup_for_preimage_pack(
                 Poseidon2SpongePreimagePackCtl {
                     clk: SPONGE.clk,
                     value,
-                    fe_addr: SPONGE.input_addr + limb_index,
-                    byte_addr: SPONGE.input_addr_padded
+                    byte_addr: SPONGE.input_addr
                         + limb_index
                             * i64::try_from(MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT)
                                 .expect("Should be < 255"),

Diff for: circuits/src/poseidon2_sponge/generation.rs

@@ -21,7 +21,6 @@ fn unroll_sponge_data<F: RichField>(row: &Row<F>) -> Vec<Poseidon2Sponge<F>> {
 
     let output_addr = poseidon2.output_addr;
     let mut input_addr = poseidon2.addr;
-    let mut input_addr_padded = poseidon2.addr;
     let mut input_len = poseidon2.len;
     for i in 0..unroll_count {
         let ops: Ops<F> = Ops {
@@ -41,10 +40,8 @@ fn unroll_sponge_data<F: RichField>(row: &Row<F>) -> Vec<Poseidon2Sponge<F>> {
             preimage: sponge_datum.preimage,
             output: sponge_datum.output,
             gen_output: sponge_datum.gen_output,
-            input_addr_padded: F::from_canonical_u32(input_addr_padded),
         });
-        input_addr_padded += u32::try_from(MozakPoseidon2::DATA_PADDING).expect("Should succeed");
-        input_addr += rate_size;
+        input_addr += u32::try_from(MozakPoseidon2::DATA_PADDING).expect("Should succeed");
         input_len -= rate_size;
     }
 

Diff for: circuits/src/poseidon2_sponge/stark.rs

@@ -96,16 +96,10 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for Poseidon2Spon
         // input
         yield_constr
             .constraint_transition(not_last_sponge * (lv.input_len - (nv.input_len + rate_scalar)));
-        // and input_addr increases by RATE
+        // and input_addr increases by RATE *
         yield_constr.constraint_transition(
-            not_last_sponge * (lv.input_addr - (nv.input_addr - rate_scalar)),
+            not_last_sponge * (lv.input_addr - (nv.input_addr - padding_scalar)),
         );
-        // and input_addr_padded increases by DATA_PADDING = PACK_CAP * RATE
-        yield_constr.constraint_transition(
-            not_last_sponge * (lv.input_addr_padded - (nv.input_addr_padded - padding_scalar)),
-        );
-        // and for init permute, input_addr = input_addr_padded
-        yield_constr.constraint(lv.ops.is_init_permute * (lv.input_addr - lv.input_addr_padded));
 
         // For each init_permute capacity bits are zero.
         (rate..state_size).for_each(|i| {
@@ -187,27 +181,12 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for Poseidon2Spon
         // length decreases by RATE, note that only actual execution row can consume
         // input
         yield_constr.constraint_transition(builder, len_check);
-        // and input_addr increases by RATE
-        let nv_input_addr_rate = builder.sub_extension(nv.input_addr, rate_ext);
+        // and input_addr increases by PADDING
+        let nv_input_addr_rate = builder.sub_extension(nv.input_addr, padding_ext);
         let input_addr_diff_rate = builder.sub_extension(lv.input_addr, nv_input_addr_rate);
         let addr_check = builder.mul_extension(not_last_sponge, input_addr_diff_rate);
         yield_constr.constraint_transition(builder, addr_check);
 
-        // and input_addr_padding increases by DATA_PADDING
-        let nv_input_addr_padding = builder.sub_extension(nv.input_addr_padded, padding_ext);
-        let input_addr_diff_padding =
-            builder.sub_extension(lv.input_addr_padded, nv_input_addr_padding);
-        let padding_addr_check = builder.mul_extension(not_last_sponge, input_addr_diff_padding);
-        yield_constr.constraint_transition(builder, padding_addr_check);
-
-        // and for init permute, input_addr = input_addr_padded
-        // yield_constr.constraint(lv.ops.is_init_permute * (lv.input_addr -
-        // lv.input_addr_padded));
-        let lv_diff_addr_padding = builder.sub_extension(lv.input_addr, lv.input_addr_padded);
-        let padding_addr_init_check =
-            builder.mul_extension(lv.ops.is_init_permute, lv_diff_addr_padding);
-        yield_constr.constraint(builder, padding_addr_init_check);
-
         let zero = builder.constant_extension(F::Extension::ZERO);
         // For each init_permute capacity bits are zero.
         (rate..state_size).for_each(|i| {