Integrate session pallet into validator-sets pallet (#440)

* remove pallet-session

* Store key shares in InSet

* integrate grandpa to vs-pallet

* integrate pallet babe

* remove pallet-session & authority discovery from runtime

* update the grandpa pallet path

* cargo update grandpa

* cargo update substrate

* Misc tweaks

Sets validators for BABE/GRANDPA in chain_spec, per Akil's realization that was
the missing piece.

* fix pr comments

* bug fix & tidy up

---------

Co-authored-by: Luke Parker <lukeparker5132@gmail.com>

substrate/client/src/serai/validator_sets.rs

@@ -44,7 +44,10 @@ impl<'a> SeraiValidatorSets<'a> {
     self.0.storage(PALLET, "CurrentSession", Some(vec![scale_value(network)])).await
   }
 
-  pub async fn participants(&self, network: NetworkId) -> Result<Option<Vec<Public>>, SeraiError> {
+  pub async fn participants(
+    &self,
+    network: NetworkId,
+  ) -> Result<Option<Vec<(Public, u64)>>, SeraiError> {
     self.0.storage(PALLET, "Participants", Some(vec![scale_value(network)])).await
   }
 

substrate/client/tests/validator_sets.rs

@@ -49,7 +49,12 @@ serai_test!(
 
     {
       let vs_serai = serai.with_current_latest_block().await.unwrap().validator_sets();
-      let participants = vs_serai.participants(set.network).await.unwrap().unwrap();
+      let participants = vs_serai.participants(set.network).await
+        .unwrap()
+        .unwrap()
+        .into_iter()
+        .map(|(k, _)| k)
+        .collect::<Vec<_>>();
       let participants_ref: &[_] = participants.as_ref();
       assert_eq!(participants_ref, [public].as_ref());
       assert_eq!(vs_serai.musig_key(set).await.unwrap().unwrap(), musig_key(set, &[public]).0);

substrate/node/src/chain_spec.rs

@@ -5,9 +5,8 @@ use sp_core::Pair as PairTrait;
 use sc_service::ChainType;
 
 use serai_runtime::{
-  primitives::*, WASM_BINARY, opaque::SessionKeys, BABE_GENESIS_EPOCH_CONFIG, RuntimeGenesisConfig,
-  SystemConfig, CoinsConfig, DexConfig, ValidatorSetsConfig, SessionConfig, BabeConfig,
-  GrandpaConfig, AuthorityDiscoveryConfig,
+  primitives::*, WASM_BINARY, BABE_GENESIS_EPOCH_CONFIG, RuntimeGenesisConfig, SystemConfig,
+  CoinsConfig, DexConfig, ValidatorSetsConfig, BabeConfig, GrandpaConfig,
 };
 
 pub type ChainSpec = sc_service::GenericChainSpec<RuntimeGenesisConfig>;
@@ -21,16 +20,7 @@ fn testnet_genesis(
   validators: &[&'static str],
   endowed_accounts: Vec<PublicKey>,
 ) -> RuntimeGenesisConfig {
-  let session_key = |name| {
-    let key = account_from_name(name);
-    (
-      key,
-      key,
-      // TODO: Properly diversify these?
-      SessionKeys { babe: key.into(), grandpa: key.into(), authority_discovery: key.into() },
-    )
-  };
-
+  let validators = validators.iter().map(|name| account_from_name(name)).collect::<Vec<_>>();
   RuntimeGenesisConfig {
     system: SystemConfig { code: wasm_binary.to_vec(), _config: PhantomData },
 
@@ -59,17 +49,17 @@ fn testnet_genesis(
           NetworkId::Monero => (NetworkId::Monero, Amount(100_000 * 10_u64.pow(8))),
         })
         .collect(),
-      participants: validators.iter().map(|name| account_from_name(name)).collect(),
+      participants: validators.clone(),
     },
-    session: SessionConfig { keys: validators.iter().map(|name| session_key(*name)).collect() },
     babe: BabeConfig {
-      authorities: vec![],
+      authorities: validators.iter().map(|validator| ((*validator).into(), 1)).collect(),
       epoch_config: Some(BABE_GENESIS_EPOCH_CONFIG),
       _config: PhantomData,
     },
-    grandpa: GrandpaConfig { authorities: vec![], _config: PhantomData },
-
-    authority_discovery: AuthorityDiscoveryConfig { keys: vec![], _config: PhantomData },
+    grandpa: GrandpaConfig {
+      authorities: validators.into_iter().map(|validator| (validator.into(), 1)).collect(),
+      _config: PhantomData,
+    },
   }
 }
 

substrate/runtime/Cargo.toml

@@ -50,7 +50,6 @@ coins-pallet = { package = "serai-coins-pallet", path = "../coins/pallet", defau
 dex-pallet = { package = "serai-dex-pallet", path = "../dex/pallet", default-features = false }
 
 validator-sets-pallet = { package = "serai-validator-sets-pallet", path = "../validator-sets/pallet", default-features = false }
-pallet-session = { git = "https://github.com/serai-dex/substrate", default-features = false }
 
 in-instructions-pallet = { package = "serai-in-instructions-pallet", path = "../in-instructions/pallet", default-features = false }
 
@@ -59,8 +58,6 @@ signals-pallet = { package = "serai-signals-pallet", path = "../signals/pallet",
 pallet-babe = { git = "https://github.com/serai-dex/substrate", default-features = false }
 pallet-grandpa = { git = "https://github.com/serai-dex/substrate", default-features = false }
 
-pallet-authority-discovery = { git = "https://github.com/serai-dex/substrate", default-features = false }
-
 frame-system-rpc-runtime-api = { git = "https://github.com/serai-dex/substrate", default-features = false }
 pallet-transaction-payment-rpc-runtime-api = { git = "https://github.com/serai-dex/substrate", default-features = false }
 
@@ -105,7 +102,6 @@ std = [
   "dex-pallet/std",
 
   "validator-sets-pallet/std",
-  "pallet-session/std",
 
   "in-instructions-pallet/std",
 
@@ -114,8 +110,6 @@ std = [
   "pallet-babe/std",
   "pallet-grandpa/std",
 
-  "pallet-authority-discovery/std",
-
   "frame-system-rpc-runtime-api/std",
   "pallet-transaction-payment-rpc-runtime-api/std",
 ]

substrate/runtime/src/lib.rs

@@ -20,7 +20,6 @@ pub use coins_pallet as coins;
 pub use dex_pallet as dex;
 
 pub use validator_sets_pallet as validator_sets;
-pub use pallet_session as session;
 
 pub use in_instructions_pallet as in_instructions;
 
@@ -29,8 +28,6 @@ pub use signals_pallet as signals;
 pub use pallet_babe as babe;
 pub use pallet_grandpa as grandpa;
 
-pub use pallet_authority_discovery as authority_discovery;
-
 // Actually used by the runtime
 use sp_core::OpaqueMetadata;
 use sp_std::prelude::*;
@@ -41,7 +38,7 @@ use sp_version::NativeVersion;
 
 use sp_runtime::{
   create_runtime_str, generic, impl_opaque_keys, KeyTypeId,
-  traits::{Convert, OpaqueKeys, BlakeTwo256, Block as BlockT},
+  traits::{Convert, BlakeTwo256, Block as BlockT},
   transaction_validity::{TransactionSource, TransactionValidity},
   ApplyExtrinsicResult, Perbill,
 };
@@ -83,13 +80,10 @@ pub mod opaque {
     pub struct SessionKeys {
       pub babe: Babe,
       pub grandpa: Grandpa,
-      pub authority_discovery: AuthorityDiscovery,
     }
   }
 }
 
-use opaque::SessionKeys;
-
 #[sp_version::runtime_version]
 pub const VERSION: RuntimeVersion = RuntimeVersion {
   spec_name: create_runtime_str!("serai"),
@@ -167,12 +161,6 @@ impl Contains<RuntimeCall> for CallFilter {
       RuntimeCall::InInstructions(call) => !matches!(call, in_instructions::Call::__Ignore(_, _)),
       RuntimeCall::Signals(call) => !matches!(call, signals::Call::__Ignore(_, _)),
 
-      RuntimeCall::Session(call) => match call {
-        session::Call::set_keys { .. } => true,
-        session::Call::purge_keys { .. } => false,
-        session::Call::__Ignore(_, _) => false,
-      },
-
       RuntimeCall::Babe(call) => match call {
         babe::Call::report_equivocation { .. } => true,
         babe::Call::report_equivocation_unsigned { .. } => true,
@@ -256,6 +244,8 @@ impl dex::Config for Runtime {
 
 impl validator_sets::Config for Runtime {
   type RuntimeEvent = RuntimeEvent;
+
+  type ShouldEndSession = Babe;
 }
 
 pub struct IdentityValidatorIdOf;
@@ -265,18 +255,6 @@ impl Convert<PublicKey, Option<PublicKey>> for IdentityValidatorIdOf {
   }
 }
 
-impl session::Config for Runtime {
-  type RuntimeEvent = RuntimeEvent;
-  type ValidatorId = PublicKey;
-  type ValidatorIdOf = IdentityValidatorIdOf;
-  type ShouldEndSession = Babe;
-  type NextSessionRotation = Babe;
-  type SessionManager = ValidatorSets;
-  type SessionHandler = <SessionKeys as OpaqueKeys>::KeyTypeIdProviders;
-  type Keys = SessionKeys;
-  type WeightInfo = session::weights::SubstrateWeight<Runtime>;
-}
-
 impl signals::Config for Runtime {
   type RuntimeEvent = RuntimeEvent;
   // 1 week
@@ -294,7 +272,7 @@ impl babe::Config for Runtime {
   type EpochDuration = ConstU64<{ 1 * DAYS }>;
   type ExpectedBlockTime = ConstU64<{ TARGET_BLOCK_TIME * 1000 }>;
   type EpochChangeTrigger = pallet_babe::ExternalTrigger;
-  type DisabledValidators = Session;
+  type DisabledValidators = ValidatorSets;
 
   type WeightInfo = ();
 
@@ -317,10 +295,6 @@ impl grandpa::Config for Runtime {
   type EquivocationReportSystem = ();
 }
 
-impl authority_discovery::Config for Runtime {
-  type MaxAuthorities = MaxAuthorities;
-}
-
 pub type Header = generic::Header<BlockNumber, BlakeTwo256>;
 pub type Block = generic::Block<Header, UncheckedExtrinsic>;
 pub type SignedExtra = (
@@ -357,7 +331,6 @@ construct_runtime!(
     Dex: dex,
 
     ValidatorSets: validator_sets,
-    Session: session,
 
     InInstructions: in_instructions,
 
@@ -365,8 +338,6 @@ construct_runtime!(
 
     Babe: babe,
     Grandpa: grandpa,
-
-    AuthorityDiscovery: authority_discovery,
   }
 );
 
@@ -569,7 +540,10 @@ sp_api::impl_runtime_apis! {
 
   impl sp_authority_discovery::AuthorityDiscoveryApi<Block> for Runtime {
     fn authorities() -> Vec<AuthorityDiscoveryId> {
-      AuthorityDiscovery::authorities()
+      Babe::authorities()
+        .into_iter()
+        .map(|(id, _)| AuthorityDiscoveryId::from(id.into_inner()))
+        .collect()
     }
   }
 }

substrate/validator-sets/pallet/Cargo.toml

@@ -21,11 +21,13 @@ sp-io = { git = "https://github.com/serai-dex/substrate", default-features = fal
 sp-std = { git = "https://github.com/serai-dex/substrate", default-features = false }
 sp-application-crypto = { git = "https://github.com/serai-dex/substrate", default-features = false }
 sp-runtime = { git = "https://github.com/serai-dex/substrate", default-features = false }
+sp-session = { git = "https://github.com/serai-dex/substrate", default-features = false }
 
 frame-system = { git = "https://github.com/serai-dex/substrate", default-features = false }
 frame-support = { git = "https://github.com/serai-dex/substrate", default-features = false }
 
-pallet-session = { git = "https://github.com/serai-dex/substrate", default-features = false }
+pallet-babe = { git = "https://github.com/serai-dex/substrate", default-features = false }
+pallet-grandpa = { git = "https://github.com/serai-dex/substrate", default-features = false }
 
 serai-primitives = { path = "../../primitives", default-features = false }
 validator-sets-primitives = { package = "serai-validator-sets-primitives", path = "../primitives", default-features = false }
@@ -42,11 +44,13 @@ std = [
   "sp-std/std",
   "sp-application-crypto/std",
   "sp-runtime/std",
+  "sp-session/std",
 
   "frame-system/std",
   "frame-support/std",
 
-  "pallet-session/std",
+  "pallet-babe/std",
+  "pallet-grandpa/std",
 
   "serai-primitives/std",
   "validator-sets-primitives/std",

substrate/validator-sets/pallet/src/lib.rs

@@ -3,14 +3,20 @@
 #[allow(deprecated, clippy::let_unit_value)] // TODO
 #[frame_support::pallet]
 pub mod pallet {
+  use super::*;
+
   use scale_info::TypeInfo;
 
   use sp_core::sr25519::{Public, Signature};
   use sp_std::{vec, vec::Vec};
   use sp_application_crypto::RuntimePublic;
+  use sp_session::ShouldEndSession;
+  use sp_runtime::traits::IsMember;
 
   use frame_system::pallet_prelude::*;
-  use frame_support::{pallet_prelude::*, StoragePrefixedMap};
+  use frame_support::{
+    pallet_prelude::*, traits::DisabledValidators, BoundedVec, WeakBoundedVec, StoragePrefixedMap,
+  };
 
   use serai_primitives::*;
   pub use validator_sets_primitives as primitives;
@@ -18,14 +24,20 @@ pub mod pallet {
 
   use coins_pallet::Pallet as Coins;
 
+  use pallet_babe::{Pallet as Babe, AuthorityId as BabeAuthorityId};
+  use pallet_grandpa::{Pallet as Grandpa, AuthorityId as GrandpaAuthorityId};
+
   #[pallet::config]
   pub trait Config:
     frame_system::Config<AccountId = Public>
     + coins_pallet::Config
-    + pallet_session::Config<ValidatorId = Public>
+    + pallet_babe::Config
+    + pallet_grandpa::Config
     + TypeInfo
   {
     type RuntimeEvent: IsType<<Self as frame_system::Config>::RuntimeEvent> + From<Event<Self>>;
+
+    type ShouldEndSession: ShouldEndSession<BlockNumberFor<Self>>;
   }
 
   #[pallet::genesis_config]
@@ -52,22 +64,18 @@ pub mod pallet {
   pub struct Pallet<T>(PhantomData<T>);
 
   /// The current session for a network.
-  ///
-  /// This does not store the current session for Serai. pallet_session handles that.
   // Uses Identity for the lookup to avoid a hash of a severely limited fixed key-space.
   #[pallet::storage]
+  #[pallet::getter(fn session)]
   pub type CurrentSession<T: Config> = StorageMap<_, Identity, NetworkId, Session, OptionQuery>;
   impl<T: Config> Pallet<T> {
-    pub fn session(network: NetworkId) -> Option<Session> {
-      if network == NetworkId::Serai {
-        Some(Session(pallet_session::Pallet::<T>::current_index()))
-      } else {
-        CurrentSession::<T>::get(network)
-      }
-    }
-
     pub fn latest_decided_session(network: NetworkId) -> Option<Session> {
-      CurrentSession::<T>::get(network)
+      let session = Self::session(network);
+      // we already decided about the next session for serai.
+      if network == NetworkId::Serai {
+        return session.map(|s| Session(s.0 + 1));
+      }
+      session
     }
   }
 
@@ -84,7 +92,7 @@ pub mod pallet {
     _,
     Identity,
     NetworkId,
-    BoundedVec<Public, ConstU32<{ MAX_KEY_SHARES_PER_SET }>>,
+    BoundedVec<(Public, u64), ConstU32<{ MAX_KEY_SHARES_PER_SET }>>,
     ValueQuery,
   >;
   /// The validators selected to be in-set, yet with the ability to perform a check for presence.
@@ -105,13 +113,8 @@ pub mod pallet {
     // current set's validators
     #[inline]
     fn in_active_serai_set(account: Public) -> bool {
-      // TODO: This is bounded O(n). Can we get O(1) via a storage lookup, like we do with InSet?
-      for validator in pallet_session::Pallet::<T>::validators() {
-        if validator == account {
-          return true;
-        }
-      }
-      false
+      // TODO: is_member is internally O(n). Update Babe to use an O(1) storage lookup?
+      Babe::<T>::is_member(&BabeAuthorityId::from(account))
     }
 
     /// Returns true if the account is included in an active set.
@@ -297,14 +300,12 @@ pub mod pallet {
   impl<T: Config> Pallet<T> {
     fn new_set(network: NetworkId) {
       // Update CurrentSession
-      let session = if network != NetworkId::Serai {
+      let session = {
         let new_session = CurrentSession::<T>::get(network)
           .map(|session| Session(session.0 + 1))
           .unwrap_or(Session(0));
         CurrentSession::<T>::set(network, Some(new_session));
         new_session
-      } else {
-        Self::session(network).unwrap_or(Session(0))
       };
 
       // Clear the current InSet
@@ -326,20 +327,26 @@ pub mod pallet {
       while key_shares < u64::from(MAX_KEY_SHARES_PER_SET) {
         let Some((key, amount)) = iter.next() else { break };
 
+        let these_key_shares = amount.0 / allocation_per_key_share;
         InSet::<T>::set(Self::in_set_key(network, key), Some(()));
-        participants.push(key);
+        participants.push((key, these_key_shares));
 
         // This can technically set key_shares to a value exceeding MAX_KEY_SHARES_PER_SET
         // Off-chain, the key shares per validator will be accordingly adjusted
-        key_shares += amount.0 / allocation_per_key_share;
+        key_shares += these_key_shares;
         total_stake += amount.0;
       }
       TotalAllocatedStake::<T>::set(network, Some(Amount(total_stake)));
 
       let set = ValidatorSet { network, session };
       Pallet::<T>::deposit_event(Event::NewSet { set });
+
+      // Only set the MuSig key for non-Serai sets, as only non-Serai sets should publish keys
       if network != NetworkId::Serai {
-        MuSigKeys::<T>::set(set, Some(musig_key(set, &participants)));
+        MuSigKeys::<T>::set(
+          set,
+          Some(musig_key(set, &participants.iter().map(|(id, _)| *id).collect::<Vec<_>>())),
+        );
       }
       Participants::<T>::set(network, participants.try_into().unwrap());
     }
@@ -372,6 +379,19 @@ pub mod pallet {
     NonExistentValidator,
   }
 
+  #[pallet::hooks]
+  impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+    fn on_initialize(n: BlockNumberFor<T>) -> Weight {
+      if T::ShouldEndSession::should_end_session(n) {
+        Self::rotate_session();
+        // TODO: set the proper weights
+        T::BlockWeights::get().max_block
+      } else {
+        Weight::zero()
+      }
+    }
+  }
+
   #[pallet::genesis_build]
   impl<T: Config> BuildGenesisConfig for GenesisConfig<T> {
     fn build(&self) {
@@ -621,9 +641,10 @@ pub mod pallet {
       for network in serai_primitives::NETWORKS {
         // If this network hasn't started sessions yet, don't start one now
         let Some(current_session) = Self::session(network) else { continue };
-        // Only spawn a NewSet if the current set was actually established with a completed
-        // handover protocol
-        if Self::handover_completed(network, current_session) {
+        // Only spawn a new set if:
+        // - This is Serai, as we need to rotate Serai upon a new session (per Babe)
+        // - The current set was actually established with a completed handover protocol
+        if (network == NetworkId::Serai) || Self::handover_completed(network, current_session) {
           Pallet::<T>::new_set(network);
         }
       }
@@ -649,6 +670,39 @@ pub mod pallet {
       }
       PendingDeallocations::<T>::take((network, session, key))
     }
+
+    fn rotate_session() {
+      let prior_serai_participants = Self::participants(NetworkId::Serai);
+      let prior_serai_session = Self::session(NetworkId::Serai).unwrap();
+
+      // TODO: T::SessionHandler::on_before_session_ending() was here.
+      // end the current serai session.
+      Self::retire_set(ValidatorSet { network: NetworkId::Serai, session: prior_serai_session });
+
+      // make a new session and get the next validator set.
+      Self::new_session();
+
+      // Update Babe and Grandpa
+      let session = prior_serai_session.0 + 1;
+      let validators = prior_serai_participants;
+      let next_validators = Self::participants(NetworkId::Serai);
+      Babe::<T>::enact_epoch_change(
+        WeakBoundedVec::force_from(
+          validators.iter().copied().map(|(id, w)| (BabeAuthorityId::from(id), w)).collect(),
+          None,
+        ),
+        WeakBoundedVec::force_from(
+          next_validators.into_iter().map(|(id, w)| (BabeAuthorityId::from(id), w)).collect(),
+          None,
+        ),
+        Some(session),
+      );
+      Grandpa::<T>::new_session(
+        true,
+        session,
+        validators.into_iter().map(|(id, w)| (GrandpaAuthorityId::from(id), w)).collect(),
+      );
+    }
   }
 
   #[pallet::call]
@@ -667,7 +721,7 @@ pub mod pallet {
       // (called by pre_dispatch) checks it
       let _ = signature;
 
-      let session = Session(pallet_session::Pallet::<T>::current_index());
+      let session = Self::session(NetworkId::Serai).unwrap();
 
       let set = ValidatorSet { session, network };
 
@@ -742,11 +796,13 @@ pub mod pallet {
       };
 
       // Don't allow the Serai set to set_keys, as they have no reason to do so
+      // This should already be covered by the lack of key in MuSigKeys, yet it doesn't hurt to be
+      // explicit
       if network == &NetworkId::Serai {
         Err(InvalidTransaction::Custom(0))?;
       }
 
-      let session = Session(pallet_session::Pallet::<T>::current_index());
+      let session = Self::session(NetworkId::Serai).unwrap();
 
       let set = ValidatorSet { session, network: *network };
       match Self::verify_signature(set, key_pair, signature) {
@@ -779,26 +835,11 @@ pub mod pallet {
     }
   }
 
-  // Call order is end_session(i - 1) -> start_session(i) -> new_session(i + 1)
-  // new_session(i + 1) is called immediately after start_session(i)
-  // then we wait until the session ends then get a call to end_session(i) and so on.
-  impl<T: Config> pallet_session::SessionManager<T::ValidatorId> for Pallet<T> {
-    fn new_session(_new_index: u32) -> Option<Vec<T::ValidatorId>> {
-      Self::new_session();
-      // TODO: Where do we return their stake?
-      Some(Self::participants(NetworkId::Serai).into())
+  impl<T: Config> DisabledValidators for Pallet<T> {
+    fn is_disabled(_: u32) -> bool {
+      // TODO
+      false
     }
-
-    fn new_session_genesis(_: u32) -> Option<Vec<T::ValidatorId>> {
-      // TODO: Because we don't call new_session here, we don't emit NewSet { Serai, session: 1 }
-      Some(Self::participants(NetworkId::Serai).into())
-    }
-
-    fn end_session(end_index: u32) {
-      Self::retire_set(ValidatorSet { network: NetworkId::Serai, session: Session(end_index) })
-    }
-
-    fn start_session(_start_index: u32) {}
   }
 }