Skip to content

Commit 2d62b18

Browse files
joyqvqjoyqvq
committed
[zklogin] add support for v2 vk
1 parent b97ada5 commit 2d62b18

5 files changed

Lines changed: 432 additions & 43 deletions

File tree

‎fastcrypto-zkp/benches/zklogin.rs‎

Lines changed: 107 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ mod zklogin_benches {
1717
use fastcrypto_zkp::bn254::zk_login::ZkLoginInputs;
1818
use fastcrypto_zkp::bn254::zk_login::JWK;
1919
use fastcrypto_zkp::bn254::zk_login::{JwkId, OIDCProvider};
20-
use fastcrypto_zkp::bn254::zk_login_api::ZkLoginEnv;
20+
use fastcrypto_zkp::bn254::zk_login_api::{ZkLoginEnv, CIRCUIT_CONFIG_V1, CIRCUIT_CONFIG_V2};
2121
use im::hashmap::HashMap as ImHashMap;
2222

2323
/// Benchmark the `fastcrypto_zkp::bn254::zk_login_api::verify_zk_login` function and it's main
@@ -32,7 +32,6 @@ mod zklogin_benches {
3232
// Test values taken from `test_verify_zk_login_google`. See the test for more details on
3333
// the values.
3434
let user_salt = "206703048842351542647799591018316385612";
35-
let max_epoch = 10;
3635
let address_seed = gen_address_seed(
3736
user_salt,
3837
"sub",
@@ -58,6 +57,7 @@ mod zklogin_benches {
5857
),
5958
content.clone(),
6059
);
60+
let max_epoch = 10;
6161
let modulus = Base64UrlUnpadded::decode_vec(&content.n)
6262
.map_err(|_| {
6363
FastCryptoError::GeneralError("Invalid Base64 encoded jwk modulus".to_string())
@@ -77,12 +77,17 @@ mod zklogin_benches {
7777
c.bench_function("verify_zk_login/calculate_all_inputs_hash", move |b| {
7878
b.iter(|| {
7979
input_clone
80-
.calculate_all_inputs_hash(&eph_pubkey_clone, &modulus_clone, max_epoch)
80+
.calculate_all_inputs_hash(
81+
&eph_pubkey_clone,
82+
&modulus_clone,
83+
max_epoch,
84+
&CIRCUIT_CONFIG_V1,
85+
)
8186
.unwrap()
8287
});
8388
});
8489
let input_hashes = input
85-
.calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch)
90+
.calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch, &CIRCUIT_CONFIG_V1)
8691
.unwrap();
8792

8893
// Benchmark the `verify_zk_login_proof_with_fixed_vk` function called by `verify_zk_login`.
@@ -92,9 +97,10 @@ mod zklogin_benches {
9297
move |b| {
9398
b.iter(|| {
9499
fastcrypto_zkp::bn254::zk_login_api::verify_zk_login_proof_with_fixed_vk(
95-
&ZkLoginEnv::Prod,
100+
&ZkLoginEnv::Test,
96101
&proof,
97102
&[input_hashes],
103+
false,
98104
)
99105
})
100106
},
@@ -105,10 +111,103 @@ mod zklogin_benches {
105111
b.iter(|| {
106112
fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
107113
&input,
108-
10,
114+
max_epoch,
115+
&eph_pubkey,
116+
&map,
117+
&ZkLoginEnv::Test,
118+
)
119+
})
120+
});
121+
}
122+
123+
/// Benchmark V2 proof verification for 8192-bit RSA keys
124+
fn verify_zk_login_v2(c: &mut Criterion) {
125+
// Test values captured from test_zklogin_v2
126+
let max_epoch = 10;
127+
let address_seed =
128+
"1930628255822123795956154519923524356793387287437090556144422698180443693114";
129+
130+
let input = ZkLoginInputs::from_json(
131+
r#"{"proofPoints":{"a":["4913491815640002925508764814861178584881454035317776104347888483537912573177","17464247119089096977765585378460061328465709176842125201639874369409917083365","1"],"b":[["13623903508208593385147109129252793918112295419570003309520868038720322470557","21609423682403605552756457705069928412495291852654002331866073641632927420027"],["21392198638402084688930318789933313022805249822640479452861513428525783839707","1188996632803951473949030842369314644349566079256879538309939741515182911983"],["1","0"]],"c":["8847019028968200963788057481027139711885570926967685201543612972187276716667","14579483098715294861159755601821797996287919909580326110060065627124968449243","1"]},"issBase64Details":{"value":"wiaXNzIjoiaHR0cHM6Ly9qd3QtdGVzdGVyLm15c3RlbmxhYnMuY29tIiw","indexMod4":2},"headerBase64":"eyJraWQiOiJzdWkta2V5LWlkLTgxOTIiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9","addressSeed":"1930628255822123795956154519923524356793387287437090556144422698180443693114"}"#,
132+
address_seed,
133+
).unwrap();
134+
135+
let kp = Ed25519KeyPair::generate(&mut StdRng::from_seed([0; 32]));
136+
let mut eph_pubkey = vec![0x00];
137+
eph_pubkey.extend(kp.public().as_ref());
138+
139+
let mut map = ImHashMap::new();
140+
let content = JWK {
141+
kty: "RSA".to_string(),
142+
e: "AQAB".to_string(),
143+
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(),
144+
alg: "RS256".to_string(),
145+
};
146+
147+
map.insert(
148+
JwkId::new(
149+
OIDCProvider::TestIssuerKey8192.get_config().iss,
150+
"sui-key-id-8192".to_string(),
151+
),
152+
content.clone(),
153+
);
154+
let modulus = Base64UrlUnpadded::decode_vec(&content.n)
155+
.map_err(|_| {
156+
FastCryptoError::GeneralError("Invalid Base64 encoded jwk modulus".to_string())
157+
})
158+
.unwrap();
159+
160+
// Benchmark the `as_arkworks` function called by `verify_zk_login`.
161+
let input_clone = input.clone();
162+
c.bench_function("verify_zk_login_v2/as_arkworks", move |b| {
163+
b.iter(|| input_clone.get_proof().as_arkworks().unwrap())
164+
});
165+
166+
// Benchmark the `calculate_all_inputs_hash` function called by `verify_zk_login`.
167+
let eph_pubkey_clone = eph_pubkey.clone();
168+
let input_clone = input.clone();
169+
let modulus_clone = modulus.clone();
170+
c.bench_function("verify_zk_login_v2/calculate_all_inputs_hash", move |b| {
171+
b.iter(|| {
172+
input_clone
173+
.calculate_all_inputs_hash(
174+
&eph_pubkey_clone,
175+
&modulus_clone,
176+
max_epoch,
177+
&CIRCUIT_CONFIG_V2,
178+
)
179+
.unwrap()
180+
});
181+
});
182+
let input_hashes = input
183+
.calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch, &CIRCUIT_CONFIG_V2)
184+
.unwrap();
185+
186+
// Benchmark the `verify_zk_login_proof_with_fixed_vk` function called by `verify_zk_login`.
187+
let proof = input.get_proof().as_arkworks().unwrap();
188+
c.bench_function(
189+
"verify_zk_login_v2/verify_zk_login_proof_with_fixed_vk",
190+
move |b| {
191+
b.iter(|| {
192+
fastcrypto_zkp::bn254::zk_login_api::verify_zk_login_proof_with_fixed_vk(
193+
&ZkLoginEnv::Test,
194+
&proof,
195+
&[input_hashes],
196+
true,
197+
)
198+
})
199+
},
200+
);
201+
202+
// Benchmark the entire `verify_zk_login` function.
203+
c.bench_function("verify_zk_login_v2", move |b| {
204+
b.iter(|| {
205+
fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
206+
&input,
207+
max_epoch,
109208
&eph_pubkey,
110209
&map,
111-
&ZkLoginEnv::Prod,
210+
&ZkLoginEnv::Test,
112211
)
113212
})
114213
});
@@ -117,7 +216,7 @@ mod zklogin_benches {
117216
criterion_group! {
118217
name = zklogin_benches;
119218
config = Criterion::default();
120-
targets = verify_zk_login,
219+
targets = verify_zk_login, verify_zk_login_v2,
121220
}
122221
}
123222

‎fastcrypto-zkp/src/bn254/unit_tests/zk_login_tests.rs‎

Lines changed: 80 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,18 +4,16 @@
44
use std::str::FromStr;
55

66
use crate::bn254::utils::{
7-
gen_address_seed, gen_address_seed_with_salt_hash, get_nonce, get_zk_login_address,
7+
gen_address_seed, gen_address_seed_with_salt_hash, get_nonce, get_proof, get_zk_login_address,
88
};
99
use crate::bn254::zk_login::big_int_array_to_bits;
1010
use crate::bn254::zk_login::bitarray_to_bytearray;
1111
use crate::bn254::zk_login::poseidon_zk_login;
12-
use crate::bn254::zk_login::OIDCProvider;
1312
use crate::bn254::zk_login::{
14-
base64_to_bitarray, convert_base, decode_base64_url, hash_ascii_str_to_field, hash_to_field,
15-
parse_jwks, trim, verify_extended_claim, Claim, JWTDetails, JwkId,
13+
base64_to_bitarray, convert_base, decode_base64_url, fetch_jwks, hash_ascii_str_to_field,
14+
hash_to_field, parse_jwks, trim, verify_extended_claim, Claim, JWTDetails, JwkId, OIDCProvider,
1615
};
17-
use crate::bn254::zk_login_api::ZkLoginEnv;
18-
use crate::bn254::zk_login_api::{verify_zk_login_id, verify_zk_login_iss, Bn254Fr};
16+
use crate::bn254::zk_login_api::{verify_zk_login_id, verify_zk_login_iss, Bn254Fr, ZkLoginEnv};
1917
use crate::bn254::{
2018
zk_login::{ZkLoginInputs, JWK},
2119
zk_login_api::verify_zk_login,
@@ -27,7 +25,7 @@ use ark_std::rand::SeedableRng;
2725
use fastcrypto::ed25519::Ed25519KeyPair;
2826
use fastcrypto::encoding::{Encoding, Hex};
2927
use fastcrypto::error::FastCryptoError;
30-
use fastcrypto::jwt_utils::JWTHeader;
28+
use fastcrypto::jwt_utils::{parse_and_validate_jwt, JWTHeader};
3129
use fastcrypto::traits::KeyPair;
3230
use im::hashmap::HashMap as ImHashMap;
3331
use num_bigint::BigUint;
@@ -652,6 +650,81 @@ fn test_alternative_iss_for_google() {
652650
assert!(invalid_res.is_err());
653651
}
654652

653+
#[tokio::test]
654+
async fn test_zklogin_v2() {
655+
let max_epoch = 10;
656+
let jwt_randomness = "100681567828351849884072155819400689117";
657+
let user_salt = "129390038577185583942388216820280642146";
658+
659+
// Generate an ephemeral key pair
660+
let kp = Ed25519KeyPair::generate(&mut StdRng::from_seed([0; 32]));
661+
let mut eph_pubkey = vec![0x00];
662+
eph_pubkey.extend(kp.public().as_ref());
663+
let kp_bigint = BigUint::from_bytes_be(&eph_pubkey).to_string();
664+
665+
// Get nonce
666+
let nonce = get_nonce(kp.public().as_ref(), max_epoch, jwt_randomness).unwrap();
667+
668+
// Get JWT from 8192-bit key endpoint
669+
let client = reqwest::Client::new();
670+
let iss = OIDCProvider::TestIssuerKey8192.get_config().iss;
671+
let response = client
672+
.post(format!(
673+
"https://jwt-tester.mystenlabs.com/8192/jwt?nonce={}&iss={}&sub={}",
674+
nonce, iss, "test"
675+
))
676+
.header("Content-Type", "application/json")
677+
.header("Content-Length", "0")
678+
.send()
679+
.await
680+
.unwrap();
681+
let jwt_response: serde_json::Value = response.json().await.unwrap();
682+
let parsed_token = jwt_response["jwt"].as_str().unwrap().to_string();
683+
684+
// Get a proof from the V2 endpoint
685+
let reader = get_proof(
686+
&parsed_token,
687+
max_epoch,
688+
jwt_randomness,
689+
&kp_bigint,
690+
user_salt,
691+
"http://185.209.177.123:8080/v2",
692+
)
693+
.await
694+
.expect("get_proof failed");
695+
696+
// Get sub and aud
697+
let (sub, aud, _) =
698+
parse_and_validate_jwt(&parsed_token).expect("parse_and_validate_jwt failed");
699+
700+
// Get the address seed
701+
let address_seed =
702+
gen_address_seed(user_salt, "sub", &sub, &aud).expect("gen_address_seed failed");
703+
704+
let zk_login_inputs =
705+
ZkLoginInputs::from_reader(reader, &address_seed).expect("from_reader failed");
706+
707+
// Fetch the 8192-bit RSA JWK
708+
let jwks_vec = fetch_jwks(&OIDCProvider::TestIssuerKey8192, &client, false)
709+
.await
710+
.unwrap();
711+
712+
let mut all_jwk = ImHashMap::new();
713+
for (jwk_id, jwk) in jwks_vec {
714+
all_jwk.insert(jwk_id, jwk);
715+
}
716+
717+
// V2 proof should verify using TEST_VERIFYING_KEY_V2
718+
let res_v2 = verify_zk_login(
719+
&zk_login_inputs,
720+
max_epoch,
721+
&eph_pubkey,
722+
&all_jwk,
723+
&ZkLoginEnv::Test,
724+
);
725+
assert!(res_v2.is_ok());
726+
}
727+
655728
#[test]
656729
fn test_base64_to_bitarray() {
657730
let input = "a";

‎fastcrypto-zkp/src/bn254/utils.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,9 @@ use std::str::FromStr;
1818
use super::zk_login::hash_ascii_str_to_field;
1919

2020
const ZK_LOGIN_AUTHENTICATOR_FLAG: u8 = 0x05;
21-
const MAX_KEY_CLAIM_NAME_LENGTH: u8 = 32;
22-
const MAX_KEY_CLAIM_VALUE_LENGTH: u8 = 115;
23-
const MAX_AUD_VALUE_LENGTH: u8 = 145;
21+
const MAX_KEY_CLAIM_NAME_LENGTH: u16 = 32;
22+
const MAX_KEY_CLAIM_VALUE_LENGTH: u16 = 115;
23+
const MAX_AUD_VALUE_LENGTH: u16 = 145;
2424

2525
/// Calculate the Sui address based on address seed and address params.
2626
pub fn get_zk_login_address(

0 commit comments

Comments
 (0)