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