Add support for loading keypairs from file

This commit is contained in:
Paul Hauner
2019-03-12 12:46:44 +11:00
parent 9de6a0c733
commit f34ae86cde
19 changed files with 264 additions and 183 deletions

View File

@@ -1 +0,0 @@
validators/

View File

@@ -16,13 +16,11 @@ use std::fs::File;
use std::iter::FromIterator;
use std::path::Path;
use std::sync::Arc;
use types::{beacon_state::BeaconStateBuilder, *};
use types::{beacon_state::BeaconStateBuilder, test_utils::generate_deterministic_keypairs, *};
mod generate_deposits;
mod load_deposits_from_file;
pub use generate_deposits::{generate_deposits_from_keypairs, generate_deterministic_keypairs};
pub use load_deposits_from_file::load_deposits_from_file;
pub use generate_deposits::generate_deposits_from_keypairs;
/// The beacon chain harness simulates a single beacon node with `validator_count` validators connected
/// to it. Each validator is provided a borrow to the beacon chain, where it may read

View File

@@ -1,35 +1,8 @@
use bls::get_withdrawal_credentials;
use int_to_bytes::int_to_bytes48;
use log::debug;
use rayon::prelude::*;
use types::*;
/// 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
}
/// Generates a `Deposit` for each keypairs
pub fn generate_deposits_from_keypairs(
keypairs: &[Keypair],

View File

@@ -1,38 +0,0 @@
use log::debug;
use serde_yaml;
use std::fs::File;
use std::path::Path;
use types::*;
pub fn load_deposits_from_file(
validator_count: usize,
keypairs_path: &Path,
deposits_path: &Path,
) -> (Vec<Keypair>, Vec<Deposit>) {
debug!("Loading keypairs from file...");
let keypairs_file = File::open(keypairs_path).unwrap();
let mut keypairs: Vec<Keypair> = serde_yaml::from_reader(&keypairs_file).unwrap();
debug!("Loading deposits from file...");
let deposits_file = File::open(deposits_path).unwrap();
let mut deposits: Vec<Deposit> = serde_yaml::from_reader(&deposits_file).unwrap();
assert!(
keypairs.len() >= validator_count,
"Unable to load {} keypairs from file ({} available)",
validator_count,
keypairs.len()
);
assert!(
deposits.len() >= validator_count,
"Unable to load {} deposits from file ({} available)",
validator_count,
deposits.len()
);
keypairs.truncate(validator_count);
deposits.truncate(validator_count);
(keypairs, deposits)
}

View File

@@ -1,11 +1,11 @@
use clap::{App, Arg, SubCommand};
use env_logger::{Builder, Env};
use prepare::prepare;
use gen_keys::gen_keys;
use run_test::run_test;
use types::ChainSpec;
mod beacon_chain_harness;
mod prepare;
mod gen_keys;
mod run_test;
mod test_case;
mod validator_harness;
@@ -55,8 +55,8 @@ fn main() {
),
)
.subcommand(
SubCommand::with_name("prepare")
.about("Builds validator YAML files for faster tests.")
SubCommand::with_name("gen_keys")
.about("Builds a file of BLS keypairs for faster tests.")
.arg(
Arg::with_name("validator_count")
.long("validator_count")
@@ -66,20 +66,12 @@ fn main() {
.required(true),
)
.arg(
Arg::with_name("genesis_time")
.long("genesis_time")
.short("t")
.value_name("GENESIS_TIME")
.help("Time for validator deposits.")
.required(true),
)
.arg(
Arg::with_name("output_dir")
.long("output_dir")
Arg::with_name("output_file")
.long("output_file")
.short("d")
.value_name("GENESIS_TIME")
.help("Output directory for generated YAML.")
.default_value("validators"),
.default_value("keypairs.raw_keypairs"),
),
)
.get_matches();
@@ -88,7 +80,7 @@ fn main() {
Builder::from_env(Env::default().default_filter_or(log_level)).init();
}
let spec = match matches.value_of("spec") {
let _spec = match matches.value_of("spec") {
Some("foundation") => ChainSpec::foundation(),
Some("few_validators") => ChainSpec::few_validators(),
_ => unreachable!(), // Has a default value, should always exist.
@@ -98,7 +90,7 @@ fn main() {
run_test(matches);
}
if let Some(matches) = matches.subcommand_matches("prepare") {
prepare(matches, &spec);
if let Some(matches) = matches.subcommand_matches("gen_keys") {
gen_keys(matches);
}
}

View File

@@ -0,0 +1,20 @@
use clap::{value_t, ArgMatches};
use log::debug;
use std::path::Path;
use types::test_utils::{generate_deterministic_keypairs, KeypairsFile};
pub fn gen_keys(matches: &ArgMatches) {
let validator_count = value_t!(matches.value_of("validator_count"), usize)
.expect("Validator count is required argument");
let output_file = matches
.value_of("output_file")
.expect("Output file has a default value.");
let keypairs = generate_deterministic_keypairs(validator_count);
debug!("Writing keypairs to file...");
let keypairs_path = Path::new(output_file);
keypairs.to_raw_file(&keypairs_path, &keypairs).unwrap();
}

View File

@@ -1,69 +0,0 @@
use crate::beacon_chain_harness::generate_deterministic_keypairs;
use bls::get_withdrawal_credentials;
use clap::{value_t, ArgMatches};
use log::debug;
use serde_yaml;
use std::path::Path;
use std::{fs, fs::File};
use types::*;
const KEYPAIRS_FILE: &str = "keypairs.yaml";
const VALIDATORS_FILE: &str = "validators.yaml";
pub fn prepare(matches: &ArgMatches, spec: &ChainSpec) {
let validator_count = value_t!(matches.value_of("validator_count"), usize)
.expect("Validator count is required argument");
let output_dir = matches
.value_of("output_dir")
.expect("Output dir has a default value.");
debug!("Created keypairs and validators, writing to file...");
fs::create_dir_all(Path::new(output_dir)).unwrap();
// Ensure that keypairs is dropped before writing validators, this provides a big memory saving
// for large validator_counts.
let validators: Vec<Validator> = {
debug!("Creating {} keypairs...", validator_count);
let keypairs = generate_deterministic_keypairs(validator_count);
debug!("Writing {} keypairs to file...", validator_count);
write_keypairs(output_dir, &keypairs);
debug!("Creating {} validators...", validator_count);
keypairs
.iter()
.map(|keypair| generate_validator(&keypair, spec))
.collect()
};
debug!("Writing {} validators to file...", validator_count);
write_validators(output_dir, &validators);
}
fn generate_validator(keypair: &Keypair, spec: &ChainSpec) -> Validator {
let withdrawal_credentials = Hash256::from_slice(&get_withdrawal_credentials(
&keypair.pk,
spec.bls_withdrawal_prefix_byte,
));
Validator {
pubkey: keypair.pk.clone(),
withdrawal_credentials,
activation_epoch: spec.far_future_epoch,
exit_epoch: spec.far_future_epoch,
withdrawable_epoch: spec.far_future_epoch,
initiated_exit: false,
slashed: false,
}
}
fn write_keypairs(output_dir: &str, keypairs: &[Keypair]) {
let keypairs_path = Path::new(output_dir).join(KEYPAIRS_FILE);
let keypairs_file = File::create(keypairs_path).unwrap();
serde_yaml::to_writer(keypairs_file, &keypairs).unwrap();
}
fn write_validators(output_dir: &str, validators: &[Validator]) {
let validators_path = Path::new(output_dir).join(VALIDATORS_FILE);
let validators_file = File::create(validators_path).unwrap();
serde_yaml::to_writer(validators_file, &validators).unwrap();
}