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        // Verify the ephemeral signature over the intent message of the transaction data.
216        self.user_signature.verify_secure(
217            intent_msg,
218            author,
219            SignatureScheme::ZkLoginAuthenticator,
220        )?;
221
222        let zklogin_circuit_mode = aux_verify_data.zklogin_circuit_mode;
223        let caching_params = self.get_caching_params(zklogin_circuit_mode);
224        let inputs_digest = caching_params.digest();
225        if zklogin_inputs_cache.is_cached(&inputs_digest) {
226            // If the zklogin inputs hits the cache, we don't need to verify the zklogin
227            // again that contains the heavy computation.
228            Ok(())
229        } else {
230            // if it is not cached, we verify the full zklogin inputs.
231            verify_zklogin_inputs_wrapper(
232                caching_params,
233                &aux_verify_data.oidc_provider_jwks,
234                &aux_verify_data.zk_login_env,
235            )
236            .map_err(|e| -> SuiError {
237                SuiErrorKind::InvalidSignature {
238                    error: e.to_string(),
239                }
240                .into()
241            })?;
242            // If it's verified ok, we cache the digest.
243            zklogin_inputs_cache.cache_digest(inputs_digest);
244            Ok(())
245        }
246    }
247}
248
249fn verify_zklogin_inputs_wrapper(
250    params: ZkLoginCachingParams,
251    all_jwk: &im::HashMap<JwkId, JWK>,
252    env: &ZkLoginEnv,
253) -> SuiResult<()> {
254    verify_zk_login(
255        &params.inputs,
256        params.max_epoch,
257        &params.extended_pk_bytes,
258        all_jwk,
259        env,
260        zklogin_circuit_mode_from_flag(params.zklogin_circuit_mode),
261    )
262    .map_err(|e| {
263        SuiErrorKind::InvalidSignature {
264            error: e.to_string(),
265        }
266        .into()
267    })
268}
269
270impl ToFromBytes for ZkLoginAuthenticator {
271    fn from_bytes(bytes: &[u8]) -> Result<Self, FastCryptoError> {
272        // The first byte matches the flag of MultiSig.
273        if bytes.first().ok_or(FastCryptoError::InvalidInput)?
274            != &SignatureScheme::ZkLoginAuthenticator.flag()
275        {
276            return Err(FastCryptoError::InvalidInput);
277        }
278        let mut zk_login: ZkLoginAuthenticator =
279            bcs::from_bytes(&bytes[1..]).map_err(|_| FastCryptoError::InvalidSignature)?;
280        zk_login.inputs.init()?;
281        Ok(zk_login)
282    }
283}
284
285impl AsRef<[u8]> for ZkLoginAuthenticator {
286    fn as_ref(&self) -> &[u8] {
287        self.bytes
288            .get_or_try_init::<_, eyre::Report>(|| {
289                let as_bytes = bcs::to_bytes(self).expect("BCS serialization should not fail");
290                let mut bytes = Vec::with_capacity(1 + as_bytes.len());
291                bytes.push(SignatureScheme::ZkLoginAuthenticator.flag());
292                bytes.extend_from_slice(as_bytes.as_slice());
293                Ok(bytes)
294            })
295            .expect("OnceCell invariant violated")
296    }
297}
298
299#[derive(Debug, Clone)]
300pub struct AddressSeed([u8; 32]);
301
302impl AddressSeed {
303    pub fn unpadded(&self) -> &[u8] {
304        let mut buf = self.0.as_slice();
305
306        while !buf.is_empty() && buf[0] == 0 {
307            buf = &buf[1..];
308        }
309
310        // If the value is '0' then just return a slice of length 1 of the final byte
311        if buf.is_empty() { &self.0[31..] } else { buf }
312    }
313
314    pub fn padded(&self) -> &[u8] {
315        &self.0
316    }
317}
318
319impl std::fmt::Display for AddressSeed {
320    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
321        let big_int = num_bigint::BigUint::from_bytes_be(&self.0);
322        let radix10 = big_int.to_str_radix(10);
323        f.write_str(&radix10)
324    }
325}
326
327#[derive(thiserror::Error, Debug)]
328pub enum AddressSeedParseError {
329    #[error("unable to parse radix10 encoded value `{0}`")]
330    Parse(#[from] num_bigint::ParseBigIntError),
331    #[error("larger than 32 bytes")]
332    TooBig,
333}
334
335impl std::str::FromStr for AddressSeed {
336    type Err = AddressSeedParseError;
337
338    fn from_str(s: &str) -> Result<Self, Self::Err> {
339        let big_int = <num_bigint::BigUint as num_traits::Num>::from_str_radix(s, 10)?;
340        let be_bytes = big_int.to_bytes_be();
341        let len = be_bytes.len();
342        let mut buf = [0; 32];
343
344        if len > 32 {
345            return Err(AddressSeedParseError::TooBig);
346        }
347
348        buf[32 - len..].copy_from_slice(&be_bytes);
349        Ok(Self(buf))
350    }
351}
352
353// AddressSeed's serialized format is as a radix10 encoded string
354impl Serialize for AddressSeed {
355    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
356    where
357        S: serde::Serializer,
358    {
359        self.to_string().serialize(serializer)
360    }
361}
362
363impl<'de> Deserialize<'de> for AddressSeed {
364    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
365    where
366        D: serde::Deserializer<'de>,
367    {
368        let s = std::borrow::Cow::<'de, str>::deserialize(deserializer)?;
369        std::str::FromStr::from_str(&s).map_err(serde::de::Error::custom)
370    }
371}
372
373#[cfg(test)]
374mod test {
375    use std::str::FromStr;
376
377    use super::AddressSeed;
378    use num_bigint::BigUint;
379    use proptest::prelude::*;
380
381    #[test]
382    fn unpadded_slice() {
383        let seed = AddressSeed([0; 32]);
384        let zero: [u8; 1] = [0];
385        assert_eq!(seed.unpadded(), zero.as_slice());
386
387        let mut seed = AddressSeed([1; 32]);
388        seed.0[0] = 0;
389        assert_eq!(seed.unpadded(), [1; 31].as_slice());
390    }
391
392    proptest! {
393        #[test]
394        fn dont_crash_on_large_inputs(
395            bytes in proptest::collection::vec(any::<u8>(), 33..1024)
396        ) {
397            let big_int = BigUint::from_bytes_be(&bytes);
398            let radix10 = big_int.to_str_radix(10);
399
400            // doesn't crash
401            let _ = AddressSeed::from_str(&radix10);
402        }
403
404        #[test]
405        fn valid_address_seeds(
406            bytes in proptest::collection::vec(any::<u8>(), 1..=32)
407        ) {
408            let big_int = BigUint::from_bytes_be(&bytes);
409            let radix10 = big_int.to_str_radix(10);
410
411            let seed = AddressSeed::from_str(&radix10).unwrap();
412            assert_eq!(radix10, seed.to_string());
413            // Ensure unpadded doesn't crash
414            seed.unpadded();
415        }
416    }
417}
sui_types/zk_login_authenticator.rs

sui_types/
zk_login_authenticator.rs
