sui_types/
zk_login_authenticator.rs

1// Copyright (c) 2021, Facebook, Inc. and its affiliates
2// Copyright (c) Mysten Labs, Inc.
3// SPDX-License-Identifier: Apache-2.0
4use crate::crypto::PublicKey;
5use crate::signature_verification::VerifiedDigestCache;
6use crate::{
7    base_types::{EpochId, SuiAddress},
8    crypto::{DefaultHash, Signature, SignatureScheme, SuiSignature},
9    digests::ZKLoginInputsDigest,
10    error::{SuiError, SuiErrorKind, SuiResult},
11    signature::{AuthenticatorTrait, VerifyParams},
12};
13use fastcrypto::{error::FastCryptoError, traits::ToFromBytes};
14use fastcrypto_zkp::bn254::zk_login::JwkId;
15use fastcrypto_zkp::bn254::zk_login::{JWK, OIDCProvider};
16use fastcrypto_zkp::bn254::zk_login_api::{ZkLoginCircuitMode, ZkLoginEnv};
17use fastcrypto_zkp::bn254::{zk_login::ZkLoginInputs, zk_login_api::verify_zk_login};
18use once_cell::sync::OnceCell;
19use schemars::JsonSchema;
20use serde::{Deserialize, Serialize};
21use shared_crypto::intent::IntentMessage;
22use std::hash::Hash;
23use std::hash::Hasher;
24use std::sync::Arc;
25#[cfg(test)]
26#[path = "unit_tests/zk_login_authenticator_test.rs"]
27mod zk_login_authenticator_test;
28
29/// An zk login authenticator with all the necessary fields.
30#[derive(Debug, Clone, JsonSchema, Serialize, Deserialize)]
31#[serde(rename_all = "camelCase")]
32pub struct ZkLoginAuthenticator {
33    pub inputs: ZkLoginInputs,
34    max_epoch: EpochId,
35    pub user_signature: Signature,
36    #[serde(skip)]
37    pub bytes: OnceCell<Vec<u8>>,
38}
39
40/// A helper struct that contains the necessary fields to calculate caching key.
41/// If the verify_zk_login() api changes, additional fields must be added here
42/// so the cache is not skipped.
43#[derive(Serialize, Deserialize)]
44struct ZkLoginCachingParams {
45    inputs: ZkLoginInputs,
46    max_epoch: EpochId,
47    extended_pk_bytes: Vec<u8>,
48    zklogin_circuit_mode: u64,
49}
50
51impl ZkLoginCachingParams {
52    fn digest(&self) -> ZKLoginInputsDigest {
53        use fastcrypto::hash::HashFunction;
54        let mut hasher = DefaultHash::default();
55        bcs::serialize_into(&mut hasher, self).expect("serde should not fail");
56        ZKLoginInputsDigest::new(hasher.finalize().into())
57    }
58}
59
60/// Map the protocol config circuit mode flag to the fastcrypto circuit mode:
61/// 0 = v1 circuit only, 1 = v2 circuit with fallback to v1 (migration mode), 2 = v2 circuit only.
62/// The flag comes from protocol config, so any other value is a config bug.
63pub fn zklogin_circuit_mode_from_flag(flag: u64) -> ZkLoginCircuitMode {
64    match flag {
65        0 => ZkLoginCircuitMode::V1Only,
66        1 => ZkLoginCircuitMode::Both,
67        2 => ZkLoginCircuitMode::V2Only,
68        _ => panic!("invalid zklogin circuit mode flag: {flag}"),
69    }
70}
71
72impl ZkLoginAuthenticator {
73    /// The caching key for zklogin signature, it is the hash of bcs bytes of
74    /// ZkLoginInputs || max_epoch || flagged_pk_bytes || zklogin_circuit_mode. If any of these
75    /// fields change, zklogin signature is re-verified without using the caching result.
76    fn get_caching_params(&self, zklogin_circuit_mode: u64) -> ZkLoginCachingParams {
77        let mut extended_pk_bytes = vec![self.user_signature.scheme().flag()];
78        extended_pk_bytes.extend(self.user_signature.public_key_bytes());
79        ZkLoginCachingParams {
80            inputs: self.inputs.clone(),
81            max_epoch: self.max_epoch,
82            extended_pk_bytes,
83            zklogin_circuit_mode,
84        }
85    }
86
87    pub fn hash_inputs(&self, zklogin_circuit_mode: u64) -> ZKLoginInputsDigest {
88        self.get_caching_params(zklogin_circuit_mode).digest()
89    }
90
91    /// Create a new [struct ZkLoginAuthenticator] with necessary fields.
92    pub fn new(inputs: ZkLoginInputs, max_epoch: EpochId, user_signature: Signature) -> Self {
93        Self {
94            inputs,
95            max_epoch,
96            user_signature,
97            bytes: OnceCell::new(),
98        }
99    }
100
101    pub fn get_pk(&self) -> SuiResult<PublicKey> {
102        PublicKey::from_zklogin_inputs(&self.inputs)
103    }
104
105    pub fn get_iss(&self) -> &str {
106        self.inputs.get_iss()
107    }
108
109    pub fn get_max_epoch(&self) -> EpochId {
110        self.max_epoch
111    }
112
113    pub fn user_signature_mut_for_testing(&mut self) -> &mut Signature {
114        &mut self.user_signature
115    }
116    pub fn max_epoch_mut_for_testing(&mut self) -> &mut EpochId {
117        &mut self.max_epoch
118    }
119    pub fn zk_login_inputs_mut_for_testing(&mut self) -> &mut ZkLoginInputs {
120        &mut self.inputs
121    }
122}
123
124/// Necessary trait for [struct SenderSignedData].
125impl PartialEq for ZkLoginAuthenticator {
126    fn eq(&self, other: &Self) -> bool {
127        self.as_ref() == other.as_ref()
128    }
129}
130
131/// Necessary trait for [struct SenderSignedData].
132impl Eq for ZkLoginAuthenticator {}
133
134/// Necessary trait for [struct SenderSignedData].
135impl Hash for ZkLoginAuthenticator {
136    fn hash<H: Hasher>(&self, state: &mut H) {
137        self.as_ref().hash(state);
138    }
139}
140
141impl AuthenticatorTrait for ZkLoginAuthenticator {
142    fn verify_user_authenticator_epoch(
143        &self,
144        epoch: EpochId,
145        max_epoch_upper_bound_delta: Option<u64>,
146    ) -> SuiResult {
147        // the checks here ensure that `current_epoch + max_epoch_upper_bound_delta >= self.max_epoch >= current_epoch`.
148        // 1. if the config for upper bound is set, ensure that the max epoch in signature is not larger than epoch + upper_bound.
149        if let Some(delta) = max_epoch_upper_bound_delta {
150            let max_epoch_upper_bound = epoch + delta;
151            if self.get_max_epoch() > max_epoch_upper_bound {
152                return Err(SuiErrorKind::InvalidSignature {
153                    error: format!(
154                        "ZKLogin max epoch too large {}, current epoch {}, max accepted: {}",
155                        self.get_max_epoch(),
156                        epoch,
157                        max_epoch_upper_bound
158                    ),
159                }
160                .into());
161            }
162        }
163        // 2. ensure that max epoch in signature is greater than the current epoch.
164        if epoch > self.get_max_epoch() {
165            return Err(SuiErrorKind::InvalidSignature {
166                error: format!(
167                    "ZKLogin expired at epoch {}, current epoch {}",
168                    self.get_max_epoch(),
169                    epoch
170                ),
171            }
172            .into());
173        }
174        Ok(())
175    }
176
177    /// Verify an intent message of a transaction with an zk login authenticator.
178    fn verify_claims<T>(
179        &self,
180        intent_msg: &IntentMessage<T>,
181        author: SuiAddress,
182        aux_verify_data: &VerifyParams,
183        zklogin_inputs_cache: Arc<VerifiedDigestCache<ZKLoginInputsDigest>>,
184    ) -> SuiResult
185    where
186        T: Serialize,
187    {
188        // Always evaluate the unpadded address derivation.
189        if author != SuiAddress::try_from_unpadded(&self.inputs)? {
190            // If the verify_legacy_zklogin_address flag is set, also evaluate the padded address derivation.
191            if !aux_verify_data.verify_legacy_zklogin_address
192                || author != SuiAddress::try_from_padded(&self.inputs)?
193            {
194                return Err(SuiErrorKind::InvalidAddress.into());
195            }
196        }
197
198        // Only when supported_providers list is not empty, we check if the provider is supported. Otherwise,
199        // we just use the JWK map to check if its supported.
200        if !aux_verify_data.supported_providers.is_empty()
201            && !aux_verify_data.supported_providers.contains(
202                &OIDCProvider::from_iss(self.inputs.get_iss()).map_err(|_| {
203                    SuiErrorKind::InvalidSignature {
204                        error: "Unknown provider".to_string(),
205                    }
206                })?,
207            )
208        {
209            return Err(SuiErrorKind::InvalidSignature {
210                error: format!("OIDC provider not supported: {}", self.inputs.get_iss()),
211            }
212            .into());
213        }
214
215        // zkLogin is classical (Groth16 over BN254); a PQ ephemeral key adds nothing.
216        match self.user_signature.scheme() {
217            SignatureScheme::ED25519 | SignatureScheme::Secp256k1 | SignatureScheme::Secp256r1 => {}
218            s => {
219                return Err(SuiErrorKind::InvalidSignature {
220                    error: format!("zkLogin ephemeral signature scheme not supported: {s}"),
221                }
222                .into());
223            }
224        }
225
226        // Verify the ephemeral signature over the intent message of the transaction data.
227        self.user_signature.verify_secure(
228            intent_msg,
229            author,
230            SignatureScheme::ZkLoginAuthenticator,
231        )?;
232
233        let zklogin_circuit_mode = aux_verify_data.zklogin_circuit_mode;
234        let caching_params = self.get_caching_params(zklogin_circuit_mode);
235        let inputs_digest = caching_params.digest();
236        if zklogin_inputs_cache.is_cached(&inputs_digest) {
237            // If the zklogin inputs hits the cache, we don't need to verify the zklogin
238            // again that contains the heavy computation.
239            Ok(())
240        } else {
241            // if it is not cached, we verify the full zklogin inputs.
242            verify_zklogin_inputs_wrapper(
243                caching_params,
244                &aux_verify_data.oidc_provider_jwks,
245                &aux_verify_data.zk_login_env,
246            )
247            .map_err(|e| -> SuiError {
248                SuiErrorKind::InvalidSignature {
249                    error: e.to_string(),
250                }
251                .into()
252            })?;
253            // If it's verified ok, we cache the digest.
254            zklogin_inputs_cache.cache_digest(inputs_digest);
255            Ok(())
256        }
257    }
258}
259
260fn verify_zklogin_inputs_wrapper(
261    params: ZkLoginCachingParams,
262    all_jwk: &imbl::HashMap<JwkId, JWK>,
263    env: &ZkLoginEnv,
264) -> SuiResult<()> {
265    verify_zk_login(
266        &params.inputs,
267        params.max_epoch,
268        &params.extended_pk_bytes,
269        all_jwk,
270        env,
271        zklogin_circuit_mode_from_flag(params.zklogin_circuit_mode),
272    )
273    .map_err(|e| {
274        SuiErrorKind::InvalidSignature {
275            error: e.to_string(),
276        }
277        .into()
278    })
279}
280
281impl ToFromBytes for ZkLoginAuthenticator {
282    fn from_bytes(bytes: &[u8]) -> Result<Self, FastCryptoError> {
283        // The first byte matches the flag of MultiSig.
284        if bytes.first().ok_or(FastCryptoError::InvalidInput)?
285            != &SignatureScheme::ZkLoginAuthenticator.flag()
286        {
287            return Err(FastCryptoError::InvalidInput);
288        }
289        let mut zk_login: ZkLoginAuthenticator =
290            bcs::from_bytes(&bytes[1..]).map_err(|_| FastCryptoError::InvalidSignature)?;
291        zk_login.inputs.init()?;
292        Ok(zk_login)
293    }
294}
295
296impl AsRef<[u8]> for ZkLoginAuthenticator {
297    fn as_ref(&self) -> &[u8] {
298        self.bytes
299            .get_or_try_init::<_, eyre::Report>(|| {
300                let as_bytes = bcs::to_bytes(self).expect("BCS serialization should not fail");
301                let mut bytes = Vec::with_capacity(1 + as_bytes.len());
302                bytes.push(SignatureScheme::ZkLoginAuthenticator.flag());
303                bytes.extend_from_slice(as_bytes.as_slice());
304                Ok(bytes)
305            })
306            .expect("OnceCell invariant violated")
307    }
308}
309
310#[derive(Debug, Clone)]
311pub struct AddressSeed([u8; 32]);
312
313impl AddressSeed {
314    pub fn unpadded(&self) -> &[u8] {
315        let mut buf = self.0.as_slice();
316
317        while !buf.is_empty() && buf[0] == 0 {
318            buf = &buf[1..];
319        }
320
321        // If the value is '0' then just return a slice of length 1 of the final byte
322        if buf.is_empty() { &self.0[31..] } else { buf }
323    }
324
325    pub fn padded(&self) -> &[u8] {
326        &self.0
327    }
328}
329
330impl std::fmt::Display for AddressSeed {
331    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332        let big_int = num_bigint::BigUint::from_bytes_be(&self.0);
333        let radix10 = big_int.to_str_radix(10);
334        f.write_str(&radix10)
335    }
336}
337
338#[derive(thiserror::Error, Debug)]
339pub enum AddressSeedParseError {
340    #[error("unable to parse radix10 encoded value `{0}`")]
341    Parse(#[from] num_bigint::ParseBigIntError),
342    #[error("larger than 32 bytes")]
343    TooBig,
344}
345
346impl std::str::FromStr for AddressSeed {
347    type Err = AddressSeedParseError;
348
349    fn from_str(s: &str) -> Result<Self, Self::Err> {
350        let big_int = <num_bigint::BigUint as num_traits::Num>::from_str_radix(s, 10)?;
351        let be_bytes = big_int.to_bytes_be();
352        let len = be_bytes.len();
353        let mut buf = [0; 32];
354
355        if len > 32 {
356            return Err(AddressSeedParseError::TooBig);
357        }
358
359        buf[32 - len..].copy_from_slice(&be_bytes);
360        Ok(Self(buf))
361    }
362}
363
364// AddressSeed's serialized format is as a radix10 encoded string
365impl Serialize for AddressSeed {
366    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
367    where
368        S: serde::Serializer,
369    {
370        self.to_string().serialize(serializer)
371    }
372}
373
374impl<'de> Deserialize<'de> for AddressSeed {
375    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
376    where
377        D: serde::Deserializer<'de>,
378    {
379        let s = std::borrow::Cow::<'de, str>::deserialize(deserializer)?;
380        std::str::FromStr::from_str(&s).map_err(serde::de::Error::custom)
381    }
382}
383
384#[cfg(test)]
385mod test {
386    use std::str::FromStr;
387
388    use super::AddressSeed;
389    use num_bigint::BigUint;
390    use proptest::prelude::*;
391
392    #[test]
393    fn unpadded_slice() {
394        let seed = AddressSeed([0; 32]);
395        let zero: [u8; 1] = [0];
396        assert_eq!(seed.unpadded(), zero.as_slice());
397
398        let mut seed = AddressSeed([1; 32]);
399        seed.0[0] = 0;
400        assert_eq!(seed.unpadded(), [1; 31].as_slice());
401    }
402
403    proptest! {
404        #[test]
405        fn dont_crash_on_large_inputs(
406            bytes in proptest::collection::vec(any::<u8>(), 33..1024)
407        ) {
408            let big_int = BigUint::from_bytes_be(&bytes);
409            let radix10 = big_int.to_str_radix(10);
410
411            // doesn't crash
412            let _ = AddressSeed::from_str(&radix10);
413        }
414
415        #[test]
416        fn valid_address_seeds(
417            bytes in proptest::collection::vec(any::<u8>(), 1..=32)
418        ) {
419            let big_int = BigUint::from_bytes_be(&bytes);
420            let radix10 = big_int.to_str_radix(10);
421
422            let seed = AddressSeed::from_str(&radix10).unwrap();
423            assert_eq!(radix10, seed.to_string());
424            // Ensure unpadded doesn't crash
425            seed.unpadded();
426        }
427    }
428}
sui_types/zk_login_authenticator.rs

sui_types/
zk_login_authenticator.rs
