1use crate::{
5    crypto::{CompressedSignature, DefaultHash, SignatureScheme},
6    digests::ZKLoginInputsDigest,
7    error::SuiErrorKind,
8    passkey_authenticator::PasskeyAuthenticator,
9    signature::{AuthenticatorTrait, GenericSignature, VerifyParams},
10    signature_verification::VerifiedDigestCache,
11    zk_login_authenticator::ZkLoginAuthenticator,
12};
13pub use enum_dispatch::enum_dispatch;
14use fastcrypto::{
15    ed25519::Ed25519PublicKey,
16    encoding::{Base64, Encoding},
17    error::FastCryptoError,
18    hash::HashFunction,
19    secp256k1::Secp256k1PublicKey,
20    secp256r1::Secp256r1PublicKey,
21    traits::{EncodeDecodeBase64, ToFromBytes, VerifyingKey},
22};
23use once_cell::sync::OnceCell;
24use schemars::JsonSchema;
25use serde::{Deserialize, Serialize};
26use serde_with::serde_as;
27use shared_crypto::intent::IntentMessage;
28use std::{
29    hash::{Hash, Hasher},
30    str::FromStr,
31    sync::Arc,
32};
33
34use crate::{
35    base_types::{EpochId, SuiAddress},
36    crypto::PublicKey,
37    error::SuiError,
38};
39
40#[cfg(test)]
41#[path = "unit_tests/multisig_tests.rs"]
42mod multisig_tests;
43
44pub type WeightUnit = u8;
45pub type ThresholdUnit = u16;
46pub type BitmapUnit = u16;
47pub const MAX_SIGNER_IN_MULTISIG: usize = 10;
48pub const MAX_BITMAP_VALUE: BitmapUnit = 0b1111111111;
49#[serde_as]
51#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)]
52pub struct MultiSig {
53    sigs: Vec<CompressedSignature>,
55    bitmap: BitmapUnit,
57    multisig_pk: MultiSigPublicKey,
59    #[serde(skip)]
61    bytes: OnceCell<Vec<u8>>,
62}
63
64impl PartialEq for MultiSig {
66    fn eq(&self, other: &Self) -> bool {
67        self.sigs == other.sigs
68            && self.bitmap == other.bitmap
69            && self.multisig_pk == other.multisig_pk
70    }
71}
72
73impl Eq for MultiSig {}
75
76impl Hash for MultiSig {
78    fn hash<H: Hasher>(&self, state: &mut H) {
79        self.as_ref().hash(state);
80    }
81}
82
83impl AuthenticatorTrait for MultiSig {
84    fn verify_user_authenticator_epoch(
85        &self,
86        epoch_id: EpochId,
87        max_epoch_upper_bound_delta: Option<u64>,
88    ) -> Result<(), SuiError> {
89        self.get_zklogin_sigs()?.iter().try_for_each(|s| {
92            s.verify_user_authenticator_epoch(epoch_id, max_epoch_upper_bound_delta)
93        })
94    }
95
96    fn verify_claims<T>(
97        &self,
98        value: &IntentMessage<T>,
99        multisig_address: SuiAddress,
100        verify_params: &VerifyParams,
101        zklogin_inputs_cache: Arc<VerifiedDigestCache<ZKLoginInputsDigest>>,
102    ) -> Result<(), SuiError>
103    where
104        T: Serialize,
105    {
106        self.multisig_pk
107            .validate()
108            .map_err(|_| SuiErrorKind::InvalidSignature {
109                error: "Invalid multisig pubkey".to_string(),
110            })?;
111
112        if SuiAddress::from(&self.multisig_pk) != multisig_address {
113            return Err(SuiErrorKind::InvalidSignature {
114                error: "Invalid address derived from pks".to_string(),
115            }
116            .into());
117        }
118
119        if !self.get_zklogin_sigs()?.is_empty() && !verify_params.accept_zklogin_in_multisig {
120            return Err(SuiErrorKind::InvalidSignature {
121                error: "zkLogin sig not supported inside multisig".to_string(),
122            }
123            .into());
124        }
125
126        if self.has_passkey_sigs() && !verify_params.accept_passkey_in_multisig {
127            return Err(SuiErrorKind::InvalidSignature {
128                error: "Passkey sig not supported inside multisig".to_string(),
129            }
130            .into());
131        }
132
133        let mut weight_sum: u16 = 0;
134        let message = bcs::to_bytes(&value).expect("Message serialization should not fail");
135        let mut hasher = DefaultHash::default();
136        hasher.update(message);
137        let digest = hasher.finalize().digest;
138        for (sig, i) in self.sigs.iter().zip(as_indices(self.bitmap)?) {
141            let (subsig_pubkey, weight) =
142                self.multisig_pk
143                    .pk_map
144                    .get(i as usize)
145                    .ok_or(SuiErrorKind::InvalidSignature {
146                        error: "Invalid public keys index".to_string(),
147                    })?;
148            let res = match sig {
149                CompressedSignature::Ed25519(s) => {
150                    if verify_params.additional_multisig_checks
151                        && !matches!(subsig_pubkey.scheme(), SignatureScheme::ED25519)
152                    {
153                        return Err(SuiErrorKind::InvalidSignature {
154                            error: format!(
155                                "Invalid sig for pk={} address={:?} error=signature/pubkey type mismatch",
156                                subsig_pubkey.encode_base64(),
157                                SuiAddress::from(subsig_pubkey)
158                            ),
159                        }.into());
160                    }
161                    let pk =
162                        Ed25519PublicKey::from_bytes(subsig_pubkey.as_ref()).map_err(|_| {
163                            SuiErrorKind::InvalidSignature {
164                                error: "Invalid ed25519 pk bytes".to_string(),
165                            }
166                        })?;
167                    pk.verify(
168                        &digest,
169                        &s.try_into().map_err(|_| SuiErrorKind::InvalidSignature {
170                            error: "Invalid ed25519 signature bytes".to_string(),
171                        })?,
172                    )
173                }
174                CompressedSignature::Secp256k1(s) => {
175                    if verify_params.additional_multisig_checks
176                        && !matches!(subsig_pubkey.scheme(), SignatureScheme::Secp256k1)
177                    {
178                        return Err(SuiErrorKind::InvalidSignature {
179                            error: format!(
180                                "Invalid sig for pk={} address={:?} error=signature/pubkey type mismatch",
181                                subsig_pubkey.encode_base64(),
182                                SuiAddress::from(subsig_pubkey)
183                            ),
184                        }.into());
185                    }
186                    let pk =
187                        Secp256k1PublicKey::from_bytes(subsig_pubkey.as_ref()).map_err(|_| {
188                            SuiErrorKind::InvalidSignature {
189                                error: "Invalid k1 pk bytes".to_string(),
190                            }
191                        })?;
192                    pk.verify(
193                        &digest,
194                        &s.try_into().map_err(|_| SuiErrorKind::InvalidSignature {
195                            error: "Invalid k1 signature bytes".to_string(),
196                        })?,
197                    )
198                }
199                CompressedSignature::Secp256r1(s) => {
200                    if verify_params.additional_multisig_checks
201                        && !matches!(subsig_pubkey.scheme(), SignatureScheme::Secp256r1)
202                    {
203                        return Err(SuiErrorKind::InvalidSignature {
204                            error: format!(
205                                "Invalid sig for pk={} address={:?} error=signature/pubkey type mismatch",
206                                subsig_pubkey.encode_base64(),
207                                SuiAddress::from(subsig_pubkey)
208                            ),
209                        }.into());
210                    }
211                    let pk =
212                        Secp256r1PublicKey::from_bytes(subsig_pubkey.as_ref()).map_err(|_| {
213                            SuiErrorKind::InvalidSignature {
214                                error: "Invalid r1 pk bytes".to_string(),
215                            }
216                        })?;
217                    pk.verify(
218                        &digest,
219                        &s.try_into().map_err(|_| SuiErrorKind::InvalidSignature {
220                            error: "Invalid r1 signature bytes".to_string(),
221                        })?,
222                    )
223                }
224                CompressedSignature::ZkLogin(z) => {
225                    if verify_params.additional_multisig_checks
226                        && !matches!(
227                            subsig_pubkey.scheme(),
228                            SignatureScheme::ZkLoginAuthenticator
229                        )
230                    {
231                        return Err(SuiErrorKind::InvalidSignature {
232                            error: format!(
233                                "Invalid sig for pk={} address={:?} error=signature/pubkey type mismatch",
234                                subsig_pubkey.encode_base64(),
235                                SuiAddress::from(subsig_pubkey)
236                            ),
237                        }.into());
238                    }
239                    let authenticator = ZkLoginAuthenticator::from_bytes(&z.0).map_err(|_| {
240                        SuiErrorKind::InvalidSignature {
241                            error: "Invalid zklogin authenticator bytes".to_string(),
242                        }
243                    })?;
244                    authenticator
245                        .verify_claims(
246                            value,
247                            SuiAddress::from(subsig_pubkey),
248                            verify_params,
249                            zklogin_inputs_cache.clone(),
250                        )
251                        .map_err(|e| FastCryptoError::GeneralError(e.to_string()))
252                }
253                CompressedSignature::Passkey(bytes) => {
254                    let authenticator =
255                        PasskeyAuthenticator::from_bytes(&bytes.0).map_err(|_| {
256                            SuiErrorKind::InvalidSignature {
257                                error: "Invalid passkey authenticator bytes".to_string(),
258                            }
259                        })?;
260                    authenticator
261                        .verify_claims(
262                            value,
263                            SuiAddress::from(subsig_pubkey),
264                            verify_params,
265                            zklogin_inputs_cache.clone(),
266                        )
267                        .map_err(|e| FastCryptoError::GeneralError(e.to_string()))
268                }
269            };
270            if res.is_ok() {
271                weight_sum += *weight as u16;
272            } else {
273                return res.map_err(|e| {
274                    SuiErrorKind::InvalidSignature {
275                        error: format!(
276                            "Invalid sig for pk={} address={:?} error={:?}",
277                            subsig_pubkey.encode_base64(),
278                            SuiAddress::from(subsig_pubkey),
279                            e.to_string()
280                        ),
281                    }
282                    .into()
283                });
284            }
285        }
286        if weight_sum >= self.multisig_pk.threshold {
287            Ok(())
288        } else {
289            Err(SuiErrorKind::InvalidSignature {
290                error: format!(
291                    "Insufficient weight={:?} threshold={:?}",
292                    weight_sum, self.multisig_pk.threshold
293                ),
294            }
295            .into())
296        }
297    }
298}
299
300pub fn as_indices(bitmap: u16) -> Result<Vec<u8>, SuiError> {
303    if bitmap > MAX_BITMAP_VALUE {
304        return Err(SuiErrorKind::InvalidSignature {
305            error: "Invalid bitmap".to_string(),
306        }
307        .into());
308    }
309    let mut res = Vec::new();
310    for i in 0..10 {
311        if bitmap & (1 << i) != 0 {
312            res.push(i as u8);
313        }
314    }
315    Ok(res)
316}
317
318impl MultiSig {
319    pub fn insecure_new(
321        sigs: Vec<CompressedSignature>,
322        bitmap: BitmapUnit,
323        multisig_pk: MultiSigPublicKey,
324    ) -> Self {
325        Self {
326            sigs,
327            bitmap,
328            multisig_pk,
329            bytes: OnceCell::new(),
330        }
331    }
332    pub fn combine(
337        full_sigs: Vec<GenericSignature>,
338        multisig_pk: MultiSigPublicKey,
339    ) -> Result<Self, SuiError> {
340        multisig_pk
341            .validate()
342            .map_err(|_| SuiErrorKind::InvalidSignature {
343                error: "Invalid multisig public key".to_string(),
344            })?;
345
346        if full_sigs.len() > multisig_pk.pk_map.len() || full_sigs.is_empty() {
347            return Err(SuiErrorKind::InvalidSignature {
348                error: "Invalid number of signatures".to_string(),
349            }
350            .into());
351        }
352        let mut bitmap = 0;
353        let mut sigs = Vec::with_capacity(full_sigs.len());
354        for s in full_sigs {
355            let pk = s.to_public_key()?;
356            let index = multisig_pk
357                .get_index(&pk)
358                .ok_or(SuiErrorKind::IncorrectSigner {
359                    error: format!("pk does not exist: {:?}", pk),
360                })?;
361            if bitmap & (1 << index) != 0 {
362                return Err(SuiErrorKind::InvalidSignature {
363                    error: "Duplicate public key".to_string(),
364                }
365                .into());
366            }
367            bitmap |= 1 << index;
368            sigs.push(s.to_compressed()?);
369        }
370
371        Ok(MultiSig {
372            sigs,
373            bitmap,
374            multisig_pk,
375            bytes: OnceCell::new(),
376        })
377    }
378
379    pub fn init_and_validate(&mut self) -> Result<Self, FastCryptoError> {
380        if self.sigs.len() > self.multisig_pk.pk_map.len()
381            || self.sigs.is_empty()
382            || self.bitmap > MAX_BITMAP_VALUE
383        {
384            return Err(FastCryptoError::InvalidInput);
385        }
386        self.multisig_pk.validate()?;
387        Ok(self.to_owned())
388    }
389
390    pub fn get_pk(&self) -> &MultiSigPublicKey {
391        &self.multisig_pk
392    }
393
394    pub fn get_sigs(&self) -> &[CompressedSignature] {
395        &self.sigs
396    }
397
398    pub fn get_bitmap(&self) -> u16 {
399        self.bitmap
400    }
401
402    pub fn get_zklogin_sigs(&self) -> Result<Vec<ZkLoginAuthenticator>, SuiError> {
403        let authenticator_as_bytes: Vec<_> = self
404            .sigs
405            .iter()
406            .filter_map(|s| match s {
407                CompressedSignature::ZkLogin(z) => Some(z),
408                _ => None,
409            })
410            .collect();
411        authenticator_as_bytes
412            .iter()
413            .map(|z| {
414                ZkLoginAuthenticator::from_bytes(&z.0).map_err(|_| {
415                    SuiErrorKind::InvalidSignature {
416                        error: "Invalid zklogin authenticator bytes".to_string(),
417                    }
418                    .into()
419                })
420            })
421            .collect()
422    }
423
424    pub fn get_indices(&self) -> Result<Vec<u8>, SuiError> {
425        as_indices(self.bitmap)
426    }
427
428    pub fn has_passkey_sigs(&self) -> bool {
429        self.sigs
430            .iter()
431            .any(|s| matches!(s, CompressedSignature::Passkey(_)))
432    }
433}
434
435impl ToFromBytes for MultiSig {
436    fn from_bytes(bytes: &[u8]) -> Result<MultiSig, FastCryptoError> {
437        if bytes.first().ok_or(FastCryptoError::InvalidInput)? != &SignatureScheme::MultiSig.flag()
439        {
440            return Err(FastCryptoError::InvalidInput);
441        }
442        let mut multisig: MultiSig =
443            bcs::from_bytes(&bytes[1..]).map_err(|_| FastCryptoError::InvalidSignature)?;
444        multisig.init_and_validate()
445    }
446}
447
448impl FromStr for MultiSig {
449    type Err = SuiError;
450
451    fn from_str(s: &str) -> Result<Self, Self::Err> {
452        let bytes = Base64::decode(s).map_err(|_| SuiErrorKind::InvalidSignature {
453            error: "Invalid base64 string".to_string(),
454        })?;
455        let sig = MultiSig::from_bytes(&bytes).map_err(|_| SuiErrorKind::InvalidSignature {
456            error: "Invalid multisig bytes".to_string(),
457        })?;
458        Ok(sig)
459    }
460}
461
462impl AsRef<[u8]> for MultiSig {
466    fn as_ref(&self) -> &[u8] {
467        self.bytes
468            .get_or_try_init::<_, eyre::Report>(|| {
469                let as_bytes = bcs::to_bytes(self).expect("BCS serialization should not fail");
470                let mut bytes = Vec::with_capacity(1 + as_bytes.len());
471                bytes.push(SignatureScheme::MultiSig.flag());
472                bytes.extend_from_slice(as_bytes.as_slice());
473                Ok(bytes)
474            })
475            .expect("OnceCell invariant violated")
476    }
477}
478
479#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
481pub struct MultiSigPublicKey {
482    pk_map: Vec<(PublicKey, WeightUnit)>,
484    threshold: ThresholdUnit,
486}
487
488impl MultiSigPublicKey {
489    pub fn insecure_new(pk_map: Vec<(PublicKey, WeightUnit)>, threshold: ThresholdUnit) -> Self {
491        Self { pk_map, threshold }
492    }
493
494    pub fn new(
495        pks: Vec<PublicKey>,
496        weights: Vec<WeightUnit>,
497        threshold: ThresholdUnit,
498    ) -> Result<Self, SuiError> {
499        if pks.is_empty()
500            || weights.is_empty()
501            || threshold == 0
502            || pks.len() != weights.len()
503            || pks.len() > MAX_SIGNER_IN_MULTISIG
504            || weights.contains(&0)
505            || weights
506                .iter()
507                .map(|w| *w as ThresholdUnit)
508                .sum::<ThresholdUnit>()
509                < threshold
510            || pks
511                .iter()
512                .enumerate()
513                .any(|(i, pk)| pks.iter().skip(i + 1).any(|other_pk| *pk == *other_pk))
514        {
515            return Err(SuiErrorKind::InvalidSignature {
516                error: "Invalid multisig public key construction".to_string(),
517            }
518            .into());
519        }
520
521        Ok(MultiSigPublicKey {
522            pk_map: pks.into_iter().zip(weights).collect(),
523            threshold,
524        })
525    }
526
527    pub fn get_index(&self, pk: &PublicKey) -> Option<u8> {
528        self.pk_map.iter().position(|x| &x.0 == pk).map(|x| x as u8)
529    }
530
531    pub fn threshold(&self) -> &ThresholdUnit {
532        &self.threshold
533    }
534
535    pub fn pubkeys(&self) -> &Vec<(PublicKey, WeightUnit)> {
536        &self.pk_map
537    }
538
539    pub fn validate(&self) -> Result<MultiSigPublicKey, FastCryptoError> {
540        let pk_map = self.pubkeys();
541        if self.threshold == 0
542            || pk_map.is_empty()
543            || pk_map.len() > MAX_SIGNER_IN_MULTISIG
544            || pk_map.iter().any(|(_pk, weight)| *weight == 0)
545            || pk_map
546                .iter()
547                .map(|(_pk, weight)| *weight as ThresholdUnit)
548                .sum::<ThresholdUnit>()
549                < self.threshold
550            || pk_map.iter().enumerate().any(|(i, (pk, _weight))| {
551                pk_map
552                    .iter()
553                    .skip(i + 1)
554                    .any(|(other_pk, _weight)| *pk == *other_pk)
555            })
556        {
557            return Err(FastCryptoError::InvalidInput);
558        }
559        Ok(self.to_owned())
560    }
561}