chore(keccak): slicker code and more stable dependencies (#45)

Move to the keccak crate
Change visibility and implementation of some structures
Add from_iv for codecs to initialize directly the sponge

---
Co-Authored-By: nougzarm <[email protected]>

Author: mmaker

Cargo.toml

@@ -27,7 +27,9 @@ num-traits = "0.2.19"
 rand = "0.8.5"
 sha3 = "0.10.8"
 thiserror = "1"
-tiny-keccak = { version = "2.0.2", features = ["fips202"] }
+keccak = "0.1.5"
+zerocopy = "0.8"
+zeroize = "1.8.1"
 
 [dev-dependencies]
 bls12_381 = "0.8.0"

src/codec.rs

@@ -24,6 +24,9 @@ pub trait Codec {
     /// Generates an empty codec that can be identified by a domain separator.
     fn new(domain_sep: &[u8]) -> Self;
 
+    /// Allows for precomputed initialization of the codec with a specific IV.
+    fn from_iv(iv: [u8; 32]) -> Self;
+
     /// Absorbs data into the codec.
     fn prover_message(&mut self, data: &[u8]);
 
@@ -58,10 +61,19 @@ where
     type Challenge = <G as Group>::Scalar;
 
     fn new(domain_sep: &[u8]) -> Self {
-        let hasher = H::new(domain_sep);
+        let iv = {
+            let mut tmp = H::new([0u8; 32]);
+            tmp.absorb(domain_sep);
+            tmp.squeeze(32).try_into().unwrap()
+        };
+
+        Self::from_iv(iv)
+    }
+
+    fn from_iv(iv: [u8; 32]) -> Self {
         Self {
-            hasher,
-            _marker: Default::default(),
+            hasher: H::new(iv),
+            _marker: core::marker::PhantomData,
         }
     }
 

src/duplex_sponge/keccak.rs

@@ -4,125 +4,93 @@
 //! It is designed to match test vectors from the original Sage implementation.
 
 use crate::duplex_sponge::DuplexSpongeInterface;
-use std::convert::TryInto;
-use tiny_keccak::keccakf;
+use zerocopy::IntoBytes;
 
-const R: usize = 136;
-const N: usize = 136 + 64;
+const RATE: usize = 136;
+const LENGTH: usize = 136 + 64;
 
 /// Low-level Keccak-f[1600] state representation.
-#[derive(Clone)]
-pub struct KeccakPermutationState {
-    pub state: [u8; 200],
-    pub rate: usize,
-    pub capacity: usize,
-}
-
-impl Default for KeccakPermutationState {
-    fn default() -> Self {
-        Self::new([0u8; 32])
-    }
-}
+#[derive(Clone, Default)]
+pub struct KeccakPermutationState([u64; LENGTH / 8]);
 
 impl KeccakPermutationState {
     pub fn new(iv: [u8; 32]) -> Self {
-        let rate = 136;
-        let mut state = [0u8; N];
-        state[rate..rate + 32].copy_from_slice(&iv);
-
-        KeccakPermutationState {
-            state,
-            rate,
-            capacity: 64,
-        }
+        let mut state = Self::default();
+        state.as_mut()[RATE..RATE + 32].copy_from_slice(&iv);
+        state
     }
 
-    fn bytes_to_flat_state(&self) -> [u64; 25] {
-        let mut flat = [0u64; 25];
-        for (i, item) in flat.iter_mut().enumerate() {
-            let start = i * 8;
-            *item = u64::from_le_bytes(self.state[start..start + 8].try_into().unwrap());
-        }
-        flat
+    pub fn permute(&mut self) {
+        keccak::f1600(&mut self.0);
     }
+}
 
-    fn flat_state_to_bytes(&mut self, flat: [u64; 25]) {
-        for (i, item) in flat.iter().enumerate() {
-            let bytes = item.to_le_bytes();
-            let start = i * 8;
-            self.state[start..start + 8].copy_from_slice(&bytes);
-        }
+impl AsRef<[u8]> for KeccakPermutationState {
+    fn as_ref(&self) -> &[u8] {
+        self.0.as_bytes()
     }
+}
 
-    pub fn permute(&mut self) {
-        let mut flat = self.bytes_to_flat_state();
-        keccakf(&mut flat);
-        self.flat_state_to_bytes(flat);
+impl AsMut<[u8]> for KeccakPermutationState {
+    fn as_mut(&mut self) -> &mut [u8] {
+        self.0.as_mut_bytes()
     }
 }
 
 /// Duplex sponge construction using Keccak-f[1600].
 #[derive(Clone)]
 pub struct KeccakDuplexSponge {
-    pub state: KeccakPermutationState,
-    pub rate: usize,
-    pub capacity: usize,
+    state: KeccakPermutationState,
     absorb_index: usize,
     squeeze_index: usize,
 }
 
 impl KeccakDuplexSponge {
-    pub fn new(iv: &[u8]) -> Self {
-        assert_eq!(iv.len(), 32);
-
-        let state = KeccakPermutationState::new(iv.try_into().unwrap());
-        let rate = R;
-        let capacity = N - R;
+    pub fn new(iv: [u8; 32]) -> Self {
+        let state = KeccakPermutationState::new(iv);
         KeccakDuplexSponge {
             state,
-            rate,
-            capacity,
             absorb_index: 0,
-            squeeze_index: rate,
+            squeeze_index: RATE,
         }
     }
 }
 
 impl DuplexSpongeInterface for KeccakDuplexSponge {
-    fn new(iv: &[u8]) -> Self {
+    fn new(iv: [u8; 32]) -> Self {
         KeccakDuplexSponge::new(iv)
     }
 
     fn absorb(&mut self, mut input: &[u8]) {
-        self.squeeze_index = self.rate;
+        self.squeeze_index = RATE;
 
         while !input.is_empty() {
-            if self.absorb_index == self.rate {
+            if self.absorb_index == RATE {
                 self.state.permute();
                 self.absorb_index = 0;
             }
 
-            let chunk_size = usize::min(self.rate - self.absorb_index, input.len());
-            let dest = &mut self.state.state[self.absorb_index..self.absorb_index + chunk_size];
+            let chunk_size = usize::min(RATE - self.absorb_index, input.len());
+            let dest = &mut self.state.as_mut()[self.absorb_index..self.absorb_index + chunk_size];
             dest.copy_from_slice(&input[..chunk_size]);
             self.absorb_index += chunk_size;
             input = &input[chunk_size..];
         }
     }
 
     fn squeeze(&mut self, mut length: usize) -> Vec<u8> {
-        self.absorb_index = self.rate;
+        self.absorb_index = RATE;
 
         let mut output = Vec::new();
         while length != 0 {
-            if self.squeeze_index == self.rate {
+            if self.squeeze_index == RATE {
                 self.state.permute();
                 self.squeeze_index = 0;
             }
 
-            let chunk_size = usize::min(self.rate - self.squeeze_index, length);
+            let chunk_size = usize::min(RATE - self.squeeze_index, length);
             output.extend_from_slice(
-                &self.state.state[self.squeeze_index..self.squeeze_index + chunk_size],
+                &self.state.as_mut()[self.squeeze_index..self.squeeze_index + chunk_size],
             );
             self.squeeze_index += chunk_size;
             length -= chunk_size;
@@ -131,3 +99,20 @@ impl DuplexSpongeInterface for KeccakDuplexSponge {
         output
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::duplex_sponge::DuplexSpongeInterface;
+
+    #[test]
+    fn test_keccak_duplex_sponge() {
+        let mut sponge = KeccakDuplexSponge::new([0u8; 32]);
+
+        let input = b"Hello, World!";
+        sponge.absorb(input);
+        let output = sponge.squeeze(64);
+
+        assert_eq!(output, hex::decode("30b74a98221dd643d0814095c212d663a67945c6a582ef8f71bd2a14607ebade3f16e5975ad13d313d9aa0aa97ad29f7df5cff249fa633d3a7ac70d8587bec90").unwrap());
+    }
+}

src/duplex_sponge/mod.rs

@@ -16,7 +16,7 @@ pub mod shake;
 /// This is the core primitive used for building cryptographic codecs.
 pub trait DuplexSpongeInterface {
     /// Creates a new sponge instance with an initialization vector.
-    fn new(iv: &[u8]) -> Self;
+    fn new(iv: [u8; 32]) -> Self;
 
     /// Absorbs input data into the sponge state.
     fn absorb(&mut self, input: &[u8]);

src/duplex_sponge/shake.rs

@@ -10,25 +10,22 @@ use sha3::{
 
 /// Duplex sponge construction using SHAKE128.
 #[derive(Clone, Debug)]
-pub struct ShakeDuplexSponge {
-    /// Internal SHAKE128 hasher state.
-    hasher: Shake128,
-}
+pub struct ShakeDuplexSponge(Shake128);
 
 impl DuplexSpongeInterface for ShakeDuplexSponge {
-    fn new(iv: &[u8]) -> Self {
+    fn new(iv: [u8; 32]) -> Self {
         let mut hasher = Shake128::default();
-        hasher.update(iv);
-        Self { hasher }
+        hasher.update(&iv);
+        Self(hasher)
     }
 
     fn absorb(&mut self, input: &[u8]) {
-        self.hasher.update(input);
+        self.0.update(input);
     }
 
     fn squeeze(&mut self, length: usize) -> Vec<u8> {
         let mut output = vec![0u8; length];
-        self.hasher.clone().finalize_xof_into(&mut output);
+        self.0.clone().finalize_xof_into(&mut output);
         output
     }
 }

src/fiat_shamir.rs

@@ -70,6 +70,14 @@ where
         }
     }
 
+    pub fn from_iv(iv: [u8; 32], instance: P) -> Self {
+        let hash_state = C::from_iv(iv);
+        Self {
+            hash_state,
+            ip: instance,
+        }
+    }
+
     /// Generates a non-interactive proof for a witness.
     ///
     /// Executes the interactive protocol steps (commit, derive challenge via hash, respond),

src/tests/spec/test_vectors.rs

@@ -22,13 +22,12 @@ type NISigmaP = NISigmaProtocol<SigmaP, Codec>;
 macro_rules! generate_ni_function {
     ($name:ident, $test_fn:ident, $($param:tt),*) => {
         #[allow(non_snake_case)]
-        fn $name(seed: &[u8], context: &[u8]) -> (Vec<Scalar>, Vec<u8>) {
+        fn $name(seed: &[u8], iv: [u8; 32]) -> (Vec<Scalar>, Vec<u8>) {
             let mut rng = TestDRNG::new(seed);
-            let (morphismp, witness) = $test_fn($(generate_ni_function!(@arg rng, $param)),*);
+            let (instance, witness) = $test_fn($(generate_ni_function!(@arg rng, $param)),*);
 
-            let protocol = SchnorrProtocolCustom(morphismp);
-            let domain_sep: Vec<u8> = context.to_vec();
-            let nizk = NISigmaP::new(&domain_sep, protocol);
+            let protocol = SchnorrProtocolCustom(instance);
+            let nizk = NISigmaP::from_iv(iv, protocol);
 
             let proof_bytes = nizk.prove_batchable(&witness, &mut rng).unwrap();
             let verified = nizk.verify_batchable(&proof_bytes).is_ok();
@@ -73,11 +72,11 @@ generate_ni_function!(
 #[test]
 fn sage_test_vectors() {
     let seed = b"hello world";
-    let context = b"yellow submarineyellow submarine";
+    let iv = *b"yellow submarineyellow submarine";
 
     let vectors = extract_vectors("src/tests/spec/allVectors.json").unwrap();
 
-    let functions: [fn(&[u8], &[u8]) -> (Vec<Scalar>, Vec<u8>); 5] = [
+    let functions: [fn(&[u8], [u8; 32]) -> (Vec<Scalar>, Vec<u8>); 5] = [
         NI_discrete_logarithm,
         NI_dleq,
         NI_pedersen_commitment,
@@ -86,10 +85,10 @@ fn sage_test_vectors() {
     ];
 
     for (i, f) in functions.iter().enumerate() {
-        let (_, proof_bytes) = f(seed, context);
+        let (_, proof_bytes) = f(seed, iv);
         assert_eq!(
-            context.to_vec(),
-            vectors[i].0,
+            iv.as_slice(),
+            vectors[i].0.as_slice(),
             "context for test vector {i} does not match"
         );
         assert_eq!(