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
No known key found for this signature in database
GPG Key ID: D362883A9218FCC6
19 changed files with 264 additions and 183 deletions

1
.gitignore vendored
View File

@ -3,3 +3,4 @@ target/
Cargo.lock
*.pk
*.sk
*.raw_keypairs

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();
}

View File

@ -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();

View File

@ -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);

View File

@ -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);

View File

@ -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);

View File

@ -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);

View File

@ -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

View File

@ -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
}

View File

@ -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);
}
}

View File

@ -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;

View File

@ -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()

View File

@ -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
}