Skip to main content

aptos_sdk/crypto/
multi_key.rs

1//! `MultiKey` signature scheme implementation.
2//!
3//! `MultiKey` enables M-of-N threshold signatures with mixed key types.
4//! Unlike `MultiEd25519`, each key can be a different signature scheme
5//! (Ed25519, Secp256k1, Secp256r1, etc.).
6
7use crate::error::{AptosError, AptosResult};
8use serde::{Deserialize, Serialize};
9use std::fmt;
10
11/// Maximum number of keys in a multi-key account.
12pub const MAX_NUM_OF_KEYS: usize = 32;
13
14// Compile-time assertion: MAX_NUM_OF_KEYS must fit in u8 for bitmap operations
15const _: () = assert!(MAX_NUM_OF_KEYS <= u8::MAX as usize);
16
17/// Minimum threshold (at least 1 signature required).
18pub const MIN_THRESHOLD: u8 = 1;
19
20/// Supported signature schemes for multi-key.
21#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
22#[repr(u8)]
23pub enum AnyPublicKeyVariant {
24    /// Ed25519 public key.
25    Ed25519 = 0,
26    /// Secp256k1 ECDSA public key.
27    Secp256k1 = 1,
28    /// Secp256r1 (P-256) ECDSA public key.
29    Secp256r1 = 2,
30    /// Keyless public key.
31    Keyless = 3,
32}
33
34impl AnyPublicKeyVariant {
35    /// Get the variant from a byte.
36    ///
37    /// # Errors
38    ///
39    /// Returns [`AptosError::InvalidPublicKey`] if the byte value is not a valid variant (0-3).
40    pub fn from_byte(byte: u8) -> AptosResult<Self> {
41        match byte {
42            0 => Ok(Self::Ed25519),
43            1 => Ok(Self::Secp256k1),
44            2 => Ok(Self::Secp256r1),
45            3 => Ok(Self::Keyless),
46            _ => Err(AptosError::InvalidPublicKey(format!(
47                "unknown public key variant: {byte}"
48            ))),
49        }
50    }
51
52    /// Get the byte representation.
53    pub fn as_byte(&self) -> u8 {
54        *self as u8
55    }
56}
57
58/// A public key that can be any supported signature scheme.
59#[derive(Clone, PartialEq, Eq)]
60pub struct AnyPublicKey {
61    /// The signature scheme variant.
62    pub variant: AnyPublicKeyVariant,
63    /// The raw public key bytes.
64    pub bytes: Vec<u8>,
65}
66
67impl AnyPublicKey {
68    /// Creates a new `AnyPublicKey`.
69    pub fn new(variant: AnyPublicKeyVariant, bytes: Vec<u8>) -> Self {
70        Self { variant, bytes }
71    }
72
73    /// Creates an Ed25519 public key.
74    #[cfg(feature = "ed25519")]
75    pub fn ed25519(public_key: &crate::crypto::Ed25519PublicKey) -> Self {
76        Self {
77            variant: AnyPublicKeyVariant::Ed25519,
78            bytes: public_key.to_bytes().to_vec(),
79        }
80    }
81
82    /// Creates a Secp256k1 public key.
83    /// Uses uncompressed format (65 bytes) as required by the Aptos protocol.
84    #[cfg(feature = "secp256k1")]
85    pub fn secp256k1(public_key: &crate::crypto::Secp256k1PublicKey) -> Self {
86        Self {
87            variant: AnyPublicKeyVariant::Secp256k1,
88            bytes: public_key.to_uncompressed_bytes(),
89        }
90    }
91
92    /// Creates a Secp256r1 public key.
93    /// Uses uncompressed format (65 bytes) as required by the Aptos protocol.
94    #[cfg(feature = "secp256r1")]
95    pub fn secp256r1(public_key: &crate::crypto::Secp256r1PublicKey) -> Self {
96        Self {
97            variant: AnyPublicKeyVariant::Secp256r1,
98            bytes: public_key.to_uncompressed_bytes(),
99        }
100    }
101
102    /// Serializes to BCS format: `variant_byte` || ULEB128(length) || bytes
103    ///
104    /// This is the correct BCS serialization format for `AnyPublicKey` used
105    /// in authentication key derivation: `SHA3-256(BCS(AnyPublicKey) || scheme_id)`
106    pub fn to_bcs_bytes(&self) -> Vec<u8> {
107        let mut result = Vec::with_capacity(1 + 1 + self.bytes.len());
108        result.push(self.variant.as_byte());
109        // BCS uses ULEB128 for vector lengths
110        result.extend(uleb128_encode(self.bytes.len()));
111        result.extend_from_slice(&self.bytes);
112        result
113    }
114
115    /// Verifies a signature against a message.
116    ///
117    /// # Errors
118    ///
119    /// This function will return an error if:
120    /// - The signature variant doesn't match the public key variant
121    /// - The public key bytes are invalid for the variant
122    /// - The signature bytes are invalid for the variant
123    /// - Signature verification fails
124    /// - Verification is not supported for the variant
125    #[allow(unused_variables)]
126    pub fn verify(&self, message: &[u8], signature: &AnySignature) -> AptosResult<()> {
127        if signature.variant != self.variant {
128            return Err(AptosError::InvalidSignature(format!(
129                "signature variant {:?} doesn't match public key variant {:?}",
130                signature.variant, self.variant
131            )));
132        }
133
134        match self.variant {
135            #[cfg(feature = "ed25519")]
136            AnyPublicKeyVariant::Ed25519 => {
137                let pk = crate::crypto::Ed25519PublicKey::from_bytes(&self.bytes)?;
138                let sig = crate::crypto::Ed25519Signature::from_bytes(&signature.bytes)?;
139                pk.verify(message, &sig)
140            }
141            #[cfg(feature = "secp256k1")]
142            AnyPublicKeyVariant::Secp256k1 => {
143                // Public key can be either compressed (33 bytes) or uncompressed (65 bytes)
144                let pk = crate::crypto::Secp256k1PublicKey::from_bytes(&self.bytes)?;
145                let sig = crate::crypto::Secp256k1Signature::from_bytes(&signature.bytes)?;
146                pk.verify(message, &sig)
147            }
148            #[cfg(feature = "secp256r1")]
149            AnyPublicKeyVariant::Secp256r1 => {
150                // Public key can be either compressed (33 bytes) or uncompressed (65 bytes)
151                let pk = crate::crypto::Secp256r1PublicKey::from_bytes(&self.bytes)?;
152                let sig = crate::crypto::Secp256r1Signature::from_bytes(&signature.bytes)?;
153                pk.verify(message, &sig)
154            }
155            #[allow(unreachable_patterns)]
156            _ => Err(AptosError::InvalidPublicKey(format!(
157                "verification not supported for variant {:?}",
158                self.variant
159            ))),
160        }
161    }
162}
163
164impl fmt::Debug for AnyPublicKey {
165    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
166        write!(
167            f,
168            "AnyPublicKey({:?}, 0x{})",
169            self.variant,
170            hex::encode(&self.bytes)
171        )
172    }
173}
174
175impl fmt::Display for AnyPublicKey {
176    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
177        write!(f, "{:?}:0x{}", self.variant, hex::encode(&self.bytes))
178    }
179}
180
181/// A signature that can be any supported signature scheme.
182#[derive(Clone, PartialEq, Eq)]
183pub struct AnySignature {
184    /// The signature scheme variant.
185    pub variant: AnyPublicKeyVariant,
186    /// The raw signature bytes.
187    pub bytes: Vec<u8>,
188}
189
190impl AnySignature {
191    /// Creates a new `AnySignature`.
192    pub fn new(variant: AnyPublicKeyVariant, bytes: Vec<u8>) -> Self {
193        Self { variant, bytes }
194    }
195
196    /// Creates an Ed25519 signature.
197    #[cfg(feature = "ed25519")]
198    pub fn ed25519(signature: &crate::crypto::Ed25519Signature) -> Self {
199        Self {
200            variant: AnyPublicKeyVariant::Ed25519,
201            bytes: signature.to_bytes().to_vec(),
202        }
203    }
204
205    /// Creates a Secp256k1 signature.
206    #[cfg(feature = "secp256k1")]
207    pub fn secp256k1(signature: &crate::crypto::Secp256k1Signature) -> Self {
208        Self {
209            variant: AnyPublicKeyVariant::Secp256k1,
210            bytes: signature.to_bytes().to_vec(),
211        }
212    }
213
214    /// Creates a Secp256r1 signature.
215    #[cfg(feature = "secp256r1")]
216    pub fn secp256r1(signature: &crate::crypto::Secp256r1Signature) -> Self {
217        Self {
218            variant: AnyPublicKeyVariant::Secp256r1,
219            bytes: signature.to_bytes().to_vec(),
220        }
221    }
222
223    /// Serializes to BCS format: `variant_byte` || ULEB128(length) || bytes
224    pub fn to_bcs_bytes(&self) -> Vec<u8> {
225        let mut result = Vec::with_capacity(1 + 1 + self.bytes.len());
226        result.push(self.variant.as_byte());
227        // BCS uses ULEB128 for vector lengths
228        result.extend(uleb128_encode(self.bytes.len()));
229        result.extend_from_slice(&self.bytes);
230        result
231    }
232}
233
234/// Encodes a value as `ULEB128` (unsigned `LEB128`).
235/// BCS uses `ULEB128` for encoding vector/sequence lengths.
236/// For typical sizes (< 128), this returns a single byte.
237#[allow(clippy::cast_possible_truncation)] // value & 0x7F is always <= 127
238#[inline]
239fn uleb128_encode(mut value: usize) -> Vec<u8> {
240    // Pre-allocate for common case: values < 128 need 1 byte, < 16384 need 2 bytes
241    let mut result = Vec::with_capacity(if value < 128 { 1 } else { 2 });
242    loop {
243        let byte = (value & 0x7F) as u8;
244        value >>= 7;
245        if value == 0 {
246            result.push(byte);
247            break;
248        }
249        result.push(byte | 0x80);
250    }
251    result
252}
253
254/// Decodes a `ULEB128` value from bytes, returning `(value, bytes_consumed)`.
255fn uleb128_decode(bytes: &[u8]) -> Option<(usize, usize)> {
256    let mut result: usize = 0;
257    let mut shift = 0;
258    for (i, &byte) in bytes.iter().enumerate() {
259        result |= ((byte & 0x7F) as usize) << shift;
260        if byte & 0x80 == 0 {
261            return Some((result, i + 1));
262        }
263        shift += 7;
264        if shift >= 64 {
265            return None; // Overflow
266        }
267    }
268    None
269}
270
271impl fmt::Debug for AnySignature {
272    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
273        write!(
274            f,
275            "AnySignature({:?}, {} bytes)",
276            self.variant,
277            self.bytes.len()
278        )
279    }
280}
281
282/// A multi-key public key supporting mixed signature schemes.
283///
284/// This allows M-of-N threshold signing where each key can be a different type.
285#[derive(Clone, PartialEq, Eq)]
286pub struct MultiKeyPublicKey {
287    /// The individual public keys.
288    public_keys: Vec<AnyPublicKey>,
289    /// The required threshold (M in M-of-N).
290    threshold: u8,
291}
292
293impl MultiKeyPublicKey {
294    /// Creates a new multi-key public key.
295    ///
296    /// # Arguments
297    ///
298    /// * `public_keys` - The individual public keys (can be mixed types)
299    /// * `threshold` - The number of signatures required (M in M-of-N)
300    ///
301    /// # Errors
302    ///
303    /// Returns [`AptosError::InvalidPublicKey`] if:
304    /// - No public keys are provided
305    /// - More than 32 public keys are provided
306    /// - Threshold is 0
307    /// - Threshold exceeds the number of keys
308    pub fn new(public_keys: Vec<AnyPublicKey>, threshold: u8) -> AptosResult<Self> {
309        if public_keys.is_empty() {
310            return Err(AptosError::InvalidPublicKey(
311                "multi-key requires at least one public key".into(),
312            ));
313        }
314        if public_keys.len() > MAX_NUM_OF_KEYS {
315            return Err(AptosError::InvalidPublicKey(format!(
316                "multi-key supports at most {} keys, got {}",
317                MAX_NUM_OF_KEYS,
318                public_keys.len()
319            )));
320        }
321        if threshold < MIN_THRESHOLD {
322            return Err(AptosError::InvalidPublicKey(
323                "threshold must be at least 1".into(),
324            ));
325        }
326        if threshold as usize > public_keys.len() {
327            return Err(AptosError::InvalidPublicKey(format!(
328                "threshold {} exceeds number of keys {}",
329                threshold,
330                public_keys.len()
331            )));
332        }
333        Ok(Self {
334            public_keys,
335            threshold,
336        })
337    }
338
339    /// Returns the number of public keys.
340    pub fn num_keys(&self) -> usize {
341        self.public_keys.len()
342    }
343
344    /// Returns the threshold.
345    pub fn threshold(&self) -> u8 {
346        self.threshold
347    }
348
349    /// Returns the individual public keys.
350    pub fn public_keys(&self) -> &[AnyPublicKey] {
351        &self.public_keys
352    }
353
354    /// Returns the key at the given index.
355    pub fn get(&self, index: usize) -> Option<&AnyPublicKey> {
356        self.public_keys.get(index)
357    }
358
359    /// Serializes to bytes for authentication key derivation.
360    ///
361    /// Format: `num_keys` || `pk1_bcs` || `pk2_bcs` || ... || threshold
362    #[allow(clippy::cast_possible_truncation)] // public_keys.len() <= MAX_NUM_OF_KEYS (32)
363    pub fn to_bytes(&self) -> Vec<u8> {
364        // Pre-allocate: 1 byte num_keys + estimated key size (avg ~35 bytes per key) + 1 byte threshold
365        let estimated_size = 2 + self.public_keys.len() * 36;
366        let mut bytes = Vec::with_capacity(estimated_size);
367
368        // Number of keys (1 byte, validated in new())
369        bytes.push(self.public_keys.len() as u8);
370
371        // Each public key in BCS format
372        for pk in &self.public_keys {
373            bytes.extend(pk.to_bcs_bytes());
374        }
375
376        // Threshold (1 byte)
377        bytes.push(self.threshold);
378
379        bytes
380    }
381
382    /// Creates from bytes.
383    ///
384    /// # Errors
385    ///
386    /// Returns [`AptosError::InvalidPublicKey`] if:
387    /// - The bytes are empty
388    /// - The number of keys is invalid (0 or > 32)
389    /// - The bytes are too short for the expected structure
390    /// - Any public key variant byte is invalid
391    /// - Any public key length or data is invalid
392    /// - The threshold is invalid
393    pub fn from_bytes(bytes: &[u8]) -> AptosResult<Self> {
394        // SECURITY: Limit individual key size to prevent DoS via large allocations
395        // Largest supported key is uncompressed secp256k1/secp256r1 at 65 bytes
396        const MAX_KEY_SIZE: usize = 128;
397
398        if bytes.is_empty() {
399            return Err(AptosError::InvalidPublicKey("empty bytes".into()));
400        }
401
402        let num_keys = bytes[0] as usize;
403        if num_keys == 0 || num_keys > MAX_NUM_OF_KEYS {
404            return Err(AptosError::InvalidPublicKey(format!(
405                "invalid number of keys: {num_keys}"
406            )));
407        }
408
409        let mut offset = 1;
410        let mut public_keys = Vec::with_capacity(num_keys);
411
412        for _ in 0..num_keys {
413            if offset >= bytes.len() {
414                return Err(AptosError::InvalidPublicKey("bytes too short".into()));
415            }
416
417            let variant = AnyPublicKeyVariant::from_byte(bytes[offset])?;
418            offset += 1;
419
420            // Decode ULEB128 length
421            let (len, len_bytes) = uleb128_decode(&bytes[offset..]).ok_or_else(|| {
422                AptosError::InvalidPublicKey("invalid ULEB128 length encoding".into())
423            })?;
424            offset += len_bytes;
425
426            if len > MAX_KEY_SIZE {
427                return Err(AptosError::InvalidPublicKey(format!(
428                    "key size {len} exceeds maximum {MAX_KEY_SIZE}"
429                )));
430            }
431
432            if offset + len > bytes.len() {
433                return Err(AptosError::InvalidPublicKey(
434                    "bytes too short for key".into(),
435                ));
436            }
437
438            let key_bytes = bytes[offset..offset + len].to_vec();
439            offset += len;
440
441            public_keys.push(AnyPublicKey::new(variant, key_bytes));
442        }
443
444        if offset >= bytes.len() {
445            return Err(AptosError::InvalidPublicKey(
446                "bytes too short for threshold".into(),
447            ));
448        }
449
450        let threshold = bytes[offset];
451
452        Self::new(public_keys, threshold)
453    }
454
455    /// Derives the account address for this multi-key public key.
456    pub fn to_address(&self) -> crate::types::AccountAddress {
457        crate::crypto::derive_address(&self.to_bytes(), crate::crypto::MULTI_KEY_SCHEME)
458    }
459
460    /// Derives the authentication key for this public key.
461    pub fn to_authentication_key(&self) -> [u8; 32] {
462        crate::crypto::derive_authentication_key(&self.to_bytes(), crate::crypto::MULTI_KEY_SCHEME)
463    }
464
465    /// Verifies a multi-key signature against a message.
466    ///
467    /// # Errors
468    ///
469    /// This function will return an error if:
470    /// - The number of signatures is less than the threshold
471    /// - Any individual signature verification fails
472    /// - A signer index is out of bounds
473    pub fn verify(&self, message: &[u8], signature: &MultiKeySignature) -> AptosResult<()> {
474        // Check that we have enough signatures
475        if signature.num_signatures() < self.threshold as usize {
476            return Err(AptosError::SignatureVerificationFailed);
477        }
478
479        // Verify each signature
480        for (index, sig) in signature.signatures() {
481            if *index as usize >= self.public_keys.len() {
482                return Err(AptosError::InvalidSignature(format!(
483                    "signer index {} out of bounds (max {})",
484                    index,
485                    self.public_keys.len() - 1
486                )));
487            }
488            let pk = &self.public_keys[*index as usize];
489            pk.verify(message, sig)?;
490        }
491
492        Ok(())
493    }
494}
495
496impl fmt::Debug for MultiKeyPublicKey {
497    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
498        write!(
499            f,
500            "MultiKeyPublicKey({}-of-{} keys)",
501            self.threshold,
502            self.public_keys.len()
503        )
504    }
505}
506
507impl fmt::Display for MultiKeyPublicKey {
508    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
509        write!(f, "0x{}", hex::encode(self.to_bytes()))
510    }
511}
512
513/// A multi-key signature containing signatures from multiple signers.
514#[derive(Clone, PartialEq, Eq)]
515pub struct MultiKeySignature {
516    /// Individual signatures with their signer index.
517    signatures: Vec<(u8, AnySignature)>,
518    /// Bitmap indicating which keys signed (little-endian, up to 4 bytes for 32 keys).
519    bitmap: [u8; 4],
520}
521
522impl MultiKeySignature {
523    /// Creates a new multi-key signature from individual signatures.
524    ///
525    /// # Arguments
526    ///
527    /// * `signatures` - Vec of (`signer_index`, signature) pairs
528    ///
529    /// # Errors
530    ///
531    /// Returns [`AptosError::InvalidSignature`] if:
532    /// - No signatures are provided
533    /// - More than 32 signatures are provided
534    /// - A signer index is out of bounds (>= 32)
535    /// - Duplicate signer indices are present
536    pub fn new(mut signatures: Vec<(u8, AnySignature)>) -> AptosResult<Self> {
537        if signatures.is_empty() {
538            return Err(AptosError::InvalidSignature(
539                "multi-key signature requires at least one signature".into(),
540            ));
541        }
542        if signatures.len() > MAX_NUM_OF_KEYS {
543            return Err(AptosError::InvalidSignature(format!(
544                "too many signatures: {} (max {})",
545                signatures.len(),
546                MAX_NUM_OF_KEYS
547            )));
548        }
549
550        // Sort by index
551        signatures.sort_by_key(|(idx, _)| *idx);
552
553        // Check for duplicates and bounds, build bitmap
554        let mut bitmap = [0u8; 4];
555        let mut last_index: Option<u8> = None;
556
557        for (index, _) in &signatures {
558            if *index as usize >= MAX_NUM_OF_KEYS {
559                return Err(AptosError::InvalidSignature(format!(
560                    "signer index {} out of bounds (max {})",
561                    index,
562                    MAX_NUM_OF_KEYS - 1
563                )));
564            }
565            if last_index == Some(*index) {
566                return Err(AptosError::InvalidSignature(format!(
567                    "duplicate signer index {index}"
568                )));
569            }
570            last_index = Some(*index);
571
572            // Set bit in bitmap
573            let byte_index = (index / 8) as usize;
574            let bit_index = index % 8;
575            bitmap[byte_index] |= 1 << bit_index;
576        }
577
578        Ok(Self { signatures, bitmap })
579    }
580
581    /// Returns the number of signatures.
582    pub fn num_signatures(&self) -> usize {
583        self.signatures.len()
584    }
585
586    /// Returns the individual signatures with their indices.
587    pub fn signatures(&self) -> &[(u8, AnySignature)] {
588        &self.signatures
589    }
590
591    /// Returns the signer bitmap.
592    pub fn bitmap(&self) -> &[u8; 4] {
593        &self.bitmap
594    }
595
596    /// Checks if a particular index signed.
597    pub fn has_signature(&self, index: u8) -> bool {
598        if index as usize >= MAX_NUM_OF_KEYS {
599            return false;
600        }
601        let byte_index = (index / 8) as usize;
602        let bit_index = index % 8;
603        (self.bitmap[byte_index] >> bit_index) & 1 == 1
604    }
605
606    /// Serializes to bytes.
607    ///
608    /// Format: `num_signatures` || `sig1_bcs` || `sig2_bcs` || ... || bitmap (4 bytes)
609    #[allow(clippy::cast_possible_truncation)] // signatures.len() <= MAX_NUM_OF_KEYS (32)
610    pub fn to_bytes(&self) -> Vec<u8> {
611        // Pre-allocate: 1 byte num_sigs + estimated sig size (avg ~66 bytes per sig) + 4 bytes bitmap
612        let estimated_size = 5 + self.signatures.len() * 68;
613        let mut bytes = Vec::with_capacity(estimated_size);
614
615        // Number of signatures (validated in new())
616        bytes.push(self.signatures.len() as u8);
617
618        // Each signature in BCS format (ordered by index)
619        for (_, sig) in &self.signatures {
620            bytes.extend(sig.to_bcs_bytes());
621        }
622
623        // Bitmap (4 bytes)
624        bytes.extend_from_slice(&self.bitmap);
625
626        bytes
627    }
628
629    /// Creates from bytes.
630    ///
631    /// # Errors
632    ///
633    /// Returns [`AptosError::InvalidSignature`] if:
634    /// - The bytes are too short (less than 5 bytes for `num_sigs` + bitmap)
635    /// - The number of signatures is invalid (0 or > 32)
636    /// - The bitmap doesn't match the number of signatures
637    /// - The bytes are too short for the expected structure
638    /// - Any signature variant byte is invalid
639    /// - Any signature length or data is invalid
640    pub fn from_bytes(bytes: &[u8]) -> AptosResult<Self> {
641        // SECURITY: Limit individual signature size to prevent DoS via large allocations
642        // Largest supported signature is ~72 bytes for ECDSA DER format
643        const MAX_SIGNATURE_SIZE: usize = 128;
644
645        if bytes.len() < 5 {
646            return Err(AptosError::InvalidSignature("bytes too short".into()));
647        }
648
649        let num_sigs = bytes[0] as usize;
650        if num_sigs == 0 || num_sigs > MAX_NUM_OF_KEYS {
651            return Err(AptosError::InvalidSignature(format!(
652                "invalid number of signatures: {num_sigs}"
653            )));
654        }
655
656        // Read bitmap from the end
657        let bitmap_start = bytes.len() - 4;
658        let mut bitmap = [0u8; 4];
659        bitmap.copy_from_slice(&bytes[bitmap_start..]);
660
661        // Parse signatures
662        let mut offset = 1;
663        let mut signatures = Vec::with_capacity(num_sigs);
664
665        // Determine signer indices from bitmap (MAX_NUM_OF_KEYS is 32, fits in u8)
666        let mut signer_indices = Vec::new();
667        #[allow(clippy::cast_possible_truncation)]
668        for bit_pos in 0..(MAX_NUM_OF_KEYS as u8) {
669            let byte_idx = (bit_pos / 8) as usize;
670            let bit_idx = bit_pos % 8;
671            if (bitmap[byte_idx] >> bit_idx) & 1 == 1 {
672                signer_indices.push(bit_pos);
673            }
674        }
675
676        if signer_indices.len() != num_sigs {
677            return Err(AptosError::InvalidSignature(
678                "bitmap doesn't match number of signatures".into(),
679            ));
680        }
681
682        for &index in &signer_indices {
683            if offset >= bitmap_start {
684                return Err(AptosError::InvalidSignature("bytes too short".into()));
685            }
686
687            let variant = AnyPublicKeyVariant::from_byte(bytes[offset])?;
688            offset += 1;
689
690            // Decode ULEB128 length
691            let (len, len_bytes) =
692                uleb128_decode(&bytes[offset..bitmap_start]).ok_or_else(|| {
693                    AptosError::InvalidSignature("invalid ULEB128 length encoding".into())
694                })?;
695            offset += len_bytes;
696
697            if len > MAX_SIGNATURE_SIZE {
698                return Err(AptosError::InvalidSignature(format!(
699                    "signature size {len} exceeds maximum {MAX_SIGNATURE_SIZE}"
700                )));
701            }
702
703            if offset + len > bitmap_start {
704                return Err(AptosError::InvalidSignature(
705                    "bytes too short for signature".into(),
706                ));
707            }
708
709            let sig_bytes = bytes[offset..offset + len].to_vec();
710            offset += len;
711
712            signatures.push((index, AnySignature::new(variant, sig_bytes)));
713        }
714
715        Ok(Self { signatures, bitmap })
716    }
717}
718
719impl fmt::Debug for MultiKeySignature {
720    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
721        write!(
722            f,
723            "MultiKeySignature({} signatures, bitmap={:?})",
724            self.signatures.len(),
725            self.bitmap
726        )
727    }
728}
729
730impl fmt::Display for MultiKeySignature {
731    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
732        write!(f, "0x{}", hex::encode(self.to_bytes()))
733    }
734}
735
736#[cfg(test)]
737mod tests {
738    use super::*;
739
740    #[test]
741    fn test_any_public_key_variant_from_byte() {
742        assert_eq!(
743            AnyPublicKeyVariant::from_byte(0).unwrap(),
744            AnyPublicKeyVariant::Ed25519
745        );
746        assert_eq!(
747            AnyPublicKeyVariant::from_byte(1).unwrap(),
748            AnyPublicKeyVariant::Secp256k1
749        );
750        assert_eq!(
751            AnyPublicKeyVariant::from_byte(2).unwrap(),
752            AnyPublicKeyVariant::Secp256r1
753        );
754        assert_eq!(
755            AnyPublicKeyVariant::from_byte(3).unwrap(),
756            AnyPublicKeyVariant::Keyless
757        );
758        assert!(AnyPublicKeyVariant::from_byte(4).is_err());
759        assert!(AnyPublicKeyVariant::from_byte(255).is_err());
760    }
761
762    #[test]
763    fn test_any_public_key_variant_as_byte() {
764        assert_eq!(AnyPublicKeyVariant::Ed25519.as_byte(), 0);
765        assert_eq!(AnyPublicKeyVariant::Secp256k1.as_byte(), 1);
766        assert_eq!(AnyPublicKeyVariant::Secp256r1.as_byte(), 2);
767        assert_eq!(AnyPublicKeyVariant::Keyless.as_byte(), 3);
768    }
769
770    #[test]
771    fn test_any_public_key_new() {
772        let pk = AnyPublicKey::new(AnyPublicKeyVariant::Ed25519, vec![0x11; 32]);
773        assert_eq!(pk.variant, AnyPublicKeyVariant::Ed25519);
774        assert_eq!(pk.bytes.len(), 32);
775        assert_eq!(pk.bytes[0], 0x11);
776    }
777
778    #[test]
779    fn test_any_public_key_to_bcs_bytes() {
780        let pk = AnyPublicKey::new(AnyPublicKeyVariant::Ed25519, vec![0xaa; 32]);
781        let bcs = pk.to_bcs_bytes();
782
783        // Format: variant_byte || ULEB128(length) || bytes
784        assert_eq!(bcs[0], 0); // Ed25519 variant
785        assert_eq!(bcs[1], 32); // ULEB128(32) = 0x20 (since 32 < 128)
786        assert_eq!(bcs[2], 0xaa); // first byte of key
787        assert_eq!(bcs.len(), 1 + 1 + 32); // variant + length + bytes
788    }
789
790    #[test]
791    fn test_any_public_key_debug() {
792        let pk = AnyPublicKey::new(AnyPublicKeyVariant::Secp256k1, vec![0xbb; 33]);
793        let debug = format!("{pk:?}");
794        assert!(debug.contains("Secp256k1"));
795        assert!(debug.contains("0x"));
796    }
797
798    #[test]
799    fn test_any_signature_new() {
800        let sig = AnySignature::new(AnyPublicKeyVariant::Ed25519, vec![0xcc; 64]);
801        assert_eq!(sig.variant, AnyPublicKeyVariant::Ed25519);
802        assert_eq!(sig.bytes.len(), 64);
803    }
804
805    #[test]
806    fn test_any_signature_to_bcs_bytes() {
807        let sig = AnySignature::new(AnyPublicKeyVariant::Ed25519, vec![0xdd; 64]);
808        let bcs = sig.to_bcs_bytes();
809
810        // Format: variant_byte || ULEB128(length) || bytes
811        assert_eq!(bcs[0], 0); // Ed25519 variant
812        assert_eq!(bcs[1], 64); // ULEB128(64) = 0x40 (since 64 < 128)
813        assert_eq!(bcs[2], 0xdd); // first byte of signature
814        assert_eq!(bcs.len(), 1 + 1 + 64); // variant + length + bytes
815    }
816
817    #[test]
818    fn test_any_signature_debug() {
819        let sig = AnySignature::new(AnyPublicKeyVariant::Secp256r1, vec![0xee; 64]);
820        let debug = format!("{sig:?}");
821        assert!(debug.contains("Secp256r1"));
822        assert!(debug.contains("64 bytes"));
823    }
824
825    #[test]
826    fn test_any_public_key_verify_mismatched_variant() {
827        let pk = AnyPublicKey::new(AnyPublicKeyVariant::Ed25519, vec![0; 32]);
828        let sig = AnySignature::new(AnyPublicKeyVariant::Secp256k1, vec![0; 64]);
829
830        let result = pk.verify(b"message", &sig);
831        assert!(result.is_err());
832        assert!(result.unwrap_err().to_string().contains("variant"));
833    }
834
835    #[test]
836    fn test_multi_key_signature_insufficient_sigs() {
837        // Empty signatures should fail
838        let result = MultiKeySignature::new(vec![]);
839        assert!(result.is_err());
840    }
841
842    #[test]
843    fn test_multi_key_signature_duplicate_indices() {
844        let sig1 = AnySignature::new(AnyPublicKeyVariant::Ed25519, vec![0; 64]);
845        let sig2 = AnySignature::new(AnyPublicKeyVariant::Ed25519, vec![1; 64]);
846
847        // Duplicate index should fail
848        let result = MultiKeySignature::new(vec![(0, sig1.clone()), (0, sig2)]);
849        assert!(result.is_err());
850    }
851
852    #[test]
853    fn test_multi_key_signature_index_out_of_range() {
854        let sig = AnySignature::new(AnyPublicKeyVariant::Ed25519, vec![0; 64]);
855
856        // Index 32 is out of range (0-31)
857        let result = MultiKeySignature::new(vec![(32, sig)]);
858        assert!(result.is_err());
859    }
860
861    #[test]
862    fn test_multi_key_signature_basic() {
863        let sig1 = AnySignature::new(AnyPublicKeyVariant::Ed25519, vec![0xaa; 64]);
864        let sig2 = AnySignature::new(AnyPublicKeyVariant::Ed25519, vec![0xbb; 64]);
865
866        let multi_sig = MultiKeySignature::new(vec![(0, sig1), (5, sig2)]).unwrap();
867
868        assert_eq!(multi_sig.num_signatures(), 2);
869        assert!(multi_sig.has_signature(0));
870        assert!(!multi_sig.has_signature(1));
871        assert!(multi_sig.has_signature(5));
872    }
873
874    #[test]
875    fn test_multi_key_signature_debug_display() {
876        let sig = AnySignature::new(AnyPublicKeyVariant::Ed25519, vec![0; 64]);
877        let multi_sig = MultiKeySignature::new(vec![(0, sig)]).unwrap();
878
879        let debug = format!("{multi_sig:?}");
880        let display = format!("{multi_sig}");
881
882        assert!(debug.contains("MultiKeySignature"));
883        assert!(display.starts_with("0x"));
884    }
885
886    #[test]
887    #[cfg(feature = "ed25519")]
888    fn test_multi_key_public_key_creation() {
889        use crate::crypto::Ed25519PrivateKey;
890
891        let keys: Vec<_> = (0..3)
892            .map(|_| AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key()))
893            .collect();
894
895        // Valid 2-of-3
896        let multi_pk = MultiKeyPublicKey::new(keys.clone(), 2).unwrap();
897        assert_eq!(multi_pk.num_keys(), 3);
898        assert_eq!(multi_pk.threshold(), 2);
899
900        // Invalid: threshold > num_keys
901        assert!(MultiKeyPublicKey::new(keys.clone(), 4).is_err());
902
903        // Invalid: threshold = 0
904        assert!(MultiKeyPublicKey::new(keys.clone(), 0).is_err());
905
906        // Invalid: empty keys
907        assert!(MultiKeyPublicKey::new(vec![], 1).is_err());
908    }
909
910    #[test]
911    #[cfg(all(feature = "ed25519", feature = "secp256k1"))]
912    fn test_multi_key_mixed_types() {
913        use crate::crypto::{Ed25519PrivateKey, Secp256k1PrivateKey};
914
915        // Create mixed key types
916        let ed_key = AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key());
917        let secp_key = AnyPublicKey::secp256k1(&Secp256k1PrivateKey::generate().public_key());
918
919        let multi_pk = MultiKeyPublicKey::new(vec![ed_key, secp_key], 2).unwrap();
920        assert_eq!(multi_pk.num_keys(), 2);
921        assert_eq!(
922            multi_pk.get(0).unwrap().variant,
923            AnyPublicKeyVariant::Ed25519
924        );
925        assert_eq!(
926            multi_pk.get(1).unwrap().variant,
927            AnyPublicKeyVariant::Secp256k1
928        );
929    }
930
931    #[test]
932    #[cfg(feature = "ed25519")]
933    fn test_multi_key_sign_verify() {
934        use crate::crypto::Ed25519PrivateKey;
935
936        let private_keys: Vec<_> = (0..3).map(|_| Ed25519PrivateKey::generate()).collect();
937        let public_keys: Vec<_> = private_keys
938            .iter()
939            .map(|k| AnyPublicKey::ed25519(&k.public_key()))
940            .collect();
941
942        let multi_pk = MultiKeyPublicKey::new(public_keys, 2).unwrap();
943        let message = b"test message";
944
945        // Sign with keys 0 and 2 (2-of-3)
946        let sig0 = AnySignature::ed25519(&private_keys[0].sign(message));
947        let sig2 = AnySignature::ed25519(&private_keys[2].sign(message));
948
949        let multi_sig = MultiKeySignature::new(vec![(0, sig0), (2, sig2)]).unwrap();
950
951        // Verify should succeed
952        assert!(multi_pk.verify(message, &multi_sig).is_ok());
953
954        // Wrong message should fail
955        assert!(multi_pk.verify(b"wrong message", &multi_sig).is_err());
956    }
957
958    #[test]
959    #[cfg(feature = "ed25519")]
960    fn test_multi_key_bytes_roundtrip() {
961        use crate::crypto::Ed25519PrivateKey;
962
963        let keys: Vec<_> = (0..3)
964            .map(|_| AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key()))
965            .collect();
966        let multi_pk = MultiKeyPublicKey::new(keys, 2).unwrap();
967
968        let bytes = multi_pk.to_bytes();
969        let restored = MultiKeyPublicKey::from_bytes(&bytes).unwrap();
970
971        assert_eq!(multi_pk.threshold(), restored.threshold());
972        assert_eq!(multi_pk.num_keys(), restored.num_keys());
973    }
974
975    #[test]
976    #[cfg(feature = "ed25519")]
977    fn test_multi_key_signature_bytes_roundtrip() {
978        use crate::crypto::Ed25519PrivateKey;
979
980        let private_keys: Vec<_> = (0..3).map(|_| Ed25519PrivateKey::generate()).collect();
981        let message = b"test";
982
983        let sig0 = AnySignature::ed25519(&private_keys[0].sign(message));
984        let sig2 = AnySignature::ed25519(&private_keys[2].sign(message));
985
986        let multi_sig = MultiKeySignature::new(vec![(0, sig0), (2, sig2)]).unwrap();
987
988        let bytes = multi_sig.to_bytes();
989        let restored = MultiKeySignature::from_bytes(&bytes).unwrap();
990
991        assert_eq!(multi_sig.num_signatures(), restored.num_signatures());
992        assert_eq!(multi_sig.bitmap(), restored.bitmap());
993    }
994
995    #[test]
996    #[cfg(feature = "ed25519")]
997    fn test_signature_bitmap() {
998        use crate::crypto::Ed25519PrivateKey;
999
1000        let private_keys: Vec<_> = (0..5).map(|_| Ed25519PrivateKey::generate()).collect();
1001        let message = b"test";
1002
1003        // Sign with indices 1, 3, 4
1004        let signatures: Vec<_> = [1, 3, 4]
1005            .iter()
1006            .map(|&i| {
1007                (
1008                    i,
1009                    AnySignature::ed25519(&private_keys[i as usize].sign(message)),
1010                )
1011            })
1012            .collect();
1013
1014        let multi_sig = MultiKeySignature::new(signatures).unwrap();
1015
1016        assert!(!multi_sig.has_signature(0));
1017        assert!(multi_sig.has_signature(1));
1018        assert!(!multi_sig.has_signature(2));
1019        assert!(multi_sig.has_signature(3));
1020        assert!(multi_sig.has_signature(4));
1021        assert!(!multi_sig.has_signature(5));
1022    }
1023
1024    #[test]
1025    fn test_multi_key_public_key_empty_keys() {
1026        let result = MultiKeyPublicKey::new(vec![], 1);
1027        assert!(result.is_err());
1028        assert!(result.unwrap_err().to_string().contains("at least one"));
1029    }
1030
1031    #[test]
1032    #[cfg(feature = "ed25519")]
1033    fn test_multi_key_public_key_threshold_zero() {
1034        use crate::crypto::Ed25519PrivateKey;
1035
1036        let keys: Vec<_> = (0..2)
1037            .map(|_| AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key()))
1038            .collect();
1039        let result = MultiKeyPublicKey::new(keys, 0);
1040        assert!(result.is_err());
1041        assert!(result.unwrap_err().to_string().contains("at least 1"));
1042    }
1043
1044    #[test]
1045    #[cfg(feature = "ed25519")]
1046    fn test_multi_key_public_key_threshold_exceeds() {
1047        use crate::crypto::Ed25519PrivateKey;
1048
1049        let keys: Vec<_> = (0..2)
1050            .map(|_| AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key()))
1051            .collect();
1052        let result = MultiKeyPublicKey::new(keys, 5);
1053        assert!(result.is_err());
1054        assert!(result.unwrap_err().to_string().contains("exceed"));
1055    }
1056
1057    #[test]
1058    #[cfg(feature = "ed25519")]
1059    fn test_multi_key_signature_empty() {
1060        let result = MultiKeySignature::new(vec![]);
1061        assert!(result.is_err());
1062        assert!(result.unwrap_err().to_string().contains("at least one"));
1063    }
1064
1065    #[test]
1066    #[cfg(feature = "ed25519")]
1067    fn test_multi_key_signature_duplicate_index() {
1068        use crate::crypto::Ed25519PrivateKey;
1069
1070        let private_key = Ed25519PrivateKey::generate();
1071        let sig = AnySignature::ed25519(&private_key.sign(b"test"));
1072
1073        let result = MultiKeySignature::new(vec![(0, sig.clone()), (0, sig)]);
1074        assert!(result.is_err());
1075        assert!(result.unwrap_err().to_string().contains("duplicate"));
1076    }
1077
1078    #[test]
1079    #[cfg(feature = "ed25519")]
1080    fn test_multi_key_public_key_accessors() {
1081        use crate::crypto::Ed25519PrivateKey;
1082
1083        let keys: Vec<_> = (0..3)
1084            .map(|_| AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key()))
1085            .collect();
1086        let multi_pk = MultiKeyPublicKey::new(keys, 2).unwrap();
1087
1088        assert_eq!(multi_pk.threshold(), 2);
1089        assert_eq!(multi_pk.num_keys(), 3);
1090        assert_eq!(multi_pk.public_keys().len(), 3);
1091    }
1092
1093    #[test]
1094    #[cfg(feature = "ed25519")]
1095    fn test_multi_key_signature_accessors() {
1096        use crate::crypto::Ed25519PrivateKey;
1097
1098        let private_key = Ed25519PrivateKey::generate();
1099        let sig0 = AnySignature::ed25519(&private_key.sign(b"test"));
1100        let sig2 = AnySignature::ed25519(&private_key.sign(b"test"));
1101
1102        let multi_sig = MultiKeySignature::new(vec![(0, sig0), (2, sig2)]).unwrap();
1103
1104        assert_eq!(multi_sig.num_signatures(), 2);
1105        assert_eq!(multi_sig.signatures().len(), 2);
1106    }
1107
1108    #[test]
1109    #[cfg(feature = "ed25519")]
1110    fn test_multi_key_public_key_debug() {
1111        use crate::crypto::Ed25519PrivateKey;
1112
1113        let keys: Vec<_> = (0..2)
1114            .map(|_| AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key()))
1115            .collect();
1116        let multi_pk = MultiKeyPublicKey::new(keys, 2).unwrap();
1117
1118        let debug = format!("{multi_pk:?}");
1119        assert!(debug.contains("MultiKeyPublicKey"));
1120    }
1121
1122    #[test]
1123    #[cfg(feature = "ed25519")]
1124    fn test_multi_key_signature_debug() {
1125        use crate::crypto::Ed25519PrivateKey;
1126
1127        let private_key = Ed25519PrivateKey::generate();
1128        let sig = AnySignature::ed25519(&private_key.sign(b"test"));
1129        let multi_sig = MultiKeySignature::new(vec![(0, sig)]).unwrap();
1130
1131        let debug = format!("{multi_sig:?}");
1132        assert!(debug.contains("MultiKeySignature"));
1133    }
1134
1135    #[test]
1136    #[cfg(feature = "ed25519")]
1137    fn test_multi_key_public_key_display() {
1138        use crate::crypto::Ed25519PrivateKey;
1139
1140        let keys: Vec<_> = (0..2)
1141            .map(|_| AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key()))
1142            .collect();
1143        let multi_pk = MultiKeyPublicKey::new(keys, 2).unwrap();
1144
1145        let display = format!("{multi_pk}");
1146        assert!(display.starts_with("0x"));
1147    }
1148
1149    #[test]
1150    #[cfg(feature = "ed25519")]
1151    fn test_multi_key_signature_display() {
1152        use crate::crypto::Ed25519PrivateKey;
1153
1154        let private_key = Ed25519PrivateKey::generate();
1155        let sig = AnySignature::ed25519(&private_key.sign(b"test"));
1156        let multi_sig = MultiKeySignature::new(vec![(0, sig)]).unwrap();
1157
1158        let display = format!("{multi_sig}");
1159        assert!(display.starts_with("0x"));
1160    }
1161
1162    #[test]
1163    #[cfg(feature = "ed25519")]
1164    fn test_multi_key_public_key_to_address() {
1165        use crate::crypto::Ed25519PrivateKey;
1166
1167        let keys: Vec<_> = (0..2)
1168            .map(|_| AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key()))
1169            .collect();
1170        let multi_pk = MultiKeyPublicKey::new(keys, 2).unwrap();
1171
1172        let address = multi_pk.to_address();
1173        assert!(!address.is_zero());
1174    }
1175
1176    #[test]
1177    fn test_any_public_key_variant_debug() {
1178        let variant = AnyPublicKeyVariant::Ed25519;
1179        let debug = format!("{variant:?}");
1180        assert!(debug.contains("Ed25519"));
1181    }
1182
1183    #[test]
1184    #[cfg(feature = "ed25519")]
1185    fn test_any_public_key_ed25519_debug() {
1186        use crate::crypto::Ed25519PrivateKey;
1187
1188        let pk = AnyPublicKey::ed25519(&Ed25519PrivateKey::generate().public_key());
1189        let debug = format!("{pk:?}");
1190        assert!(debug.contains("Ed25519"));
1191    }
1192
1193    #[test]
1194    #[cfg(feature = "ed25519")]
1195    fn test_any_signature_ed25519_debug() {
1196        use crate::crypto::Ed25519PrivateKey;
1197
1198        let private_key = Ed25519PrivateKey::generate();
1199        let sig = AnySignature::ed25519(&private_key.sign(b"test"));
1200        let debug = format!("{sig:?}");
1201        assert!(debug.contains("Ed25519"));
1202    }
1203
1204    #[test]
1205    #[cfg(feature = "ed25519")]
1206    fn test_multi_key_insufficient_signatures() {
1207        use crate::crypto::Ed25519PrivateKey;
1208
1209        let private_keys: Vec<_> = (0..3).map(|_| Ed25519PrivateKey::generate()).collect();
1210        let public_keys: Vec<_> = private_keys
1211            .iter()
1212            .map(|k| AnyPublicKey::ed25519(&k.public_key()))
1213            .collect();
1214
1215        let multi_pk = MultiKeyPublicKey::new(public_keys, 2).unwrap();
1216        let message = b"test message";
1217
1218        // Only provide 1 signature when threshold is 2
1219        let sig0 = AnySignature::ed25519(&private_keys[0].sign(message));
1220        let multi_sig = MultiKeySignature::new(vec![(0, sig0)]).unwrap();
1221
1222        // Verify should fail
1223        let result = multi_pk.verify(message, &multi_sig);
1224        assert!(result.is_err());
1225    }
1226}