Apply consensus committee changes at epoch boundaries

The most important changes in this commit are in `deposit.sol`.
Everything else is pretty much just supporting the new
interfaces in the deposit contract. The key feature is that
we now keep track of the current 'epoch' and only apply changes
to the consensus committee two epochs after they are made.

There are a few other changes to the deposit contract in this
commit too:
* The initial stakers are now passed to the contract constructor,
rather than relying on the `setStake` hack.
* `leader()` was removed since it isn't used anywhere and it
doesn't make sense to have both `leader` and `leaderAtView`.
* Delayed withdrawals are now supported - We keep a circular
buffer of pending withdrawals for each validator. 2 weeks after
unstaking, the validator can make another call to withdraw the
funds from the deposit contract.

Author: JamesHinshelwood

z2/src/converter.rs

@@ -319,47 +319,6 @@ pub async fn convert_persistence(
         }
     }
 
-    // Add stake for this validator. For now, we just assume they've always had 64 ZIL staked.
-    // This assumptions will need to change for the actual testnet and mainnet launches, where we cannot invent ZIL
-    // out of thin air (like we do below).
-    let data = contracts::deposit::SET_STAKE.encode_input(&[
-        Token::Bytes(secret_key.node_public_key().as_bytes()),
-        Token::Bytes(
-            secret_key
-                .to_libp2p_keypair()
-                .public()
-                .to_peer_id()
-                .to_bytes(),
-        ),
-        Token::Address(ethabi::Address::from_low_u64_be(1)),
-        Token::Uint((64 * 10u128.pow(18)).into()),
-    ])?;
-    let (
-        ResultAndState {
-            result,
-            state: result_state,
-        },
-        ..,
-    ) = state.apply_transaction_evm(
-        Address::ZERO,
-        Some(contract_addr::DEPOSIT),
-        0,
-        node_config.consensus.eth_block_gas_limit,
-        0,
-        data,
-        None,
-        BlockHeader::default(),
-        inspector::noop(),
-        BaseFeeCheck::Ignore,
-    )?;
-    if !result.is_success() {
-        return Err(anyhow!("setting stake failed: {result:?}"));
-    }
-    state.apply_delta_evm(&result_state)?;
-
-    // Flush any pending changes to db
-    let _ = state.root_hash()?;
-
     if !convert_blocks {
         println!("Accounts converted. Skipping blocks.");
         return Ok(());

z2/src/setup.rs

@@ -9,7 +9,6 @@ use alloy::{
 };
 use anyhow::{anyhow, Context, Result};
 use k256::ecdsa::SigningKey;
-use libp2p::PeerId;
 use serde::{Deserialize, Serialize};
 use serde_yaml;
 use tera::Tera;
@@ -26,9 +25,8 @@ use zilliqa::{
         local_address_default, max_blocks_in_flight_default,
         minimum_time_left_for_empty_block_default, scilla_address_default,
         scilla_ext_libs_path_default, scilla_stdlib_dir_default, state_rpc_limit_default,
-        total_native_token_supply_default, Amount, ConsensusConfig,
+        total_native_token_supply_default, Amount, ConsensusConfig, GenesisDeposit,
     },
-    crypto::NodePublicKey,
     transaction::EvmGas,
 };
 
@@ -281,7 +279,7 @@ impl Setup {
 
     pub async fn generate_config(&self) -> Result<()> {
         // The genesis deposits.
-        let mut genesis_deposits: Vec<(NodePublicKey, PeerId, Amount, Address)> = Vec::new();
+        let mut genesis_deposits: Vec<GenesisDeposit> = Vec::new();
         for (node, desc) in self.config.shape.nodes.iter() {
             if desc.is_validator {
                 let data = self
@@ -291,12 +289,13 @@ impl Setup {
                     .ok_or(anyhow!("no node data for {node}"))?;
                 // Better have a genesis deposit.
                 let secret_key = SecretKey::from_hex(&data.secret_key)?;
-                genesis_deposits.push((
-                    secret_key.node_public_key(),
-                    secret_key.to_libp2p_keypair().public().to_peer_id(),
-                    GENESIS_DEPOSIT.into(),
-                    data.address,
-                ))
+                genesis_deposits.push(GenesisDeposit {
+                    public_key: secret_key.node_public_key(),
+                    peer_id: secret_key.to_libp2p_keypair().public().to_peer_id(),
+                    stake: GENESIS_DEPOSIT.into(),
+                    reward_address: data.address,
+                    control_address: data.address,
+                });
             }
         }
 

zilliqa-macros/src/test.rs

@@ -7,6 +7,7 @@ use syn::{parse::Parser, ItemFn};
 pub(crate) fn test_macro(args: TokenStream, item: TokenStream) -> TokenStream {
     let mut restrict_concurrency = false;
     let mut do_checkpoints = false;
+    let mut blocks_per_epoch = 10;
 
     let parsed_args =
         match syn::punctuated::Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated
@@ -35,6 +36,41 @@ pub(crate) fn test_macro(args: TokenStream, item: TokenStream) -> TokenStream {
                     }
                 }
             }
+            syn::Meta::NameValue(a) => {
+                let Some(name) = a.path.get_ident() else {
+                    return token_stream_with_error(
+                        args,
+                        syn::Error::new_spanned(a.path, "Attribute parameter must be ident"),
+                    );
+                };
+                match name.to_string().as_str() {
+                    "blocks_per_epoch" => {
+                        let syn::Expr::Lit(syn::ExprLit {
+                            lit: syn::Lit::Int(val),
+                            ..
+                        }) = a.value
+                        else {
+                            return token_stream_with_error(
+                                args,
+                                syn::Error::new_spanned(
+                                    a.value,
+                                    "Attribute parameter value must be an int",
+                                ),
+                            );
+                        };
+                        match val.base10_parse::<u64>() {
+                            Ok(val) => blocks_per_epoch = val,
+                            Err(e) => return token_stream_with_error(args, e),
+                        };
+                    }
+                    _ => {
+                        return token_stream_with_error(
+                            args,
+                            syn::Error::new_spanned(a, "Unknown attribute"),
+                        )
+                    }
+                }
+            }
             // Can match syn::Meta::Namevalue(a) here for args with params, e.g. to set node_count
             // in the future
             other => {
@@ -199,7 +235,7 @@ pub(crate) fn test_macro(args: TokenStream, item: TokenStream) -> TokenStream {
                     async move {
                         let mut rng = <rand_chacha::ChaCha8Rng as rand_core::SeedableRng>::seed_from_u64(seed);
                         let network = crate::Network::new(std::sync::Arc::new(std::sync::Mutex::new(rng)), 4, seed, format!("http://{addr}"),
-                                                          scilla_stdlib_dir.to_string(), #do_checkpoints);
+                                                          scilla_stdlib_dir.to_string(), #do_checkpoints, #blocks_per_epoch);
 
                         // Call the original test function, wrapped in `catch_unwind` so we can detect the panic.
                         let result = futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(

zilliqa/src/cfg.rs

@@ -233,7 +233,7 @@ pub struct ConsensusConfig {
     /// The initially staked deposits in the deposit contract at genesis, composed of
     /// (public key, peerId, amount, reward address) tuples.
     #[serde(default)]
-    pub genesis_deposits: Vec<(NodePublicKey, PeerId, Amount, Address)>,
+    pub genesis_deposits: Vec<GenesisDeposit>,
     /// Accounts that will be pre-funded at genesis.
     #[serde(default)]
     pub genesis_accounts: Vec<(Address, Amount)>,
@@ -281,6 +281,16 @@ pub struct ConsensusConfig {
     pub total_native_token_supply: Amount,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct GenesisDeposit {
+    pub public_key: NodePublicKey,
+    pub peer_id: PeerId,
+    pub stake: Amount,
+    pub reward_address: Address,
+    pub control_address: Address,
+}
+
 pub fn consensus_timeout_default() -> Duration {
     Duration::from_secs(5)
 }

zilliqa/src/consensus.rs

@@ -950,6 +950,7 @@ impl Consensus {
             );
             return Ok(None);
         }
+
         let committee = self.state.get_stakers_at_block(&block)?;
 
         // verify the sender's signature on block_hash
@@ -2007,42 +2008,40 @@ impl Consensus {
             }
         }
 
-        if self.block_is_first_in_epoch(block.number()) && !block.is_genesis() {
-            // TODO: handle epochs (#1140)
-
-            if self.config.do_checkpoints
-                && self.epoch_is_checkpoint(self.epoch_number(block.number()))
-            {
-                if let Some(checkpoint_path) = self.db.get_checkpoint_dir()? {
-                    let parent = self
-                        .db
-                        .get_block_by_hash(block.parent_hash())?
-                        .ok_or(anyhow!(
-                            "Trying to checkpoint block, but we don't have its parent"
+        if self.block_is_first_in_epoch(block.number())
+            && !block.is_genesis()
+            && self.config.do_checkpoints
+            && self.epoch_is_checkpoint(self.epoch_number(block.number()))
+        {
+            if let Some(checkpoint_path) = self.db.get_checkpoint_dir()? {
+                let parent = self
+                    .db
+                    .get_block_by_hash(block.parent_hash())?
+                    .ok_or(anyhow!(
+                        "Trying to checkpoint block, but we don't have its parent"
+                    ))?;
+                let transactions: Vec<SignedTransaction> = block
+                    .transactions
+                    .iter()
+                    .map(|txn_hash| {
+                        let tx = self.db.get_transaction(txn_hash)?.ok_or(anyhow!(
+                            "failed to fetch transaction {} for checkpoint parent {}",
+                            txn_hash,
+                            parent.hash()
                         ))?;
-                    let transactions: Vec<SignedTransaction> = block
-                        .transactions
-                        .iter()
-                        .map(|txn_hash| {
-                            let tx = self.db.get_transaction(txn_hash)?.ok_or(anyhow!(
-                                "failed to fetch transaction {} for checkpoint parent {}",
-                                txn_hash,
-                                parent.hash()
-                            ))?;
-                            Ok::<_, anyhow::Error>(tx)
-                        })
-                        .collect::<Result<Vec<SignedTransaction>>>()?;
-
-                    self.message_sender.send_message_to_coordinator(
-                        InternalMessage::ExportBlockCheckpoint(
-                            Box::new(block),
-                            transactions,
-                            Box::new(parent),
-                            self.db.state_trie()?.clone(),
-                            checkpoint_path,
-                        ),
-                    )?;
-                }
+                        Ok::<_, anyhow::Error>(tx)
+                    })
+                    .collect::<Result<Vec<SignedTransaction>>>()?;
+
+                self.message_sender.send_message_to_coordinator(
+                    InternalMessage::ExportBlockCheckpoint(
+                        Box::new(block),
+                        transactions,
+                        Box::new(parent),
+                        self.db.state_trie()?.clone(),
+                        checkpoint_path,
+                    ),
+                )?;
             }
         }
 
@@ -2463,7 +2462,11 @@ impl Consensus {
             return None;
         };
 
-        let public_key = state_at.leader(view).unwrap();
+        let executed_block = BlockHeader {
+            number: block.header.number + 1,
+            ..Default::default()
+        };
+        let public_key = state_at.leader(view, executed_block).unwrap();
         let peer_id = state_at.get_peer_id(public_key).unwrap().unwrap();
 
         Some(Validator {

zilliqa/src/contracts/compiled.json

@@ -14,24 +14,36 @@ pub mod deposit {
     pub static BYTECODE: Lazy<Vec<u8>> = Lazy::new(|| CONTRACT.bytecode.clone());
     pub static LEADER_AT_VIEW: Lazy<Function> =
         Lazy::new(|| CONTRACT.abi.function("leaderAtView").unwrap().clone());
-    pub static TEMP_REMOVE_STAKER: Lazy<Function> =
-        Lazy::new(|| CONTRACT.abi.function("tempRemoveStaker").unwrap().clone());
     pub static DEPOSIT: Lazy<Function> =
         Lazy::new(|| CONTRACT.abi.function("deposit").unwrap().clone());
-    pub static SET_STAKE: Lazy<Function> =
-        Lazy::new(|| CONTRACT.abi.function("setStake").unwrap().clone());
+    pub static DEPOSIT_TOPUP: Lazy<Function> =
+        Lazy::new(|| CONTRACT.abi.function("depositTopup").unwrap().clone());
+    pub static UNSTAKE: Lazy<Function> =
+        Lazy::new(|| CONTRACT.abi.function("unstake").unwrap().clone());
+    pub static WITHDRAW_FUNDS: Lazy<Function> =
+        Lazy::new(|| CONTRACT.abi.function("withdrawFunds").unwrap().clone());
+    pub static EJECT: Lazy<Function> =
+        Lazy::new(|| CONTRACT.abi.function("eject").unwrap().clone());
+    pub static CURRENT_EPOCH: Lazy<Function> =
+        Lazy::new(|| CONTRACT.abi.function("currentEpoch").unwrap().clone());
     pub static GET_STAKE: Lazy<Function> =
         Lazy::new(|| CONTRACT.abi.function("getStake").unwrap().clone());
     pub static GET_REWARD_ADDRESS: Lazy<Function> =
         Lazy::new(|| CONTRACT.abi.function("getRewardAddress").unwrap().clone());
     pub static GET_PEER_ID: Lazy<Function> =
         Lazy::new(|| CONTRACT.abi.function("getPeerId").unwrap().clone());
+    pub static GET_STAKERS_MAP: Lazy<Function> =
+        Lazy::new(|| CONTRACT.abi.function("getStakersMap").unwrap().clone());
     pub static GET_STAKERS: Lazy<Function> =
         Lazy::new(|| CONTRACT.abi.function("getStakers").unwrap().clone());
     pub static TOTAL_STAKE: Lazy<Function> =
         Lazy::new(|| CONTRACT.abi.function("totalStake").unwrap().clone());
     pub static MIN_DEPOSIT: Lazy<Function> =
-        Lazy::new(|| CONTRACT.abi.function("_minimumStake").unwrap().clone());
+        Lazy::new(|| CONTRACT.abi.function("minimumStake").unwrap().clone());
+    pub static COMMITTEE: Lazy<Function> =
+        Lazy::new(|| CONTRACT.abi.function("committee").unwrap().clone());
+    pub static INTERNAL_COMMITTEE: Lazy<Function> =
+        Lazy::new(|| CONTRACT.abi.function("internalCommittee").unwrap().clone());
 }
 
 pub mod shard {

zilliqa/src/contracts/deposit.sol

@@ -422,9 +422,7 @@ impl State {
                 self.apply_delta_evm(&state)?;
                 Ok(addr)
             }
-            ExecutionResult::Success { .. } => {
-                Err(anyhow!("deployment did not create a transaction"))
-            }
+            ExecutionResult::Success { .. } => Err(anyhow!("deployment did not create a contract")),
             ExecutionResult::Revert { .. } => Err(anyhow!("deployment reverted")),
             ExecutionResult::Halt { reason, .. } => Err(anyhow!("deployment halted: {reason:?}")),
         }
@@ -715,15 +713,15 @@ impl State {
         Ok(())
     }
 
-    pub fn leader(&self, view: u64) -> Result<NodePublicKey> {
+    pub fn leader(&self, view: u64, current_block: BlockHeader) -> Result<NodePublicKey> {
         let data = contracts::deposit::LEADER_AT_VIEW.encode_input(&[Token::Uint(view.into())])?;
 
         let leader = self.call_contract(
             Address::ZERO,
             Some(contract_addr::DEPOSIT),
             data,
             0,
-            BlockHeader::default(),
+            current_block,
         )?;
 
         NodePublicKey::from_bytes(
@@ -754,7 +752,7 @@ impl State {
             0,
             // The current block is not accessed when the native balance is read, so we just pass in some
             // dummy values.
-            BlockHeader::default(),
+            BlockHeader::default(), // TODO
         )?;
 
         let stakers = contracts::deposit::GET_STAKERS
@@ -770,6 +768,22 @@ impl State {
             .collect()
     }
 
+    pub fn committee(&self) -> Result<()> {
+        let data = contracts::deposit::COMMITTEE.encode_input(&[])?;
+
+        let committee = self.call_contract(
+            Address::ZERO,
+            Some(contract_addr::DEPOSIT),
+            data,
+            0,
+            BlockHeader::default(),
+        )?;
+        let committee = contracts::deposit::COMMITTEE.decode_output(&committee)?;
+        info!("committee: {committee:?}");
+
+        Ok(())
+    }
+
     pub fn get_stake(&self, public_key: NodePublicKey) -> Result<Option<NonZeroU128>> {
         let data =
             contracts::deposit::GET_STAKE.encode_input(&[Token::Bytes(public_key.as_bytes())])?;
@@ -781,7 +795,7 @@ impl State {
             0,
             // The current block is not accessed when the native balance is read, so we just pass in some
             // dummy values.
-            BlockHeader::default(),
+            BlockHeader::default(), // TODO
         )?;
 
         let stake = NonZeroU128::new(U256::from_be_slice(&stake).to());
@@ -844,7 +858,7 @@ impl State {
             0,
             // The current block is not accessed when the native balance is read, so we just pass in some
             // dummy values.
-            BlockHeader::default(),
+            BlockHeader::default(), // TODO
         )?;
 
         let amount = contracts::deposit::TOTAL_STAKE.decode_output(&return_value)?[0]

zilliqa/src/contracts/mod.rs

@@ -11,7 +11,6 @@ use alloy::{
 use anyhow::{anyhow, Result};
 use eth_trie::{EthTrie as PatriciaTrie, Trie};
 use ethabi::Token;
-use revm::primitives::ResultAndState;
 use serde::{Deserialize, Serialize};
 use sha3::{Digest, Keccak256};
 
@@ -20,9 +19,7 @@ use crate::{
     cfg::{Amount, NodeConfig, ScillaExtLibsPath},
     contracts, crypto,
     db::TrieStorage,
-    exec::BaseFeeCheck,
-    inspector,
-    message::{BlockHeader, MAX_COMMITTEE_SIZE},
+    message::MAX_COMMITTEE_SIZE,
     node::ChainId,
     scilla::{ParamValue, Scilla, Transition},
     serde_util::vec_param_value,
@@ -130,12 +127,13 @@ impl State {
                     .fold(0, |acc, item: &(Address, Amount)| acc + item.1 .0),
             )
             .expect("Genesis accounts sum to more than total native token supply")
-            .checked_sub(config.consensus.genesis_deposits.iter().fold(
-                0,
-                |acc, item: &(crypto::NodePublicKey, libp2p::PeerId, Amount, Address)| {
-                    acc + item.2 .0
-                },
-            ))
+            .checked_sub(
+                config
+                    .consensus
+                    .genesis_deposits
+                    .iter()
+                    .fold(0, |acc, item| acc + item.stake.0),
+            )
             .expect(
                 "Genesis accounts + genesis deposits sum to more than total native token supply",
             );
@@ -151,46 +149,69 @@ impl State {
             })?;
         }
 
+        let initial_stakers: Vec<_> = config
+            .consensus
+            .genesis_deposits
+            .into_iter()
+            .map(|deposit| {
+                Token::Tuple(vec![
+                    Token::Bytes(deposit.public_key.as_bytes()),
+                    Token::Bytes(deposit.peer_id.to_bytes()),
+                    Token::Address(ethabi::Address::from(deposit.reward_address.into_array())),
+                    Token::Address(ethabi::Address::from(deposit.control_address.into_array())),
+                    Token::Uint((*deposit.stake).into()),
+                ])
+            })
+            .collect();
         let deposit_data = contracts::deposit::CONSTRUCTOR.encode_input(
             contracts::deposit::BYTECODE.to_vec(),
             &[
                 Token::Uint((*config.consensus.minimum_stake).into()),
                 Token::Uint(MAX_COMMITTEE_SIZE.into()),
+                Token::Uint(config.consensus.blocks_per_epoch.into()),
+                Token::Array(initial_stakers),
             ],
         )?;
-
         state.force_deploy_contract_evm(deposit_data, Some(contract_addr::DEPOSIT))?;
 
-        for (pub_key, peer_id, stake, reward_address) in config.consensus.genesis_deposits {
-            let data = contracts::deposit::SET_STAKE.encode_input(&[
-                Token::Bytes(pub_key.as_bytes()),
-                Token::Bytes(peer_id.to_bytes()),
-                Token::Address(ethabi::Address::from(reward_address.into_array())),
-                Token::Uint((*stake).into()),
-            ])?;
-            let (
-                ResultAndState {
-                    result,
-                    state: result_state,
-                },
-                ..,
-            ) = state.apply_transaction_evm(
-                Address::ZERO,
-                Some(contract_addr::DEPOSIT),
-                0,
-                config.consensus.eth_block_gas_limit,
-                0,
-                data,
-                None,
-                BlockHeader::default(),
-                inspector::noop(),
-                BaseFeeCheck::Ignore,
-            )?;
-            if !result.is_success() {
-                return Err(anyhow!("setting stake failed: {result:?}"));
-            }
-            state.apply_delta_evm(&result_state)?;
-        }
+        //for GenesisDeposit {
+        //    public_key,
+        //    peer_id,
+        //    stake,
+        //    reward_address,
+        //    control_address,
+        //} in config.consensus.genesis_deposits
+        //{
+        //    let data = contracts::deposit::SET_STAKE.encode_input(&[
+        //        Token::Bytes(public_key.as_bytes()),
+        //        Token::Bytes(peer_id.to_bytes()),
+        //        Token::Address(ethabi::Address::from(reward_address.into_array())),
+        //        Token::Address(ethabi::Address::from(control_address.into_array())),
+        //        Token::Uint((*stake).into()),
+        //    ])?;
+        //    let (
+        //        ResultAndState {
+        //            result,
+        //            state: result_state,
+        //        },
+        //        ..,
+        //    ) = state.apply_transaction_evm(
+        //        Address::ZERO,
+        //        Some(contract_addr::DEPOSIT),
+        //        0,
+        //        config.consensus.eth_block_gas_limit,
+        //        0,
+        //        data,
+        //        None,
+        //        BlockHeader::default(),
+        //        inspector::noop(),
+        //        BaseFeeCheck::Ignore,
+        //    )?;
+        //    if !result.is_success() {
+        //        return Err(anyhow!("setting stake failed: {result:?}"));
+        //    }
+        //    state.apply_delta_evm(&result_state)?;
+        //}
 
         Ok(state)
     }

zilliqa/src/exec.rs

@@ -152,6 +152,7 @@ async fn create_shard(
         network.scilla_address.clone(),
         network.scilla_stdlib_dir.clone(),
         false,
+        2,
     );
     let shard_wallet = shard_network.genesis_wallet().await;
 
@@ -658,7 +659,7 @@ async fn zero_account_per_block_balance_updates(mut network: Network) {
         .consensus
         .genesis_deposits
         .clone();
-    let total_staked: u128 = genesis_deposits[0].2 .0 * 4;
+    let total_staked: u128 = genesis_deposits[0].stake.0 * 4;
 
     // Zero account balance plus genesis account plus initial stakes should equal total_native_token_supply
     let zero_account_balance: u128 = wallet

zilliqa/src/state.rs

@@ -69,9 +69,9 @@ use zilliqa::{
         max_blocks_in_flight_default, minimum_time_left_for_empty_block_default,
         scilla_address_default, scilla_ext_libs_path_default, scilla_stdlib_dir_default,
         state_rpc_limit_default, total_native_token_supply_default, Amount, Checkpoint,
-        ConsensusConfig, NodeConfig,
+        ConsensusConfig, GenesisDeposit, NodeConfig,
     },
-    crypto::{NodePublicKey, SecretKey, TransactionPublicKey},
+    crypto::{SecretKey, TransactionPublicKey},
     db,
     message::{ExternalMessage, InternalMessage},
     node::{Node, RequestId},
@@ -209,7 +209,7 @@ struct TestNode {
 }
 
 struct Network {
-    pub genesis_deposits: Vec<(NodePublicKey, PeerId, Amount, Address)>,
+    pub genesis_deposits: Vec<GenesisDeposit>,
     /// Child shards.
     pub children: HashMap<u64, Network>,
     pub shard_id: u64,
@@ -238,9 +238,12 @@ struct Network {
     scilla_address: String,
     scilla_stdlib_dir: String,
     do_checkpoints: bool,
+    blocks_per_epoch: u64,
 }
 
 impl Network {
+    // This is only used in the zilliqa_macros::test macro. Consider refactoring this to a builder
+    // or removing entirely (and calling new_shard there)?
     /// Create a main shard network with reasonable defaults.
     pub fn new(
         rng: Arc<Mutex<ChaCha8Rng>>,
@@ -249,6 +252,7 @@ impl Network {
         scilla_address: String,
         scilla_stdlib_dir: String,
         do_checkpoints: bool,
+        blocks_per_epoch: u64,
     ) -> Network {
         Self::new_shard(
             rng,
@@ -260,6 +264,7 @@ impl Network {
             scilla_address,
             scilla_stdlib_dir,
             do_checkpoints,
+            blocks_per_epoch,
         )
     }
 
@@ -274,6 +279,7 @@ impl Network {
         scilla_address: String,
         scilla_stdlib_dir: String,
         do_checkpoints: bool,
+        blocks_per_epoch: u64,
     ) -> Network {
         let mut signing_keys = keys.unwrap_or_else(|| {
             (0..nodes)
@@ -296,13 +302,13 @@ impl Network {
         let stake = 32_000_000_000_000_000_000u128;
         let genesis_deposits: Vec<_> = keys
             .iter()
-            .map(|k| {
-                (
-                    k.0.node_public_key(),
-                    k.0.to_libp2p_keypair().public().to_peer_id(),
-                    stake.into(),
-                    TransactionPublicKey::Ecdsa(*k.1.verifying_key(), true).into_addr(),
-                )
+            .map(|k| GenesisDeposit {
+                public_key: k.0.node_public_key(),
+                peer_id: k.0.to_libp2p_keypair().public().to_peer_id(),
+                stake: stake.into(),
+                reward_address: TransactionPublicKey::Ecdsa(*k.1.verifying_key(), true).into_addr(),
+                control_address: TransactionPublicKey::Ecdsa(*k.1.verifying_key(), true)
+                    .into_addr(),
             })
             .collect();
 
@@ -326,7 +332,7 @@ impl Network {
                 eth_block_gas_limit: EvmGas(84000000),
                 gas_price: 4_761_904_800_000u128.into(),
                 main_shard_id: None,
-                blocks_per_epoch: 10,
+                blocks_per_epoch,
                 epochs_per_checkpoint: 1,
                 total_native_token_supply: total_native_token_supply_default(),
             },
@@ -398,6 +404,7 @@ impl Network {
             genesis_key,
             scilla_address,
             do_checkpoints,
+            blocks_per_epoch,
             scilla_stdlib_dir,
         }
     }
@@ -444,7 +451,7 @@ impl Network {
                 main_shard_id: None,
                 minimum_time_left_for_empty_block: minimum_time_left_for_empty_block_default(),
                 scilla_address: scilla_address_default(),
-                blocks_per_epoch: 10,
+                blocks_per_epoch: self.blocks_per_epoch,
                 epochs_per_checkpoint: 1,
                 scilla_stdlib_dir: scilla_stdlib_dir_default(),
                 scilla_ext_libs_path: scilla_ext_libs_path_default(),
@@ -493,13 +500,13 @@ impl Network {
         let stake = 32_000_000_000_000_000_000u128;
         let genesis_deposits: Vec<_> = keys
             .iter()
-            .map(|k| {
-                (
-                    k.0.node_public_key(),
-                    k.0.to_libp2p_keypair().public().to_peer_id(),
-                    stake.into(),
-                    TransactionPublicKey::Ecdsa(*k.1.verifying_key(), true).into_addr(),
-                )
+            .map(|k| GenesisDeposit {
+                public_key: k.0.node_public_key(),
+                peer_id: k.0.to_libp2p_keypair().public().to_peer_id(),
+                stake: stake.into(),
+                reward_address: TransactionPublicKey::Ecdsa(*k.1.verifying_key(), true).into_addr(),
+                control_address: TransactionPublicKey::Ecdsa(*k.1.verifying_key(), true)
+                    .into_addr(),
             })
             .collect();
 
@@ -551,7 +558,7 @@ impl Network {
                         minimum_time_left_for_empty_block:
                             minimum_time_left_for_empty_block_default(),
                         scilla_address: scilla_address_default(),
-                        blocks_per_epoch: 10,
+                        blocks_per_epoch: self.blocks_per_epoch,
                         epochs_per_checkpoint: 1,
                         scilla_stdlib_dir: scilla_stdlib_dir_default(),
                         scilla_ext_libs_path: scilla_ext_libs_path_default(),
@@ -867,6 +874,7 @@ impl Network {
                                     self.scilla_address.clone(),
                                     self.scilla_stdlib_dir.clone(),
                                     self.do_checkpoints,
+                                    self.blocks_per_epoch,
                                 ),
                             );
                         }

zilliqa/tests/it/consensus.rs

@@ -14,13 +14,13 @@ use zilliqa::{
         failed_request_sleep_duration_default, json_rpc_port_default, max_blocks_in_flight_default,
         minimum_time_left_for_empty_block_default, scilla_address_default,
         scilla_ext_libs_path_default, scilla_stdlib_dir_default, state_rpc_limit_default,
-        total_native_token_supply_default, Checkpoint,
+        total_native_token_supply_default, Checkpoint, ConsensusConfig, NodeConfig,
     },
     crypto::{Hash, SecretKey},
     transaction::EvmGas,
 };
 
-use crate::{deploy_contract, ConsensusConfig, Network, NewNodeOptions, NodeConfig, TestNode};
+use crate::{deploy_contract, Network, NewNodeOptions, TestNode};
 
 #[zilliqa_macros::test]
 async fn block_and_tx_data_persistence(mut network: Network) {
@@ -75,6 +75,7 @@ async fn block_and_tx_data_persistence(mut network: Network) {
     let block_with_tx = inner.get_block(receipt.block_hash).unwrap().unwrap();
     let last_block = inner.get_block(last_number).unwrap().unwrap();
     let tx = inner.get_transaction_by_hash(hash).unwrap().unwrap();
+
     // sanity check
     assert_eq!(tx.hash, hash);
     assert_eq!(block_with_tx.transactions.len(), 1);

zilliqa/tests/it/main.rs

@@ -9,7 +9,7 @@ use ethers::{
     types::{BlockId, BlockNumber, TransactionRequest},
 };
 use libp2p::PeerId;
-use primitive_types::H160;
+use primitive_types::{H160, H256};
 use rand::Rng;
 use tracing::{info, trace};
 use zilliqa::{contracts, crypto::NodePublicKey, state::contract_addr};
@@ -36,17 +36,21 @@ async fn check_miner_got_reward(
 
 async fn deposit_stake(
     network: &mut Network,
-    wallet: &SignerMiddleware<Provider<LocalRpcClient>, LocalWallet>,
+    control_wallet: &SignerMiddleware<Provider<LocalRpcClient>, LocalWallet>,
     key: NodePublicKey,
     peer_id: PeerId,
     stake: u128,
     reward_address: H160,
     pop: blsful::ProofOfPossession<Bls12381G2Impl>,
-) {
+) -> H256 {
     // Transfer the new validator enough ZIL to stake.
     let tx = TransactionRequest::pay(reward_address, stake);
-    let hash = wallet.send_transaction(tx, None).await.unwrap().tx_hash();
-    network.run_until_receipt(wallet, hash, 80).await;
+    let hash = control_wallet
+        .send_transaction(tx, None)
+        .await
+        .unwrap()
+        .tx_hash();
+    network.run_until_receipt(control_wallet, hash, 80).await;
 
     // Stake the new validator's funds.
     let tx = TransactionRequest::new()
@@ -62,20 +66,77 @@ async fn deposit_stake(
                 ])
                 .unwrap(),
         );
-    let hash = wallet.send_transaction(tx, None).await.unwrap().tx_hash();
-    network.run_until_receipt(wallet, hash, 80).await;
+    let hash = control_wallet
+        .send_transaction(tx, None)
+        .await
+        .unwrap()
+        .tx_hash();
+    network.run_until_receipt(control_wallet, hash, 80).await;
+    hash
+}
+
+async fn current_epoch(
+    wallet: &SignerMiddleware<Provider<LocalRpcClient>, LocalWallet>,
+    block: Option<u64>,
+) -> u64 {
+    let tx = TransactionRequest::new()
+        .to(H160(contract_addr::DEPOSIT.into_array()))
+        .data(contracts::deposit::CURRENT_EPOCH.encode_input(&[]).unwrap());
+    let response = wallet
+        .call(&tx.into(), block.map(|b| b.into()))
+        .await
+        .unwrap();
+    let epoch = contracts::deposit::CURRENT_EPOCH
+        .decode_output(&response)
+        .unwrap()
+        .remove(0)
+        .into_uint()
+        .unwrap()
+        .as_u64();
+    let current_block = block.unwrap_or(wallet.get_block_number().await.unwrap().as_u64());
+
+    // Sanity check that epochs are calculated correctly (assuming `blocks_per_epoch = 2`).
+    assert_eq!(epoch, current_block / 2);
+
+    epoch
 }
 
-async fn remove_staker(network: &mut Network, wallet: &Wallet, key: NodePublicKey) {
+async fn unstake_amount(network: &mut Network, control_wallet: &Wallet, amount: u128) -> H256 {
     let tx = TransactionRequest::new()
         .to(H160(contract_addr::DEPOSIT.into_array()))
         .data(
-            contracts::deposit::TEMP_REMOVE_STAKER
-                .encode_input(&[Token::Bytes(key.as_bytes())])
+            contracts::deposit::UNSTAKE
+                .encode_input(&[Token::Uint(amount.into())])
+                .unwrap(),
+        )
+        .gas(10000000); // TODO: Why needed?
+    let hash = control_wallet
+        .send_transaction(tx, None)
+        .await
+        .unwrap()
+        .tx_hash();
+    let receipt = network.run_until_receipt(control_wallet, hash, 100).await;
+    assert_eq!(receipt.status.unwrap().as_u64(), 1);
+    hash
+}
+
+async fn get_stake(wallet: &Wallet, staker: &NodePublicKey) -> u128 {
+    let tx = TransactionRequest::new()
+        .to(H160(contract_addr::DEPOSIT.into_array()))
+        .data(
+            contracts::deposit::GET_STAKE
+                .encode_input(&[Token::Bytes(staker.as_bytes())])
                 .unwrap(),
         );
-    let hash = wallet.send_transaction(tx, None).await.unwrap().tx_hash();
-    network.run_until_receipt(wallet, hash, 50).await;
+    let output = wallet.call(&tx.into(), None).await.unwrap();
+
+    contracts::deposit::GET_STAKE
+        .decode_output(&output)
+        .unwrap()[0]
+        .clone()
+        .into_uint()
+        .unwrap()
+        .as_u128()
 }
 
 async fn get_stakers(
@@ -145,10 +206,16 @@ async fn rewards_are_sent_to_reward_address_of_proposer(mut network: Network) {
     check_miner_got_reward(&wallet, 1).await;
 }
 
-#[zilliqa_macros::test]
+#[zilliqa_macros::test(blocks_per_epoch = 2)]
 async fn validators_can_join_and_become_proposer(mut network: Network) {
     let wallet = network.genesis_wallet().await;
 
+    // randomise the current epoch state and current leader
+    let blocks_to_prerun = network.rng.lock().unwrap().gen_range(0..8);
+    network
+        .run_until_block(&wallet, blocks_to_prerun.into(), 100)
+        .await;
+
     let index = network.add_node();
     let new_validator_key = network.get_node_raw(index).secret_key;
     let reward_address = H160::random_using(&mut network.rng.lock().unwrap().deref_mut());
@@ -157,22 +224,50 @@ async fn validators_can_join_and_become_proposer(mut network: Network) {
     assert_eq!(stakers.len(), 4);
     assert!(!stakers.contains(&new_validator_key.node_public_key()));
 
-    let pop = new_validator_key.pop_prove();
-
-    deposit_stake(
+    let deposit_hash = deposit_stake(
         &mut network,
         &wallet,
         new_validator_key.node_public_key(),
         new_validator_key.to_libp2p_keypair().public().to_peer_id(),
         32 * 10u128.pow(18),
         reward_address,
-        pop,
+        new_validator_key.pop_prove(),
     )
     .await;
 
-    let stakers = get_stakers(&wallet).await;
-    assert_eq!(stakers.len(), 5);
-    assert!(stakers.contains(&new_validator_key.node_public_key()));
+    let deposit_block = wallet
+        .get_transaction_receipt(deposit_hash)
+        .await
+        .unwrap()
+        .unwrap()
+        .block_number
+        .unwrap()
+        .as_u64();
+
+    // The new validator should become part of the committee exactly two epochs after the one in which the deposit was
+    // made.
+    let deposit_epoch = current_epoch(&wallet, Some(deposit_block)).await;
+    network
+        .run_until_async(
+            || async {
+                let should_be_in_committee =
+                    current_epoch(&wallet, None).await == deposit_epoch + 2;
+
+                let stakers = get_stakers(&wallet).await;
+                if !should_be_in_committee {
+                    assert_eq!(stakers.len(), 4);
+                    assert!(!stakers.contains(&new_validator_key.node_public_key()));
+                    false // Keep running
+                } else {
+                    assert_eq!(stakers.len(), 5);
+                    assert!(stakers.contains(&new_validator_key.node_public_key()));
+                    true
+                }
+            },
+            200,
+        )
+        .await
+        .unwrap();
 
     // Check the new validator eventually gets to be a block proposer.
     network
@@ -268,45 +363,85 @@ async fn block_proposers_are_selected_proportionally_to_their_stake(mut network:
     );
 }
 
-#[zilliqa_macros::test]
-async fn validators_can_leave(mut network: Network) {
-    let genesis_wallet = network.genesis_wallet().await;
+#[zilliqa_macros::test(blocks_per_epoch = 2)]
+async fn validators_can_unstake(mut network: Network) {
+    let wallet = network.genesis_wallet().await;
 
-    let blocks_to_prerun = network.rng.lock().unwrap().gen_range(0..5);
+    // randomise the current epoch state and current leader
+    let blocks_to_prerun = network.rng.lock().unwrap().gen_range(0..8);
     network
-        .run_until_block(&genesis_wallet, blocks_to_prerun.into(), 100)
+        .run_until_block(&wallet, blocks_to_prerun.into(), 100)
         .await;
 
-    let validator_to_remove = network.random_index();
-    let validator_sending_removal = network.random_index();
-
-    let key_to_remove = network.get_node_raw(validator_to_remove).secret_key;
-    let wallet_sending_removal = network
-        .wallet_from_key(
-            network
-                .get_node_raw(validator_sending_removal)
-                .onchain_key
-                .clone(),
-        )
+    let validator_idx = network.random_index();
+    let validator_blskey = network
+        .get_node_raw(validator_idx)
+        .secret_key
+        .node_public_key();
+    let validator_control_wallet = network
+        .wallet_from_key(network.get_node_raw(validator_idx).onchain_key.clone())
         .await;
-    fund_wallet(&mut network, &genesis_wallet, &wallet_sending_removal).await;
+    fund_wallet(&mut network, &wallet, &validator_control_wallet).await;
 
-    let stakers = get_stakers(&genesis_wallet).await;
+    let stakers = get_stakers(&wallet).await;
     assert_eq!(stakers.len(), 4);
-    assert!(stakers.contains(&key_to_remove.node_public_key()));
-
-    remove_staker(
-        &mut network,
-        &wallet_sending_removal,
-        key_to_remove.node_public_key(),
-    )
-    .await;
+    assert!(stakers.contains(&validator_blskey));
 
-    let stakers = get_stakers(&genesis_wallet).await;
-    assert_eq!(stakers.len(), 3);
-    assert!(!stakers.contains(&key_to_remove.node_public_key()));
+    // unstake validator's entire stake
+    let stake = get_stake(&wallet, &validator_blskey).await;
+    let unstake_hash = unstake_amount(&mut network, &validator_control_wallet, stake).await;
+    let unstake_block = wallet
+        .get_transaction_receipt(unstake_hash)
+        .await
+        .unwrap()
+        .unwrap()
+        .block_number
+        .unwrap()
+        .as_u64();
+
+    // The validator should leave the committee exactly two epochs after the one in which the deposit was made.
+    let unstake_epoch = current_epoch(&wallet, Some(unstake_block)).await;
+    network
+        .run_until_async(
+            || async {
+                let should_be_in_committee =
+                    current_epoch(&wallet, None).await != unstake_epoch + 2;
+
+                let stakers = get_stakers(&wallet).await;
+                if should_be_in_committee {
+                    assert_eq!(stakers.len(), 4);
+                    assert!(stakers.contains(&validator_blskey));
+                    false // Keep running
+                } else {
+                    assert_eq!(stakers.len(), 3);
+                    assert!(!stakers.contains(&validator_blskey));
+                    true
+                }
+            },
+            200,
+        )
+        .await
+        .unwrap();
 
+    // ensure network still runs well
     network
-        .run_until_block(&genesis_wallet, (blocks_to_prerun + 10).into(), 500)
-        .await;
+        .run_until_async(
+            || async {
+                let stakers = get_stakers(&wallet).await;
+                assert_eq!(stakers.len(), 3);
+                assert!(!stakers.contains(&validator_blskey));
+                wallet.get_block_number().await.unwrap().as_u64() >= unstake_block + 15
+            },
+            1000,
+        )
+        .await
+        .unwrap();
 }
+
+// TODO: Tests for:
+// * partial unstaking staying above the minimum
+// * partial unstaking under the minimum (should fail)
+// * increase stake
+// * updating staker details (reward address)
+// * disallow access to callers other than the controlAddress
+// * withdraw stake after 2 weeks (exercise the circular buffer logic or test it separately if difficult to do here)