From 2d2c6d66a9d2687eb770ceec6d4cf11c7187b5e9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Michele=20Orr=C3=B9?= <m@orru.net>
Date: Sun, 10 Nov 2024 17:38:14 +0100
Subject: [PATCH] Fix audit section 3.2 and test with nimue-poseidon.

---
 nimue-poseidon/src/tests.rs        | 32 ++++++++++++++++++++++++++++++
 nimue/src/plugins/ark/common.rs    | 32 +++++++++++++++---------------
 nimue/src/plugins/ark/iopattern.rs |  2 +-
 3 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/nimue-poseidon/src/tests.rs b/nimue-poseidon/src/tests.rs
index b4f80d6..08979b9 100644
--- a/nimue-poseidon/src/tests.rs
+++ b/nimue-poseidon/src/tests.rs
@@ -1,5 +1,6 @@
 use nimue::hash::sponge::Sponge;
 
+#[allow(unused)]
 fn test_vector<H: Sponge>(input: &[H::U], output: &[H::U])
 where
     H::U: PartialEq + std::fmt::Debug,
@@ -10,6 +11,37 @@ where
     assert_eq!(hash.as_ref(), output);
 }
 
+#[cfg(feature = "bls12-381")]
+#[test]
+fn test_squeeze_bytes_from_algebraic_hash() {
+    use nimue::ByteChallenges;
+
+    type F = ark_bls12_381::Fr;
+    type H = crate::bls12_381::Poseidonx5_255_3;
+
+    let io = nimue::IOPattern::<H, F>::new("test").absorb(1, "in");
+    let io = <nimue::IOPattern<H, F> as nimue::plugins::ark::ByteIOPattern>::challenge_bytes(
+        io, 2048, "out",
+    );
+    let mut merlin = io.to_merlin();
+    merlin.add_units(&[F::from(0x42)]).unwrap();
+
+    let mut merlin_challenges = [0u8; 2048];
+    merlin.fill_challenge_bytes(&mut merlin_challenges).unwrap();
+
+    let mut arthur = io.to_arthur(merlin.transcript());
+    // write the unit to an throw-away array
+    arthur.fill_next_units(&mut [F::from(0)]).unwrap();
+    let arthur_challenges: [u8; 2048] = arthur.challenge_bytes().unwrap();
+
+    assert_eq!(merlin_challenges, arthur_challenges);
+    let frequencies = (0u8..=255)
+        .map(|i| merlin_challenges.iter().filter(|&&x| x == i).count())
+        .collect::<Vec<_>>();
+    // each element should appear roughly 8 times on average. Checking we're not too far from that.
+    assert!(frequencies.iter().all(|&x| x < 32 && x > 0), "This array should have random bytes but hasn't: {:?}", frequencies);
+}
+
 #[cfg(feature = "bls12-381")]
 #[test]
 fn test_poseidon_bls12_381() {
diff --git a/nimue/src/plugins/ark/common.rs b/nimue/src/plugins/ark/common.rs
index df7bd91..61d47db 100644
--- a/nimue/src/plugins/ark/common.rs
+++ b/nimue/src/plugins/ark/common.rs
@@ -1,7 +1,7 @@
 use std::io;
 
 use ark_ec::{AffineRepr, CurveGroup};
-use ark_ff::{Field, Fp, FpConfig, PrimeField};
+use ark_ff::{BigInteger, Field, Fp, FpConfig, PrimeField};
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, SerializationError};
 use rand::{CryptoRng, RngCore};
 
@@ -234,12 +234,9 @@ where
                 crate::plugins::random_bytes_in_random_modp(Fp::<C, N>::MODULUS),
                 output.len(),
             );
-            let len = crate::plugins::bytes_modp(Fp::<C, N>::MODULUS_BIT_SIZE);
             let mut tmp = [Fp::from(0); 1];
-            let mut buf = vec![0u8; len];
             self.fill_challenge_units(&mut tmp)?;
-            tmp[0].serialize_compressed(&mut buf).unwrap();
-
+            let buf = tmp[0].into_bigint().to_bytes_le();
             output[..len_good].copy_from_slice(&buf[..len_good]);
 
             // recursively fill the rest of the buffer
@@ -255,17 +252,20 @@ where
     H: DuplexHash<Fp<C, N>>,
 {
     fn fill_challenge_bytes(&mut self, output: &mut [u8]) -> Result<(), IOPatternError> {
-        let len_good = usize::min(
-            crate::plugins::random_bytes_in_random_modp(Fp::<C, N>::MODULUS),
-            output.len(),
-        );
-        let len = crate::plugins::bytes_modp(Fp::<C, N>::MODULUS_BIT_SIZE);
-        let mut tmp = [Fp::from(0); 1];
-        let mut buf = vec![0u8; len];
-        self.fill_challenge_units(&mut tmp)?;
-        tmp[0].serialize_compressed(&mut buf).unwrap();
+        if output == &[] {
+            Ok(())
+        } else {
+            let len_good = usize::min(
+                crate::plugins::random_bytes_in_random_modp(Fp::<C, N>::MODULUS),
+                output.len(),
+            );
+            let mut tmp = [Fp::from(0); 1];
+            self.fill_challenge_units(&mut tmp)?;
+            let buf = tmp[0].into_bigint().to_bytes_le();
+            output[..len_good].copy_from_slice(&buf[..len_good]);
 
-        output[..len_good].copy_from_slice(&buf[..len_good]);
-        Ok(())
+            // recursively fill the rest of the buffer
+            self.fill_challenge_bytes(&mut output[len_good..])
+        }
     }
 }
diff --git a/nimue/src/plugins/ark/iopattern.rs b/nimue/src/plugins/ark/iopattern.rs
index f26d879..089e810 100644
--- a/nimue/src/plugins/ark/iopattern.rs
+++ b/nimue/src/plugins/ark/iopattern.rs
@@ -53,7 +53,7 @@ where
     }
 
     fn challenge_bytes(self, count: usize, label: &str) -> Self {
-        let n = bytes_uniform_modp(Fp::<C, N>::MODULUS_BIT_SIZE);
+        let n = crate::plugins::random_bits_in_random_modp(Fp::<C, N>::MODULUS) / 8;
         self.squeeze((count + n - 1) / n, label)
     }
 }