feat: Parse and expand authority public keys once and reuse (#121)

* Parse and expand authority public keys once and reuse * Key derivation in single place + add TODO validation issue Co-authored-by: George Danezis <[email protected]>
MystenLabs · Jan 5, 2022 · 5a482af · 5a482af · github-actions · Jan 5, 2022
1 parent e29c107
commit 5a482af
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 38 deletions.
diff --git a/fastpay/src/bench.rs b/fastpay/src/bench.rs
@@ -109,10 +109,7 @@ impl ClientServerBenchmark {
         for _ in 0..self.committee_size {
             keys.push(get_key_pair());
         }
-        let committee = Committee {
-            voting_rights: keys.iter().map(|(k, _)| (*k, 1)).collect(),
-            total_votes: self.committee_size,
-        };
+        let committee = Committee::new(keys.iter().map(|(k, _)| (*k, 1)).collect());
 
         // Pick an authority and create state.
         let (public_auth0, secret_auth0) = keys.pop().unwrap();

diff --git a/fastpay/src/client.rs b/fastpay/src/client.rs
@@ -130,10 +130,7 @@ fn make_benchmark_certificates_from_orders_and_server_configs(
         let server_config = AuthorityServerConfig::read(file).expect("Fail to read server config");
         keys.push((server_config.authority.address, server_config.key));
     }
-    let committee = Committee {
-        voting_rights: keys.iter().map(|(k, _)| (*k, 1)).collect(),
-        total_votes: keys.len(),
-    };
+    let committee = Committee::new(keys.iter().map(|(k, _)| (*k, 1)).collect());
     assert!(
         keys.len() >= committee.quorum_threshold(),
         "Not enough server configs were provided with --server-configs"

diff --git a/fastx_types/src/base_types.rs b/fastx_types/src/base_types.rs
@@ -1,10 +1,11 @@
 // Copyright (c) Facebook, Inc. and its affiliates.
 // SPDX-License-Identifier: Apache-2.0
 use crate::error::FastPayError;
+use std::collections::HashMap;
 use std::convert::{TryFrom, TryInto};
 
 use ed25519_dalek as dalek;
-use ed25519_dalek::{Digest, Signer, Verifier};
+use ed25519_dalek::{Digest, PublicKey, Signer, Verifier};
 use move_core_types::account_address::AccountAddress;
 use move_core_types::ident_str;
 use move_core_types::identifier::IdentStr;
@@ -44,6 +45,18 @@ impl PublicKeyBytes {
     pub fn to_vec(&self) -> Vec<u8> {
         self.0.to_vec()
     }
+
+    pub fn to_public_key(&self) -> Result<PublicKey, FastPayError> {
+        // TODO(https://github.com/MystenLabs/fastnft/issues/101): Do better key validation
+        // to ensure the bytes represent a point on the curve.
+        PublicKey::from_bytes(self.as_ref()).map_err(|_| FastPayError::InvalidAuthenticator)
+    }
+}
+
+impl AsRef<[u8]> for PublicKeyBytes {
+    fn as_ref(&self) -> &[u8] {
+        &self.0[..]
+    }
 }
 
 // TODO(https://github.com/MystenLabs/fastnft/issues/101): more robust key validation
@@ -425,31 +438,25 @@ impl Signature {
         Signature(signature)
     }
 
-    fn check_internal<T>(
-        &self,
-        value: &T,
-        author: FastPayAddress,
-    ) -> Result<(), dalek::SignatureError>
+    pub fn check<T>(&self, value: &T, author: FastPayAddress) -> Result<(), FastPayError>
     where
         T: Signable<Vec<u8>>,
     {
         let mut message = Vec::new();
         value.write(&mut message);
-        let public_key = dalek::PublicKey::from_bytes(&author.0)?;
-        public_key.verify(&message, &self.0)
-    }
-
-    pub fn check<T>(&self, value: &T, author: FastPayAddress) -> Result<(), FastPayError>
-    where
-        T: Signable<Vec<u8>>,
-    {
-        self.check_internal(value, author)
+        let public_key = author.to_public_key()?;
+        public_key
+            .verify(&message, &self.0)
             .map_err(|error| FastPayError::InvalidSignature {
                 error: format!("{}", error),
             })
     }
 
-    fn verify_batch_internal<'a, T, I>(value: &'a T, votes: I) -> Result<(), dalek::SignatureError>
+    pub fn verify_batch<'a, T, I>(
+        value: &'a T,
+        votes: I,
+        key_cache: &HashMap<PublicKeyBytes, PublicKey>,
+    ) -> Result<(), FastPayError>
     where
         T: Signable<Vec<u8>>,
         I: IntoIterator<Item = &'a (FastPayAddress, Signature)>,
@@ -462,17 +469,12 @@ impl Signature {
         for (addr, sig) in votes.into_iter() {
             messages.push(&msg);
             signatures.push(sig.0);
-            public_keys.push(dalek::PublicKey::from_bytes(&addr.0)?);
+            match key_cache.get(addr) {
+                Some(v) => public_keys.push(*v),
+                None => public_keys.push(addr.to_public_key()?),
+            }
         }
-        dalek::verify_batch(&messages[..], &signatures[..], &public_keys[..])
-    }
-
-    pub fn verify_batch<'a, T, I>(value: &'a T, votes: I) -> Result<(), FastPayError>
-    where
-        T: Signable<Vec<u8>>,
-        I: IntoIterator<Item = &'a (FastPayAddress, Signature)>,
-    {
-        Signature::verify_batch_internal(value, votes).map_err(|error| {
+        dalek::verify_batch(&messages[..], &signatures[..], &public_keys[..]).map_err(|error| {
             FastPayError::InvalidSignature {
                 error: format!("{}", error),
             }

diff --git a/fastx_types/src/committee.rs b/fastx_types/src/committee.rs
@@ -2,20 +2,27 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::base_types::*;
-use std::collections::BTreeMap;
+use ed25519_dalek::PublicKey;
+use std::collections::{BTreeMap, HashMap};
 
-#[derive(Eq, PartialEq, Clone, Hash, Debug)]
+#[derive(Eq, PartialEq, Clone, Debug)]
 pub struct Committee {
     pub voting_rights: BTreeMap<AuthorityName, usize>,
     pub total_votes: usize,
+    pub expanded_keys: HashMap<AuthorityName, PublicKey>,
 }
 
 impl Committee {
     pub fn new(voting_rights: BTreeMap<AuthorityName, usize>) -> Self {
-        let total_votes = voting_rights.iter().fold(0, |sum, (_, votes)| sum + *votes);
+        let total_votes = voting_rights.iter().map(|(_, votes)| votes).sum();
+        let expanded_keys: HashMap<_, _> = voting_rights
+            .iter()
+            .map(|(addr, _)| (*addr, addr.to_public_key().expect("Invalid Authority Key")))
+            .collect();
         Committee {
             voting_rights,
             total_votes,
+            expanded_keys,
         }
     }
 

diff --git a/fastx_types/src/messages.rs b/fastx_types/src/messages.rs
@@ -404,6 +404,7 @@ impl CertifiedOrder {
         Signature::verify_batch(
             &self.order.kind,
             std::iter::once(&inner_sig).chain(&self.signatures),
+            &committee.expanded_keys,
         )
     }
 }

diff --git a/fastx_types/src/unit_tests/serialize_tests.rs b/fastx_types/src/unit_tests/serialize_tests.rs
@@ -312,10 +312,17 @@ fn test_time_cert() {
         signatures: Vec::new(),
     };
 
+    use ed25519_dalek::PublicKey;
+    use std::collections::HashMap;
+    let mut cache = HashMap::new();
     for _ in 0..7 {
         let (authority_name, authority_key) = get_key_pair();
         let sig = Signature::new(&cert.order.kind, &authority_key);
         cert.signatures.push((authority_name, sig));
+        cache.insert(
+            authority_name,
+            PublicKey::from_bytes(authority_name.as_ref()).expect("No problem parsing key."),
+        );
     }
 
     let mut buf = Vec::new();
@@ -330,7 +337,7 @@ fn test_time_cert() {
     let mut buf2 = buf.as_slice();
     for _ in 0..count {
         if let SerializedMessage::Cert(cert) = deserialize_message(&mut buf2).unwrap() {
-            Signature::verify_batch(&cert.order.kind, &cert.signatures).unwrap();
+            Signature::verify_batch(&cert.order.kind, &cert.signatures, &cache).unwrap();
         }
     }
     assert!(deserialize_message(buf2).is_err());