1use crate::{
5    crypto::{CompressedSignature, SignatureScheme},
6    digests::ZKLoginInputsDigest,
7    error::SuiErrorKind,
8    multisig::{MultiSig, MultiSigPublicKey},
9    signature::{AuthenticatorTrait, GenericSignature, VerifyParams},
10    signature_verification::VerifiedDigestCache,
11    sui_serde::SuiBitmap,
12};
13pub use enum_dispatch::enum_dispatch;
14use fastcrypto::{
15    encoding::Base64,
16    error::FastCryptoError,
17    traits::{EncodeDecodeBase64, ToFromBytes},
18};
19use once_cell::sync::OnceCell;
20use roaring::RoaringBitmap;
21use schemars::JsonSchema;
22use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeSeq};
23use serde_with::serde_as;
24use shared_crypto::intent::IntentMessage;
25use std::{
26    hash::{Hash, Hasher},
27    sync::Arc,
28};
29
30use crate::{
31    base_types::{EpochId, SuiAddress},
32    crypto::PublicKey,
33    error::SuiError,
34};
35
36pub type WeightUnit = u8;
37pub type ThresholdUnit = u16;
38pub const MAX_SIGNER_IN_MULTISIG: usize = 10;
39
40#[serde_as]
43#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)]
44pub struct MultiSigLegacy {
45    sigs: Vec<CompressedSignature>,
47    #[schemars(with = "Base64")]
49    #[serde_as(as = "SuiBitmap")]
50    bitmap: RoaringBitmap,
51    multisig_pk: MultiSigPublicKeyLegacy,
53    #[serde(skip)]
55    bytes: OnceCell<Vec<u8>>,
56}
57
58impl AsRef<[u8]> for MultiSigLegacy {
62    fn as_ref(&self) -> &[u8] {
63        self.bytes
64            .get_or_try_init::<_, eyre::Report>(|| {
65                let as_bytes = bcs::to_bytes(self).expect("BCS serialization should not fail");
66                let mut bytes = Vec::with_capacity(1 + as_bytes.len());
67                bytes.push(SignatureScheme::MultiSig.flag());
68                bytes.extend_from_slice(as_bytes.as_slice());
69                Ok(bytes)
70            })
71            .expect("OnceCell invariant violated")
72    }
73}
74
75impl PartialEq for MultiSigLegacy {
77    fn eq(&self, other: &Self) -> bool {
78        self.sigs == other.sigs
79            && self.bitmap == other.bitmap
80            && self.multisig_pk == other.multisig_pk
81    }
82}
83
84impl Eq for MultiSigLegacy {}
86
87impl Hash for MultiSigLegacy {
89    fn hash<H: Hasher>(&self, state: &mut H) {
90        self.as_ref().hash(state);
91    }
92}
93
94impl AuthenticatorTrait for MultiSigLegacy {
95    fn verify_user_authenticator_epoch(
96        &self,
97        epoch_id: EpochId,
98        max_epoch_upper_bound_delta: Option<u64>,
99    ) -> Result<(), SuiError> {
100        let multisig: MultiSig =
101            self.clone()
102                .try_into()
103                .map_err(|_| SuiErrorKind::InvalidSignature {
104                    error: "Invalid legacy multisig".to_string(),
105                })?;
106        multisig.verify_user_authenticator_epoch(epoch_id, max_epoch_upper_bound_delta)
107    }
108
109    fn verify_claims<T>(
110        &self,
111        value: &IntentMessage<T>,
112        author: SuiAddress,
113        aux_verify_data: &VerifyParams,
114        zklogin_inputs_cache: Arc<VerifiedDigestCache<ZKLoginInputsDigest>>,
115    ) -> Result<(), SuiError>
116    where
117        T: Serialize,
118    {
119        let multisig: MultiSig =
120            self.clone()
121                .try_into()
122                .map_err(|_| SuiErrorKind::InvalidSignature {
123                    error: "Invalid legacy multisig".to_string(),
124                })?;
125        multisig.verify_claims(value, author, aux_verify_data, zklogin_inputs_cache)
126    }
127}
128
129impl TryFrom<MultiSigLegacy> for MultiSig {
130    type Error = FastCryptoError;
131
132    fn try_from(multisig: MultiSigLegacy) -> Result<Self, Self::Error> {
133        MultiSig::insecure_new(
134            multisig.clone().sigs,
135            bitmap_to_u16(multisig.clone().bitmap)?,
136            multisig.multisig_pk.try_into()?,
137        )
138        .init_and_validate()
139    }
140}
141
142impl TryFrom<MultiSigPublicKeyLegacy> for MultiSigPublicKey {
143    type Error = FastCryptoError;
144    fn try_from(multisig: MultiSigPublicKeyLegacy) -> Result<Self, Self::Error> {
145        let multisig_pk_legacy =
146            MultiSigPublicKey::insecure_new(multisig.pk_map, multisig.threshold).validate()?;
147        Ok(multisig_pk_legacy)
148    }
149}
150
151pub fn bitmap_to_u16(roaring: RoaringBitmap) -> Result<u16, FastCryptoError> {
153    let indices: Vec<u32> = roaring.into_iter().collect();
154    let mut val = 0;
155    for i in indices {
156        if i >= 10 {
157            return Err(FastCryptoError::InvalidInput);
158        }
159        val |= 1 << i as u8;
160    }
161    Ok(val)
162}
163
164impl MultiSigLegacy {
165    pub fn combine(
167        full_sigs: Vec<GenericSignature>,
168        multisig_pk: MultiSigPublicKeyLegacy,
169    ) -> Result<Self, SuiError> {
170        multisig_pk
171            .validate()
172            .map_err(|_| SuiErrorKind::InvalidSignature {
173                error: "Invalid multisig public key".to_string(),
174            })?;
175
176        if full_sigs.len() > multisig_pk.pk_map.len() || full_sigs.is_empty() {
177            return Err(SuiErrorKind::InvalidSignature {
178                error: "Invalid number of signatures".to_string(),
179            }
180            .into());
181        }
182        let mut bitmap = RoaringBitmap::new();
183        let mut sigs = Vec::with_capacity(full_sigs.len());
184        for s in full_sigs {
185            let pk = s.to_public_key()?;
186            let inserted = bitmap.insert(multisig_pk.get_index(&pk).ok_or(
187                SuiErrorKind::IncorrectSigner {
188                    error: format!("pk does not exist: {:?}", pk),
189                },
190            )?);
191            if !inserted {
192                return Err(SuiErrorKind::InvalidSignature {
193                    error: "Duplicate signature".to_string(),
194                }
195                .into());
196            }
197            sigs.push(s.to_compressed()?);
198        }
199        Ok(MultiSigLegacy {
200            sigs,
201            bitmap,
202            multisig_pk,
203            bytes: OnceCell::new(),
204        })
205    }
206
207    pub fn validate(&self) -> Result<(), FastCryptoError> {
208        if bitmap_to_u16(self.bitmap.clone()).is_err()
209            || self.sigs.len() > self.multisig_pk.pk_map.len()
210            || self.sigs.is_empty()
211        {
212            return Err(FastCryptoError::InvalidInput);
213        }
214        self.multisig_pk.validate()?;
215        Ok(())
216    }
217
218    pub fn get_pk(&self) -> &MultiSigPublicKeyLegacy {
219        &self.multisig_pk
220    }
221
222    pub fn get_sigs(&self) -> &[CompressedSignature] {
223        &self.sigs
224    }
225
226    pub fn get_bitmap(&self) -> &RoaringBitmap {
227        &self.bitmap
228    }
229}
230
231impl ToFromBytes for MultiSigLegacy {
232    fn from_bytes(bytes: &[u8]) -> Result<MultiSigLegacy, FastCryptoError> {
233        if bytes.first().ok_or(FastCryptoError::InvalidInput)? != &SignatureScheme::MultiSig.flag()
235        {
236            return Err(FastCryptoError::InvalidInput);
237        }
238        let multisig: MultiSigLegacy =
239            bcs::from_bytes(&bytes[1..]).map_err(|_| FastCryptoError::InvalidInput)?;
240        multisig.validate()?;
241        Ok(multisig)
242    }
243}
244
245#[derive(Debug, Clone, PartialEq, Eq, JsonSchema, Serialize, Deserialize)]
248pub struct MultiSigPublicKeyLegacy {
249    #[serde(serialize_with = "serialize_pk_map")]
251    #[serde(deserialize_with = "deserialize_pk_map")]
252    pk_map: Vec<(PublicKey, WeightUnit)>,
253    threshold: ThresholdUnit,
255}
256
257fn serialize_pk_map<S>(pk_map: &[(PublicKey, WeightUnit)], serializer: S) -> Result<S::Ok, S::Error>
259where
260    S: Serializer,
261{
262    let pk_weight_arr: Vec<(String, WeightUnit)> = pk_map
263        .iter()
264        .map(|(pk, w)| (pk.encode_base64(), *w))
265        .collect();
266
267    let mut seq = serializer.serialize_seq(Some(pk_weight_arr.len()))?;
268    for (pk_string, w) in pk_weight_arr {
269        seq.serialize_element(&(pk_string, w))?;
270    }
271    seq.end()
272}
273
274fn deserialize_pk_map<'de, D>(deserializer: D) -> Result<Vec<(PublicKey, WeightUnit)>, D::Error>
276where
277    D: Deserializer<'de>,
278{
279    use serde::de::Error;
280    let pk_weight_arr: Vec<(String, WeightUnit)> = Vec::deserialize(deserializer)?;
281    pk_weight_arr
282        .into_iter()
283        .map(|(s, w)| {
284            let pk = <PublicKey as EncodeDecodeBase64>::decode_base64(&s)
285                .map_err(|e| Error::custom(e.to_string()))?;
286            Ok((pk, w))
287        })
288        .collect()
289}
290
291impl MultiSigPublicKeyLegacy {
292    pub fn new(
293        pks: Vec<PublicKey>,
294        weights: Vec<WeightUnit>,
295        threshold: ThresholdUnit,
296    ) -> Result<Self, SuiError> {
297        if pks.is_empty()
298            || weights.is_empty()
299            || threshold == 0
300            || pks.len() != weights.len()
301            || pks.len() > MAX_SIGNER_IN_MULTISIG
302            || weights.contains(&0)
303            || weights
304                .iter()
305                .map(|w| *w as ThresholdUnit)
306                .sum::<ThresholdUnit>()
307                < threshold
308        {
309            return Err(SuiErrorKind::InvalidSignature {
310                error: "Invalid multisig public key construction".to_string(),
311            }
312            .into());
313        }
314        Ok(MultiSigPublicKeyLegacy {
315            pk_map: pks.into_iter().zip(weights).collect(),
316            threshold,
317        })
318    }
319
320    pub fn get_index(&self, pk: &PublicKey) -> Option<u32> {
321        self.pk_map
322            .iter()
323            .position(|x| &x.0 == pk)
324            .map(|x| x as u32)
325    }
326
327    pub fn threshold(&self) -> &ThresholdUnit {
328        &self.threshold
329    }
330
331    pub fn pubkeys(&self) -> &[(PublicKey, WeightUnit)] {
332        &self.pk_map
333    }
334
335    pub fn validate(&self) -> Result<Self, FastCryptoError> {
336        let multisig: MultiSigPublicKey = self.clone().try_into()?;
337        multisig.validate()?;
338        Ok(self.clone())
339    }
340}