bug fixes

substrate/client/tests/common/in_instructions.rs

@@ -31,7 +31,7 @@ pub async fn provide_batch(serai: &Serai, batch: Batch) -> [u8; 32] {
     keys
   } else {
     let keys = KeyPair(pair.public(), vec![].try_into().unwrap());
-    set_keys(serai, set, keys.clone()).await;
+    set_keys(serai, set, keys.clone(), &[insecure_pair_from_name("Alice")]).await;
     keys
   };
   assert_eq!(keys.0, pair.public());

substrate/client/tests/common/validator_sets.rs

@@ -14,7 +14,6 @@ use frost::dkg::musig::musig;
 use schnorrkel::Schnorrkel;
 
 use serai_client::{
-  primitives::insecure_pair_from_name,
   validator_sets::{
     primitives::{ValidatorSet, KeyPair, musig_context, set_keys_message},
     ValidatorSetsEvent,
@@ -25,26 +24,40 @@ use serai_client::{
 use crate::common::tx::publish_tx;
 
 #[allow(dead_code)]
-pub async fn set_keys(serai: &Serai, set: ValidatorSet, key_pair: KeyPair) -> [u8; 32] {
-  let pair = insecure_pair_from_name("Alice");
-  let public = pair.public();
+pub async fn set_keys(
+  serai: &Serai,
+  set: ValidatorSet,
+  key_pair: KeyPair,
+  pairs: &[Pair],
+) -> [u8; 32] {
+  let mut pub_keys = vec![];
+  for pair in pairs {
+    let public_key =
+      <Ristretto as Ciphersuite>::read_G::<&[u8]>(&mut pair.public().0.as_ref()).unwrap();
+    pub_keys.push(public_key);
+  }
 
-  let public_key = <Ristretto as Ciphersuite>::read_G::<&[u8]>(&mut public.0.as_ref()).unwrap();
-  let secret_key = <Ristretto as Ciphersuite>::read_F::<&[u8]>(
-    &mut pair.as_ref().secret.to_bytes()[.. 32].as_ref(),
-  )
-  .unwrap();
-  assert_eq!(Ristretto::generator() * secret_key, public_key);
-  let threshold_keys =
-    musig::<Ristretto>(&musig_context(set), &Zeroizing::new(secret_key), &[public_key]).unwrap();
+  let mut threshold_keys = vec![];
+  for i in 0 .. pairs.len() {
+    let secret_key = <Ristretto as Ciphersuite>::read_F::<&[u8]>(
+      &mut pairs[i].as_ref().secret.to_bytes()[.. 32].as_ref(),
+    )
+    .unwrap();
+    assert_eq!(Ristretto::generator() * secret_key, pub_keys[i]);
+
+    threshold_keys.push(
+      musig::<Ristretto>(&musig_context(set), &Zeroizing::new(secret_key), &pub_keys).unwrap(),
+    );
+  }
+
+  let mut musig_keys = HashMap::new();
+  for tk in threshold_keys {
+    musig_keys.insert(tk.params().i(), tk.into());
+  }
 
   let sig = frost::tests::sign_without_caching(
     &mut OsRng,
-    frost::tests::algorithm_machines(
-      &mut OsRng,
-      &Schnorrkel::new(b"substrate"),
-      &HashMap::from([(threshold_keys.params().i(), threshold_keys.into())]),
-    ),
+    frost::tests::algorithm_machines(&mut OsRng, &Schnorrkel::new(b"substrate"), &musig_keys),
     &set_keys_message(&set, &[], &key_pair),
   );
 

substrate/client/tests/validator_sets.rs

@@ -1,34 +1,71 @@
 use rand_core::{RngCore, OsRng};
 
-use sp_core::{sr25519::Public, Pair};
+use sp_core::{
+  sr25519::{Public, Pair},
+  Pair as PairTrait,
+};
 
 use serai_client::{
-  primitives::{NETWORKS, NetworkId, insecure_pair_from_name},
+  primitives::{NETWORKS, NetworkId, BlockHash, insecure_pair_from_name},
   validator_sets::{
     primitives::{Session, ValidatorSet, KeyPair},
     ValidatorSetsEvent,
   },
+  in_instructions::{
+    primitives::{Batch, SignedBatch, batch_message},
+    SeraiInInstructions,
+  },
   Amount, Serai,
 };
 
 mod common;
-use common::validator_sets::{set_keys, allocate_stake, deallocate_stake};
+use common::{
+  tx::publish_tx,
+  validator_sets::{allocate_stake, deallocate_stake, set_keys},
+};
+
+fn get_random_key_pair() -> KeyPair {
+  let mut ristretto_key = [0; 32];
+  OsRng.fill_bytes(&mut ristretto_key);
+  let mut external_key = vec![0; 33];
+  OsRng.fill_bytes(&mut external_key);
+  KeyPair(Public(ristretto_key), external_key.try_into().unwrap())
+}
+
+async fn get_correct_pairs(serai: &Serai, network: NetworkId, accounts: &[Pair]) -> Vec<Pair> {
+  // retrieve the current session validators so that we know the order of the keys
+  // that is necessary for the correct musig signature.
+  let validators = serai
+    .as_of_latest_finalized_block()
+    .await
+    .unwrap()
+    .validator_sets()
+    .active_network_validators(network)
+    .await
+    .unwrap();
+
+  // collect the pairs of the validators
+  let mut pairs = vec![];
+  for v in validators {
+    let p = accounts.iter().find(|pair| pair.public() == v).unwrap().clone();
+    pairs.push(p);
+  }
+
+  pairs
+}
 
 serai_test!(
   set_keys_test: (|serai: Serai| async move {
     let network = NetworkId::Bitcoin;
     let set = ValidatorSet { session: Session(0), network };
 
-    let public = insecure_pair_from_name("Alice").public();
+    let pair = insecure_pair_from_name("Alice");
+    let public = pair.public();
 
     // Neither of these keys are validated
     // The external key is infeasible to validate on-chain, the Ristretto key is feasible
     // TODO: Should the Ristretto key be validated?
-    let mut ristretto_key = [0; 32];
-    OsRng.fill_bytes(&mut ristretto_key);
-    let mut external_key = vec![0; 33];
-    OsRng.fill_bytes(&mut external_key);
-    let key_pair = KeyPair(Public(ristretto_key), external_key.try_into().unwrap());
+    let key_pair = get_random_key_pair();
 
     // Make sure the genesis is as expected
     assert_eq!(
@@ -60,7 +97,7 @@ serai_test!(
       assert_eq!(participants_ref, [public].as_ref());
     }
 
-    let block = set_keys(&serai, set, key_pair.clone()).await;
+    let block = set_keys(&serai, set, key_pair.clone(), &[pair]).await;
 
     // While the set_keys function should handle this, it's beneficial to
     // independently test it
@@ -126,7 +163,7 @@ async fn validator_set_rotation() {
       let alice_rpc = format!("http://{}:{}", alice_rpc.0, alice_rpc.1);
 
       // Sleep for some time
-      tokio::time::sleep(core::time::Duration::from_secs(20)).await;
+      tokio::time::sleep(tokio::time::Duration::from_secs(20)).await;
       let serai = Serai::new(alice_rpc.clone()).await.unwrap();
 
       // Make sure the genesis is as expected
@@ -147,11 +184,13 @@ async fn validator_set_rotation() {
       );
 
       // 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");
+      let accounts = vec![
+        insecure_pair_from_name("Alice"),
+        insecure_pair_from_name("Bob"),
+        insecure_pair_from_name("Charlie"),
+        insecure_pair_from_name("Dave"),
+        insecure_pair_from_name("Eve"),
+      ];
 
       // amounts for single key share per network
       let key_shares = HashMap::from([
@@ -162,8 +201,9 @@ async fn validator_set_rotation() {
       ]);
 
       // genesis participants per network
+      #[allow(clippy::redundant_closure_for_method_calls)]
       let default_participants =
-        vec![pair1.public(), pair2.public(), pair3.public(), pair4.public()];
+        accounts[.. 4].to_vec().iter().map(|pair| pair.public()).collect::<Vec<_>>();
       let mut participants = HashMap::from([
         (NetworkId::Serai, default_participants.clone()),
         (NetworkId::Bitcoin, default_participants.clone()),
@@ -179,26 +219,80 @@ async fn validator_set_rotation() {
         participants.sort();
         verify_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();
-        verify_session_and_active_validators(
+        // add 1 participant
+        let last_participant = accounts[4].clone();
+        let hash = allocate_stake(
           &serai,
           network,
-          get_active_session(&serai, network, hash).await,
-          participants,
+          key_shares[&network],
+          &last_participant,
+          i.try_into().unwrap(),
         )
         .await;
+        participants.push(last_participant.public());
+        // the session at which set changes becomes active
+        let active_session = get_active_session(&serai, network, hash).await;
+
+        // set the keys if it is an external set
+        if network != NetworkId::Serai {
+          let set = ValidatorSet { session: Session(0), network };
+          let key_pair = get_random_key_pair();
+          let pairs = get_correct_pairs(&serai, network, &accounts).await;
+          set_keys(&serai, set, key_pair, &pairs).await;
+        }
+
+        // verify
+        participants.sort();
+        verify_session_and_active_validators(&serai, network, active_session, 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());
+        // remove 1 participant
+        // TODO: this participant can be selected at random
+        let participant_to_remove = accounts[1].clone();
+        let hash = deallocate_stake(
+          &serai,
+          network,
+          key_shares[&network],
+          &participant_to_remove,
+          i.try_into().unwrap(),
+        )
+        .await;
+        participants.swap_remove(
+          participants.iter().position(|k| *k == participant_to_remove.public()).unwrap(),
+        );
         let active_session = get_active_session(&serai, network, hash).await;
+
+        if network != NetworkId::Serai {
+          // set the keys if it is an external set
+          let set = ValidatorSet { session: Session(1), network };
+
+          // we need the whole substrate key pair to sign the bath
+          let (substrate_pair, key_pair) = {
+            let pair = insecure_pair_from_name("session-1-key-pair");
+            let public = pair.public();
+
+            let mut external_key = vec![0; 33];
+            OsRng.fill_bytes(&mut external_key);
+
+            (pair, KeyPair(public, external_key.try_into().unwrap()))
+          };
+          let pairs = get_correct_pairs(&serai, network, &accounts).await;
+          set_keys(&serai, set, key_pair, &pairs).await;
+
+          // provide a batch to retire the previous set
+          let mut block_hash = BlockHash([0; 32]);
+          OsRng.fill_bytes(&mut block_hash.0);
+          let batch = Batch { network, id: 0, block: block_hash, instructions: vec![] };
+          publish_tx(
+            &serai,
+            &SeraiInInstructions::execute_batch(SignedBatch {
+              batch: batch.clone(),
+              signature: substrate_pair.sign(&batch_message(&batch)),
+            }),
+          )
+          .await;
+        }
+
+        // verify
         participants.sort();
         verify_session_and_active_validators(&serai, network, active_session, participants).await;
 
@@ -210,8 +304,8 @@ async fn validator_set_rotation() {
           .validator_sets()
           .pending_deallocations(
             network,
-            pair2.public(),
-            Session(u32::try_from(active_session + 1).unwrap()),
+            participant_to_remove.public(),
+            Session(active_session + 1),
           )
           .await
           .unwrap();
@@ -221,42 +315,39 @@ async fn validator_set_rotation() {
     .await;
 }
 
-async fn epoch_for_block(serai: &Serai, block: [u8; 32]) -> u64 {
-  let epoch: String = serai
-    .call("state_call", ["BabeApi_current_epoch".to_string(), String::new(), hex::encode(block)])
-    .await
-    .unwrap();
-  <u64 as scale::Decode>::decode(
-    &mut hex::decode(epoch.strip_prefix("0x").unwrap()).unwrap().as_slice(),
-  )
-  .unwrap()
+async fn session_for_block(serai: &Serai, block: [u8; 32], network: NetworkId) -> u32 {
+  serai.as_of(block).validator_sets().session(network).await.unwrap().unwrap().0
 }
 
 async fn verify_session_and_active_validators(
   serai: &Serai,
   network: NetworkId,
-  session: u64,
+  session: u32,
   participants: &[Public],
 ) {
-  // wait untill the epoch block finalizes
-  let block = loop {
-    let mut block = serai.latest_finalized_block_hash().await.unwrap();
-    if epoch_for_block(serai, block).await < session {
-      // Sleep a block
-      tokio::time::sleep(tokio::time::Duration::from_secs(6)).await;
-      continue;
-    }
-    while epoch_for_block(serai, block).await > session {
-      block = serai.block(block).await.unwrap().unwrap().header.parent_hash.0;
+  // wait until the active session. This wait should be max 30 secs since the epoch time.
+  let block = tokio::time::timeout(tokio::time::Duration::from_secs(2 * 60), async move {
+    loop {
+      let mut block = serai.latest_finalized_block_hash().await.unwrap();
+      if session_for_block(serai, block, network).await < session {
+        // Sleep a block
+        tokio::time::sleep(tokio::time::Duration::from_secs(6)).await;
+        continue;
+      }
+      while session_for_block(serai, block, network).await > session {
+        block = serai.block(block).await.unwrap().unwrap().header.parent_hash.0;
+      }
+      assert_eq!(session_for_block(serai, block, network).await, session);
+      break block;
     }
-    assert_eq!(epoch_for_block(serai, block).await, session);
-    break block;
-  };
+  })
+  .await
+  .unwrap();
   let serai_for_block = serai.as_of(block);
 
   // verify session
   let s = serai_for_block.validator_sets().session(network).await.unwrap().unwrap();
-  assert_eq!(u64::from(s.0), session);
+  assert_eq!(s.0, session);
 
   // verify participants
   let mut validators =
@@ -279,14 +370,14 @@ async fn verify_session_and_active_validators(
   // TODO: verify key shares as well?
 }
 
-async fn get_active_session(serai: &Serai, network: NetworkId, hash: [u8; 32]) -> u64 {
-  let epoch = epoch_for_block(serai, hash).await;
+async fn get_active_session(serai: &Serai, network: NetworkId, hash: [u8; 32]) -> u32 {
+  let session = session_for_block(serai, hash, network).await;
 
   // 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
+    session + 2
   } else {
-    epoch + 1
+    session + 1
   }
 }