fix(zk): Poseidon2 hash fixes (#2668)

Cargo.lock

@@ -4107,6 +4107,7 @@ dependencies = [
  "logos-blockchain-key-management-system-keys",
  "logos-blockchain-pol",
  "logos-blockchain-poq",
+ "logos-blockchain-poseidon2",
  "logos-blockchain-utils",
  "num-bigint",
  "serde",

blend/crypto/src/merkle.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 
 use lb_groth16::{fr_from_bytes_unchecked, fr_to_bytes};
 use lb_poq::{CORE_MERKLE_TREE_HEIGHT, CorePathAndSelectors};
+use lb_poseidon2::Digest;
 use rs_merkle_tree::{Node, stores::MemoryStore, tree::MerkleProof};
 use thiserror::Error;
 
@@ -31,14 +32,13 @@ struct InnerTreeZkHasher;
 
 impl rs_merkle_tree::hasher::Hasher for InnerTreeZkHasher {
     fn hash(&self, left: &Node, right: &Node) -> Node {
-        let mut hasher = ZkHasher::new();
-        hasher.compress(&[
-            // We use `unchecked` because we control the inputs, and poseidon hasher is guaranteed
-            // to always output valid `Fr` points.
+        let hash = <ZkHasher as Digest>::compress(&[
+            // We use `unchecked` because we control the inputs, and poseidon hasher is
+            // guaranteed to always output valid `Fr` points.
             fr_from_bytes_unchecked(left.as_ref()),
             fr_from_bytes_unchecked(right.as_ref()),
         ]);
-        fr_to_bytes(&hasher.finalize()).into()
+        fr_to_bytes(&hash).into()
     }
 }
 

blend/proofs/Cargo.toml

@@ -20,6 +20,7 @@ lb-blend-crypto = { workspace = true }
 lb-groth16      = { workspace = true }
 lb-pol          = { workspace = true }
 lb-poq          = { workspace = true }
+lb-poseidon2    = { workspace = true }
 lb-utils        = { workspace = true }
 num-bigint      = { workspace = true }
 serde           = { workspace = true }

blend/proofs/src/lib.rs

@@ -1,5 +1,6 @@
 use lb_blend_crypto::{ZkHash, ZkHasher};
 pub use lb_poq::CorePathAndSelectors;
+use lb_poseidon2::Digest;
 
 pub mod quota;
 pub mod selection;
@@ -13,9 +14,10 @@ where
     T: AsRef<[ZkHash]>,
 {
     fn hash(&self) -> ZkHash {
-        let mut hasher = ZkHasher::new();
-        hasher.update(self.as_ref());
-        hasher.finalize()
+        // let mut hasher = ZkHasher::new();
+        // hasher.update(self.as_ref());
+        // hasher.finalize();
+        <ZkHasher as Digest>::digest(self.as_ref())
     }
 }
 
@@ -25,16 +27,12 @@ trait ZkCompressExt {
 
 impl ZkCompressExt for [ZkHash; 2] {
     fn compress(&self) -> ZkHash {
-        let mut hasher = ZkHasher::new();
-        hasher.compress(self);
-        hasher.finalize()
+        <ZkHasher as Digest>::compress(self)
     }
 }
 
 impl ZkCompressExt for &[ZkHash; 2] {
     fn compress(&self) -> ZkHash {
-        let mut hasher = ZkHasher::new();
-        hasher.compress(self);
-        hasher.finalize()
+        <ZkHasher as Digest>::compress(self)
     }
 }

core/src/mantle/ops/leader_claim.rs

@@ -2,7 +2,7 @@ use std::sync::LazyLock;
 
 use lb_groth16::{fr_from_bytes, fr_to_bytes, serde::serde_fr};
 use lb_key_management_system_keys::keys::ZkPublicKey;
-use lb_poseidon2::{Fr, ZkHash};
+use lb_poseidon2::{Digest, Fr, ZkHash};
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
@@ -117,9 +117,10 @@ impl From<VoucherNullifier> for Fr {
 impl VoucherNullifier {
     #[must_use]
     pub fn from_secret(voucher_secret: VoucherSecret) -> Self {
-        let mut hash = ZkHasher::new();
-        hash.compress(&[*VOUCHER_NF, voucher_secret.into()]);
-        hash.finalize().into()
+        Self(<ZkHasher as Digest>::compress(&[
+            *VOUCHER_NF,
+            voucher_secret.into(),
+        ]))
     }
 }
 
@@ -137,9 +138,10 @@ impl VoucherCm {
 
     #[must_use]
     pub fn from_secret(voucher_secret: VoucherSecret) -> Self {
-        let mut hash = ZkHasher::new();
-        hash.compress(&[*REWARD_VOUCHER, voucher_secret.into()]);
-        hash.finalize().into()
+        Self(<ZkHasher as Digest>::compress(&[
+            *REWARD_VOUCHER,
+            voucher_secret.into(),
+        ]))
     }
 }
 

zk/poseidon2/src/hasher.rs

@@ -28,7 +28,8 @@ impl Poseidon2Hasher {
         );
     }
 
-    pub fn update(&mut self, input: &[Fr]) {
+    fn update(&mut self, input: &[Fr]) {
+        assert!(!input.is_empty());
         for fr in input {
             self.update_one(fr);
         }
@@ -36,7 +37,7 @@ impl Poseidon2Hasher {
     }
 
     /// Only use `compress` before `finalize` for poseidon2 without SAFE padding
-    pub fn compress(&mut self, inputs: &[Fr; 2]) {
+    fn compress(&mut self, inputs: &[Fr; 2]) {
         self.state[0] += inputs[0];
         self.state[1] += inputs[1];
         Poseidon2Bn254::permute_mut::<jf_poseidon2::constants::bn254::Poseidon2ParamsBn3, 3>(
@@ -44,7 +45,7 @@ impl Poseidon2Hasher {
         );
     }
 
-    pub const fn finalize(self) -> Fr {
+    const fn finalize(self) -> Fr {
         self.state[0]
     }
 }