Add support for loading keypairs from file

eth2/fork_choice/tests/tests.rs

@@ -224,7 +224,7 @@ fn setup_inital_state(
 
     let spec = ChainSpec::foundation();
 
-    let state_builder = TestingBeaconStateBuilder::new(no_validators, &spec);
+    let state_builder = TestingBeaconStateBuilder::new(no_validators, None, &spec);
     let (state, _keypairs) = state_builder.build();
 
     let state_root = state.canonical_root();

eth2/state_processing/benches/bench_block_processing.rs

@@ -9,14 +9,19 @@ use state_processing::{
         verify_block_signature,
     },
 };
+use std::path::Path;
 use types::test_utils::{TestingBeaconBlockBuilder, TestingBeaconStateBuilder};
 use types::*;
 
 /// Run the benchmarking suite on a foundation spec with 16,384 validators.
-pub fn bench_block_processing_n_validators(c: &mut Criterion, validator_count: usize) {
+pub fn bench_block_processing_n_validators(
+    c: &mut Criterion,
+    validator_count: usize,
+    keypair_file: Option<&Path>,
+) {
     let spec = ChainSpec::foundation();
 
-    let (mut state, keypairs) = build_state(validator_count, &spec);
+    let (mut state, keypairs) = build_state(validator_count, keypair_file, &spec);
     let block = build_block(&mut state, &keypairs, &spec);
 
     assert_eq!(
@@ -79,8 +84,12 @@ pub fn bench_block_processing_n_validators(c: &mut Criterion, validator_count: u
     );
 }
 
-fn build_state(validator_count: usize, spec: &ChainSpec) -> (BeaconState, Vec<Keypair>) {
-    let mut builder = TestingBeaconStateBuilder::new(validator_count, &spec);
+fn build_state(
+    validator_count: usize,
+    keypair_file: Option<&Path>,
+    spec: &ChainSpec,
+) -> (BeaconState, Vec<Keypair>) {
+    let mut builder = TestingBeaconStateBuilder::new(validator_count, keypair_file, &spec);
 
     // Set the state to be just before an epoch transition.
     let target_slot = (spec.genesis_epoch + 4).end_slot(spec.slots_per_epoch);

eth2/state_processing/benches/bench_epoch_processing.rs

@@ -10,6 +10,7 @@ use state_processing::{
         update_latest_slashed_balances,
     },
 };
+use std::path::Path;
 use types::test_utils::TestingBeaconStateBuilder;
 use types::{validator_registry::get_active_validator_indices, *};
 
@@ -17,10 +18,14 @@ pub const BENCHING_SAMPLE_SIZE: usize = 10;
 pub const SMALL_BENCHING_SAMPLE_SIZE: usize = 10;
 
 /// Run the benchmarking suite on a foundation spec with 16,384 validators.
-pub fn bench_epoch_processing_n_validators(c: &mut Criterion, validator_count: usize) {
+pub fn bench_epoch_processing_n_validators(
+    c: &mut Criterion,
+    validator_count: usize,
+    keypair_file: Option<&Path>,
+) {
     let spec = ChainSpec::foundation();
 
-    let mut builder = TestingBeaconStateBuilder::new(validator_count, &spec);
+    let mut builder = TestingBeaconStateBuilder::new(validator_count, keypair_file, &spec);
 
     // Set the state to be just before an epoch transition.
     let target_slot = (spec.genesis_epoch + 4).end_slot(spec.slots_per_epoch);

eth2/state_processing/benches/benches.rs

@@ -1,5 +1,9 @@
+use criterion::Benchmark;
 use criterion::Criterion;
 use criterion::{criterion_group, criterion_main};
+use std::path::Path;
+use types::test_utils::TestingBeaconStateBuilder;
+use types::*;
 
 mod bench_block_processing;
 mod bench_epoch_processing;
@@ -7,9 +11,47 @@ mod bench_epoch_processing;
 pub const VALIDATOR_COUNT: usize = 300_032;
 
 pub fn state_processing(c: &mut Criterion) {
-    bench_block_processing::bench_block_processing_n_validators(c, VALIDATOR_COUNT);
-    bench_epoch_processing::bench_epoch_processing_n_validators(c, VALIDATOR_COUNT);
+    bench_block_processing::bench_block_processing_n_validators(c, VALIDATOR_COUNT, None);
+    bench_epoch_processing::bench_epoch_processing_n_validators(c, VALIDATOR_COUNT, None);
 }
 
-criterion_group!(benches, state_processing,);
+pub fn key_loading(c: &mut Criterion) {
+    let validator_count = 1000;
+
+    c.bench(
+        &format!("{}_validators", validator_count),
+        Benchmark::new("generated", move |b| {
+            b.iter_batched(
+                || (),
+                |_| TestingBeaconStateBuilder::new(validator_count, None, &ChainSpec::foundation()),
+                criterion::BatchSize::SmallInput,
+            )
+        })
+        .sample_size(10),
+    );
+
+    // Note: path needs to be relative to where cargo is executed from.
+    let keypair_file =
+        Path::new("../../beacon_node/beacon_chain/test_harness/keypairs.raw_keypairs");
+    c.bench(
+        &format!("{}_validators", validator_count),
+        Benchmark::new("from_file", move |b| {
+            b.iter_batched(
+                || (),
+                |_| {
+                    TestingBeaconStateBuilder::new(
+                        validator_count,
+                        Some(&keypair_file),
+                        &ChainSpec::foundation(),
+                    )
+                },
+                criterion::BatchSize::SmallInput,
+            )
+        })
+        .sample_size(10),
+    );
+}
+
+// criterion_group!(benches, state_processing, key_loading);
+criterion_group!(benches, key_loading);
 criterion_main!(benches);

eth2/state_processing/src/per_epoch_processing/tests.rs

@@ -10,7 +10,7 @@ fn runs_without_error() {
 
     let spec = ChainSpec::few_validators();
 
-    let mut builder = TestingBeaconStateBuilder::new(8, &spec);
+    let mut builder = TestingBeaconStateBuilder::new(8, None, &spec);
 
     let target_slot = (spec.genesis_epoch + 4).end_slot(spec.slots_per_epoch);
     builder.teleport_to_slot(target_slot, &spec);

eth2/types/src/beacon_state/tests.rs

@@ -13,7 +13,7 @@ pub fn get_attestation_participants_consistency() {
     let mut rng = XorShiftRng::from_seed([42; 16]);
 
     let spec = ChainSpec::few_validators();
-    let builder = TestingBeaconStateBuilder::new(8, &spec);
+    let builder = TestingBeaconStateBuilder::new(8, None, &spec);
     let (mut state, _keypairs) = builder.build();
 
     state

eth2/types/src/test_utils/generate_deterministic_keypairs.rs

@@ -0,0 +1,30 @@
+use crate::*;
+use int_to_bytes::int_to_bytes48;
+use log::debug;
+use rayon::prelude::*;
+
+/// Generates `validator_count` keypairs where the secret key is the index of the
+/// validator.
+///
+/// For example, the first validator has a secret key of `int_to_bytes48(1)`, the second has
+/// `int_to_bytes48(2)` and so on. (We skip `0` as it generates a weird looking public key and is
+/// probably invalid).
+pub fn generate_deterministic_keypairs(validator_count: usize) -> Vec<Keypair> {
+    debug!(
+        "Generating {} deterministic validator keypairs...",
+        validator_count
+    );
+
+    let keypairs: Vec<Keypair> = (0..validator_count)
+        .collect::<Vec<usize>>()
+        .par_iter()
+        .map(|&i| {
+            let secret = int_to_bytes48(i as u64 + 1);
+            let sk = SecretKey::from_bytes(&secret).unwrap();
+            let pk = PublicKey::from_secret_key(&sk);
+            Keypair { sk, pk }
+        })
+        .collect();
+
+    keypairs
+}

eth2/types/src/test_utils/keypairs_file.rs

@@ -0,0 +1,113 @@
+use crate::*;
+use std::fs::File;
+use std::io::{Error, ErrorKind, Read, Write};
+use std::path::Path;
+
+pub const PUBLIC_KEY_BYTES_LEN: usize = 48;
+pub const SECRET_KEY_BYTES_LEN: usize = 48;
+
+pub const BATCH_SIZE: usize = 1_000; // ~15MB
+
+pub const KEYPAIR_BYTES_LEN: usize = PUBLIC_KEY_BYTES_LEN + SECRET_KEY_BYTES_LEN;
+pub const BATCH_BYTE_LEN: usize = KEYPAIR_BYTES_LEN * BATCH_SIZE;
+
+pub trait KeypairsFile {
+    fn to_raw_file(&self, path: &Path, keypairs: &[Keypair]) -> Result<(), Error>;
+    fn from_raw_file(path: &Path, count: usize) -> Result<Vec<Keypair>, Error>;
+}
+
+impl KeypairsFile for Vec<Keypair> {
+    fn to_raw_file(&self, path: &Path, keypairs: &[Keypair]) -> Result<(), Error> {
+        let mut keypairs_file = File::create(path)?;
+
+        for keypair_batch in keypairs.chunks(BATCH_SIZE) {
+            let mut buf = Vec::with_capacity(BATCH_BYTE_LEN);
+
+            for keypair in keypair_batch {
+                buf.append(&mut keypair.sk.as_raw().as_bytes());
+                buf.append(&mut keypair.pk.as_raw().as_bytes());
+            }
+
+            keypairs_file.write_all(&buf)?;
+        }
+
+        Ok(())
+    }
+
+    fn from_raw_file(path: &Path, count: usize) -> Result<Vec<Keypair>, Error> {
+        let mut keypairs_file = File::open(path)?;
+
+        let mut keypairs = Vec::with_capacity(count);
+
+        let indices: Vec<usize> = (0..count).collect();
+
+        for batch in indices.chunks(BATCH_SIZE) {
+            let mut buf = vec![0; batch.len() * KEYPAIR_BYTES_LEN];
+            keypairs_file.read_exact(&mut buf)?;
+
+            for (i, _) in batch.iter().enumerate() {
+                let sk_start = i * KEYPAIR_BYTES_LEN;
+                let sk_end = sk_start + SECRET_KEY_BYTES_LEN;
+                let sk = SecretKey::from_bytes(&buf[sk_start..sk_end])
+                    .map_err(|_| Error::new(ErrorKind::Other, "Invalid SecretKey bytes"))?;
+
+                let pk_start = sk_end;
+                let pk_end = pk_start + PUBLIC_KEY_BYTES_LEN;
+                let pk = PublicKey::from_bytes(&buf[pk_start..pk_end])
+                    .map_err(|_| Error::new(ErrorKind::Other, "Invalid PublicKey bytes"))?;
+
+                keypairs.push(Keypair { sk, pk });
+            }
+        }
+
+        Ok(keypairs)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rand::{distributions::Alphanumeric, thread_rng, Rng};
+    use rayon::prelude::*;
+    use std::fs::remove_file;
+
+    fn random_keypairs(n: usize) -> Vec<Keypair> {
+        (0..n).into_par_iter().map(|_| Keypair::random()).collect()
+    }
+
+    fn random_tmp_file() -> String {
+        let mut rng = thread_rng();
+
+        rng.sample_iter(&Alphanumeric).take(7).collect()
+    }
+
+    #[test]
+    #[ignore]
+    fn read_write_consistency_small_batch() {
+        let num_keypairs = 10;
+        let keypairs = random_keypairs(num_keypairs);
+
+        let keypairs_path = Path::new("/tmp").join(random_tmp_file());
+        keypairs.to_raw_file(&keypairs_path, &keypairs).unwrap();
+
+        let decoded = Vec::from_raw_file(&keypairs_path, num_keypairs).unwrap();
+        remove_file(keypairs_path).unwrap();
+
+        assert_eq!(keypairs, decoded);
+    }
+
+    #[test]
+    #[ignore]
+    fn read_write_consistency_big_batch() {
+        let num_keypairs = BATCH_SIZE + 1;
+        let keypairs = random_keypairs(num_keypairs);
+
+        let keypairs_path = Path::new("/tmp").join(random_tmp_file());
+        keypairs.to_raw_file(&keypairs_path, &keypairs).unwrap();
+
+        let decoded = Vec::from_raw_file(&keypairs_path, num_keypairs).unwrap();
+        remove_file(keypairs_path).unwrap();
+
+        assert_eq!(keypairs, decoded);
+    }
+}

eth2/types/src/test_utils/mod.rs

@@ -1,3 +1,5 @@
+mod generate_deterministic_keypairs;
+mod keypairs_file;
 mod test_random;
 mod testing_attestation_builder;
 mod testing_beacon_block_builder;
@@ -6,6 +8,8 @@ mod testing_deposit_builder;
 mod testing_transfer_builder;
 mod testing_voluntary_exit_builder;
 
+pub use generate_deterministic_keypairs::generate_deterministic_keypairs;
+pub use keypairs_file::KeypairsFile;
 pub use rand::{prng::XorShiftRng, SeedableRng};
 pub use test_random::TestRandom;
 pub use testing_attestation_builder::TestingAttestationBuilder;

eth2/types/src/test_utils/testing_beacon_state_builder.rs

@@ -1,8 +1,9 @@
+use super::{generate_deterministic_keypairs, KeypairsFile};
 use crate::beacon_state::BeaconStateBuilder;
 use crate::*;
 use bls::get_withdrawal_credentials;
-use int_to_bytes::int_to_bytes48;
 use rayon::prelude::*;
+use std::path::Path;
 
 pub struct TestingBeaconStateBuilder {
     state: BeaconState,
@@ -10,17 +11,11 @@ pub struct TestingBeaconStateBuilder {
 }
 
 impl TestingBeaconStateBuilder {
-    pub fn new(validator_count: usize, spec: &ChainSpec) -> Self {
-        let keypairs: Vec<Keypair> = (0..validator_count)
-            .collect::<Vec<usize>>()
-            .par_iter()
-            .map(|&i| {
-                let secret = int_to_bytes48(i as u64 + 1);
-                let sk = SecretKey::from_bytes(&secret).unwrap();
-                let pk = PublicKey::from_secret_key(&sk);
-                Keypair { sk, pk }
-            })
-            .collect();
+    pub fn new(validator_count: usize, keypairs_path: Option<&Path>, spec: &ChainSpec) -> Self {
+        let keypairs = match keypairs_path {
+            None => generate_deterministic_keypairs(validator_count),
+            Some(path) => Vec::from_raw_file(path, validator_count).unwrap(),
+        };
 
         let validators = keypairs
             .par_iter()

eth2/utils/bls/src/public_key.rs

@@ -1,6 +1,6 @@
 use super::serde_vistors::HexVisitor;
 use super::SecretKey;
-use bls_aggregates::PublicKey as RawPublicKey;
+use bls_aggregates::{DecodeError as BlsDecodeError, PublicKey as RawPublicKey};
 use hex::encode as hex_encode;
 use serde::de::{Deserialize, Deserializer};
 use serde::ser::{Serialize, Serializer};
@@ -22,7 +22,14 @@ impl PublicKey {
         PublicKey(RawPublicKey::from_secret_key(secret_key.as_raw()))
     }
 
-    /// Returns the underlying signature.
+    /// Instantiate a PublicKey from existing bytes.
+    ///
+    /// Note: this is _not_ SSZ decoding.
+    pub fn from_bytes(bytes: &[u8]) -> Result<PublicKey, BlsDecodeError> {
+        Ok(Self(RawPublicKey::from_bytes(bytes)?))
+    }
+
+    /// Returns the underlying public key.
     pub fn as_raw(&self) -> &RawPublicKey {
         &self.0
     }