Improve DX for loading validator keys from file

beacon_node/beacon_chain/test_harness/src/bin.rs

@@ -2,6 +2,8 @@ use clap::{App, Arg, SubCommand};
 use env_logger::{Builder, Env};
 use gen_keys::gen_keys;
 use run_test::run_test;
+use std::fs;
+use types::test_utils::keypairs_path;
 use types::ChainSpec;
 
 mod beacon_chain_harness;
@@ -13,6 +15,10 @@ mod validator_harness;
 use validator_harness::ValidatorHarness;
 
 fn main() {
+    let validator_file_path = keypairs_path();
+
+    fs::create_dir(validator_file_path.parent().unwrap()).unwrap();
+
     let matches = App::new("Lighthouse Test Harness Runner")
         .version("0.0.1")
         .author("Sigma Prime <contact@sigmaprime.io>")
@@ -71,7 +77,7 @@ fn main() {
                         .short("d")
                         .value_name("GENESIS_TIME")
                         .help("Output directory for generated YAML.")
-                        .default_value("keypairs.raw_keypairs"),
+                        .default_value(validator_file_path.to_str().unwrap()),
                 ),
         )
         .get_matches();

eth2/fork_choice/tests/tests.rs

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

eth2/state_processing/benches/bench_block_processing.rs

@@ -9,19 +9,14 @@ 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(
+pub fn bench_block_processing_n_validators(c: &mut Criterion, validator_count: usize) {
-    c: &mut Criterion,
-    validator_count: usize,
-    keypair_file: Option<&Path>,
-) {
     let spec = ChainSpec::foundation();
 
-    let (mut state, keypairs) = build_state(validator_count, keypair_file, &spec);
+    let (mut state, keypairs) = build_state(validator_count, &spec);
     let block = build_block(&mut state, &keypairs, &spec);
 
     assert_eq!(
@@ -84,12 +79,9 @@ pub fn bench_block_processing_n_validators(
     );
 }
 
-fn build_state(
+fn build_state(validator_count: usize, spec: &ChainSpec) -> (BeaconState, Vec<Keypair>) {
-    validator_count: usize,
+    let mut builder =
-    keypair_file: Option<&Path>,
+        TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(validator_count, &spec);
-    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,7 +10,6 @@ use state_processing::{
         update_latest_slashed_balances,
     },
 };
-use std::path::Path;
 use types::test_utils::TestingBeaconStateBuilder;
 use types::{validator_registry::get_active_validator_indices, *};
 
@@ -18,14 +17,11 @@ 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(
+pub fn bench_epoch_processing_n_validators(c: &mut Criterion, validator_count: usize) {
-    c: &mut Criterion,
-    validator_count: usize,
-    keypair_file: Option<&Path>,
-) {
     let spec = ChainSpec::foundation();
 
-    let mut builder = TestingBeaconStateBuilder::new(validator_count, keypair_file, &spec);
+    let mut builder =
+        TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(validator_count, &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,7 +1,6 @@
 use criterion::Benchmark;
 use criterion::Criterion;
 use criterion::{criterion_group, criterion_main};
-use std::path::Path;
 use types::test_utils::TestingBeaconStateBuilder;
 use types::*;
 
@@ -11,8 +10,8 @@ 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, None);
+    bench_block_processing::bench_block_processing_n_validators(c, VALIDATOR_COUNT);
-    bench_epoch_processing::bench_epoch_processing_n_validators(c, VALIDATOR_COUNT, None);
+    bench_epoch_processing::bench_epoch_processing_n_validators(c, VALIDATOR_COUNT);
 }
 
 pub fn key_loading(c: &mut Criterion) {
@@ -23,7 +22,12 @@ pub fn key_loading(c: &mut Criterion) {
         Benchmark::new("generated", move |b| {
             b.iter_batched(
                 || (),
-                |_| TestingBeaconStateBuilder::new(validator_count, None, &ChainSpec::foundation()),
+                |_| {
+                    TestingBeaconStateBuilder::from_deterministic_keypairs(
+                        validator_count,
+                        &ChainSpec::foundation(),
+                    )
+                },
                 criterion::BatchSize::SmallInput,
             )
         })
@@ -31,17 +35,14 @@ pub fn key_loading(c: &mut Criterion) {
     );
 
     // 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(
+                    TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(
                         validator_count,
-                        Some(&keypair_file),
                         &ChainSpec::foundation(),
                     )
                 },

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, None, &spec);
+    let mut builder = TestingBeaconStateBuilder::from_deterministic_keypairs(8, &spec);
 
     let target_slot = (spec.genesis_epoch + 4).end_slot(spec.slots_per_epoch);
     builder.teleport_to_slot(target_slot, &spec);

eth2/types/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2018"
 [dependencies]
 bls = { path = "../utils/bls" }
 boolean-bitfield = { path = "../utils/boolean-bitfield" }
+dirs = "1.0"
 ethereum-types = "0.5"
 hashing = { path = "../utils/hashing" }
 honey-badger-split =  { path = "../utils/honey-badger-split" }

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, None, &spec);
+    let builder = TestingBeaconStateBuilder::from_deterministic_keypairs(8, &spec);
     let (mut state, _keypairs) = builder.build();
 
     state

eth2/types/src/test_utils/keypairs_file.rs

@@ -1,4 +1,5 @@
 use crate::*;
+use rayon::prelude::*;
 use std::fs::File;
 use std::io::{Error, ErrorKind, Read, Write};
 use std::path::Path;
@@ -45,19 +46,27 @@ impl KeypairsFile for Vec<Keypair> {
             let mut buf = vec![0; batch.len() * KEYPAIR_BYTES_LEN];
             keypairs_file.read_exact(&mut buf)?;
 
-            for (i, _) in batch.iter().enumerate() {
+            let mut keypair_batch = batch
-                let sk_start = i * KEYPAIR_BYTES_LEN;
+                .par_iter()
-                let sk_end = sk_start + SECRET_KEY_BYTES_LEN;
+                .enumerate()
-                let sk = SecretKey::from_bytes(&buf[sk_start..sk_end])
+                .map(|(i, _)| {
-                    .map_err(|_| Error::new(ErrorKind::Other, "Invalid SecretKey bytes"))?;
+                    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"))
+                        .unwrap();
 
-                let pk_start = sk_end;
+                    let pk_start = sk_end;
-                let pk_end = pk_start + PUBLIC_KEY_BYTES_LEN;
+                    let pk_end = pk_start + PUBLIC_KEY_BYTES_LEN;
-                let pk = PublicKey::from_bytes(&buf[pk_start..pk_end])
+                    let pk = PublicKey::from_bytes(&buf[pk_start..pk_end])
-                    .map_err(|_| Error::new(ErrorKind::Other, "Invalid PublicKey bytes"))?;
+                        .map_err(|_| Error::new(ErrorKind::Other, "Invalid PublicKey bytes"))
+                        .unwrap();
 
-                keypairs.push(Keypair { sk, pk });
+                    Keypair { sk, pk }
-            }
+                })
+                .collect();
+
+            keypairs.append(&mut keypair_batch);
         }
 
         Ok(keypairs)
@@ -68,7 +77,6 @@ impl KeypairsFile for Vec<Keypair> {
 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> {

eth2/types/src/test_utils/mod.rs

@@ -14,7 +14,7 @@ pub use rand::{prng::XorShiftRng, SeedableRng};
 pub use test_random::TestRandom;
 pub use testing_attestation_builder::TestingAttestationBuilder;
 pub use testing_beacon_block_builder::TestingBeaconBlockBuilder;
-pub use testing_beacon_state_builder::TestingBeaconStateBuilder;
+pub use testing_beacon_state_builder::{keypairs_path, TestingBeaconStateBuilder};
 pub use testing_deposit_builder::TestingDepositBuilder;
 pub use testing_transfer_builder::TestingTransferBuilder;
 pub use testing_voluntary_exit_builder::TestingVoluntaryExitBuilder;

eth2/types/src/test_utils/testing_beacon_state_builder.rs

@@ -2,8 +2,22 @@ use super::{generate_deterministic_keypairs, KeypairsFile};
 use crate::beacon_state::BeaconStateBuilder;
 use crate::*;
 use bls::get_withdrawal_credentials;
+use dirs;
 use rayon::prelude::*;
-use std::path::Path;
+use std::path::{Path, PathBuf};
+
+pub const KEYPAIRS_FILE: &str = "keypairs.raw_keypairs";
+
+/// Returns the directory where the generated keypairs should be stored.
+///
+/// It is either `$HOME/.lighthouse/keypairs.raw_keypairs` or, if `$HOME` is not available,
+/// `./keypairs.raw_keypairs`.
+pub fn keypairs_path() -> PathBuf {
+    let dir = dirs::home_dir()
+        .and_then(|home| Some(home.join(".lighthouse")))
+        .unwrap_or_else(|| PathBuf::from(""));
+    dir.join(KEYPAIRS_FILE)
+}
 
 pub struct TestingBeaconStateBuilder {
     state: BeaconState,
@@ -11,11 +25,52 @@ pub struct TestingBeaconStateBuilder {
 }
 
 impl TestingBeaconStateBuilder {
-    pub fn new(validator_count: usize, keypairs_path: Option<&Path>, spec: &ChainSpec) -> Self {
+    /// Attempts to load validators from a file in the `CARGO_MANIFEST_DIR`. If the file is
-        let keypairs = match keypairs_path {
+    /// unavailable, it generates the keys at runtime.
-            None => generate_deterministic_keypairs(validator_count),
+    ///
-            Some(path) => Vec::from_raw_file(path, validator_count).unwrap(),
+    /// If the `CARGO_MANIFEST_DIR` environment variable is not set, the local directory is used.
-        };
+    ///
+    /// See the `Self::from_keypairs_file` method for more info.
+    ///
+    /// # Panics
+    ///
+    /// If the file does not contain enough keypairs or is invalid.
+    pub fn from_default_keypairs_file_if_exists(validator_count: usize, spec: &ChainSpec) -> Self {
+        let dir = dirs::home_dir()
+            .and_then(|home| Some(home.join(".lighthouse")))
+            .unwrap_or_else(|| PathBuf::from(""));
+        let file = dir.join(KEYPAIRS_FILE);
+
+        if file.exists() {
+            TestingBeaconStateBuilder::from_keypairs_file(validator_count, &file, spec)
+        } else {
+            TestingBeaconStateBuilder::from_deterministic_keypairs(validator_count, spec)
+        }
+    }
+
+    /// Loads the initial validator keypairs from a file on disk.
+    ///
+    /// Loading keypairs from file is ~10x faster than generating them. Use the `gen_keys` command
+    /// on the  `test_harness` binary to generate the keys. In the `test_harness` dir, run `cargo
+    /// run -- gen_keys -h` for help.
+    ///
+    /// # Panics
+    ///
+    /// If the file does not exist, is invalid or does not contain enough keypairs.
+    pub fn from_keypairs_file(validator_count: usize, path: &Path, spec: &ChainSpec) -> Self {
+        let keypairs = Vec::from_raw_file(path, validator_count).unwrap();
+        TestingBeaconStateBuilder::from_keypairs(keypairs, spec)
+    }
+
+    /// Generates the validator keypairs deterministically.
+    pub fn from_deterministic_keypairs(validator_count: usize, spec: &ChainSpec) -> Self {
+        let keypairs = generate_deterministic_keypairs(validator_count);
+        TestingBeaconStateBuilder::from_keypairs(keypairs, spec)
+    }
+
+    /// Creates the builder from an existing set of keypairs.
+    pub fn from_keypairs(keypairs: Vec<Keypair>, spec: &ChainSpec) -> Self {
+        let validator_count = keypairs.len();
 
         let validators = keypairs
             .par_iter()
@@ -61,10 +116,15 @@ impl TestingBeaconStateBuilder {
         }
     }
 
+    /// Consume the builder and return the `BeaconState` and the keypairs for each validator.
     pub fn build(self) -> (BeaconState, Vec<Keypair>) {
         (self.state, self.keypairs)
     }
 
+    /// Ensures that the state returned from `Self::build(..)` has all caches pre-built.
+    ///
+    /// Note: this performs the build when called. Ensure that no changes are made that would
+    /// invalidate this cache.
     pub fn build_caches(&mut self, spec: &ChainSpec) -> Result<(), BeaconStateError> {
         let state = &mut self.state;
 
@@ -147,6 +207,9 @@ impl TestingBeaconStateBuilder {
     }
 }
 
+/// Maps a committee to a `PendingAttestation`.
+///
+/// The committee will be signed by all validators in the committee.
 fn committee_to_pending_attestation(
     state: &BeaconState,
     committee: &[usize],