Add validator set rotation test for the node side

.gitignore

@@ -1,5 +1,6 @@
 target
 Dockerfile
+Dockerfile.fast-epoch
 !orchestration/runtime/Dockerfile
 .test-logs
 

orchestration/src/serai.rs

@@ -5,6 +5,8 @@ use crate::{Network, Os, mimalloc, os, build_serai_service, write_dockerfile};
 pub fn serai(orchestration_path: &Path, network: Network) {
   // Always builds in release for performance reasons
   let setup = mimalloc(Os::Debian).to_string() + &build_serai_service(true, "", "serai-node");
+  let setup_fast_epoch =
+    mimalloc(Os::Debian).to_string() + &build_serai_service(true, "fast-epoch", "serai-node");
 
   // TODO: Review the ports exposed here
   let run_serai = format!(
@@ -24,10 +26,16 @@ CMD ["/run.sh"]
 
   let run = os(Os::Debian, "", "serai") + &run_serai;
   let res = setup + &run;
+  let res_fast_epoch = setup_fast_epoch + &run;
 
   let mut serai_path = orchestration_path.to_path_buf();
   serai_path.push("serai");
+
+  let mut serai_fast_epoch_path = serai_path.clone();
+
   serai_path.push("Dockerfile");
+  serai_fast_epoch_path.push("Dockerfile.fast-epoch");
 
   write_dockerfile(serai_path, &res);
+  write_dockerfile(serai_fast_epoch_path, &res_fast_epoch);
 }

substrate/client/src/serai/mod.rs

@@ -3,7 +3,7 @@ use thiserror::Error;
 use async_lock::RwLock;
 use simple_request::{hyper, Request, Client};
 
-use scale::{Encode, Decode, Compact};
+use scale::{decode_from_bytes, Compact, Decode, Encode};
 use serde::{Serialize, Deserialize, de::DeserializeOwned};
 
 pub use sp_core::{
@@ -195,6 +195,19 @@ impl Serai {
     Ok(())
   }
 
+  pub async fn active_network_validators(
+    &self,
+    network: NetworkId,
+  ) -> Result<Vec<Public>, SeraiError> {
+    let hash: String = self
+      .call("state_call", ["SeraiRuntimeApi_validators".to_string(), hex::encode(network.encode())])
+      .await?;
+    let bytes = Self::hex_decode(hash)?;
+    let r = decode_from_bytes::<Vec<Public>>(bytes.into())
+      .map_err(|e| SeraiError::ErrorInResponse(e.to_string()))?;
+    Ok(r)
+  }
+
   pub async fn latest_finalized_block_hash(&self) -> Result<[u8; 32], SeraiError> {
     let hash: String = self.call("chain_getFinalizedHead", ()).await?;
     Self::hex_decode(hash)?.try_into().map_err(|_| {

substrate/client/src/serai/validator_sets.rs

@@ -143,6 +143,22 @@ impl<'a> SeraiValidatorSets<'a> {
       .await
   }
 
+  pub async fn pending_deallocations(
+    &self,
+    network: NetworkId,
+    account: Public,
+    session: Session,
+  ) -> Result<Option<Amount>, SeraiError> {
+    self
+      .0
+      .storage(
+        PALLET,
+        "PendingDeallocations",
+        (sp_core::hashing::blake2_128(&(network, account).encode()), (network, account, session)),
+      )
+      .await
+  }
+
   // TODO: Store these separately since we almost never need both at once?
   pub async fn keys(&self, set: ValidatorSet) -> Result<Option<KeyPair>, SeraiError> {
     self.0.storage(PALLET, "Keys", (sp_core::hashing::twox_64(&set.encode()), set)).await
@@ -169,6 +185,14 @@ impl<'a> SeraiValidatorSets<'a> {
     }))
   }
 
+  pub fn allocate(network: NetworkId, amount: Amount) -> serai_abi::Call {
+    serai_abi::Call::ValidatorSets(serai_abi::validator_sets::Call::allocate { network, amount })
+  }
+
+  pub fn deallocate(network: NetworkId, amount: Amount) -> serai_abi::Call {
+    serai_abi::Call::ValidatorSets(serai_abi::validator_sets::Call::deallocate { network, amount })
+  }
+
   pub fn report_slashes(
     network: NetworkId,
     slashes: sp_runtime::BoundedVec<

substrate/client/tests/common/validator_sets.rs

@@ -1,9 +1,13 @@
 use std::collections::HashMap;
 
+use serai_abi::primitives::NetworkId;
 use zeroize::Zeroizing;
 use rand_core::OsRng;
 
-use sp_core::{Pair, sr25519::Signature};
+use sp_core::{
+  sr25519::{Pair, Signature},
+  Pair as PairTrait,
+};
 
 use ciphersuite::{Ciphersuite, Ristretto};
 use frost::dkg::musig::musig;
@@ -15,7 +19,7 @@ use serai_client::{
     primitives::{ValidatorSet, KeyPair, musig_context, set_keys_message},
     ValidatorSetsEvent,
   },
-  SeraiValidatorSets, Serai,
+  Amount, Serai, SeraiValidatorSets,
 };
 
 use crate::common::tx::publish_tx;
@@ -59,3 +63,29 @@ pub async fn set_keys(serai: &Serai, set: ValidatorSet, key_pair: KeyPair) -> [u
 
   block
 }
+
+#[allow(dead_code)]
+pub async fn allocate_stake(
+  serai: &Serai,
+  network: NetworkId,
+  amount: Amount,
+  pair: &Pair,
+  nonce: u32,
+) -> [u8; 32] {
+  // get the call
+  let tx = serai.sign(&pair, SeraiValidatorSets::allocate(network, amount), nonce, 0);
+  publish_tx(serai, &tx).await
+}
+
+#[allow(dead_code)]
+pub async fn deallocate_stake(
+  serai: &Serai,
+  network: NetworkId,
+  amount: Amount,
+  pair: &Pair,
+  nonce: u32,
+) -> [u8; 32] {
+  // get the call
+  let tx = serai.sign(&pair, SeraiValidatorSets::deallocate(network, amount), nonce, 0);
+  publish_tx(serai, &tx).await
+}

substrate/client/tests/validator_sets.rs

@@ -8,11 +8,13 @@ use serai_client::{
     primitives::{Session, ValidatorSet, KeyPair},
     ValidatorSetsEvent,
   },
-  Serai,
+  Amount, Serai,
 };
 
 mod common;
-use common::validator_sets::set_keys;
+use common::validator_sets::{set_keys, allocate_stake, deallocate_stake};
+
+const EPOCH_INTERVAL: u64 = 5;
 
 serai_test!(
   set_keys_test: (|serai: Serai| async move {
@@ -73,3 +75,198 @@ serai_test!(
     assert_eq!(serai.keys(set).await.unwrap(), Some(key_pair));
   })
 );
+
+#[tokio::test]
+async fn validator_set_rotation() {
+  use dockertest::{
+    PullPolicy, StartPolicy, LogOptions, LogAction, LogPolicy, LogSource, Image,
+    TestBodySpecification, DockerTest,
+  };
+  use std::collections::HashMap;
+
+  serai_docker_tests::build("serai-fast-epoch".to_string());
+
+  let handle = |name| format!("serai_client-serai_node-{name}");
+  let composition = |name| {
+    TestBodySpecification::with_image(
+      Image::with_repository("serai-dev-serai-fast-epoch").pull_policy(PullPolicy::Never),
+    )
+    .replace_cmd(vec![
+      "serai-node".to_string(),
+      "--unsafe-rpc-external".to_string(),
+      "--rpc-cors".to_string(),
+      "all".to_string(),
+      "--chain".to_string(),
+      "local".to_string(),
+      format!("--{name}"),
+    ])
+    .replace_env(HashMap::from([("RUST_LOG=runtime".to_string(), "debug".to_string())]))
+    .set_publish_all_ports(true)
+    .set_handle(handle(name))
+    .set_start_policy(StartPolicy::Strict)
+    .set_log_options(Some(LogOptions {
+      action: LogAction::Forward,
+      policy: LogPolicy::Always,
+      source: LogSource::Both,
+    }))
+  };
+
+  let mut test = DockerTest::new().with_network(dockertest::Network::Isolated);
+  test.provide_container(composition("alice"));
+  test.provide_container(composition("bob"));
+  test.provide_container(composition("charlie"));
+  test.provide_container(composition("dave"));
+  test.provide_container(composition("eve"));
+  test
+    .run_async(|ops| async move {
+      // Sleep until the Substrate RPC starts
+      let alice = handle("alice");
+      let alice_rpc = ops.handle(&alice).host_port(9944).unwrap();
+      let alice_rpc = format!("http://{}:{}", alice_rpc.0, alice_rpc.1);
+
+      // Sleep for a minute
+      tokio::time::sleep(core::time::Duration::from_secs(20)).await;
+      let serai = Serai::new(alice_rpc.clone()).await.unwrap();
+
+      // Make sure the genesis is as expected
+      assert_eq!(
+        serai
+          .as_of(serai.finalized_block_by_number(0).await.unwrap().unwrap().hash())
+          .validator_sets()
+          .new_set_events()
+          .await
+          .unwrap(),
+        NETWORKS
+          .iter()
+          .copied()
+          .map(|network| ValidatorSetsEvent::NewSet {
+            set: ValidatorSet { session: Session(0), network }
+          })
+          .collect::<Vec<_>>(),
+      );
+
+      // genesis accounts
+      let pair1 = insecure_pair_from_name("Alice");
+      let pair2 = insecure_pair_from_name("Bob");
+      let pair3 = insecure_pair_from_name("Charlie");
+      let pair4 = insecure_pair_from_name("Dave");
+      let pair5 = insecure_pair_from_name("Eve");
+
+      // amounts for single key share per network
+      let key_shares = HashMap::from([
+        (NetworkId::Serai, Amount(50_000 * 10_u64.pow(8))),
+        (NetworkId::Bitcoin, Amount(1_000_000 * 10_u64.pow(8))),
+        (NetworkId::Monero, Amount(100_000 * 10_u64.pow(8))),
+        (NetworkId::Ethereum, Amount(1_000_000 * 10_u64.pow(8))),
+      ]);
+
+      // genesis participants per network
+      let default_participants =
+        vec![pair1.public(), pair2.public(), pair3.public(), pair4.public()];
+      let mut participants = HashMap::from([
+        (NetworkId::Serai, default_participants.clone()),
+        (NetworkId::Bitcoin, default_participants.clone()),
+        (NetworkId::Monero, default_participants.clone()),
+        (NetworkId::Ethereum, default_participants),
+      ]);
+
+      // test the set rotation
+      for (i, network) in NETWORKS.into_iter().enumerate() {
+        let participants = participants.get_mut(&network).unwrap();
+
+        // we start the chain with 4 default participants that has a single key share each
+        participants.sort();
+        verfiy_session_and_active_validators(&serai, network, 0, &participants).await;
+
+        // add 1 participant & verify
+        let hash =
+          allocate_stake(&serai, network, key_shares[&network], &pair5, i.try_into().unwrap())
+            .await;
+        participants.push(pair5.public());
+        participants.sort();
+        verfiy_session_and_active_validators(
+          &serai,
+          network,
+          get_active_session(&serai, network, hash).await,
+          &participants,
+        )
+        .await;
+
+        // remove 1 participant & verify
+        let hash =
+          deallocate_stake(&serai, network, key_shares[&network], &pair2, i.try_into().unwrap())
+            .await;
+        participants.swap_remove(participants.iter().position(|k| *k == pair2.public()).unwrap());
+        let active_session = get_active_session(&serai, network, hash).await;
+        participants.sort();
+        verfiy_session_and_active_validators(&serai, network, active_session, &participants).await;
+
+        // check pending deallocations
+        let pending = serai
+          .as_of_latest_finalized_block()
+          .await
+          .unwrap()
+          .validator_sets()
+          .pending_deallocations(
+            network,
+            pair2.public(),
+            Session(u32::try_from(active_session + 1).unwrap()),
+          )
+          .await
+          .unwrap();
+        assert_eq!(pending, Some(key_shares[&network]));
+      }
+    })
+    .await;
+}
+
+async fn verfiy_session_and_active_validators(
+  serai: &Serai,
+  network: NetworkId,
+  session: u64,
+  participants: &[Public],
+) {
+  // wait untill the epoch block finalized
+  let epoch_block = (session * EPOCH_INTERVAL) + 1;
+  while serai.finalized_block_by_number(epoch_block).await.unwrap().is_none() {
+    // sleep 1 block
+    tokio::time::sleep(tokio::time::Duration::from_secs(6)).await;
+  }
+  let serai_for_block =
+    serai.as_of(serai.finalized_block_by_number(epoch_block).await.unwrap().unwrap().hash());
+
+  // verfiy session
+  let s = serai_for_block.validator_sets().session(network).await.unwrap().unwrap();
+  assert_eq!(u64::from(s.0), session);
+
+  // verify participants
+  let mut validators = serai.active_network_validators(network).await.unwrap();
+  validators.sort();
+  assert_eq!(validators, participants);
+
+  // make sure finalization continues as usual after the changes
+  tokio::time::timeout(tokio::time::Duration::from_secs(60), async move {
+    let mut finalized_block = serai.latest_finalized_block().await.unwrap().header.number;
+    while finalized_block <= epoch_block + 2 {
+      tokio::time::sleep(tokio::time::Duration::from_secs(6)).await;
+      finalized_block = serai.latest_finalized_block().await.unwrap().header.number;
+    }
+  })
+  .await
+  .unwrap();
+
+  // TODO: verfiy key shares as well?
+}
+
+async fn get_active_session(serai: &Serai, network: NetworkId, hash: [u8; 32]) -> u64 {
+  let block_number = serai.block(hash).await.unwrap().unwrap().header.number;
+  let epoch = block_number / EPOCH_INTERVAL;
+
+  // changes should be active in the next session
+  if network == NetworkId::Serai {
+    // it takes 1 extra session for serai net to make the changes active.
+    epoch + 2
+  } else {
+    epoch + 1
+  }
+}

substrate/node/Cargo.toml

@@ -75,6 +75,7 @@ substrate-build-script-utils = { git = "https://github.com/serai-dex/substrate"
 
 [features]
 default = []
+fast-epoch = ["serai-runtime/fast-epoch"]
 runtime-benchmarks = [
   "frame-benchmarking/runtime-benchmarks",
 

substrate/runtime/Cargo.toml

@@ -124,6 +124,8 @@ std = [
   "pallet-transaction-payment-rpc-runtime-api/std",
 ]
 
+fast-epoch = []
+
 runtime-benchmarks = [
   "sp-runtime/runtime-benchmarks",