Skip to main content

sui_crypto/zklogin/
mod.rs

1use std::collections::HashMap;
2
3use crate::SignatureError;
4use poseidon::POSEIDON;
5use signature::Verifier;
6use sui_sdk_types::Jwk;
7use sui_sdk_types::JwkId;
8use sui_sdk_types::UserSignature;
9use sui_sdk_types::ZkLoginAuthenticator;
10use verify::CircuitVersion;
11
12mod poseidon;
13mod verify;
14
15#[cfg(test)]
16mod tests;
17
18/// Which zkLogin circuit versions to accept during proof verification.
19#[derive(Debug, Copy, Clone, Eq, PartialEq, Default)]
20pub enum ZkLoginCircuitMode {
21    /// Accept proofs against the v1 circuit only; v2 proofs are rejected.
22    #[default]
23    V1Only,
24    /// Try the v2 circuit first, then fall back to v1. Migration phase.
25    Both,
26    /// Accept proofs against the v2 circuit only; v1 proofs are rejected.
27    V2Only,
28}
29
30#[derive(Debug, Clone, PartialEq)]
31pub struct ZkloginVerifier {
32    proof_verifying_keys: HashMap<CircuitVersion, verify::VerifyingKey>,
33    circuit_mode: ZkLoginCircuitMode,
34    jwks: HashMap<JwkId, Jwk>,
35}
36
37impl ZkloginVerifier {
38    fn new(proof_verifying_keys: HashMap<CircuitVersion, verify::VerifyingKey>) -> Self {
39        Self {
40            proof_verifying_keys,
41            circuit_mode: Default::default(),
42            jwks: Default::default(),
43        }
44    }
45
46    pub fn new_mainnet() -> Self {
47        Self::new(
48            [(CircuitVersion::V1, verify::VerifyingKey::new_mainnet())]
49                .into_iter()
50                .collect(),
51        )
52    }
53
54    pub fn new_dev() -> Self {
55        Self::new(
56            [CircuitVersion::V1, CircuitVersion::V2]
57                .into_iter()
58                .map(|version| (version, verify::VerifyingKey::new_dev_for(version)))
59                .collect(),
60        )
61    }
62
63    pub fn jwks(&self) -> &HashMap<JwkId, Jwk> {
64        &self.jwks
65    }
66
67    pub fn jwks_mut(&mut self) -> &mut HashMap<JwkId, Jwk> {
68        &mut self.jwks
69    }
70
71    pub fn circuit_mode(&self) -> ZkLoginCircuitMode {
72        self.circuit_mode
73    }
74
75    /// Set which circuit mode version from sui.
76    pub fn set_circuit_mode(&mut self, circuit_mode: ZkLoginCircuitMode) {
77        self.circuit_mode = circuit_mode;
78    }
79}
80
81impl Verifier<ZkLoginAuthenticator> for ZkloginVerifier {
82    fn verify(
83        &self,
84        message: &[u8],
85        signature: &ZkLoginAuthenticator,
86    ) -> Result<(), SignatureError> {
87        // 1. check that we have a valid corrisponding Jwk
88        let jwk_id = signature.inputs.jwk_id();
89        let jwk = self.jwks.get(jwk_id).ok_or_else(|| {
90            SignatureError::from_source(format!(
91                "unable to find corrisponding jwk with id '{:?}' for provided authenticator",
92                jwk_id,
93            ))
94        })?;
95
96        // 2. verify that the provided SimpleSignature is valid
97        crate::simple::SimpleVerifier.verify(message, &signature.signature)?;
98
99        // 3. verify groth16 proof against the circuit versions allowed by the configured mode
100        let verify_proof = |version: CircuitVersion| {
101            self.proof_verifying_keys
102                .get(&version)
103                .ok_or_else(|| {
104                    SignatureError::from_source(format!("no verifying key for circuit {version:?}"))
105                })?
106                .verify_zklogin(
107                    jwk,
108                    &signature.inputs,
109                    &signature.signature,
110                    signature.max_epoch,
111                    version,
112                )
113        };
114        match self.circuit_mode {
115            ZkLoginCircuitMode::V1Only => verify_proof(CircuitVersion::V1),
116            ZkLoginCircuitMode::V2Only => verify_proof(CircuitVersion::V2),
117            ZkLoginCircuitMode::Both => {
118                verify_proof(CircuitVersion::V2).or_else(|_| verify_proof(CircuitVersion::V1))
119            }
120        }
121    }
122}
123
124impl Verifier<UserSignature> for ZkloginVerifier {
125    fn verify(&self, message: &[u8], signature: &UserSignature) -> Result<(), SignatureError> {
126        let UserSignature::ZkLogin(zklogin_authenticator) = signature else {
127            return Err(SignatureError::from_source("not a zklogin signature"));
128        };
129
130        self.verify(message, zklogin_authenticator.as_ref())
131    }
132}