sui_crypto/zklogin/poseidon/
mod.rs

1//! Poseidon Hash implementation using ark-ff
2//!
3//! This module is vendored from <https://github.com/arnaucube/poseidon-ark> at commit
4//! 6d2487aa1308d9d3860a2b724c485d73095c1c68 with a few minor changes on top.
5
6#![allow(clippy::needless_range_loop)]
7
8use ark_bn254::Fr;
9use ark_ff::fields::Field;
10use ark_std::Zero;
11use ark_std::str::FromStr;
12use core::ops::AddAssign;
13use core::ops::MulAssign;
14
15mod constants;
16
17pub static POSEIDON: std::sync::LazyLock<Poseidon> = std::sync::LazyLock::new(Poseidon::new);
18
19#[derive(Debug)]
20struct Constants {
21    pub c: Vec<Vec<Fr>>,
22    pub m: Vec<Vec<Vec<Fr>>>,
23    pub n_rounds_f: usize,
24    pub n_rounds_p: Vec<usize>,
25}
26
27fn load_constants() -> Constants {
28    let (c_str, m_str) = constants::constants();
29    let mut c: Vec<Vec<Fr>> = Vec::new();
30    for i in 0..c_str.len() {
31        let mut cci: Vec<Fr> = Vec::new();
32        for j in 0..c_str[i].len() {
33            let b: Fr = Fr::from_str(c_str[i][j]).unwrap();
34            cci.push(b);
35        }
36        c.push(cci);
37    }
38    let mut m: Vec<Vec<Vec<Fr>>> = Vec::new();
39    for i in 0..m_str.len() {
40        let mut mi: Vec<Vec<Fr>> = Vec::new();
41        for j in 0..m_str[i].len() {
42            let mut mij: Vec<Fr> = Vec::new();
43            for k in 0..m_str[i][j].len() {
44                let b: Fr = Fr::from_str(m_str[i][j][k]).unwrap();
45                mij.push(b);
46            }
47            mi.push(mij);
48        }
49        m.push(mi);
50    }
51    Constants {
52        c,
53        m,
54        n_rounds_f: 8,
55        n_rounds_p: vec![
56            56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68,
57        ],
58    }
59}
60
61pub struct Poseidon {
62    constants: Constants,
63}
64
65impl Poseidon {
66    pub fn new() -> Poseidon {
67        Poseidon {
68            constants: load_constants(),
69        }
70    }
71
72    fn ark(&self, state: &mut [Fr], c: &[Fr], it: usize) {
73        for i in 0..state.len() {
74            state[i].add_assign(&c[it + i]);
75        }
76    }
77
78    fn sbox(&self, n_rounds_f: usize, n_rounds_p: usize, state: &mut [Fr], i: usize) {
79        if i < n_rounds_f / 2 || i >= n_rounds_f / 2 + n_rounds_p {
80            for j in 0..state.len() {
81                let aux = state[j];
82                state[j] = state[j].square();
83                state[j] = state[j].square();
84                state[j].mul_assign(&aux);
85            }
86        } else {
87            let aux = state[0];
88            state[0] = state[0].square();
89            state[0] = state[0].square();
90            state[0].mul_assign(&aux);
91        }
92    }
93
94    fn mix(&self, state: &[Fr], m: &[Vec<Fr>]) -> Vec<Fr> {
95        let mut new_state: Vec<Fr> = Vec::new();
96        for i in 0..state.len() {
97            new_state.push(Fr::zero());
98            for j in 0..state.len() {
99                let mut mij = m[i][j];
100                mij.mul_assign(&state[j]);
101                new_state[i].add_assign(&mij);
102            }
103        }
104        new_state.clone()
105    }
106
107    /// Support up to 64 inputs.
108    ///
109    /// Equivalent to fastcrypto's `poseidon_zk_login` (which is backed by the `neptune` crate;
110    /// this implementation is vendored from `poseidon-ark` instead — both match circomlib's
111    /// Poseidon over BN254).
112    pub fn hash(&self, inp: &[Fr]) -> Result<Fr, String> {
113        let max_direct_inputs = self.constants.n_rounds_p.len();
114        if inp.len() > max_direct_inputs * 4 {
115            return Err("Wrong inputs length".to_string());
116        }
117        if inp.len() > max_direct_inputs {
118            let chunks = inp
119                .chunks(max_direct_inputs)
120                .map(|chunk| self.hash(chunk))
121                .collect::<Result<Vec<_>, _>>()?;
122            return self.hash(&chunks);
123        }
124
125        let t = inp.len() + 1;
126        if inp.is_empty() || inp.len() > self.constants.n_rounds_p.len() {
127            return Err("Wrong inputs length".to_string());
128        }
129        let n_rounds_f = self.constants.n_rounds_f;
130        let n_rounds_p = self.constants.n_rounds_p[t - 2];
131
132        let mut state = vec![Fr::zero(); t];
133        state[1..].clone_from_slice(inp);
134
135        for i in 0..(n_rounds_f + n_rounds_p) {
136            self.ark(&mut state, &self.constants.c[t - 2], i * t);
137            self.sbox(n_rounds_f, n_rounds_p, &mut state, i);
138            state = self.mix(&state, &self.constants.m[t - 2]);
139        }
140
141        Ok(state[0])
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[cfg(test)]
150    #[cfg(target_arch = "wasm32")]
151    use wasm_bindgen_test::wasm_bindgen_test as test;
152
153    #[test]
154    fn test_load_constants() {
155        let cons = load_constants();
156        assert_eq!(
157            cons.c[0][0].to_string(),
158            "4417881134626180770308697923359573201005643519861877412381846989312604493735"
159        );
160        assert_eq!(
161            cons.c[cons.c.len() - 1][0].to_string(),
162            "21579410516734741630578831791708254656585702717204712919233299001262271512412"
163        );
164        assert_eq!(
165            cons.m[0][0][0].to_string(),
166            "2910766817845651019878574839501801340070030115151021261302834310722729507541"
167        );
168        assert_eq!(
169            cons.m[cons.m.len() - 1][0][0].to_string(),
170            "11497693837059016825308731789443585196852778517742143582474723527597064448312"
171        );
172    }
173
174    #[test]
175    fn test_hash() {
176        let b0: Fr = Fr::from_str("0").unwrap();
177        let b1: Fr = Fr::from_str("1").unwrap();
178        let b2: Fr = Fr::from_str("2").unwrap();
179        let b3: Fr = Fr::from_str("3").unwrap();
180        let b4: Fr = Fr::from_str("4").unwrap();
181        let b5: Fr = Fr::from_str("5").unwrap();
182        let b6: Fr = Fr::from_str("6").unwrap();
183        let b7: Fr = Fr::from_str("7").unwrap();
184        let b8: Fr = Fr::from_str("8").unwrap();
185        let b9: Fr = Fr::from_str("9").unwrap();
186        let b10: Fr = Fr::from_str("10").unwrap();
187        let b11: Fr = Fr::from_str("11").unwrap();
188        let b12: Fr = Fr::from_str("12").unwrap();
189        let b13: Fr = Fr::from_str("13").unwrap();
190        let b14: Fr = Fr::from_str("14").unwrap();
191        let b15: Fr = Fr::from_str("15").unwrap();
192        let b16: Fr = Fr::from_str("16").unwrap();
193
194        let poseidon = Poseidon::new();
195
196        let big_arr: Vec<Fr> = vec![b1];
197        let h = poseidon.hash(&big_arr).unwrap();
198        assert_eq!(
199            h.to_string(),
200            "18586133768512220936620570745912940619677854269274689475585506675881198879027"
201        );
202
203        let big_arr: Vec<Fr> = vec![b1, b2];
204        let h = poseidon.hash(&big_arr).unwrap();
205        assert_eq!(
206            h.to_string(),
207            "7853200120776062878684798364095072458815029376092732009249414926327459813530"
208        );
209
210        let big_arr: Vec<Fr> = vec![b1, b2, b0, b0, b0];
211        let h = poseidon.hash(&big_arr).unwrap();
212        assert_eq!(
213            h.to_string(),
214            "1018317224307729531995786483840663576608797660851238720571059489595066344487"
215        );
216
217        let big_arr: Vec<Fr> = vec![b1, b2, b0, b0, b0, b0];
218        let h = poseidon.hash(&big_arr).unwrap();
219        assert_eq!(
220            h.to_string(),
221            "15336558801450556532856248569924170992202208561737609669134139141992924267169"
222        );
223
224        let big_arr: Vec<Fr> = vec![b3, b4, b0, b0, b0];
225        let h = poseidon.hash(&big_arr).unwrap();
226        assert_eq!(
227            h.to_string(),
228            "5811595552068139067952687508729883632420015185677766880877743348592482390548"
229        );
230
231        let big_arr: Vec<Fr> = vec![b3, b4, b0, b0, b0, b0];
232        let h = poseidon.hash(&big_arr).unwrap();
233        assert_eq!(
234            h.to_string(),
235            "12263118664590987767234828103155242843640892839966517009184493198782366909018"
236        );
237
238        let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6];
239        let h = poseidon.hash(&big_arr).unwrap();
240        assert_eq!(
241            h.to_string(),
242            "20400040500897583745843009878988256314335038853985262692600694741116813247201"
243        );
244
245        let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14];
246        let h = poseidon.hash(&big_arr).unwrap();
247        assert_eq!(
248            h.to_string(),
249            "8354478399926161176778659061636406690034081872658507739535256090879947077494"
250        );
251
252        let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6, b7, b8, b9, b0, b0, b0, b0, b0];
253        let h = poseidon.hash(&big_arr).unwrap();
254        assert_eq!(
255            h.to_string(),
256            "5540388656744764564518487011617040650780060800286365721923524861648744699539"
257        );
258
259        let big_arr: Vec<Fr> = vec![
260            b1, b2, b3, b4, b5, b6, b7, b8, b9, b0, b0, b0, b0, b0, b0, b0,
261        ];
262        let h = poseidon.hash(&big_arr).unwrap();
263        assert_eq!(
264            h.to_string(),
265            "11882816200654282475720830292386643970958445617880627439994635298904836126497"
266        );
267
268        let big_arr: Vec<Fr> = vec![
269            b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14, b15, b16,
270        ];
271        let h = poseidon.hash(&big_arr).unwrap();
272        assert_eq!(
273            h.to_string(),
274            "9989051620750914585850546081941653841776809718687451684622678807385399211877"
275        );
276    }
277    #[test]
278    fn test_wrong_inputs() {
279        let poseidon = Poseidon::new();
280
281        poseidon.hash(&[]).expect_err("Wrong inputs length");
282        let big_arr: Vec<Fr> = vec![Fr::from_str("1").unwrap(); 65];
283        poseidon.hash(&big_arr).expect_err("Wrong inputs length");
284    }
285
286    // 17 to 64 inputs are hashed as a 16-ary merkle tree, matching the zkLogin circuit.
287    #[test]
288    fn test_merkle_tree_inputs() {
289        let poseidon = Poseidon::new();
290
291        let arr: Vec<Fr> = (0..34)
292            .map(|i| Fr::from_str(&i.to_string()).unwrap())
293            .collect();
294        let expected = poseidon
295            .hash(&[
296                poseidon.hash(&arr[0..16]).unwrap(),
297                poseidon.hash(&arr[16..32]).unwrap(),
298                poseidon.hash(&arr[32..34]).unwrap(),
299            ])
300            .unwrap();
301        assert_eq!(poseidon.hash(&arr).unwrap(), expected);
302    }
303}
sui_crypto/zklogin/poseidon/mod.rs

sui_crypto/zklogin/poseidon/
mod.rs
