Skip to content

Commit fc8c0a1

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

6 files changed

Lines changed: 434 additions & 44 deletions

File tree

‎fastcrypto-zkp/benches/zklogin.rs‎

Lines changed: 109 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@ 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::{
21+
ZkLoginEnv, CIRCUIT_CONFIG_V1, CIRCUIT_CONFIG_V2,
22+
};
2123
use im::hashmap::HashMap as ImHashMap;
2224

2325
/// Benchmark the `fastcrypto_zkp::bn254::zk_login_api::verify_zk_login` function and it's main
@@ -32,7 +34,6 @@ mod zklogin_benches {
3234
// Test values taken from `test_verify_zk_login_google`. See the test for more details on
3335
// the values.
3436
let user_salt = "206703048842351542647799591018316385612";
35-
let max_epoch = 10;
3637
let address_seed = gen_address_seed(
3738
user_salt,
3839
"sub",
@@ -58,6 +59,7 @@ mod zklogin_benches {
5859
),
5960
content.clone(),
6061
);
62+
let max_epoch = 10;
6163
let modulus = Base64UrlUnpadded::decode_vec(&content.n)
6264
.map_err(|_| {
6365
FastCryptoError::GeneralError("Invalid Base64 encoded jwk modulus".to_string())
@@ -77,12 +79,17 @@ mod zklogin_benches {
7779
c.bench_function("verify_zk_login/calculate_all_inputs_hash", move |b| {
7880
b.iter(|| {
7981
input_clone
80-
.calculate_all_inputs_hash(&eph_pubkey_clone, &modulus_clone, max_epoch)
82+
.calculate_all_inputs_hash(
83+
&eph_pubkey_clone,
84+
&modulus_clone,
85+
max_epoch,
86+
&CIRCUIT_CONFIG_V1,
87+
)
8188
.unwrap()
8289
});
8390
});
8491
let input_hashes = input
85-
.calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch)
92+
.calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch, &CIRCUIT_CONFIG_V1)
8693
.unwrap();
8794

8895
// Benchmark the `verify_zk_login_proof_with_fixed_vk` function called by `verify_zk_login`.
@@ -92,9 +99,10 @@ mod zklogin_benches {
9299
move |b| {
93100
b.iter(|| {
94101
fastcrypto_zkp::bn254::zk_login_api::verify_zk_login_proof_with_fixed_vk(
95-
&ZkLoginEnv::Prod,
102+
&ZkLoginEnv::Test,
96103
&proof,
97104
&[input_hashes],
105+
false,
98106
)
99107
})
100108
},
@@ -105,10 +113,103 @@ mod zklogin_benches {
105113
b.iter(|| {
106114
fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
107115
&input,
108-
10,
116+
max_epoch,
117+
&eph_pubkey,
118+
&map,
119+
&ZkLoginEnv::Test,
120+
)
121+
})
122+
});
123+
}
124+
125+
/// Benchmark V2 proof verification for 8192-bit RSA keys
126+
fn verify_zk_login_v2(c: &mut Criterion) {
127+
// Test values captured from test_zklogin_v2
128+
let max_epoch = 10;
129+
let address_seed =
130+
"1930628255822123795956154519923524356793387287437090556144422698180443693114";
131+
132+
let input = ZkLoginInputs::from_json(
133+
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"}"#,
134+
address_seed,
135+
).unwrap();
136+
137+
let kp = Ed25519KeyPair::generate(&mut StdRng::from_seed([0; 32]));
138+
let mut eph_pubkey = vec![0x00];
139+
eph_pubkey.extend(kp.public().as_ref());
140+
141+
let mut map = ImHashMap::new();
142+
let content = JWK {
143+
kty: "RSA".to_string(),
144+
e: "AQAB".to_string(),
145+
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(),
146+
alg: "RS256".to_string(),
147+
};
148+
149+
map.insert(
150+
JwkId::new(
151+
OIDCProvider::TestIssuerKey8192.get_config().iss,
152+
"sui-key-id-8192".to_string(),
153+
),
154+
content.clone(),
155+
);
156+
let modulus = Base64UrlUnpadded::decode_vec(&content.n)
157+
.map_err(|_| {
158+
FastCryptoError::GeneralError("Invalid Base64 encoded jwk modulus".to_string())
159+
})
160+
.unwrap();
161+
162+
// Benchmark the `as_arkworks` function called by `verify_zk_login`.
163+
let input_clone = input.clone();
164+
c.bench_function("verify_zk_login_v2/as_arkworks", move |b| {
165+
b.iter(|| input_clone.get_proof().as_arkworks().unwrap())
166+
});
167+
168+
// Benchmark the `calculate_all_inputs_hash` function called by `verify_zk_login`.
169+
let eph_pubkey_clone = eph_pubkey.clone();
170+
let input_clone = input.clone();
171+
let modulus_clone = modulus.clone();
172+
c.bench_function("verify_zk_login_v2/calculate_all_inputs_hash", move |b| {
173+
b.iter(|| {
174+
input_clone
175+
.calculate_all_inputs_hash(
176+
&eph_pubkey_clone,
177+
&modulus_clone,
178+
max_epoch,
179+
&CIRCUIT_CONFIG_V2,
180+
)
181+
.unwrap()
182+
});
183+
});
184+
let input_hashes = input
185+
.calculate_all_inputs_hash(&eph_pubkey, &modulus, max_epoch, &CIRCUIT_CONFIG_V2)
186+
.unwrap();
187+
188+
// Benchmark the `verify_zk_login_proof_with_fixed_vk` function called by `verify_zk_login`.
189+
let proof = input.get_proof().as_arkworks().unwrap();
190+
c.bench_function(
191+
"verify_zk_login_v2/verify_zk_login_proof_with_fixed_vk",
192+
move |b| {
193+
b.iter(|| {
194+
fastcrypto_zkp::bn254::zk_login_api::verify_zk_login_proof_with_fixed_vk(
195+
&ZkLoginEnv::Test,
196+
&proof,
197+
&[input_hashes],
198+
true,
199+
)
200+
})
201+
},
202+
);
203+
204+
// Benchmark the entire `verify_zk_login` function.
205+
c.bench_function("verify_zk_login_v2", move |b| {
206+
b.iter(|| {
207+
fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
208+
&input,
209+
max_epoch,
109210
&eph_pubkey,
110211
&map,
111-
&ZkLoginEnv::Prod,
212+
&ZkLoginEnv::Test,
112213
)
113214
})
114215
});
@@ -117,7 +218,7 @@ mod zklogin_benches {
117218
criterion_group! {
118219
name = zklogin_benches;
119220
config = Criterion::default();
120-
targets = verify_zk_login,
221+
targets = verify_zk_login, verify_zk_login_v2,
121222
}
122223
}
123224

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,8 @@ async fn get_test_inputs(parsed_token: &str) -> (u64, Vec<u8>, ZkLoginInputs) {
272272
let (sub, aud, _) = parse_and_validate_jwt(parsed_token).unwrap();
273273
// Get the address seed.
274274
let address_seed = gen_address_seed(user_salt, "sub", &sub, &aud).unwrap();
275-
let zk_login_inputs = ZkLoginInputs::from_reader(reader, &address_seed).unwrap();
275+
let zk_login_inputs =
276+
ZkLoginInputs::from_reader(reader, &address_seed).unwrap();
276277
(max_epoch, eph_pubkey, zk_login_inputs)
277278
}
278279

@@ -345,7 +346,8 @@ async fn test_end_to_end_test_issuer(test_input: TestInputStruct) {
345346
// Get the address seed.
346347
let address_seed = gen_address_seed(&user_salt, "sub", &sub, &aud).unwrap();
347348
let zk_login_inputs =
348-
ZkLoginInputs::from_reader(reader, &address_seed.to_string()).unwrap();
349+
ZkLoginInputs::from_reader(reader, &address_seed.to_string())
350+
.unwrap();
349351

350352
// Make a map of jwk ids to jwks just for Microsoft.
351353
let iss = zk_login_inputs.get_iss();

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

Lines changed: 84 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,18 +4,19 @@
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,
15+
OIDCProvider,
16+
};
17+
use crate::bn254::zk_login_api::{
18+
verify_zk_login_id, verify_zk_login_iss, Bn254Fr, ZkLoginEnv,
1619
};
17-
use crate::bn254::zk_login_api::ZkLoginEnv;
18-
use crate::bn254::zk_login_api::{verify_zk_login_id, verify_zk_login_iss, Bn254Fr};
1920
use crate::bn254::{
2021
zk_login::{ZkLoginInputs, JWK},
2122
zk_login_api::verify_zk_login,
@@ -27,7 +28,7 @@ use ark_std::rand::SeedableRng;
2728
use fastcrypto::ed25519::Ed25519KeyPair;
2829
use fastcrypto::encoding::{Encoding, Hex};
2930
use fastcrypto::error::FastCryptoError;
30-
use fastcrypto::jwt_utils::JWTHeader;
31+
use fastcrypto::jwt_utils::{parse_and_validate_jwt, JWTHeader};
3132
use fastcrypto::traits::KeyPair;
3233
use im::hashmap::HashMap as ImHashMap;
3334
use num_bigint::BigUint;
@@ -652,6 +653,82 @@ fn test_alternative_iss_for_google() {
652653
assert!(invalid_res.is_err());
653654
}
654655

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