[BUG] Poseidon sponge does not permute when squeezing

This bug was reported to @Pratyush in January 2024 through a mutual partner. Since I haven't received a response about the preferred disclosure method, I'm documenting this issue to keep track of it.

Executive Summary

In the squeeze_internal function of the implementation of the Poseidon sponge, a non-necessary condition if is present, leading not to permute when squeezing self.parameters.rate elements with rate_start_index > 0.

Environment

Compiler Version: cargo 1.73.0 (9c4383fb5 2023-08-26)
Distro Version: Ubuntu 23.04

Steps to Reproduce

Cargo.toml:

    [package]
    name = "bad_if_squeeze_poseidon"
    version = "0.1.0"
    edition = "2021"

    # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

    [dependencies]
    ark-bls12-377 = "0.4.0"
    ark-crypto-primitives = { version = "0.4.0", features = ["sponge"]}

main.rs:

    use ark_crypto_primitives::sponge::{poseidon::{PoseidonSponge, PoseidonConfig,find_poseidon_ark_and_mds}, CryptographicSponge};
    use ark_crypto_primitives::sponge::FieldBasedCryptographicSponge;
    use ark_bls12_377::Fq as Fq;

    fn main() {

        let full_rounds = 8;
        let partial_rounds = 31;
        let alpha = 17;
        let rate = 2;
        let capacity = 1;

        let (ark, mds) = find_poseidon_ark_and_mds::<Fq>(
            377,
            rate,
            full_rounds as u64,
            partial_rounds as u64,
            0
        );

        let sponge_param = PoseidonConfig::new(full_rounds, partial_rounds, alpha, mds, ark, rate, capacity);

        let mut sponge = PoseidonSponge::new(&sponge_param);

        // 1 is the only integer > 0 but < rate = 2
        let x = sponge.squeeze_native_field_elements(1);
        println!("First element squeezed: {x:?}");

        let x = sponge.squeeze_native_field_elements(rate);
        println!("Second and third elements squeezed: {x:?}");
    }

If we remove the if conditions that we talk about below, we can see that the third element changes. And adding some println() in the base code after the calls to the permute function, we can see that we don't use this function enough.

Root Cause Analysis

The problem is line 173 in poseidon/constraints.rs, and line 176 in poseidon/mod.rs.

Let's remember that output_remaining.len() is equal to self.parameters.rate, and rate_start_index > 0. The condition rate_start_index + output_remaining.len() <= self.parameters.rate is therefore not respected, and we can access the previously cited lines. If the output.len() is different from self.parameters.rate, then we permute, but in our case, both are equal, so we don't permute, and we enter in the first if loop, since, now, the condition rate_start_index + output_remaining.len() <= self.parameters.rate is respected. Inside this loop, we recover the missing element to squeeze, but still without permute. Finally, we have squeezed a certain number of elements without permuting the internal state.

Recommendations

Remove the if output_remaining.len() != self.parameters.rate conditions.

References

Poseidon Paper

arkworks-rs / crypto-primitives