1use crate::{
10 ed25519::{
11 Ed25519PrivateKey, Ed25519PublicKey, Ed25519Signature, ED25519_PRIVATE_KEY_LENGTH,
12 ED25519_PUBLIC_KEY_LENGTH, ED25519_SIGNATURE_LENGTH,
13 },
14 hash::{CryptoHash, CryptoHasher},
15 traits::*,
16};
17use anyhow::{anyhow, Result};
18use aptos_crypto_derive::{DeserializeKey, SerializeKey, SilentDebug, SilentDisplay};
19use core::convert::TryFrom;
20use rand::Rng;
21use serde::Serialize;
22use std::{convert::TryInto, fmt};
23
24pub const MAX_NUM_OF_KEYS: usize = 32;
26pub const BITMAP_NUM_OF_BYTES: usize = 4;
28
29#[derive(DeserializeKey, Eq, PartialEq, SilentDisplay, SilentDebug, SerializeKey)]
31pub struct MultiEd25519PrivateKey {
32 private_keys: Vec<Ed25519PrivateKey>,
33 threshold: u8,
34}
35
36#[cfg(feature = "assert-private-keys-not-cloneable")]
37static_assertions::assert_not_impl_any!(MultiEd25519PrivateKey: Clone);
38
39#[derive(Clone, DeserializeKey, Eq, PartialEq, SerializeKey)]
41pub struct MultiEd25519PublicKey {
42 public_keys: Vec<Ed25519PublicKey>,
43 threshold: u8,
44}
45
46#[derive(Clone, DeserializeKey, Eq, PartialEq, SerializeKey)]
52pub struct MultiEd25519Signature {
53 signatures: Vec<Ed25519Signature>,
54 bitmap: [u8; BITMAP_NUM_OF_BYTES],
55}
56
57impl MultiEd25519PrivateKey {
58 pub fn new(
60 private_keys: Vec<Ed25519PrivateKey>,
61 threshold: u8,
62 ) -> std::result::Result<Self, CryptoMaterialError> {
63 let num_of_private_keys = private_keys.len();
64 if threshold == 0 || num_of_private_keys < threshold as usize {
65 Err(CryptoMaterialError::ValidationError)
66 } else if num_of_private_keys > MAX_NUM_OF_KEYS {
67 Err(CryptoMaterialError::WrongLengthError)
68 } else {
69 Ok(MultiEd25519PrivateKey {
70 private_keys,
71 threshold,
72 })
73 }
74 }
75
76 pub fn to_bytes(&self) -> Vec<u8> {
78 to_bytes(&self.private_keys, self.threshold)
79 }
80}
81
82impl MultiEd25519PublicKey {
83 pub fn new(
89 public_keys: Vec<Ed25519PublicKey>,
90 threshold: u8,
91 ) -> std::result::Result<Self, CryptoMaterialError> {
92 let num_of_public_keys = public_keys.len();
93 if threshold == 0 || num_of_public_keys < threshold as usize {
94 Err(CryptoMaterialError::ValidationError)
95 } else if num_of_public_keys > MAX_NUM_OF_KEYS {
96 Err(CryptoMaterialError::WrongLengthError)
97 } else {
98 Ok(MultiEd25519PublicKey {
99 public_keys,
100 threshold,
101 })
102 }
103 }
104
105 pub fn public_keys(&self) -> &Vec<Ed25519PublicKey> {
107 &self.public_keys
108 }
109
110 pub fn threshold(&self) -> &u8 {
112 &self.threshold
113 }
114
115 pub fn to_bytes(&self) -> Vec<u8> {
117 to_bytes(&self.public_keys, self.threshold)
118 }
119}
120
121impl From<&Ed25519PrivateKey> for MultiEd25519PrivateKey {
127 fn from(ed_private_key: &Ed25519PrivateKey) -> Self {
128 MultiEd25519PrivateKey {
129 private_keys: vec![Ed25519PrivateKey::try_from(&ed_private_key.to_bytes()[..]).unwrap()],
130 threshold: 1u8,
131 }
132 }
133}
134
135impl PrivateKey for MultiEd25519PrivateKey {
136 type PublicKeyMaterial = MultiEd25519PublicKey;
137}
138
139impl SigningKey for MultiEd25519PrivateKey {
140 type VerifyingKeyMaterial = MultiEd25519PublicKey;
141 type SignatureMaterial = MultiEd25519Signature;
142
143 fn sign<T: CryptoHash + Serialize>(
146 &self,
147 message: &T,
148 ) -> Result<MultiEd25519Signature, CryptoMaterialError> {
149 let mut bitmap = [0u8; BITMAP_NUM_OF_BYTES];
150 let mut signatures: Vec<Ed25519Signature> = vec![];
151
152 for (i, private_key) in self
153 .private_keys
154 .iter()
155 .take(self.threshold as usize)
156 .enumerate()
157 {
158 bitmap_set_bit(&mut bitmap, i);
159 signatures.push(private_key.sign(message)?);
160 }
161
162 Ok(MultiEd25519Signature { signatures, bitmap })
163 }
164
165 #[cfg(any(test, feature = "fuzzing"))]
166 fn sign_arbitrary_message(&self, message: &[u8]) -> MultiEd25519Signature {
167 let mut signatures: Vec<Ed25519Signature> = Vec::with_capacity(self.threshold as usize);
168 let mut bitmap = [0u8; BITMAP_NUM_OF_BYTES];
169 for (i, private_key) in self
170 .private_keys
171 .iter()
172 .take(self.threshold as usize)
173 .enumerate()
174 {
175 bitmap_set_bit(&mut bitmap, i);
176 signatures.push(private_key.sign_arbitrary_message(message));
177 }
178
179 MultiEd25519Signature { signatures, bitmap }
180 }
181}
182
183impl Uniform for MultiEd25519PrivateKey {
185 fn generate<R>(rng: &mut R) -> Self
186 where
187 R: ::rand::RngCore + ::rand::CryptoRng,
188 {
189 let num_of_keys = rng.gen_range(1, MAX_NUM_OF_KEYS + 1);
190 let mut private_keys: Vec<Ed25519PrivateKey> = Vec::with_capacity(num_of_keys);
191 for _ in 0..num_of_keys {
192 private_keys.push(
193 Ed25519PrivateKey::try_from(
194 &ed25519_dalek::SecretKey::generate(rng).to_bytes()[..],
195 )
196 .unwrap(),
197 );
198 }
199 let threshold = rng.gen_range(1, num_of_keys + 1) as u8;
200 MultiEd25519PrivateKey {
201 private_keys,
202 threshold,
203 }
204 }
205}
206
207impl TryFrom<&[u8]> for MultiEd25519PrivateKey {
208 type Error = CryptoMaterialError;
209
210 fn try_from(bytes: &[u8]) -> std::result::Result<MultiEd25519PrivateKey, CryptoMaterialError> {
212 if bytes.is_empty() {
213 return Err(CryptoMaterialError::WrongLengthError);
214 }
215 let threshold = check_and_get_threshold(bytes, ED25519_PRIVATE_KEY_LENGTH)?;
216
217 let private_keys: Result<Vec<Ed25519PrivateKey>, _> = bytes
218 .chunks_exact(ED25519_PRIVATE_KEY_LENGTH)
219 .map(Ed25519PrivateKey::try_from)
220 .collect();
221
222 private_keys.map(|private_keys| MultiEd25519PrivateKey {
223 private_keys,
224 threshold,
225 })
226 }
227}
228
229impl Length for MultiEd25519PrivateKey {
230 fn length(&self) -> usize {
231 self.private_keys.len() * ED25519_PRIVATE_KEY_LENGTH + 1
232 }
233}
234
235impl ValidCryptoMaterial for MultiEd25519PrivateKey {
236 fn to_bytes(&self) -> Vec<u8> {
237 self.to_bytes()
238 }
239}
240
241impl Genesis for MultiEd25519PrivateKey {
242 fn genesis() -> Self {
243 let mut buf = [0u8; ED25519_PRIVATE_KEY_LENGTH];
244 buf[ED25519_PRIVATE_KEY_LENGTH - 1] = 1u8;
245 MultiEd25519PrivateKey {
246 private_keys: vec![Ed25519PrivateKey::try_from(buf.as_ref()).unwrap()],
247 threshold: 1u8,
248 }
249 }
250}
251
252impl From<Ed25519PublicKey> for MultiEd25519PublicKey {
258 fn from(ed_public_key: Ed25519PublicKey) -> Self {
259 MultiEd25519PublicKey {
260 public_keys: vec![ed_public_key],
261 threshold: 1u8,
262 }
263 }
264}
265
266impl From<&MultiEd25519PrivateKey> for MultiEd25519PublicKey {
268 fn from(private_key: &MultiEd25519PrivateKey) -> Self {
269 let public_keys = private_key
270 .private_keys
271 .iter()
272 .map(PrivateKey::public_key)
273 .collect();
274 MultiEd25519PublicKey {
275 public_keys,
276 threshold: private_key.threshold,
277 }
278 }
279}
280
281impl PublicKey for MultiEd25519PublicKey {
283 type PrivateKeyMaterial = MultiEd25519PrivateKey;
284}
285
286#[allow(clippy::derive_hash_xor_eq)]
287impl std::hash::Hash for MultiEd25519PublicKey {
288 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
289 let encoded_pubkey = self.to_bytes();
290 state.write(&encoded_pubkey);
291 }
292}
293
294impl TryFrom<&[u8]> for MultiEd25519PublicKey {
295 type Error = CryptoMaterialError;
296
297 fn try_from(bytes: &[u8]) -> std::result::Result<MultiEd25519PublicKey, CryptoMaterialError> {
301 if bytes.is_empty() {
302 return Err(CryptoMaterialError::WrongLengthError);
303 }
304 let threshold = check_and_get_threshold(bytes, ED25519_PUBLIC_KEY_LENGTH)?;
305 let public_keys: Result<Vec<Ed25519PublicKey>, _> = bytes
306 .chunks_exact(ED25519_PUBLIC_KEY_LENGTH)
307 .map(Ed25519PublicKey::try_from)
308 .collect();
309 public_keys.map(|public_keys| MultiEd25519PublicKey {
310 public_keys,
311 threshold,
312 })
313 }
314}
315
316impl VerifyingKey for MultiEd25519PublicKey {
318 type SigningKeyMaterial = MultiEd25519PrivateKey;
319 type SignatureMaterial = MultiEd25519Signature;
320}
321
322impl fmt::Display for MultiEd25519PublicKey {
323 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
324 write!(f, "{}", hex::encode(&self.to_bytes()))
325 }
326}
327
328impl fmt::Debug for MultiEd25519PublicKey {
329 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
330 write!(f, "MultiEd25519PublicKey({})", self)
331 }
332}
333
334impl Length for MultiEd25519PublicKey {
335 fn length(&self) -> usize {
336 self.public_keys.len() * ED25519_PUBLIC_KEY_LENGTH + 1
337 }
338}
339
340impl ValidCryptoMaterial for MultiEd25519PublicKey {
341 fn to_bytes(&self) -> Vec<u8> {
342 self.to_bytes()
343 }
344}
345
346impl MultiEd25519Signature {
347 pub fn new(
349 signatures: Vec<(Ed25519Signature, u8)>,
350 ) -> std::result::Result<Self, CryptoMaterialError> {
351 let num_of_sigs = signatures.len();
352 if num_of_sigs == 0 || num_of_sigs > MAX_NUM_OF_KEYS {
353 return Err(CryptoMaterialError::ValidationError);
354 }
355
356 let mut sorted_signatures = signatures;
357 sorted_signatures.sort_by(|a, b| a.1.cmp(&b.1));
358
359 let mut bitmap = [0u8; BITMAP_NUM_OF_BYTES];
360
361 let (sigs, indexes): (Vec<_>, Vec<_>) = sorted_signatures.into_iter().unzip();
363 for i in indexes {
364 if i < MAX_NUM_OF_KEYS as u8 {
366 if bitmap_get_bit(bitmap, i as usize) {
368 return Err(CryptoMaterialError::BitVecError(
369 "Duplicate signature index".to_string(),
370 ));
371 } else {
372 bitmap_set_bit(&mut bitmap, i as usize);
373 }
374 } else {
375 return Err(CryptoMaterialError::BitVecError(
376 "Signature index is out of range".to_string(),
377 ));
378 }
379 }
380 Ok(MultiEd25519Signature {
381 signatures: sigs,
382 bitmap,
383 })
384 }
385
386 pub fn new_with_signatures_and_bitmap(
388 signatures: Vec<Ed25519Signature>,
389 bitmap: [u8; BITMAP_NUM_OF_BYTES],
390 ) -> Self {
391 Self { signatures, bitmap }
392 }
393
394 pub fn signatures(&self) -> &Vec<Ed25519Signature> {
396 &self.signatures
397 }
398
399 pub fn bitmap(&self) -> &[u8; BITMAP_NUM_OF_BYTES] {
401 &self.bitmap
402 }
403
404 pub fn to_bytes(&self) -> Vec<u8> {
406 let mut bytes: Vec<u8> = self
407 .signatures
408 .iter()
409 .flat_map(|sig| sig.to_bytes().to_vec())
410 .collect();
411 bytes.extend(&self.bitmap[..]);
412 bytes
413 }
414}
415
416impl TryFrom<&[u8]> for MultiEd25519Signature {
421 type Error = CryptoMaterialError;
422
423 fn try_from(bytes: &[u8]) -> std::result::Result<MultiEd25519Signature, CryptoMaterialError> {
426 let length = bytes.len();
427 let bitmap_num_of_bytes = length % ED25519_SIGNATURE_LENGTH;
428 let num_of_sigs = length / ED25519_SIGNATURE_LENGTH;
429
430 if num_of_sigs == 0
431 || num_of_sigs > MAX_NUM_OF_KEYS
432 || bitmap_num_of_bytes != BITMAP_NUM_OF_BYTES
433 {
434 return Err(CryptoMaterialError::WrongLengthError);
435 }
436
437 let bitmap = match bytes[length - BITMAP_NUM_OF_BYTES..].try_into() {
438 Ok(bitmap) => bitmap,
439 Err(_) => return Err(CryptoMaterialError::DeserializationError),
440 };
441 if bitmap_count_ones(bitmap) != num_of_sigs as u32 {
442 return Err(CryptoMaterialError::DeserializationError);
443 }
444
445 let signatures: Result<Vec<Ed25519Signature>, _> = bytes
446 .chunks_exact(ED25519_SIGNATURE_LENGTH)
447 .map(Ed25519Signature::try_from)
448 .collect();
449 signatures.map(|signatures| MultiEd25519Signature { signatures, bitmap })
450 }
451}
452
453impl Length for MultiEd25519Signature {
454 fn length(&self) -> usize {
455 self.signatures.len() * ED25519_SIGNATURE_LENGTH + BITMAP_NUM_OF_BYTES
456 }
457}
458
459#[allow(clippy::derive_hash_xor_eq)]
460impl std::hash::Hash for MultiEd25519Signature {
461 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
462 let encoded_signature = self.to_bytes();
463 state.write(&encoded_signature);
464 }
465}
466
467impl fmt::Display for MultiEd25519Signature {
468 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
469 write!(f, "{}", hex::encode(&self.to_bytes()[..]))
470 }
471}
472
473impl fmt::Debug for MultiEd25519Signature {
474 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
475 write!(f, "MultiEd25519Signature({})", self)
476 }
477}
478
479impl ValidCryptoMaterial for MultiEd25519Signature {
480 fn to_bytes(&self) -> Vec<u8> {
481 self.to_bytes()
482 }
483}
484
485impl Signature for MultiEd25519Signature {
486 type VerifyingKeyMaterial = MultiEd25519PublicKey;
487 type SigningKeyMaterial = MultiEd25519PrivateKey;
488
489 fn verify<T: CryptoHash + Serialize>(
490 &self,
491 message: &T,
492 public_key: &MultiEd25519PublicKey,
493 ) -> Result<()> {
494 let mut bytes = <T as CryptoHash>::Hasher::seed().to_vec();
497 bcs::serialize_into(&mut bytes, &message)
498 .map_err(|_| CryptoMaterialError::SerializationError)?;
499 Self::verify_arbitrary_msg(self, &bytes, public_key)
500 }
501
502 fn verify_arbitrary_msg(
506 &self,
507 message: &[u8],
508 public_key: &MultiEd25519PublicKey,
509 ) -> Result<()> {
510 match bitmap_last_set_bit(self.bitmap) {
513 Some(last_bit) if last_bit as usize <= public_key.public_keys.len() => (),
514 _ => {
515 return Err(anyhow!(
516 "{}",
517 CryptoMaterialError::BitVecError("Signature index is out of range".to_string())
518 ))
519 }
520 };
521 if bitmap_count_ones(self.bitmap) < public_key.threshold as u32 {
522 return Err(anyhow!(
523 "{}",
524 CryptoMaterialError::BitVecError(
525 "Not enough signatures to meet the threshold".to_string()
526 )
527 ));
528 }
529 let mut bitmap_index = 0;
530 for sig in &self.signatures {
532 while !bitmap_get_bit(self.bitmap, bitmap_index) {
533 bitmap_index += 1;
534 }
535 sig.verify_arbitrary_msg(message, &public_key.public_keys[bitmap_index as usize])?;
536 bitmap_index += 1;
537 }
538 Ok(())
539 }
540
541 fn to_bytes(&self) -> Vec<u8> {
542 self.to_bytes()
543 }
544}
545
546impl From<Ed25519Signature> for MultiEd25519Signature {
547 fn from(ed_signature: Ed25519Signature) -> Self {
548 MultiEd25519Signature {
549 signatures: vec![ed_signature],
550 bitmap: [0b1000_0000u8, 0u8, 0u8, 0u8],
552 }
553 }
554}
555
556fn to_bytes<T: ValidCryptoMaterial>(keys: &[T], threshold: u8) -> Vec<u8> {
562 let mut bytes: Vec<u8> = keys
563 .iter()
564 .flat_map(ValidCryptoMaterial::to_bytes)
565 .collect();
566 bytes.push(threshold);
567 bytes
568}
569
570fn check_and_get_threshold(
572 bytes: &[u8],
573 key_size: usize,
574) -> std::result::Result<u8, CryptoMaterialError> {
575 let payload_length = bytes.len();
576 if bytes.is_empty() {
577 return Err(CryptoMaterialError::WrongLengthError);
578 }
579 let threshold_num_of_bytes = payload_length % key_size;
580 let num_of_keys = payload_length / key_size;
581 let threshold_byte = bytes[bytes.len() - 1];
582
583 if num_of_keys == 0 || num_of_keys > MAX_NUM_OF_KEYS || threshold_num_of_bytes != 1 {
584 Err(CryptoMaterialError::WrongLengthError)
585 } else if threshold_byte == 0 || threshold_byte > num_of_keys as u8 {
586 Err(CryptoMaterialError::ValidationError)
587 } else {
588 Ok(threshold_byte)
589 }
590}
591
592fn bitmap_set_bit(input: &mut [u8; BITMAP_NUM_OF_BYTES], index: usize) {
593 let bucket = index / 8;
594 let bucket_pos = index - (bucket * 8);
596 input[bucket] |= 128 >> bucket_pos as u8;
597}
598
599fn bitmap_get_bit(input: [u8; BITMAP_NUM_OF_BYTES], index: usize) -> bool {
601 let bucket = index / 8;
602 let bucket_pos = index - (bucket * 8);
604 (input[bucket] & (128 >> bucket_pos as u8)) != 0
605}
606
607fn bitmap_count_ones(input: [u8; BITMAP_NUM_OF_BYTES]) -> u32 {
609 input.iter().map(|a| a.count_ones()).sum()
610}
611
612fn bitmap_last_set_bit(input: [u8; BITMAP_NUM_OF_BYTES]) -> Option<u8> {
614 input
615 .iter()
616 .rev()
617 .enumerate()
618 .find(|(_, byte)| byte != &&0u8)
619 .map(|(i, byte)| (8 * (BITMAP_NUM_OF_BYTES - i) - byte.trailing_zeros() as usize - 1) as u8)
620}
621
622#[test]
623fn bitmap_tests() {
624 let mut bitmap = [0b0100_0000u8, 0b1111_1111u8, 0u8, 0b1000_0000u8];
625 assert!(!bitmap_get_bit(bitmap, 0));
626 assert!(bitmap_get_bit(bitmap, 1));
627 for i in 8..16 {
628 assert!(bitmap_get_bit(bitmap, i));
629 }
630 for i in 16..24 {
631 assert!(!bitmap_get_bit(bitmap, i));
632 }
633 assert!(bitmap_get_bit(bitmap, 24));
634 assert!(!bitmap_get_bit(bitmap, 31));
635 assert_eq!(bitmap_last_set_bit(bitmap), Some(24));
636
637 bitmap_set_bit(&mut bitmap, 30);
638 assert!(bitmap_get_bit(bitmap, 30));
639 assert_eq!(bitmap_last_set_bit(bitmap), Some(30));
640}