From 14e858a0a547bb38150b252fb60281fcb2e2a221 Mon Sep 17 00:00:00 2001 From: Joy Wang <108701016+joyqvq@users.noreply.github.com> Date: Mon, 30 Mar 2026 20:57:36 -0400 Subject: [PATCH] [zklogin] add support for v2 vk --- Cargo.lock | 1 + fastcrypto-zkp/Cargo.toml | 3 + fastcrypto-zkp/benches/zklogin.rs | 222 +++++++++++++- .../bn254/unit_tests/zk_login_e2e_tests.rs | 11 + .../src/bn254/unit_tests/zk_login_tests.rs | 99 ++++++- fastcrypto-zkp/src/bn254/utils.rs | 6 +- fastcrypto-zkp/src/bn254/zk_login.rs | 137 +++++++-- fastcrypto-zkp/src/bn254/zk_login_api.rs | 271 +++++++++++++++++- 8 files changed, 684 insertions(+), 66 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 156981f9a2..87e69604cb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1566,6 +1566,7 @@ dependencies = [ "serde_json", "test-strategy", "tokio", + "tracing", "typenum", ] diff --git a/fastcrypto-zkp/Cargo.toml b/fastcrypto-zkp/Cargo.toml index d0a11dd7e0..5a2227d897 100644 --- a/fastcrypto-zkp/Cargo.toml +++ b/fastcrypto-zkp/Cargo.toml @@ -15,6 +15,7 @@ harness = false [[bench]] name = "zklogin" harness = false +required-features = ["test-utils"] [[bench]] name = "poseidon" @@ -37,6 +38,7 @@ schemars = "0.8.10" serde = { version = "1.0.152", features = ["derive"] } serde_json = "1.0.93" once_cell = "1.16" +tracing = "0.1" im = "15" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } neptune = { version = "13.0.0", default-features = false } @@ -63,3 +65,4 @@ proptest = "1.1.0" [features] e2e = [] +test-utils = [] diff --git a/fastcrypto-zkp/benches/zklogin.rs b/fastcrypto-zkp/benches/zklogin.rs index 33f36fbf94..5d2f7fc293 100644 --- a/fastcrypto-zkp/benches/zklogin.rs +++ b/fastcrypto-zkp/benches/zklogin.rs @@ -8,16 +8,17 @@ mod zklogin_benches { use ark_std::rand::rngs::StdRng; use ark_std::rand::SeedableRng; - use criterion::Criterion; + use criterion::{BatchSize, Criterion}; use fastcrypto::ed25519::Ed25519KeyPair; use fastcrypto::error::FastCryptoError; use fastcrypto::rsa::{Base64UrlUnpadded, Encoding}; use fastcrypto::traits::KeyPair; use fastcrypto_zkp::bn254::utils::gen_address_seed; + use fastcrypto_zkp::bn254::zk_login::clear_cache_for_testing; 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::ZkLoginEnv; + use fastcrypto_zkp::bn254::zk_login_api::{CircuitVersion, ZkLoginEnv}; use im::hashmap::HashMap as ImHashMap; /// Benchmark the `fastcrypto_zkp::bn254::zk_login_api::verify_zk_login` function and it's main @@ -70,19 +71,45 @@ mod zklogin_benches { b.iter(|| input_clone.get_proof().as_arkworks().unwrap()) }); - // Benchmark the `calculate_all_inputs_hash` function called by `verify_zk_login`. + // Benchmark `calculate_all_inputs_hash` with a warm modulus-hash cache. let eph_pubkey_clone = eph_pubkey.clone(); let input_clone = input.clone(); let modulus_clone = modulus.clone(); - c.bench_function("verify_zk_login/calculate_all_inputs_hash", move |b| { + c.bench_function("verify_zk_login/calculate_all_inputs_hash/warm", move |b| { b.iter(|| { input_clone - .calculate_all_inputs_hash(&eph_pubkey_clone, &modulus_clone, max_epoch) + .calculate_all_inputs_hash( + &eph_pubkey_clone, + &modulus_clone, + max_epoch, + CircuitVersion::V1, + ) .unwrap() }); }); + + // Benchmark `calculate_all_inputs_hash` without cache so modulus-hash is computed every time. + let eph_pubkey_clone = eph_pubkey.clone(); + let input_clone = input.clone(); + let modulus_clone = modulus.clone(); + c.bench_function("verify_zk_login/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::V1, + ) + .unwrap() + }, + BatchSize::PerIteration, + ) + }); let input_hashes = input - .calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch) + .calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch, CircuitVersion::V1) .unwrap(); // Benchmark the `verify_zk_login_proof_with_fixed_vk` function called by `verify_zk_login`. @@ -92,32 +119,197 @@ mod zklogin_benches { move |b| { b.iter(|| { fastcrypto_zkp::bn254::zk_login_api::verify_zk_login_proof_with_fixed_vk( - &ZkLoginEnv::Prod, + &ZkLoginEnv::Test, + &proof, + &[input_hashes], + false, + ) + }) + }, + ); + + // Benchmark the entire `verify_zk_login` function with cache hit. + let input_warm = input.clone(); + let eph_warm = eph_pubkey.clone(); + let map_warm = map.clone(); + c.bench_function("verify_zk_login/warm", move |b| { + b.iter(|| { + fastcrypto_zkp::bn254::zk_login_api::verify_zk_login( + &input_warm, + max_epoch, + &eph_warm, + &map_warm, + &ZkLoginEnv::Test, + false, + ) + }) + }); + + // Benchmark `verify_zk_login` without cache. + c.bench_function("verify_zk_login/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, + false, + ) + }, + BatchSize::PerIteration, + ) + }); + } + + /// Benchmark V2 proof verification for 8192-bit RSA keys + fn verify_zk_login_v2(c: &mut Criterion) { + // Test values captured from test_zklogin_v2 + let max_epoch = 10; + let address_seed = + "1930628255822123795956154519923524356793387287437090556144422698180443693114"; + + let input = ZkLoginInputs::from_json( + r#"{"proofPoints":{"a":["14689542003205233165405871678578194233366368919811845258469718564746750013486","10012512891965046352621498096824179167030362165334951051528928536774758249772","1"],"b":[["12799698195028240196670878042975160414182451048782851217337031435532971817054","14558393469956339622214059893728330668015496468906843425269621870546568985673"],["9037880852443368589593489366180527669270225536480864658815735187820691374421","1238746346960939376770226388493818602061028587465272966849066262147126575607"],["1","0"]],"c":["3959688241581692257982179641980114002760624005031864644765944361132259645035","14237516486394033153908845927080955059437896911853766838881954835146028699428","1"]},"issBase64Details":{"value":"wiaXNzIjoiaHR0cHM6Ly9qd3QtdGVzdGVyLm15c3RlbmxhYnMuY29tIiw","indexMod4":2},"headerBase64":"eyJraWQiOiJzdWkta2V5LWlkLTgxOTIiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9","addressSeed":"1930628255822123795956154519923524356793387287437090556144422698180443693114"}"#, + address_seed, + ).unwrap(); + + let kp = Ed25519KeyPair::generate(&mut StdRng::from_seed([0; 32])); + let mut eph_pubkey = vec![0x00]; + eph_pubkey.extend(kp.public().as_ref()); + + let mut map = ImHashMap::new(); + let content = JWK { + kty: "RSA".to_string(), + e: "AQAB".to_string(), + n: "lViYJOuLB6EZenCimgyWrwOH_QBEkCZxSIEfcQgP5MrZkRlohbrTAN1YpXGRaqugp9A4mRzCmi9ddXscpRBSsLefdPJJLG8lQZ2qrw6X2-6HD5kDFd6-K7JZS-_GOEfr5xGEDm8_MS_SorbmneKspL0n4MPYWH8qke4OBFCwL6WzGBU9rqDuvhYmafmkvVvOtHIqekBxNrCud7Spv43BHdiBM0V-jUquuNM3oK97i_GVLjGfwrGRpR3tK4nva_ryiHh9Ajs68If7-ZhIoLJ05lRsHJJpqsloiEqlCZwhge9zEMnNkoaIzdQr-xLy0GPnr5W0gikjlSGYiInfx9ITADwK3W33xdOB7npM7lqJY73Njbuw8hBQicU8t0M0gvvWfmh1KDeA5IqffZgue-ka9Jj1nrYmZtd0JimQpPDUiGbLv69gQJZcLVQWf9z6mVC4gNm8VU2OafssnolrvNndC3wIm8AgqzVzn_DIOcMQdhIe8jTF3hu1_6R4Id3KoA5Hb3uI2H86-8RjhSG2wKb3zi44yKSmxEDhzl7i450PQX64JK4ftv5jb9vSw5unpikmVvGlGsuvrqWFuWKBcrcXLgyar8pGvRO8fR9ifDHSj-D2fBiLnhK0-iqsJeU8XnfJhUvKxSjXejwsoQeLqlgq9-PgCDP3dE61fkqGpJ1UZjZ44Q9Vh4YLCPAO6oX8btXSkwreuP5m0UtWgFsc-ynWbt6NYS7JlsMtJNWybM4_auqRdil_cPMwFsUgjocztGLeG304YH-GehmyBJyGKuDIiXL9RfLoZ35jKawrWJb4UqckKWV5kOKeXsXdKtMw96ABFumcnhrzxAsqwshS5a2lT8P7Cdd9g3T1JXI7JM1AnJU9_gPXmJoc3yEFNf-JxEf00URoy2xUusyyxYdTswLJp3NQP4VjrAGwnsp7gHKC-V-mJ21FpQCHsV0JQ-1x-E3du9hkpsjTtGkffetEsV8k9enbkudox7WIlsnPcA8y7aY4lnaBqLLSzaj2GOf4KTN4cRpcPzOmSvgcVVYYQXDjRw45X86P1WJG8UDl6Wkl044tAdQRuIxW8QVzBFWWxeXcoagOBKn1_DV0RKUX9Ud4LLauy81rUNfoAcnolz9nippTBEZA_4OOBvXhdngCYaoZyjAkmYdPhKIkghGhKoVVKiEJ1Ua6nUr3zB9WFlTO9lODeV9h0tgKGtKGu3UBeaRCQSMv9gZK-eGIpcqjsqK_rEf4htdDZUBzfOJ0VtCiFYUUBPiuJNuIf9xQGVDE7qZufK1irvGug8jvWSWzB4pGLP75PnPH7B9axnXrxssaIR90Y3Vr9ih_ptzcfNrwD_wiGHUTy698FHu2fXp51HbSEQ".to_string(), + alg: "RS256".to_string(), + }; + + map.insert( + JwkId::new( + OIDCProvider::TestIssuerKey8192.get_config().iss, + "sui-key-id-8192".to_string(), + ), + content.clone(), + ); + let modulus = Base64UrlUnpadded::decode_vec(&content.n) + .map_err(|_| { + FastCryptoError::GeneralError("Invalid Base64 encoded jwk modulus".to_string()) + }) + .unwrap(); + + // Benchmark the `as_arkworks` function called by `verify_zk_login`. + let input_clone = input.clone(); + c.bench_function("verify_zk_login_v2/as_arkworks", move |b| { + b.iter(|| input_clone.get_proof().as_arkworks().unwrap()) + }); + + // Benchmark `calculate_all_inputs_hash` with a WARM modulus-hash 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() + }); + }, + ); + + // 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(); + + // Benchmark the `verify_zk_login_proof_with_fixed_vk` function called by `verify_zk_login`. + let proof = input.get_proof().as_arkworks().unwrap(); + c.bench_function( + "verify_zk_login_v2/verify_zk_login_proof_with_fixed_vk", + move |b| { + b.iter(|| { + fastcrypto_zkp::bn254::zk_login_api::verify_zk_login_proof_with_fixed_vk( + &ZkLoginEnv::Test, &proof, &[input_hashes], + true, ) }) }, ); - // Benchmark the entire `verify_zk_login` function. - c.bench_function("verify_zk_login", move |b| { + // 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| { b.iter(|| { fastcrypto_zkp::bn254::zk_login_api::verify_zk_login( - &input, - 10, - &eph_pubkey, - &map, - &ZkLoginEnv::Prod, + &input_warm, + max_epoch, + &eph_warm, + &map_warm, + &ZkLoginEnv::Test, + true, ) }) }); + + // 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, + true, + ) + }, + BatchSize::PerIteration, + ) + }); } criterion_group! { name = zklogin_benches; config = Criterion::default(); - targets = verify_zk_login, + targets = verify_zk_login, verify_zk_login_v2, } } diff --git a/fastcrypto-zkp/src/bn254/unit_tests/zk_login_e2e_tests.rs b/fastcrypto-zkp/src/bn254/unit_tests/zk_login_e2e_tests.rs index c3e8c0b314..2a45546c88 100644 --- a/fastcrypto-zkp/src/bn254/unit_tests/zk_login_e2e_tests.rs +++ b/fastcrypto-zkp/src/bn254/unit_tests/zk_login_e2e_tests.rs @@ -46,6 +46,7 @@ async fn test_end_to_end_twitch() { &eph_pubkey, &map, &ZkLoginEnv::Test, + false, ); assert!(res.is_ok()); @@ -56,6 +57,7 @@ async fn test_end_to_end_twitch() { &eph_pubkey, &map, &ZkLoginEnv::Prod, + false, ); assert!(res_prod.is_err()); } @@ -88,6 +90,7 @@ async fn test_end_to_end_kakao() { &eph_pubkey, &map, &ZkLoginEnv::Test, + false, ); assert!(res.is_ok()); @@ -98,6 +101,7 @@ async fn test_end_to_end_kakao() { &eph_pubkey, &map, &ZkLoginEnv::Prod, + false, ); assert!(res_prod.is_err()); } @@ -129,6 +133,7 @@ async fn test_end_to_end_apple() { &eph_pubkey, &map, &ZkLoginEnv::Test, + false, ); assert!(res.is_ok()); @@ -139,6 +144,7 @@ async fn test_end_to_end_apple() { &eph_pubkey, &map, &ZkLoginEnv::Prod, + false, ); assert!(res_prod.is_err()); } @@ -170,6 +176,7 @@ async fn test_end_to_end_slack() { &eph_pubkey, &map, &ZkLoginEnv::Test, + false, ); assert!(res.is_ok()); @@ -180,6 +187,7 @@ async fn test_end_to_end_slack() { &eph_pubkey, &map, &ZkLoginEnv::Prod, + false, ); assert!(res_prod.is_err()); } @@ -229,6 +237,7 @@ async fn test_end_to_end_all_providers() { &eph_pubkey, &map, &ZkLoginEnv::Test, + false, ); assert!(res.is_ok()); @@ -239,6 +248,7 @@ async fn test_end_to_end_all_providers() { &eph_pubkey, &map, &ZkLoginEnv::Prod, + false, ); assert!(res_prod.is_err()); } @@ -363,6 +373,7 @@ async fn test_end_to_end_test_issuer(test_input: TestInputStruct) { &eph_pk_bytes, &map, &ZkLoginEnv::Test, + false, ); assert!(res.is_ok()); } diff --git a/fastcrypto-zkp/src/bn254/unit_tests/zk_login_tests.rs b/fastcrypto-zkp/src/bn254/unit_tests/zk_login_tests.rs index d9290a3873..842ba402bc 100644 --- a/fastcrypto-zkp/src/bn254/unit_tests/zk_login_tests.rs +++ b/fastcrypto-zkp/src/bn254/unit_tests/zk_login_tests.rs @@ -4,18 +4,16 @@ use std::str::FromStr; use crate::bn254::utils::{ - gen_address_seed, gen_address_seed_with_salt_hash, get_nonce, get_zk_login_address, + gen_address_seed, gen_address_seed_with_salt_hash, get_nonce, get_proof, get_zk_login_address, }; use crate::bn254::zk_login::big_int_array_to_bits; use crate::bn254::zk_login::bitarray_to_bytearray; use crate::bn254::zk_login::poseidon_zk_login; -use crate::bn254::zk_login::OIDCProvider; use crate::bn254::zk_login::{ - base64_to_bitarray, convert_base, decode_base64_url, hash_ascii_str_to_field, hash_to_field, - parse_jwks, trim, verify_extended_claim, Claim, JWTDetails, JwkId, + base64_to_bitarray, convert_base, decode_base64_url, fetch_jwks, hash_ascii_str_to_field, + hash_to_field, parse_jwks, trim, verify_extended_claim, Claim, JWTDetails, JwkId, OIDCProvider, }; -use crate::bn254::zk_login_api::ZkLoginEnv; -use crate::bn254::zk_login_api::{verify_zk_login_id, verify_zk_login_iss, Bn254Fr}; +use crate::bn254::zk_login_api::{verify_zk_login_id, verify_zk_login_iss, Bn254Fr, ZkLoginEnv}; use crate::bn254::{ zk_login::{ZkLoginInputs, JWK}, zk_login_api::verify_zk_login, @@ -27,7 +25,7 @@ use ark_std::rand::SeedableRng; use fastcrypto::ed25519::Ed25519KeyPair; use fastcrypto::encoding::{Encoding, Hex}; use fastcrypto::error::FastCryptoError; -use fastcrypto::jwt_utils::JWTHeader; +use fastcrypto::jwt_utils::{parse_and_validate_jwt, JWTHeader}; use fastcrypto::traits::KeyPair; use im::hashmap::HashMap as ImHashMap; use num_bigint::BigUint; @@ -166,7 +164,14 @@ async fn test_verify_zk_login_google() { ), content, ); - let res = verify_zk_login(&zk_login_inputs, 10, &eph_pubkey, &map, &ZkLoginEnv::Prod); + let res = verify_zk_login( + &zk_login_inputs, + 10, + &eph_pubkey, + &map, + &ZkLoginEnv::Prod, + false, + ); assert!(res.is_ok()); } @@ -639,6 +644,7 @@ fn test_alternative_iss_for_google() { &eph_pubkey_bytes, &all_jwk, &ZkLoginEnv::Test, + false, ); assert!(res.is_ok()); @@ -648,10 +654,87 @@ fn test_alternative_iss_for_google() { &eph_pubkey_bytes, &all_jwk, &ZkLoginEnv::Test, + false, ); assert!(invalid_res.is_err()); } +#[tokio::test] +async fn test_zklogin_v2() { + let max_epoch = 10; + let jwt_randomness = "100681567828351849884072155819400689117"; + let user_salt = "129390038577185583942388216820280642146"; + + // Generate an ephemeral key pair + let kp = Ed25519KeyPair::generate(&mut StdRng::from_seed([0; 32])); + let mut eph_pubkey = vec![0x00]; + eph_pubkey.extend(kp.public().as_ref()); + let kp_bigint = BigUint::from_bytes_be(&eph_pubkey).to_string(); + + // Get nonce + let nonce = get_nonce(&eph_pubkey, max_epoch, jwt_randomness).unwrap(); + + // Get JWT from 8192-bit key endpoint + let client = reqwest::Client::new(); + let iss = OIDCProvider::TestIssuerKey8192.get_config().iss; + let response = client + .post(format!( + "https://jwt-tester.mystenlabs.com/8192/jwt?nonce={}&iss={}&sub={}", + nonce, iss, "test" + )) + .header("Content-Type", "application/json") + .header("Content-Length", "0") + .send() + .await + .unwrap(); + let jwt_response: serde_json::Value = response.json().await.unwrap(); + let parsed_token = jwt_response["jwt"].as_str().unwrap().to_string(); + + // Get a proof from the V2 endpoint + let reader = get_proof( + &parsed_token, + max_epoch, + jwt_randomness, + &kp_bigint, + user_salt, + "https://prover-dev-v2.mystenlabs.com/v1", + ) + .await + .expect("get_proof failed"); + + // Get sub and aud + let (sub, aud, _) = + parse_and_validate_jwt(&parsed_token).expect("parse_and_validate_jwt failed"); + + // Get the address seed + let address_seed = + gen_address_seed(user_salt, "sub", &sub, &aud).expect("gen_address_seed failed"); + + let zk_login_inputs = + ZkLoginInputs::from_reader(reader, &address_seed).expect("from_reader failed"); + + // Fetch the 8192-bit RSA JWK + let jwks_vec = fetch_jwks(&OIDCProvider::TestIssuerKey8192, &client, false) + .await + .unwrap(); + + let mut all_jwk = ImHashMap::new(); + for (jwk_id, jwk) in jwks_vec { + all_jwk.insert(jwk_id, jwk); + } + + // V2 proof should verify using INSECURE_VERIFYING_KEY_V2 + let res_v2 = verify_zk_login( + &zk_login_inputs, + max_epoch, + &eph_pubkey, + &all_jwk, + &ZkLoginEnv::Test, + true, + ); + assert!(res_v2.is_ok()); +} + #[test] fn test_base64_to_bitarray() { let input = "a"; diff --git a/fastcrypto-zkp/src/bn254/utils.rs b/fastcrypto-zkp/src/bn254/utils.rs index 090aa6bdb1..c45ab37342 100644 --- a/fastcrypto-zkp/src/bn254/utils.rs +++ b/fastcrypto-zkp/src/bn254/utils.rs @@ -18,9 +18,9 @@ use std::str::FromStr; use super::zk_login::hash_ascii_str_to_field; const ZK_LOGIN_AUTHENTICATOR_FLAG: u8 = 0x05; -const MAX_KEY_CLAIM_NAME_LENGTH: u8 = 32; -const MAX_KEY_CLAIM_VALUE_LENGTH: u8 = 115; -const MAX_AUD_VALUE_LENGTH: u8 = 145; +const MAX_KEY_CLAIM_NAME_LENGTH: u16 = 32; +const MAX_KEY_CLAIM_VALUE_LENGTH: u16 = 115; +const MAX_AUD_VALUE_LENGTH: u16 = 145; /// Calculate the Sui address based on address seed and address params. pub fn get_zk_login_address( diff --git a/fastcrypto-zkp/src/bn254/zk_login.rs b/fastcrypto-zkp/src/bn254/zk_login.rs index 0635402ded..8007bac047 100644 --- a/fastcrypto-zkp/src/bn254/zk_login.rs +++ b/fastcrypto-zkp/src/bn254/zk_login.rs @@ -7,6 +7,7 @@ use serde_json::Value; use super::utils::split_to_two_frs; use crate::bn254::poseidon::poseidon_merkle_tree; +use crate::bn254::zk_login_api::CircuitVersion; use crate::bn254::FieldElement; use crate::zk_login_utils::{ g1_affine_from_str_projective, g2_affine_from_str_projective, Bn254FrElement, CircomG1, @@ -20,13 +21,47 @@ pub use ark_serialize::{CanonicalDeserialize, CanonicalSerialize}; use fastcrypto::error::FastCryptoError; use itertools::Itertools; use num_bigint::BigUint; +use once_cell::sync::Lazy; use regex::Regex; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use std::cmp::Ordering::{Equal, Greater, Less}; +use std::collections::HashMap; use std::error::Error; use std::fmt::Display; use std::str::FromStr; +use std::sync::RwLock; + +/// Key for the modulus hash cache: (modulus bytes, max_rsa_bits). +type ModulusHashKey = (Vec, u16); + +/// JWKs rotate occasionally, so caching by (modulus bytes, max_rsa_bits) avoids recomputing +/// bit-packing + poseidon hash on every verification. +static MODULUS_HASH_CACHE: Lazy>> = + Lazy::new(|| RwLock::new(HashMap::new())); + +fn cached_modulus_hash(modulus: &[u8], max_rsa_bits: u16) -> Result { + if let Some(f) = MODULUS_HASH_CACHE + .read() + .ok() + .and_then(|m| m.get(&(modulus.to_vec(), max_rsa_bits)).copied()) + { + return Ok(f); + } + let f = hash_to_field(&[BigUint::from_bytes_be(modulus)], max_rsa_bits, PACK_WIDTH)?; + if let Ok(mut m) = MODULUS_HASH_CACHE.write() { + m.insert((modulus.to_vec(), max_rsa_bits), f); + } + Ok(f) +} + +/// Clear the modulus hash cache for testing only benchmark. +#[cfg(any(test, feature = "test-utils"))] +pub fn clear_cache_for_testing() { + if let Ok(mut m) = MODULUS_HASH_CACHE.write() { + m.clear(); + } +} #[cfg(test)] #[path = "unit_tests/zk_login_tests.rs"] @@ -37,12 +72,9 @@ mod zk_login_tests; #[path = "unit_tests/zk_login_e2e_tests.rs"] mod zk_login_e2e_tests; -const MAX_HEADER_LEN: u8 = 248; const PACK_WIDTH: u8 = 248; const ISS: &str = "iss"; const BASE64_URL_CHARSET: &str = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"; -const MAX_EXT_ISS_LEN: u8 = 165; -const MAX_ISS_LEN_B64: u8 = 4 * (1 + MAX_EXT_ISS_LEN / 3); /// Key to identify a JWK, consists of iss and kid. #[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize, PartialOrd, Ord)] @@ -112,6 +144,8 @@ pub enum OIDCProvider { Credenza3, /// This is a test issuer that will return a JWT non-interactively. TestIssuer, + /// Test issuer for key 8192. + TestIssuerKey8192, /// https://oauth2.playtron.one/.well-known/jwks.json Playtron, /// https://auth.3dos.io/.well-known/openid-configuration @@ -140,6 +174,7 @@ impl FromStr for OIDCProvider { "Apple" => Ok(Self::Apple), "Slack" => Ok(Self::Slack), "TestIssuer" => Ok(Self::TestIssuer), + "TestIssuerKey8192" => Ok(Self::TestIssuerKey8192), "Microsoft" => Ok(Self::Microsoft), "KarrierOne" => Ok(Self::KarrierOne), "Credenza3" => Ok(Self::Credenza3), @@ -177,6 +212,7 @@ impl Display for OIDCProvider { Self::Apple => write!(f, "Apple"), Self::Slack => write!(f, "Slack"), Self::TestIssuer => write!(f, "TestIssuer"), + Self::TestIssuerKey8192 => write!(f, "TestIssuerKey8192"), Self::Microsoft => write!(f, "Microsoft"), Self::KarrierOne => write!(f, "KarrierOne"), Self::Credenza3 => write!(f, "Credenza3"), @@ -241,9 +277,13 @@ impl OIDCProvider { "https://accounts.credenza3.com/jwks", ), OIDCProvider::TestIssuer => ProviderConfig::new( - "https://oauth.sui.io", + "https://oauth.sui.io", "https://jwt-tester.mystenlabs.com/.well-known/jwks.json", ), + OIDCProvider::TestIssuerKey8192 => ProviderConfig::new( + "https://jwt-tester.mystenlabs.com", + "https://jwt-tester.mystenlabs.com/8192/jwks.json", + ), OIDCProvider::Playtron => ProviderConfig::new( "https://oauth2.playtron.one", "https://oauth2.playtron.one/.well-known/jwks.json", @@ -285,6 +325,7 @@ impl OIDCProvider { "https://appleid.apple.com" => Ok(Self::Apple), "https://slack.com" => Ok(Self::Slack), "https://oauth.sui.io" => Ok(Self::TestIssuer), + "https://jwt-tester.mystenlabs.com" => Ok(Self::TestIssuerKey8192), "https://accounts.karrier.one/" => Ok(Self::KarrierOne), "https://accounts.credenza3.com" => Ok(Self::Credenza3), "https://oauth2.playtron.one" => Ok(Self::Playtron), @@ -518,7 +559,7 @@ impl ZkLoginInputs { Self::from_reader(reader, address_seed) } - /// Initialize ZkLoginInputs from the + /// Initialize ZkLoginInputs from the reader. pub fn from_reader( reader: ZkLoginInputsReader, address_seed: &str, @@ -566,32 +607,64 @@ impl ZkLoginInputs { eph_pk_bytes: &[u8], modulus: &[u8], max_epoch: u64, + version: CircuitVersion, ) -> Result { - if self.header_base64.len() > MAX_HEADER_LEN as usize { + let config = version.config(); + if self.header_base64.len() > config.max_header_len_b64 as usize { return Err(FastCryptoError::GeneralError("Header too long".to_string())); } let addr_seed = (&self.address_seed).into(); let (first, second) = split_to_two_frs(eph_pk_bytes)?; - let max_epoch_f = (&Bn254FrElement::from_str(&max_epoch.to_string())?).into(); - let index_mod_4_f = - (&Bn254FrElement::from_str(&self.iss_base64_details.index_mod_4.to_string())?).into(); - - let iss_base64_f = - hash_ascii_str_to_field(&self.iss_base64_details.value, MAX_ISS_LEN_B64)?; - let header_f = hash_ascii_str_to_field(&self.header_base64, MAX_HEADER_LEN)?; - let modulus_f = hash_to_field(&[BigUint::from_bytes_be(modulus)], 2048, PACK_WIDTH)?; - poseidon_zk_login(&[ - first, - second, - addr_seed, - max_epoch_f, - iss_base64_f, - index_mod_4_f, - header_f, - modulus_f, - ]) + let header_f = hash_ascii_str_to_field(&self.header_base64, config.max_header_len_b64)?; + let modulus_f = cached_modulus_hash(modulus, config.max_rsa_bits)?; + + match version { + CircuitVersion::V1 => { + let index_mod_4_f = + (&Bn254FrElement::from_str(&self.iss_base64_details.index_mod_4.to_string())?) + .into(); + let iss_base64_f = + hash_ascii_str_to_field(&self.iss_base64_details.value, config.max_iss_len)?; + poseidon_zk_login(&[ + first, + second, + addr_seed, + max_epoch_f, + iss_base64_f, + index_mod_4_f, + header_f, + modulus_f, + ]) + } + CircuitVersion::V2 => { + let iss_f = self.hash_iss_decoded(config.max_iss_len)?; + let rsa_num_bits = BigUint::from_bytes_be(modulus).bits(); + let rsa_num_bits_f = (&Bn254FrElement::from_str(&rsa_num_bits.to_string())?).into(); + poseidon_zk_login(&[ + first, + second, + addr_seed, + max_epoch_f, + iss_f, + header_f, + modulus_f, + rsa_num_bits_f, + ]) + } + } + } + + /// Hash the v2 circuit's `iss_F`: the *decoded* extended iss claim (e.g. `,"iss":"https://...",`), + /// obtained by base64-decoding `iss_base64_details` at its `index_mod_4` offset. This differs + /// from v1, which hashes the raw base64 value directly. + fn hash_iss_decoded(&self, max_len: u16) -> FastCryptoResult { + let ext_iss = decode_base64_url( + &self.iss_base64_details.value, + &self.iss_base64_details.index_mod_4, + )?; + hash_ascii_str_to_field(&ext_iss, max_len) } } /// The struct for zk login proof. @@ -729,12 +802,12 @@ fn bitarray_to_bytearray(bits: &[u8]) -> FastCryptoResult> { } /// Pads a stream of bytes and maps it to a field element -pub fn hash_ascii_str_to_field(str: &str, max_size: u8) -> Result { +pub fn hash_ascii_str_to_field(str: &str, max_size: u16) -> Result { let str_padded = str_to_padded_char_codes(str, max_size)?; hash_to_field(&str_padded, 8, PACK_WIDTH) } -fn str_to_padded_char_codes(str: &str, max_len: u8) -> Result, FastCryptoError> { +fn str_to_padded_char_codes(str: &str, max_len: u16) -> Result, FastCryptoError> { let arr: Vec = str .chars() .map(|c| BigUint::from_slice(&([c as u32]))) @@ -742,7 +815,7 @@ fn str_to_padded_char_codes(str: &str, max_len: u8) -> Result, Fast pad_with_zeroes(arr, max_len) } -fn pad_with_zeroes(in_arr: Vec, out_count: u8) -> Result, FastCryptoError> { +fn pad_with_zeroes(in_arr: Vec, out_count: u16) -> Result, FastCryptoError> { if in_arr.len() > out_count as usize { return Err(FastCryptoError::GeneralError("in_arr too long".to_string())); } @@ -808,11 +881,13 @@ fn big_int_array_to_bits(integers: &[BigUint], intended_size: usize) -> FastCryp /// Calculate the poseidon hash of the field element inputs. If there are no inputs, return an error. /// If input length is <= 16, calculate H(inputs), if it is <= 32, calculate H(H(inputs[0..16]), -/// H(inputs[16..])), otherwise return an error. +/// H(inputs[16..])), if it is <= 48, calculate H(H(inputs[0..16]), H(inputs[16..32]), H(inputs[32..])), +/// if it is <= 64, calculate H(H(inputs[0..16]), H(inputs[16..32]), H(inputs[32..48]), H(inputs[48..])), +/// otherwise return an error. /// /// This functions must be equivalent with the one found in the zk_login circuit. pub(crate) fn poseidon_zk_login(inputs: &[Bn254Fr]) -> FastCryptoResult { - if inputs.is_empty() || inputs.len() > 32 { + if inputs.is_empty() || inputs.len() > 64 { return Err(FastCryptoError::InputLengthWrong(inputs.len())); } poseidon_merkle_tree(&inputs.iter().map(|x| FieldElement(*x)).collect_vec()).map(|x| x.0) @@ -823,5 +898,7 @@ fn test_poseidon_zk_login_input_sizes() { assert!(poseidon_zk_login(&[]).is_err()); assert!(poseidon_zk_login(&[Bn254Fr::from_str("123").unwrap(); 1]).is_ok()); assert!(poseidon_zk_login(&[Bn254Fr::from_str("123").unwrap(); 32]).is_ok()); - assert!(poseidon_zk_login(&[Bn254Fr::from_str("123").unwrap(); 33]).is_err()); + assert!(poseidon_zk_login(&[Bn254Fr::from_str("123").unwrap(); 33]).is_ok()); + assert!(poseidon_zk_login(&[Bn254Fr::from_str("123").unwrap(); 64]).is_ok()); + assert!(poseidon_zk_login(&[Bn254Fr::from_str("123").unwrap(); 65]).is_err()); } diff --git a/fastcrypto-zkp/src/bn254/zk_login_api.rs b/fastcrypto-zkp/src/bn254/zk_login_api.rs index 31d518345f..84579ad289 100644 --- a/fastcrypto-zkp/src/bn254/zk_login_api.rs +++ b/fastcrypto-zkp/src/bn254/zk_login_api.rs @@ -36,6 +36,12 @@ static GLOBAL_VERIFYING_KEY: Lazy> = Lazy::new(globa /// Corresponding to proofs generated from prover-dev. Used in devnet or other envs. static INSECURE_VERIFYING_KEY: Lazy> = Lazy::new(insecure_pvk); +/// V2 verifying key for production (mainnet/testnet). +static GLOBAL_VERIFYING_KEY_V2: Lazy> = Lazy::new(global_pvk_v2); + +/// V2 verifying key for test environments. +static INSECURE_VERIFYING_KEY_V2: Lazy> = Lazy::new(insecure_pvk_v2); + /// Load a fixed verifying key from zkLogin.vkey output. This is based on a local setup and should not use in production. fn insecure_pvk() -> PreparedVerifyingKey { // Convert the Circom G1/G2/GT to arkworks G1/G2/GT @@ -314,13 +320,162 @@ fn global_pvk() -> PreparedVerifyingKey { PreparedVerifyingKey::from(vk) } -/// Entry point for the ZkLogin API. +/// TODO: Replace with the actual V2 production verifying key from ceremony. +/// Currently uses the V1 production key as a placeholder so that V2 verification +/// fails gracefully and falls back to V1 in verify_zk_login. +fn global_pvk_v2() -> PreparedVerifyingKey { + global_pvk() +} + +/// Load a fixed verifying key for V2 (insecure/test). Based on zklogin-circuits v2-main branch artifacts/dev/zkLogin.vkey (finalized v2 circuit). +fn insecure_pvk_v2() -> PreparedVerifyingKey { + // Convert the Circom G1/G2/GT to arkworks G1/G2/GT + let vk_alpha_1 = g1_affine_from_str_projective(&vec![ + Bn254FqElement::from_str( + "20491192805390485299153009773594534940189261866228447918068658471970481763042", + ) + .unwrap(), + Bn254FqElement::from_str( + "9383485363053290200918347156157836566562967994039712273449902621266178545958", + ) + .unwrap(), + Bn254FqElement::from_str("1").unwrap(), + ]) + .unwrap(); + let vk_beta_2 = g2_affine_from_str_projective(&vec![ + vec![ + Bn254FqElement::from_str( + "6375614351688725206403948262868962793625744043794305715222011528459656738731", + ) + .unwrap(), + Bn254FqElement::from_str( + "4252822878758300859123897981450591353533073413197771768651442665752259397132", + ) + .unwrap(), + ], + vec![ + Bn254FqElement::from_str( + "10505242626370262277552901082094356697409835680220590971873171140371331206856", + ) + .unwrap(), + Bn254FqElement::from_str( + "21847035105528745403288232691147584728191162732299865338377159692350059136679", + ) + .unwrap(), + ], + vec![ + Bn254FqElement::from_str("1").unwrap(), + Bn254FqElement::from_str("0").unwrap(), + ], + ]) + .unwrap(); + let vk_gamma_2 = g2_affine_from_str_projective(&vec![ + vec![ + Bn254FqElement::from_str( + "10857046999023057135944570762232829481370756359578518086990519993285655852781", + ) + .unwrap(), + Bn254FqElement::from_str( + "11559732032986387107991004021392285783925812861821192530917403151452391805634", + ) + .unwrap(), + ], + vec![ + Bn254FqElement::from_str( + "8495653923123431417604973247489272438418190587263600148770280649306958101930", + ) + .unwrap(), + Bn254FqElement::from_str( + "4082367875863433681332203403145435568316851327593401208105741076214120093531", + ) + .unwrap(), + ], + vec![ + Bn254FqElement::from_str("1").unwrap(), + Bn254FqElement::from_str("0").unwrap(), + ], + ]) + .unwrap(); + let vk_delta_2 = g2_affine_from_str_projective(&vec![ + vec![ + Bn254FqElement::from_str( + "10857046999023057135944570762232829481370756359578518086990519993285655852781", + ) + .unwrap(), + Bn254FqElement::from_str( + "11559732032986387107991004021392285783925812861821192530917403151452391805634", + ) + .unwrap(), + ], + vec![ + Bn254FqElement::from_str( + "8495653923123431417604973247489272438418190587263600148770280649306958101930", + ) + .unwrap(), + Bn254FqElement::from_str( + "4082367875863433681332203403145435568316851327593401208105741076214120093531", + ) + .unwrap(), + ], + vec![ + Bn254FqElement::from_str("1").unwrap(), + Bn254FqElement::from_str("0").unwrap(), + ], + ]) + .unwrap(); + + // Create a vector of G1Affine elements from the IC + let mut vk_gamma_abc_g1 = Vec::new(); + for e in [ + vec![ + Bn254FqElement::from_str( + "7494019946711010262111829149650251402606293243816947408234427261113297488209", + ) + .unwrap(), + Bn254FqElement::from_str( + "15277037309547978131462479061393755632702482465205499508808087107719455935457", + ) + .unwrap(), + Bn254FqElement::from_str("1").unwrap(), + ], + vec![ + Bn254FqElement::from_str( + "11337299590256247592846534480065240790944417381563410255189124427241414159992", + ) + .unwrap(), + Bn254FqElement::from_str( + "10513214342948426742825684049934611922239382370302051503558098869981509341761", + ) + .unwrap(), + Bn254FqElement::from_str("1").unwrap(), + ], + ] { + let g1 = g1_affine_from_str_projective(&e).unwrap(); + vk_gamma_abc_g1.push(g1); + } + + let vk = VerifyingKey { + alpha_g1: vk_alpha_1, + beta_g2: vk_beta_2, + gamma_g2: vk_gamma_2, + delta_g2: vk_delta_2, + gamma_abc_g1: vk_gamma_abc_g1, + }; + + // Convert the verifying key into the prepared form. + PreparedVerifyingKey::from(vk) +} + +/// Verify a zkLogin proof. If `enable_zklogin_v2` is set, verify against the v2 circuit first and +/// fall back to v1; otherwise verify against v1 only. This flag is expected to be wired from the +/// `enable_zklogin_auth_v2` protocol config so v2 can be enabled per-network during rollout. pub fn verify_zk_login( input: &ZkLoginInputs, max_epoch: u64, eph_pubkey_bytes: &[u8], all_jwk: &ImHashMap, env: &ZkLoginEnv, + enable_zklogin_v2: bool, ) -> Result<(), FastCryptoError> { // Load the expected JWK based on (iss, kid). let (iss, kid) = (input.get_iss().to_string(), input.get_kid().to_string()); @@ -335,12 +490,54 @@ pub fn verify_zk_login( FastCryptoError::GeneralError("Invalid Base64 encoded jwk modulus".to_string()) })?; - // Calculate all inputs hash and passed to the verification function. - match verify_zk_login_proof_with_fixed_vk( - env, - &input.get_proof().as_arkworks()?, - &[input.calculate_all_inputs_hash(eph_pubkey_bytes, &modulus, max_epoch)?], - ) { + let proof = input.get_proof().as_arkworks()?; + + // When v2 is enabled, try verifying with the v2 circuit + VK first, then fall back to v1. + if enable_zklogin_v2 { + match input.calculate_all_inputs_hash( + eph_pubkey_bytes, + &modulus, + max_epoch, + CircuitVersion::V2, + ) { + Ok(all_inputs_hash) => { + match verify_zk_login_proof_with_fixed_vk(env, &proof, &[all_inputs_hash], true) { + Ok(true) => return Ok(()), + Ok(false) => { + tracing::debug!( + "[zkLogin] v2 verify returned false (env={:?}, iss={}), falling back to v1", + env, + iss + ); + } + Err(e) => { + tracing::debug!( + "[zkLogin] v2 verify failed (env={:?}, iss={}): {:?}, falling back to v1", + env, + iss, + e + ); + } + } + } + Err(e) => { + tracing::debug!( + "[zkLogin] v2 hash computation failed (env={:?}, iss={}): {:?}, falling back to v1", + env, + iss, + e + ); + } + } + } + + let all_inputs_hash = input.calculate_all_inputs_hash( + eph_pubkey_bytes, + &modulus, + max_epoch, + CircuitVersion::V1, + )?; + match verify_zk_login_proof_with_fixed_vk(env, &proof, &[all_inputs_hash], false) { Ok(true) => Ok(()), Ok(false) | Err(_) => Err(FastCryptoError::GeneralError( "Groth16 proof verify failed".to_string(), @@ -348,16 +545,70 @@ pub fn verify_zk_login( } } +/// Circuit version, which determines the layout of the `all_inputs_hash` public input. +#[derive(Debug, Copy, Clone, Eq, PartialEq)] +pub enum CircuitVersion { + /// V1 circuit: the iss claim is folded in as its *base64* value and `index_mod_4` is folded + /// into the hash. There is no `rsa_num_bits` field. + V1, + /// V2 circuit: the iss claim is folded in as its *decoded* extended claim, `index_mod_4` is + /// dropped, and the actual RSA modulus bit length (`rsa_num_bits`) is folded in instead. + V2, +} + +const V1_MAX_EXT_ISS_LEN: u16 = 165; +const V2_MAX_EXT_ISS_LEN: u16 = 186; + +impl CircuitVersion { + /// Circuit-specific parameters (max field lengths) for this version. + pub fn config(&self) -> CircuitConfig { + match self { + CircuitVersion::V1 => CircuitConfig { + max_header_len_b64: 248, + max_iss_len: 4 * (1 + V1_MAX_EXT_ISS_LEN / 3), + max_rsa_bits: 2048, + }, + CircuitVersion::V2 => CircuitConfig { + max_header_len_b64: 279, + max_iss_len: V2_MAX_EXT_ISS_LEN, + max_rsa_bits: 8192, + }, + } + } +} + +/// Circuit-specific parameters, obtained via [CircuitVersion::config]. +#[derive(Debug, Copy, Clone)] +pub struct CircuitConfig { + /// Maximum header length in base64. + pub max_header_len_b64: u16, + /// Maximum length passed to the iss hash. For [CircuitVersion::V1] this is the base64 length + /// of the iss claim; for [CircuitVersion::V2] it is the byte length of the *decoded* extended + /// iss claim (the circuit's `maxExtIssLength`). + pub max_iss_len: u16, + /// Maximum RSA key size in bits. The modulus is padded to this size before hashing. + pub max_rsa_bits: u16, +} + /// Verify a proof against its public inputs using the fixed verifying key. pub fn verify_zk_login_proof_with_fixed_vk( usage: &ZkLoginEnv, proof: &Proof, public_inputs: &[Bn254Fr], + v2: bool, ) -> Result { - let vk = match usage { - ZkLoginEnv::Prod => &GLOBAL_VERIFYING_KEY, - ZkLoginEnv::Test => &INSECURE_VERIFYING_KEY, + let vk = if v2 { + match usage { + ZkLoginEnv::Prod => &GLOBAL_VERIFYING_KEY_V2, + ZkLoginEnv::Test => &INSECURE_VERIFYING_KEY_V2, + } + } else { + match usage { + ZkLoginEnv::Prod => &GLOBAL_VERIFYING_KEY, + ZkLoginEnv::Test => &INSECURE_VERIFYING_KEY, + } }; + Groth16::::verify_with_processed_vk(vk, public_inputs, proof) .map_err(|e| FastCryptoError::GeneralError(e.to_string())) }