[zklogin] use sha2 for public inputs hash, update vk

fastcrypto-zkp/benches/zklogin.rs

@@ -18,7 +18,7 @@ mod zklogin_benches {
     use fastcrypto_zkp::bn254::zk_login::ZkLoginInputs;
     use fastcrypto_zkp::bn254::zk_login::JWK;
     use fastcrypto_zkp::bn254::zk_login::{JwkId, OIDCProvider};
-    use fastcrypto_zkp::bn254::zk_login_api::{CircuitVersion, ZkLoginCircuitMode, ZkLoginEnv};
+    use fastcrypto_zkp::bn254::zk_login_api::{CircuitVersion, ZkLoginEnv};
     use imbl::hashmap::HashMap as ImHashMap;
 
     /// Benchmark the `fastcrypto_zkp::bn254::zk_login_api::verify_zk_login` function and it's main
@@ -122,7 +122,7 @@ mod zklogin_benches {
                         &ZkLoginEnv::Test,
                         &proof,
                         &[input_hashes],
-                        false,
+                        CircuitVersion::V1,
                     )
                 })
             },
@@ -140,7 +140,7 @@ mod zklogin_benches {
                     &eph_warm,
                     &map_warm,
                     &ZkLoginEnv::Test,
-                    ZkLoginCircuitMode::V1Only,
+                    CircuitVersion::V1,
                 )
             })
         });
@@ -156,7 +156,7 @@ mod zklogin_benches {
                         &eph_pubkey,
                         &map,
                         &ZkLoginEnv::Test,
-                        ZkLoginCircuitMode::V1Only,
+                        CircuitVersion::V1,
                     )
                 },
                 BatchSize::PerIteration,
@@ -207,49 +207,23 @@ mod zklogin_benches {
             b.iter(|| input_clone.get_proof().as_arkworks().unwrap())
         });
 
-        // Benchmark `calculate_all_inputs_hash` with a WARM modulus-hash cache.
+        // Benchmark the SHA-256 public-input calculation. V2 does not use the V1 modulus cache.
         let eph_pubkey_clone = eph_pubkey.clone();
         let input_clone = input.clone();
         let modulus_clone = modulus.clone();
-        c.bench_function(
-            "verify_zk_login_v2/calculate_all_inputs_hash/warm",
-            move |b| {
-                b.iter(|| {
-                    input_clone
-                        .calculate_all_inputs_hash(
-                            &eph_pubkey_clone,
-                            &modulus_clone,
-                            max_epoch,
-                            CircuitVersion::V2,
-                        )
-                        .unwrap()
-                });
-            },
-        );
+        c.bench_function("verify_zk_login_v2/calculate_all_inputs_hash", move |b| {
+            b.iter(|| {
+                input_clone
+                    .calculate_all_inputs_hash(
+                        &eph_pubkey_clone,
+                        &modulus_clone,
+                        max_epoch,
+                        CircuitVersion::V2,
+                    )
+                    .unwrap()
+            });
+        });
 
-        // Benchmark `calculate_all_inputs_hash` with a COLD cache (cleared each iteration).
-        let eph_pubkey_clone = eph_pubkey.clone();
-        let input_clone = input.clone();
-        let modulus_clone = modulus.clone();
-        c.bench_function(
-            "verify_zk_login_v2/calculate_all_inputs_hash/cold",
-            move |b| {
-                b.iter_batched(
-                    clear_cache_for_testing,
-                    |_| {
-                        input_clone
-                            .calculate_all_inputs_hash(
-                                &eph_pubkey_clone,
-                                &modulus_clone,
-                                max_epoch,
-                                CircuitVersion::V2,
-                            )
-                            .unwrap()
-                    },
-                    BatchSize::PerIteration,
-                )
-            },
-        );
         let input_hashes = input
             .calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch, CircuitVersion::V2)
             .unwrap();
@@ -264,46 +238,25 @@ mod zklogin_benches {
                         &ZkLoginEnv::Test,
                         &proof,
                         &[input_hashes],
-                        true,
+                        CircuitVersion::V2,
                     )
                 })
             },
         );
 
-        // Benchmark the entire `verify_zk_login` function (warm: modulus hash cache hit).
-        let input_warm = input.clone();
-        let eph_warm = eph_pubkey.clone();
-        let map_warm = map.clone();
-        c.bench_function("verify_zk_login_v2/warm", move |b| {
+        // Benchmark the entire V2 verification path.
+        c.bench_function("verify_zk_login_v2", move |b| {
             b.iter(|| {
                 fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
-                    &input_warm,
+                    &input,
                     max_epoch,
-                    &eph_warm,
-                    &map_warm,
+                    &eph_pubkey,
+                    &map,
                     &ZkLoginEnv::Test,
-                    ZkLoginCircuitMode::Both,
+                    CircuitVersion::V2,
                 )
             })
         });
-
-        // Benchmark `verify_zk_login` on a cold cache (first call after a JWK refresh).
-        c.bench_function("verify_zk_login_v2/cold", move |b| {
-            b.iter_batched(
-                clear_cache_for_testing,
-                |_| {
-                    fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
-                        &input,
-                        max_epoch,
-                        &eph_pubkey,
-                        &map,
-                        &ZkLoginEnv::Test,
-                        ZkLoginCircuitMode::Both,
-                    )
-                },
-                BatchSize::PerIteration,
-            )
-        });
     }
 
     criterion_group! {

fastcrypto-zkp/src/bn254/mod.rs

@@ -19,6 +19,8 @@ pub mod verifier;
 /// Poseidon hash function over BN254
 pub mod poseidon;
 
+mod sha256_transcripts;
+
 /// Zk login structs and utilities
 pub mod zk_login;
 

fastcrypto-zkp/src/bn254/sha256_transcripts.rs

@@ -0,0 +1,65 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use fastcrypto::error::{FastCryptoError, FastCryptoResult};
+use fastcrypto::hash::{HashFunction, Sha256};
+
+pub(crate) fn append_fixed_width_be(
+    output: &mut Vec<u8>,
+    value: &[u8],
+    width: usize,
+) -> FastCryptoResult<()> {
+    let value = value
+        .iter()
+        .position(|byte| *byte != 0)
+        .map_or(&[][..], |first_nonzero| &value[first_nonzero..]);
+    let padding_len = width
+        .checked_sub(value.len())
+        .ok_or(FastCryptoError::InvalidInput)?;
+    let padded_len = output
+        .len()
+        .checked_add(padding_len)
+        .ok_or(FastCryptoError::InvalidInput)?;
+    output.resize(padded_len, 0);
+    output.extend_from_slice(value);
+    Ok(())
+}
+
+pub(crate) fn append_length_prefixed_bytes(
+    output: &mut Vec<u8>,
+    value: &[u8],
+    max_len: u16,
+) -> FastCryptoResult<()> {
+    let value_len = u16::try_from(value.len()).map_err(|_| FastCryptoError::InvalidInput)?;
+    if value_len == 0 || value_len > max_len {
+        return Err(FastCryptoError::InvalidInput);
+    }
+    output.extend_from_slice(&value_len.to_be_bytes());
+    output.extend_from_slice(value);
+    Ok(())
+}
+
+pub(crate) fn append_length_prefixed_padded_bytes(
+    output: &mut Vec<u8>,
+    value: &[u8],
+    max_len: usize,
+) -> FastCryptoResult<()> {
+    let value_len = u16::try_from(value.len()).map_err(|_| FastCryptoError::InvalidInput)?;
+    let padding_len = max_len
+        .checked_sub(value.len())
+        .ok_or(FastCryptoError::InvalidInput)?;
+    output.extend_from_slice(&value_len.to_be_bytes());
+    output.extend_from_slice(value);
+    let padded_len = output
+        .len()
+        .checked_add(padding_len)
+        .ok_or(FastCryptoError::InvalidInput)?;
+    output.resize(padded_len, 0);
+    Ok(())
+}
+
+pub(crate) fn sha256_low_253(input: &[u8]) -> [u8; 32] {
+    let mut digest: [u8; 32] = Sha256::digest(input).into();
+    digest[0] &= 0x1f;
+    digest
+}