Merge branch 'develop' into HEAD

substrate/client/Cargo.toml

@@ -36,7 +36,7 @@ async-lock = "3"
 
 simple-request = { path = "../../common/request", version = "0.1", optional = true }
 
-bitcoin = { version = "0.31", optional = true }
+bitcoin = { version = "0.32", optional = true }
 
 ciphersuite = { path = "../../crypto/ciphersuite", version = "0.4", optional = true }
 monero-serai = { path = "../../coins/monero", version = "0.1.4-alpha", optional = true }

substrate/client/src/networks/bitcoin.rs

@@ -6,38 +6,46 @@ use bitcoin::{
   hashes::{Hash as HashTrait, hash160::Hash},
   PubkeyHash, ScriptHash,
   network::Network,
-  WitnessVersion, WitnessProgram,
-  address::{Error, Payload, NetworkChecked, Address as BAddressGeneric},
+  WitnessVersion, WitnessProgram, ScriptBuf,
+  address::{AddressType, NetworkChecked, Address as BAddress},
 };
 
-type BAddress = BAddressGeneric<NetworkChecked>;
-
 #[derive(Clone, Eq, Debug)]
-pub struct Address(BAddress);
+pub struct Address(ScriptBuf);
 
 impl PartialEq for Address {
   fn eq(&self, other: &Self) -> bool {
-    // Since Serai defines the Bitcoin-address specification as a variant of the payload alone,
-    // define equivalency as the payload alone
-    self.0.payload() == other.0.payload()
+    // Since Serai defines the Bitcoin-address specification as a variant of the script alone,
+    // define equivalency as the script alone
+    self.0 == other.0
+  }
+}
+
+impl From<Address> for ScriptBuf {
+  fn from(addr: Address) -> ScriptBuf {
+    addr.0
   }
 }
 
 impl FromStr for Address {
-  type Err = Error;
-  fn from_str(str: &str) -> Result<Address, Error> {
+  type Err = ();
+  fn from_str(str: &str) -> Result<Address, ()> {
     Address::new(
-      BAddressGeneric::from_str(str)
-        .map_err(|_| Error::UnrecognizedScript)?
-        .require_network(Network::Bitcoin)?,
+      BAddress::from_str(str)
+        .map_err(|_| ())?
+        .require_network(Network::Bitcoin)
+        .map_err(|_| ())?
+        .script_pubkey(),
     )
-    .ok_or(Error::UnrecognizedScript)
+    .ok_or(())
   }
 }
 
 impl fmt::Display for Address {
   fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-    self.0.fmt(f)
+    BAddress::<NetworkChecked>::from_script(&self.0, Network::Bitcoin)
+      .map_err(|_| fmt::Error)?
+      .fmt(f)
   }
 }
 
@@ -54,55 +62,52 @@ enum EncodedAddress {
 impl TryFrom<Vec<u8>> for Address {
   type Error = ();
   fn try_from(data: Vec<u8>) -> Result<Address, ()> {
-    Ok(Address(BAddress::new(
-      Network::Bitcoin,
-      match EncodedAddress::decode(&mut data.as_ref()).map_err(|_| ())? {
-        EncodedAddress::P2PKH(hash) => {
-          Payload::PubkeyHash(PubkeyHash::from_raw_hash(Hash::from_byte_array(hash)))
-        }
-        EncodedAddress::P2SH(hash) => {
-          Payload::ScriptHash(ScriptHash::from_raw_hash(Hash::from_byte_array(hash)))
-        }
-        EncodedAddress::P2WPKH(hash) => {
-          Payload::WitnessProgram(WitnessProgram::new(WitnessVersion::V0, hash).unwrap())
-        }
-        EncodedAddress::P2WSH(hash) => {
-          Payload::WitnessProgram(WitnessProgram::new(WitnessVersion::V0, hash).unwrap())
-        }
-        EncodedAddress::P2TR(key) => {
-          Payload::WitnessProgram(WitnessProgram::new(WitnessVersion::V1, key).unwrap())
-        }
-      },
-    )))
+    Ok(Address(match EncodedAddress::decode(&mut data.as_ref()).map_err(|_| ())? {
+      EncodedAddress::P2PKH(hash) => {
+        ScriptBuf::new_p2pkh(&PubkeyHash::from_raw_hash(Hash::from_byte_array(hash)))
+      }
+      EncodedAddress::P2SH(hash) => {
+        ScriptBuf::new_p2sh(&ScriptHash::from_raw_hash(Hash::from_byte_array(hash)))
+      }
+      EncodedAddress::P2WPKH(hash) => {
+        ScriptBuf::new_witness_program(&WitnessProgram::new(WitnessVersion::V0, &hash).unwrap())
+      }
+      EncodedAddress::P2WSH(hash) => {
+        ScriptBuf::new_witness_program(&WitnessProgram::new(WitnessVersion::V0, &hash).unwrap())
+      }
+      EncodedAddress::P2TR(key) => {
+        ScriptBuf::new_witness_program(&WitnessProgram::new(WitnessVersion::V1, &key).unwrap())
+      }
+    }))
   }
 }
 
 fn try_to_vec(addr: &Address) -> Result<Vec<u8>, ()> {
+  let parsed_addr =
+    BAddress::<NetworkChecked>::from_script(&addr.0, Network::Bitcoin).map_err(|_| ())?;
   Ok(
-    (match addr.0.payload() {
-      Payload::PubkeyHash(hash) => EncodedAddress::P2PKH(*hash.as_raw_hash().as_byte_array()),
-      Payload::ScriptHash(hash) => EncodedAddress::P2SH(*hash.as_raw_hash().as_byte_array()),
-      Payload::WitnessProgram(program) => match program.version() {
-        WitnessVersion::V0 => {
-          let program = program.program();
-          if program.len() == 20 {
-            let mut buf = [0; 20];
-            buf.copy_from_slice(program.as_ref());
-            EncodedAddress::P2WPKH(buf)
-          } else if program.len() == 32 {
-            let mut buf = [0; 32];
-            buf.copy_from_slice(program.as_ref());
-            EncodedAddress::P2WSH(buf)
-          } else {
-            Err(())?
-          }
-        }
-        WitnessVersion::V1 => {
-          let program_ref: &[u8] = program.program().as_ref();
-          EncodedAddress::P2TR(program_ref.try_into().map_err(|_| ())?)
-        }
-        _ => Err(())?,
-      },
+    (match parsed_addr.address_type() {
+      Some(AddressType::P2pkh) => {
+        EncodedAddress::P2PKH(*parsed_addr.pubkey_hash().unwrap().as_raw_hash().as_byte_array())
+      }
+      Some(AddressType::P2sh) => {
+        EncodedAddress::P2SH(*parsed_addr.script_hash().unwrap().as_raw_hash().as_byte_array())
+      }
+      Some(AddressType::P2wpkh) => {
+        let program = parsed_addr.witness_program().ok_or(())?;
+        let program = program.program().as_bytes();
+        EncodedAddress::P2WPKH(program.try_into().map_err(|_| ())?)
+      }
+      Some(AddressType::P2wsh) => {
+        let program = parsed_addr.witness_program().ok_or(())?;
+        let program = program.program().as_bytes();
+        EncodedAddress::P2WSH(program.try_into().map_err(|_| ())?)
+      }
+      Some(AddressType::P2tr) => {
+        let program = parsed_addr.witness_program().ok_or(())?;
+        let program = program.program().as_bytes();
+        EncodedAddress::P2TR(program.try_into().map_err(|_| ())?)
+      }
       _ => Err(())?,
     })
     .encode(),
@@ -116,20 +121,8 @@ impl From<Address> for Vec<u8> {
   }
 }
 
-impl From<Address> for BAddress {
-  fn from(addr: Address) -> BAddress {
-    addr.0
-  }
-}
-
-impl AsRef<BAddress> for Address {
-  fn as_ref(&self) -> &BAddress {
-    &self.0
-  }
-}
-
 impl Address {
-  pub fn new(address: BAddress) -> Option<Self> {
+  pub fn new(address: ScriptBuf) -> Option<Self> {
     let res = Self(address);
     if try_to_vec(&res).is_ok() {
       return Some(res);

substrate/client/src/serai/mod.rs

@@ -4,7 +4,7 @@ use thiserror::Error;
 use async_lock::RwLock;
 use simple_request::{hyper, Request, Client};
 
-use scale::{Compact, Decode, Encode};
+use scale::{Decode, Encode};
 use serde::{Serialize, Deserialize, de::DeserializeOwned};
 
 pub use sp_core::{
@@ -48,8 +48,8 @@ impl Block {
   /// Returns the time of this block, set by its producer, in milliseconds since the epoch.
   pub fn time(&self) -> Result<u64, SeraiError> {
     for transaction in &self.transactions {
-      if let Call::Timestamp(timestamp::Call::set { now }) = &transaction.call {
-        return Ok(u64::from(*now));
+      if let Call::Timestamp(timestamp::Call::set { now }) = transaction.call() {
+        return Ok(*now);
       }
     }
     Err(SeraiError::InvalidNode("no time was present in block".to_string()))
@@ -167,15 +167,14 @@ impl Serai {
   }
 
   fn unsigned(call: Call) -> Transaction {
-    Transaction { call, signature: None }
+    Transaction::new(call, None)
   }
 
   pub fn sign(&self, signer: &Pair, call: Call, nonce: u32, tip: u64) -> Transaction {
     const SPEC_VERSION: u32 = 1;
     const TX_VERSION: u32 = 1;
 
-    let extra =
-      Extra { era: sp_runtime::generic::Era::Immortal, nonce: Compact(nonce), tip: Compact(tip) };
+    let extra = Extra { era: sp_runtime::generic::Era::Immortal, nonce, tip };
     let signature_payload = (
       &call,
       &extra,
@@ -189,7 +188,7 @@ impl Serai {
       .encode();
     let signature = signer.sign(&signature_payload);
 
-    Transaction { call, signature: Some((signer.public().into(), signature, extra)) }
+    Transaction::new(call, Some((signer.public().into(), signature, extra)))
   }
 
   pub async fn publish(&self, tx: &Transaction) -> Result<(), SeraiError> {
@@ -202,7 +201,7 @@ impl Serai {
 
   // TODO: move this into substrate/client/src/validator_sets.rs
   async fn active_network_validators(&self, network: NetworkId) -> Result<Vec<Public>, SeraiError> {
-    let hash: String = self
+    let validators: String = self
       .call("state_call", ["SeraiRuntimeApi_validators".to_string(), hex::encode(network.encode())])
       .await?;
     let bytes = hex_decode(hash)
@@ -378,7 +377,10 @@ impl<'a> TemporalSerai<'a> {
     let res = hex_decode(res)
       .map_err(|_| SeraiError::InvalidNode("expected hex from node wasn't hex".to_string()))?;
     Ok(Some(R::decode(&mut res.as_slice()).map_err(|_| {
-      SeraiError::InvalidRuntime("different type present at storage location".to_string())
+      SeraiError::InvalidRuntime(format!(
+        "different type present at storage location, raw value: {}",
+        hex::encode(res)
+      ))
     })?))
   }
 

substrate/client/src/serai/validator_sets.rs

@@ -180,7 +180,10 @@ impl<'a> SeraiValidatorSets<'a> {
 
   pub fn set_keys(
     network: NetworkId,
-    removed_participants: Vec<SeraiAddress>,
+    removed_participants: sp_runtime::BoundedVec<
+      SeraiAddress,
+      sp_core::ConstU32<{ primitives::MAX_KEY_SHARES_PER_SET / 3 }>,
+    >,
     key_pair: KeyPair,
     signature: Signature,
   ) -> Transaction {

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,33 +24,52 @@ 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),
   );
 
   // Set the key pair
   let block = publish_tx(
     serai,
-    &SeraiValidatorSets::set_keys(set.network, vec![], key_pair.clone(), Signature(sig.to_bytes())),
+    &SeraiValidatorSets::set_keys(
+      set.network,
+      vec![].try_into().unwrap(),
+      key_pair.clone(),
+      Signature(sig.to_bytes()),
+    ),
   )
   .await;
 

substrate/client/tests/validator_sets.rs

@@ -1,36 +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},
+};
 
-const EPOCH_INTERVAL: u64 = 300;
+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_ordered_keys(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!(
@@ -62,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
@@ -149,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([
@@ -164,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()),
@@ -181,28 +219,83 @@ 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 activation_session = get_session_at_which_changes_activate(&serai, network, hash).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;
+        // 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_ordered_keys(&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;
+        verify_session_and_active_validators(&serai, network, activation_session, participants)
+          .await;
+
+        // remove 1 participant
+        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 activation_session = get_session_at_which_changes_activate(&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 batch
+          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_ordered_keys(&serai, network, &accounts).await;
+          set_keys(&serai, set, key_pair, &pairs).await;
+
+          // provide a batch to complete the handover and 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, activation_session, participants)
+          .await;
 
         // check pending deallocations
         let pending = serai
@@ -212,8 +305,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(activation_session + 1),
           )
           .await
           .unwrap();
@@ -223,24 +316,39 @@ async fn validator_set_rotation() {
     .await;
 }
 
+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 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());
+  // wait until the active session. This wait should be max 30 secs since the epoch time.
+  let block = tokio::time::timeout(core::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(core::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;
+    }
+  })
+  .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 =
@@ -249,10 +357,11 @@ async fn verify_session_and_active_validators(
   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 current_finalized_block = serai.latest_finalized_block().await.unwrap().header.number;
+  tokio::time::timeout(core::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;
+    while finalized_block <= current_finalized_block + 2 {
+      tokio::time::sleep(core::time::Duration::from_secs(6)).await;
       finalized_block = serai.latest_finalized_block().await.unwrap().header.number;
     }
   })
@@ -262,15 +371,18 @@ 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 block_number = serai.block(hash).await.unwrap().unwrap().header.number;
-  let epoch = block_number / EPOCH_INTERVAL;
+async fn get_session_at_which_changes_activate(
+  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
   }
 }