From 00d6bffc0d91bae681d1298c901028f7e80a7ae2 Mon Sep 17 00:00:00 2001 From: Dan Cline <6798349+Rjected@users.noreply.github.com> Date: Fri, 12 Jan 2024 06:43:28 -0500 Subject: [PATCH] feat: add GenesisAllocator for easier genesis alloc creation (#6021) --- bin/reth/src/builder/mod.rs | 1 + crates/consensus/beacon/src/engine/mod.rs | 23 ++- crates/primitives/src/genesis.rs | 200 +++++++++++++++++++++- crates/primitives/src/lib.rs | 2 +- 4 files changed, 210 insertions(+), 16 deletions(-) diff --git a/bin/reth/src/builder/mod.rs b/bin/reth/src/builder/mod.rs index e6b2385afc..657615ab48 100644 --- a/bin/reth/src/builder/mod.rs +++ b/bin/reth/src/builder/mod.rs @@ -1452,6 +1452,7 @@ mod tests { // spawn_test_node takes roughly 1 second per node, so this test takes ~4 seconds let num_nodes = 4; + // this reserves instances 3-6 let starting_instance = 3; let mut handles = Vec::new(); for i in 0..num_nodes { diff --git a/crates/consensus/beacon/src/engine/mod.rs b/crates/consensus/beacon/src/engine/mod.rs index fe759d3cb1..adf00530aa 100644 --- a/crates/consensus/beacon/src/engine/mod.rs +++ b/crates/consensus/beacon/src/engine/mod.rs @@ -2345,11 +2345,8 @@ mod tests { mod new_payload { use super::*; - use reth_interfaces::test_utils::{ - generators, - generators::{generate_keys, random_block}, - }; - use reth_primitives::{public_key_to_address, Genesis, GenesisAccount, Hardfork, U256}; + use reth_interfaces::test_utils::{generators, generators::random_block}; + use reth_primitives::{Genesis, GenesisAllocator, Hardfork, U256}; use reth_provider::test_utils::blocks::BlockChainTestData; #[tokio::test] @@ -2453,14 +2450,14 @@ mod tests { #[tokio::test] async fn simple_validate_block() { let mut rng = generators::rng(); - let genesis_keys = generate_keys(&mut rng, 16); - let amount = 1000000000000000000u64; - let alloc = genesis_keys.iter().map(|pair| { - ( - public_key_to_address(pair.public_key()), - GenesisAccount::default().with_balance(U256::from(amount)), - ) - }); + let amount = U256::from(1000000000000000000u64); + let mut allocator = GenesisAllocator::default().with_rng(&mut rng); + for _ in 0..16 { + // add 16 new accounts + allocator.new_funded_account(amount); + } + + let alloc = allocator.build(); let genesis = Genesis::default().extend_accounts(alloc); diff --git a/crates/primitives/src/genesis.rs b/crates/primitives/src/genesis.rs index 404e9e75b6..e81f247510 100644 --- a/crates/primitives/src/genesis.rs +++ b/crates/primitives/src/genesis.rs @@ -1,5 +1,5 @@ use crate::{ - keccak256, + keccak256, public_key_to_address, serde_helper::{ json_u256::{deserialize_json_ttd_opt, deserialize_json_u256}, num::{u64_hex_or_decimal, u64_hex_or_decimal_opt}, @@ -7,8 +7,12 @@ use crate::{ }, Account, Address, Bytes, B256, U256, }; +use secp256k1::{ + rand::{thread_rng, RngCore}, + KeyPair, Secp256k1, +}; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; +use std::collections::{hash_map::Entry, HashMap}; /// The genesis block specification. #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] @@ -185,6 +189,198 @@ impl From for Account { } } +/// Helper trait that encapsulates [RngCore], and [Debug](std::fmt::Debug) to get around rules for +/// auto traits (Opt-in built-in traits). +trait RngDebug: RngCore + std::fmt::Debug {} +impl RngDebug for T where T: RngCore + std::fmt::Debug {} + +/// This helps create a custom genesis alloc by making it easy to add funded accounts with known +/// signers to the genesis block. +/// +/// # Example +/// ``` +/// # use reth_primitives::{GenesisAllocator, Address, U256, hex, Bytes}; +/// # use std::str::FromStr; +/// let mut allocator = GenesisAllocator::default(); +/// +/// // This will add a genesis account to the alloc builder, with the provided balance. The +/// // signer for the account will be returned. +/// let (_signer, _addr) = allocator.new_funded_account(U256::from(100_000_000_000_000_000u128)); +/// +/// // You can also provide code for the account. +/// let code = Bytes::from_str("0x1234").unwrap(); +/// let (_second_signer, _second_addr) = +/// allocator.new_funded_account_with_code(U256::from(100_000_000_000_000_000u128), code); +/// +/// // You can also add an account with a specific address. +/// // This will not return a signer, since the address is provided by the user and the signer +/// // may be unknown. +/// let addr = "0Ac1dF02185025F65202660F8167210A80dD5086".parse::
().unwrap(); +/// allocator.add_funded_account_with_address(addr, U256::from(100_000_000_000_000_000u128)); +/// +/// // Once you're done adding accounts, you can build the alloc. +/// let alloc = allocator.build(); +/// ``` +#[derive(Debug)] +pub struct GenesisAllocator<'a> { + /// The genesis alloc to be built. + alloc: HashMap, + /// The rng to use for generating key pairs. + rng: Box, +} + +impl<'a> GenesisAllocator<'a> { + /// Initialize a new alloc builder with the provided rng. + pub fn new_with_rng(rng: &'a mut R) -> Self + where + R: RngCore + std::fmt::Debug, + { + Self { alloc: HashMap::default(), rng: Box::new(rng) } + } + + /// Use the provided rng for generating key pairs. + pub fn with_rng(mut self, rng: &'a mut R) -> Self + where + R: RngCore + std::fmt::Debug, + { + self.rng = Box::new(rng); + self + } + + /// Add a funded account to the genesis alloc. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_funded_account(&mut self, balance: U256) -> (KeyPair, Address) { + let secp = Secp256k1::new(); + let pair = KeyPair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert(address, GenesisAccount::default().with_balance(balance)); + + (pair, address) + } + + /// Add a funded account to the genesis alloc with the provided code. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_funded_account_with_code( + &mut self, + balance: U256, + code: Bytes, + ) -> (KeyPair, Address) { + let secp = Secp256k1::new(); + let pair = KeyPair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc + .insert(address, GenesisAccount::default().with_balance(balance).with_code(Some(code))); + + (pair, address) + } + + /// Adds a funded account to the genesis alloc with the provided storage. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_funded_account_with_storage( + &mut self, + balance: U256, + storage: HashMap, + ) -> (KeyPair, Address) { + let secp = Secp256k1::new(); + let pair = KeyPair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert( + address, + GenesisAccount::default().with_balance(balance).with_storage(Some(storage)), + ); + + (pair, address) + } + + /// Adds an account with code and storage to the genesis alloc. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_account_with_code_and_storage( + &mut self, + code: Bytes, + storage: HashMap, + ) -> (KeyPair, Address) { + let secp = Secp256k1::new(); + let pair = KeyPair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert( + address, + GenesisAccount::default().with_code(Some(code)).with_storage(Some(storage)), + ); + + (pair, address) + } + + /// Adds an account with code to the genesis alloc. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_account_with_code(&mut self, code: Bytes) -> (KeyPair, Address) { + let secp = Secp256k1::new(); + let pair = KeyPair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert(address, GenesisAccount::default().with_code(Some(code))); + + (pair, address) + } + + /// Add a funded account to the genesis alloc with the provided address. + /// + /// Neither the key pair nor the account will be returned, since the address is provided by the + /// user and the signer may be unknown. + pub fn add_funded_account_with_address(&mut self, address: Address, balance: U256) { + self.alloc.insert(address, GenesisAccount::default().with_balance(balance)); + } + + /// Adds the given [GenesisAccount] to the genesis alloc. + /// + /// Returns the key pair for the account and the account's address. + pub fn add_account(&mut self, account: GenesisAccount) -> Address { + let secp = Secp256k1::new(); + let pair = KeyPair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert(address, account); + + address + } + + /// Gets the account for the provided address. + /// + /// If it does not exist, this returns `None`. + pub fn get_account(&self, address: &Address) -> Option<&GenesisAccount> { + self.alloc.get(address) + } + + /// Gets a mutable version of the account for the provided address, if it exists. + pub fn get_account_mut(&mut self, address: &Address) -> Option<&mut GenesisAccount> { + self.alloc.get_mut(address) + } + + /// Gets an [Entry] for the provided address. + pub fn account_entry(&mut self, address: Address) -> Entry<'_, Address, GenesisAccount> { + self.alloc.entry(address) + } + + /// Build the genesis alloc. + pub fn build(self) -> HashMap { + self.alloc + } +} + +impl Default for GenesisAllocator<'_> { + fn default() -> Self { + Self { alloc: HashMap::default(), rng: Box::new(thread_rng()) } + } +} + /// Defines core blockchain settings per block. /// /// Tailors unique settings for each network based on its genesis block. diff --git a/crates/primitives/src/lib.rs b/crates/primitives/src/lib.rs index 12117fecc8..291d7122de 100644 --- a/crates/primitives/src/lib.rs +++ b/crates/primitives/src/lib.rs @@ -59,7 +59,7 @@ pub use constants::{ KECCAK_EMPTY, MAINNET_GENESIS_HASH, SEPOLIA_GENESIS_HASH, }; pub use error::{GotExpected, GotExpectedBoxed}; -pub use genesis::{ChainConfig, Genesis, GenesisAccount}; +pub use genesis::{ChainConfig, Genesis, GenesisAccount, GenesisAllocator}; pub use header::{Header, HeadersDirection, SealedHeader}; pub use integer_list::IntegerList; pub use log::{logs_bloom, Log};