Add signals pallet

Resolves #353

Implements code such that:

- 80% of validators (by stake) must be in favor of a signal for the network to
  be
- 80% of networks (by stake) must be in favor of a signal for it to be locked
  in
- After a signal has been locked in for two weeks, the network halts

The intention is to:

1) Not allow validators to unilaterally declare new consensus rules.

No method of declaring new consensus rules is provided by this pallet. Solely a
way to deprecate the current rules, with a signaled for successor. All nodes
must then individually decide whether or not to download and run a new node
which has new rules, and if so, which rules.

2) Not place blobs on chain.

Even if they'd be reproducible, it's just a lot of data to chuck on the
blockchain.

Cargo.lock

@@ -8566,6 +8566,7 @@ dependencies = [
  "serai-coins-pallet",
  "serai-in-instructions-pallet",
  "serai-primitives",
+ "serai-signals-pallet",
  "serai-staking-pallet",
  "serai-validator-sets-pallet",
  "sp-api",
@@ -8584,6 +8585,20 @@ dependencies = [
  "substrate-wasm-builder",
 ]
 
+[[package]]
+name = "serai-signals-pallet"
+version = "0.1.0"
+dependencies = [
+ "frame-support",
+ "frame-system",
+ "parity-scale-codec",
+ "scale-info",
+ "serai-primitives",
+ "serai-validator-sets-pallet",
+ "sp-core",
+ "sp-io",
+]
+
 [[package]]
 name = "serai-staking-pallet"
 version = "0.1.0"

Cargo.toml

@@ -48,6 +48,8 @@ members = [
 
   "substrate/staking/pallet",
 
+  "substrate/signals/pallet",
+
   "substrate/runtime",
   "substrate/node",
 

deny.toml

@@ -62,6 +62,8 @@ exceptions = [
 
   { allow = ["AGPL-3.0"], name = "serai-staking-pallet" },
 
+  { allow = ["AGPL-3.0"], name = "serai-signals-pallet" },
+
   { allow = ["AGPL-3.0"], name = "serai-runtime" },
   { allow = ["AGPL-3.0"], name = "serai-node" },
 

substrate/in-instructions/pallet/src/lib.rs

@@ -84,7 +84,10 @@ pub mod pallet {
   fn keys_for_network<T: Config>(
     network: NetworkId,
   ) -> Result<(Session, Option<Public>, Option<Public>), InvalidTransaction> {
-    let session = ValidatorSets::<T>::session(network);
+    // If there's no session set, and therefore no keys set, then this must be an invalid signature
+    let Some(session) = ValidatorSets::<T>::session(network) else {
+      Err(InvalidTransaction::BadProof)?
+    };
     let mut set = ValidatorSet { session, network };
     let latest = ValidatorSets::<T>::keys(set).map(|keys| keys.0);
     let prior = if set.session.0 != 0 {
@@ -93,7 +96,6 @@ pub mod pallet {
     } else {
       None
     };
-    // If there's no keys set, then this must be an invalid signature
     if prior.is_none() && latest.is_none() {
       Err(InvalidTransaction::BadProof)?;
     }

substrate/runtime/Cargo.toml

@@ -48,8 +48,10 @@ pallet-transaction-payment = { git = "https://github.com/serai-dex/substrate", d
 coins-pallet = { package = "serai-coins-pallet", path = "../coins/pallet", default-features = false }
 in-instructions-pallet = { package = "serai-in-instructions-pallet", path = "../in-instructions/pallet", default-features = false }
 
-staking-pallet = { package = "serai-staking-pallet", path = "../staking/pallet", default-features = false }
 validator-sets-pallet = { package = "serai-validator-sets-pallet", path = "../validator-sets/pallet", default-features = false }
+staking-pallet = { package = "serai-staking-pallet", path = "../staking/pallet", default-features = false }
+
+signals-pallet = { package = "serai-signals-pallet", path = "../signals/pallet", 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 }
@@ -100,8 +102,10 @@ std = [
   "coins-pallet/std",
   "in-instructions-pallet/std",
 
-  "staking-pallet/std",
   "validator-sets-pallet/std",
+  "staking-pallet/std",
+
+  "signals-pallet/std",
 
   "pallet-session/std",
   "pallet-babe/std",

substrate/runtime/src/lib.rs

@@ -22,6 +22,8 @@ pub use in_instructions_pallet as in_instructions;
 pub use staking_pallet as staking;
 pub use validator_sets_pallet as validator_sets;
 
+pub use signals_pallet as signals;
+
 pub use pallet_session as session;
 pub use pallet_babe as babe;
 pub use pallet_grandpa as grandpa;
@@ -46,7 +48,7 @@ use sp_runtime::{
 use primitives::{PublicKey, SeraiAddress, AccountLookup, Signature, SubstrateAmount};
 
 use support::{
-  traits::{ConstU8, ConstU64, Contains},
+  traits::{ConstU8, ConstU32, ConstU64, Contains},
   weights::{
     constants::{RocksDbWeight, WEIGHT_REF_TIME_PER_SECOND},
     IdentityFee, Weight,
@@ -232,11 +234,18 @@ impl in_instructions::Config for Runtime {
   type RuntimeEvent = RuntimeEvent;
 }
 
-impl staking::Config for Runtime {}
-
 impl validator_sets::Config for Runtime {
   type RuntimeEvent = RuntimeEvent;
 }
+impl staking::Config for Runtime {}
+
+impl signals::Config for Runtime {
+  type RuntimeEvent = RuntimeEvent;
+  // 1 week
+  type ValidityDuration = ConstU32<{ (7 * 24 * 60 * 60) / (TARGET_BLOCK_TIME as u32) }>;
+  // 2 weeks
+  type LockInDuration = ConstU32<{ (2 * 7 * 24 * 60 * 60) / (TARGET_BLOCK_TIME as u32) }>;
+}
 
 pub struct IdentityValidatorIdOf;
 impl Convert<PublicKey, Option<PublicKey>> for IdentityValidatorIdOf {
@@ -324,9 +333,10 @@ construct_runtime!(
     InInstructions: in_instructions,
 
     ValidatorSets: validator_sets,
-
     Staking: staking,
 
+    Signals: signals,
+
     Session: session,
     Babe: babe,
     Grandpa: grandpa,

substrate/signals/pallet/Cargo.toml

@@ -0,0 +1,47 @@
+[package]
+name = "serai-signals-pallet"
+version = "0.1.0"
+description = "Signals pallet"
+license = "AGPL-3.0-only"
+repository = "https://github.com/serai-dex/serai/tree/develop/substrate/signals/pallet"
+authors = ["Luke Parker <lukeparker5132@gmail.com>"]
+edition = "2021"
+
+[package.metadata.docs.rs]
+all-features = true
+rustdoc-args = ["--cfg", "docsrs"]
+
+[dependencies]
+scale = { package = "parity-scale-codec", version = "3", default-features = false, features = ["derive"] }
+scale-info = { version = "2", default-features = false, features = ["derive"] }
+
+sp-core = { git = "https://github.com/serai-dex/substrate", default-features = false }
+sp-io = { 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 }
+
+serai-primitives = { path = "../../primitives", default-features = false }
+validator-sets-pallet = { package = "serai-validator-sets-pallet", path = "../../validator-sets/pallet", default-features = false }
+
+[features]
+std = [
+  "scale/std",
+  "scale-info/std",
+
+  "sp-core/std",
+  "sp-io/std",
+
+  "frame-system/std",
+  "frame-support/std",
+
+  "serai-primitives/std",
+  "validator-sets-pallet/std",
+]
+
+runtime-benchmarks = [
+  "frame-system/runtime-benchmarks",
+  "frame-support/runtime-benchmarks",
+]
+
+default = ["std"]

substrate/signals/pallet/LICENSE

@@ -0,0 +1,15 @@
+AGPL-3.0-only license
+
+Copyright (c) 2023 Luke Parker
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License Version 3 as
+published by the Free Software Foundation.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program. If not, see <http://www.gnu.org/licenses/>.

substrate/signals/pallet/src/lib.rs

@@ -0,0 +1,381 @@
+#![cfg_attr(not(feature = "std"), no_std)]
+
+#[allow(deprecated, clippy::let_unit_value)] // TODO
+#[frame_support::pallet]
+pub mod pallet {
+  use scale_info::TypeInfo;
+
+  use sp_core::sr25519::Public;
+  use sp_io::hashing::blake2_256;
+
+  use frame_system::pallet_prelude::*;
+  use frame_support::pallet_prelude::*;
+
+  use serai_primitives::*;
+  use validator_sets_pallet::{primitives::ValidatorSet, Config as VsConfig, Pallet as VsPallet};
+
+  #[pallet::config]
+  pub trait Config: frame_system::Config<AccountId = Public> + VsConfig + TypeInfo {
+    type RuntimeEvent: IsType<<Self as frame_system::Config>::RuntimeEvent> + From<Event<Self>>;
+
+    type ValidityDuration: Get<u32>;
+    type LockInDuration: Get<u32>;
+  }
+
+  #[pallet::pallet]
+  pub struct Pallet<T>(PhantomData<T>);
+
+  #[derive(Clone, PartialEq, Eq, Encode, Decode, TypeInfo, MaxEncodedLen)]
+  struct RegisteredSignal<T: Config> {
+    signal: [u8; 32],
+    registrant: T::AccountId,
+    registed_at: BlockNumberFor<T>,
+  }
+
+  #[pallet::storage]
+  type RegisteredSignals<T: Config> =
+    StorageMap<_, Blake2_128Concat, [u8; 32], RegisteredSignal<T>, OptionQuery>;
+
+  #[pallet::storage]
+  pub type Favors<T: Config> = StorageDoubleMap<
+    _,
+    Blake2_128Concat,
+    ([u8; 32], NetworkId),
+    Blake2_128Concat,
+    T::AccountId,
+    (),
+    OptionQuery,
+  >;
+
+  #[pallet::storage]
+  pub type SetsInFavor<T: Config> =
+    StorageMap<_, Blake2_128Concat, ([u8; 32], ValidatorSet), (), OptionQuery>;
+
+  #[pallet::storage]
+  pub type LockedInSignal<T: Config> = StorageValue<_, ([u8; 32], BlockNumberFor<T>), OptionQuery>;
+
+  #[pallet::event]
+  #[pallet::generate_deposit(pub(super) fn deposit_event)]
+  pub enum Event<T: Config> {
+    SignalRegistered { signal_id: [u8; 32], signal: [u8; 32], registrant: T::AccountId },
+    SignalRevoked { signal_id: [u8; 32] },
+    SignalFavored { signal_id: [u8; 32], by: T::AccountId, for_network: NetworkId },
+    SetInFavor { signal_id: [u8; 32], set: ValidatorSet },
+    SignalLockedIn { signal_id: [u8; 32] },
+    SetNoLongerInFavor { signal_id: [u8; 32], set: ValidatorSet },
+    FavorRevoked { signal_id: [u8; 32], by: T::AccountId, for_network: NetworkId },
+    AgainstSignal { signal_id: [u8; 32], who: T::AccountId, for_network: NetworkId },
+  }
+
+  #[pallet::error]
+  pub enum Error<T> {
+    SignalLockedIn,
+    SignalAlreadyRegistered,
+    NotSignalRegistrant,
+    NonExistantSignal,
+    ExpiredSignal,
+    NotValidator,
+    RevokingNonExistantFavor,
+  }
+
+  // 80% threshold
+  const REQUIREMENT_NUMERATOR: u64 = 4;
+  const REQUIREMENT_DIVISOR: u64 = 5;
+
+  impl<T: Config> Pallet<T> {
+    // Returns true if this network's current set is in favor of the signal.
+    //
+    // Must only be called for networks which have a set decided.
+    fn tally_for_network(signal_id: [u8; 32], network: NetworkId) -> Result<bool, Error<T>> {
+      let this_network_session = VsPallet::<T>::latest_decided_session(network).unwrap();
+      let this_set = ValidatorSet { network, session: this_network_session };
+
+      // This is a bounded O(n) (which is still acceptable) due to the infeasibility of caching
+      // here
+      // TODO: Make caching feasible? Do a first-pass with cache then actual pass before
+      // execution?
+      let mut iter = Favors::<T>::iter_prefix_values((signal_id, network));
+      let mut needed_favor = (VsPallet::<T>::total_allocated_stake(network).unwrap().0 *
+        REQUIREMENT_NUMERATOR)
+        .div_ceil(REQUIREMENT_DIVISOR);
+      while iter.next().is_some() && (needed_favor != 0) {
+        let item_key = iter.last_raw_key();
+        // `.len() - 32` is safe because AccountId is bound to being Public, which is 32 bytes
+        let account = T::AccountId::decode(&mut &item_key[(item_key.len() - 32) ..]).unwrap();
+        if VsPallet::<T>::in_latest_decided_set(network, account) {
+          // This call uses the current allocation, not the allocation at the time of set
+          // decision
+          // This is deemed safe due to the validator-set pallet's deallocation scheduling
+          // unwrap is safe due to being in the latest decided set
+          needed_favor =
+            needed_favor.saturating_sub(VsPallet::<T>::allocation((network, account)).unwrap().0);
+        }
+      }
+
+      if needed_favor == 0 {
+        // Set the set as in favor until someone triggers a re-tally
+        //
+        // Since a re-tally is an extra step we can't assume will occur, this effectively means a
+        // network in favor across any point in its Session is in favor for its entire Session
+        // While a malicious actor could increase their stake, favor a signal, then deallocate,
+        // this is largely prevented by deallocation scheduling
+        //
+        // At any given point, only just under 50% of a set can be immediately deallocated
+        // (if each validator has just under two key shares, they can deallocate the entire amount
+        // above a single key share)
+        //
+        // This means that if a signal has a 67% adoption threshold, and someone executes this
+        // attack, they still have a majority of the allocated stake (though less of a majority
+        // than desired)
+        //
+        // With the 80% threshold, removing 39.9% creates a 40.1% to 20% ratio, which is still
+        // the BFT threshold of 67%
+        if !SetsInFavor::<T>::contains_key((signal_id, this_set)) {
+          SetsInFavor::<T>::set((signal_id, this_set), Some(()));
+          Self::deposit_event(Event::SetInFavor { signal_id, set: this_set });
+        }
+        Ok(true)
+      } else {
+        if SetsInFavor::<T>::contains_key((signal_id, this_set)) {
+          // This should no longer be under the current tally
+          SetsInFavor::<T>::remove((signal_id, this_set));
+          Self::deposit_event(Event::SetNoLongerInFavor { signal_id, set: this_set });
+        }
+        Ok(false)
+      }
+    }
+
+    fn tally_for_all_networks(signal_id: [u8; 32]) -> Result<bool, Error<T>> {
+      let mut total_in_favor_stake = 0;
+      let mut total_allocated_stake = 0;
+      for network in serai_primitives::NETWORKS {
+        let Some(latest_decided_session) = VsPallet::<T>::latest_decided_session(network) else {
+          continue;
+        };
+        // If it has a session, it should have a total allocated stake value
+        let network_stake = VsPallet::<T>::total_allocated_stake(network).unwrap();
+        if SetsInFavor::<T>::contains_key((
+          signal_id,
+          ValidatorSet { network, session: latest_decided_session },
+        )) {
+          total_in_favor_stake += network_stake.0;
+        }
+        total_allocated_stake += network_stake.0;
+      }
+
+      Ok(
+        total_in_favor_stake >=
+          (total_allocated_stake * REQUIREMENT_NUMERATOR).div_ceil(REQUIREMENT_DIVISOR),
+      )
+    }
+
+    fn revoke_favor_internal(
+      account: T::AccountId,
+      signal_id: [u8; 32],
+      for_network: NetworkId,
+    ) -> DispatchResult {
+      if !Favors::<T>::contains_key((signal_id, for_network), account) {
+        Err::<(), _>(Error::<T>::RevokingNonExistantFavor)?;
+      }
+      Favors::<T>::remove((signal_id, for_network), account);
+      Self::deposit_event(Event::<T>::FavorRevoked { signal_id, by: account, for_network });
+      // tally_for_network assumes the network is active, which is implied by having prior set a
+      // favor for it
+      // Technically, this tally may make the network in favor and justify re-tallying for all
+      // networks
+      // Its assumed not to
+      Self::tally_for_network(signal_id, for_network)?;
+      Ok(())
+    }
+  }
+
+  #[pallet::call]
+  impl<T: Config> Pallet<T> {
+    #[pallet::call_index(0)]
+    #[pallet::weight(0)]
+    pub fn register_signal(origin: OriginFor<T>, signal: [u8; 32]) -> DispatchResult {
+      if LockedInSignal::<T>::exists() {
+        Err::<(), _>(Error::<T>::SignalLockedIn)?;
+      }
+
+      let account = ensure_signed(origin)?;
+
+      // Bind the signal ID to the proposer
+      // This prevents a malicious actor from frontrunning a proposal, causing them to be the
+      // registrant, just to cancel it later
+      let mut signal_preimage = account.encode();
+      signal_preimage.extend(signal);
+      let signal_id = blake2_256(&signal_preimage);
+
+      if RegisteredSignals::<T>::get(signal_id).is_some() {
+        Err::<(), _>(Error::<T>::SignalAlreadyRegistered)?;
+      }
+      RegisteredSignals::<T>::set(
+        signal_id,
+        Some(RegisteredSignal {
+          signal,
+          registrant: account,
+          registed_at: frame_system::Pallet::<T>::block_number(),
+        }),
+      );
+      Self::deposit_event(Event::<T>::SignalRegistered { signal_id, signal, registrant: account });
+      Ok(())
+    }
+
+    #[pallet::call_index(1)]
+    #[pallet::weight(0)]
+    pub fn revoke_signal(origin: OriginFor<T>, signal_id: [u8; 32]) -> DispatchResult {
+      let account = ensure_signed(origin)?;
+      let Some(registered_signal) = RegisteredSignals::<T>::get(signal_id) else {
+        return Err::<(), _>(Error::<T>::NonExistantSignal.into());
+      };
+      if account != registered_signal.registrant {
+        Err::<(), _>(Error::<T>::NotSignalRegistrant)?;
+      }
+      RegisteredSignals::<T>::remove(signal_id);
+
+      // If this signal was locked in, remove it
+      // This lets a post-lock-in discovered fault be prevented from going live without
+      // intervention by all node runners
+      if LockedInSignal::<T>::get().map(|(signal_id, _block_number)| signal_id) == Some(signal_id) {
+        LockedInSignal::<T>::kill();
+      }
+
+      Self::deposit_event(Event::<T>::SignalRevoked { signal_id });
+      Ok(())
+    }
+
+    #[pallet::call_index(2)]
+    #[pallet::weight(0)]
+    pub fn favor(
+      origin: OriginFor<T>,
+      signal_id: [u8; 32],
+      for_network: NetworkId,
+    ) -> DispatchResult {
+      if LockedInSignal::<T>::exists() {
+        Err::<(), _>(Error::<T>::SignalLockedIn)?;
+      }
+
+      let account = ensure_signed(origin)?;
+      let Some(registered_signal) = RegisteredSignals::<T>::get(signal_id) else {
+        return Err::<(), _>(Error::<T>::NonExistantSignal.into());
+      };
+      // Check the signal isn't out of date
+      if (registered_signal.registed_at + T::ValidityDuration::get().into()) <
+        frame_system::Pallet::<T>::block_number()
+      {
+        Err::<(), _>(Error::<T>::ExpiredSignal)?;
+      }
+
+      // Check the signer is a validator
+      // Technically, in the case of Serai, this will check they're planned to be in the next set,
+      // not that they are in the current set
+      // This is a practical requirement due to the lack of tracking historical allocations, and
+      // fine for the purposes here
+      if !VsPallet::<T>::in_latest_decided_set(for_network, account) {
+        Err::<(), _>(Error::<T>::NotValidator)?;
+      }
+
+      // Set them as in-favor
+      // Doesn't error if they already voted in order to let any validator trigger a re-tally
+      if !Favors::<T>::contains_key((signal_id, for_network), account) {
+        Favors::<T>::set((signal_id, for_network), account, Some(()));
+        Self::deposit_event(Event::SignalFavored { signal_id, by: account, for_network });
+      }
+
+      // Check if the network is in favor
+      // tally_for_network expects the network to be active, which is implied by being in the
+      // latest decided set
+      let network_in_favor = Self::tally_for_network(signal_id, for_network)?;
+
+      // If this network is in favor, check if enough networks are
+      // We could optimize this by only running the following code when the network is *newly* in
+      // favor
+      // Re-running the following code ensures that if networks' allocated stakes change relative
+      // to each other, any new votes will cause a re-tally
+      if network_in_favor {
+        // If enough are, lock in the signal
+        if Self::tally_for_all_networks(signal_id)? {
+          LockedInSignal::<T>::set(Some((
+            signal_id,
+            frame_system::Pallet::<T>::block_number() + T::LockInDuration::get().into(),
+          )));
+          Self::deposit_event(Event::SignalLockedIn { signal_id });
+        }
+      }
+
+      Ok(())
+    }
+
+    /// Revoke favor into an abstaining position.
+    #[pallet::call_index(3)]
+    #[pallet::weight(0)]
+    pub fn revoke_favor(
+      origin: OriginFor<T>,
+      signal_id: [u8; 32],
+      for_network: NetworkId,
+    ) -> DispatchResult {
+      if LockedInSignal::<T>::exists() {
+        Err::<(), _>(Error::<T>::SignalLockedIn)?;
+      }
+
+      // Doesn't check the signal exists due to later checking the favor exists
+      // While the signal may have been revoked, making this pointless, it's not worth the storage
+      // read on every call to check
+      // Since revoke will re-tally, this does technically mean a network will become in-favor of a
+      // revoked signal. Since revoke won't re-tally for all networks/lock-in, this is also fine
+
+      Self::revoke_favor_internal(ensure_signed(origin)?, signal_id, for_network)
+    }
+
+    /// Emit an event standing against the signal.
+    ///
+    /// If the origin is currently in favor of the signal, their favor will be revoked.
+    #[pallet::call_index(4)]
+    #[pallet::weight(0)]
+    pub fn stand_against(
+      origin: OriginFor<T>,
+      signal_id: [u8; 32],
+      for_network: NetworkId,
+    ) -> DispatchResult {
+      if LockedInSignal::<T>::exists() {
+        Err::<(), _>(Error::<T>::SignalLockedIn)?;
+      }
+
+      let account = ensure_signed(origin)?;
+      // If currently in favor, revoke the favor
+      if Favors::<T>::contains_key((signal_id, for_network), account) {
+        Self::revoke_favor_internal(account, signal_id, for_network)?;
+      } else {
+        // Check this Signal exists (which would've been implied by Favors for it existing)
+        if RegisteredSignals::<T>::get(signal_id).is_none() {
+          Err::<(), _>(Error::<T>::NonExistantSignal)?;
+        }
+      }
+
+      // Emit an event that we're against the signal
+      // No actual effects happen besides this
+      Self::deposit_event(Event::<T>::AgainstSignal { signal_id, who: account, for_network });
+      Ok(())
+    }
+  }
+
+  #[pallet::hooks]
+  impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+    fn on_initialize(current_number: BlockNumberFor<T>) -> Weight {
+      // If this is the block at which a locked-in signal has been set for long enough, panic
+      // This will prevent this block from executing and halt the chain
+      if let Some((signal, block_number)) = LockedInSignal::<T>::get() {
+        if block_number == current_number {
+          panic!(
+            "locked-in signal {} has been set for too long",
+            sp_core::hexdisplay::HexDisplay::from(&signal)
+          );
+        }
+      }
+      Weight::zero() // TODO
+    }
+  }
+}
+
+pub use pallet::*;

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

@@ -55,13 +55,17 @@ pub mod pallet {
   #[pallet::storage]
   pub type CurrentSession<T: Config> = StorageMap<_, Identity, NetworkId, Session, OptionQuery>;
   impl<T: Config> Pallet<T> {
-    pub fn session(network: NetworkId) -> Session {
+    pub fn session(network: NetworkId) -> Option<Session> {
       if network == NetworkId::Serai {
-        Session(pallet_session::Pallet::<T>::current_index())
+        Some(Session(pallet_session::Pallet::<T>::current_index()))
       } else {
-        CurrentSession::<T>::get(network).unwrap()
+        CurrentSession::<T>::get(network)
       }
     }
+
+    pub fn latest_decided_session(network: NetworkId) -> Option<Session> {
+      CurrentSession::<T>::get(network)
+    }
   }
 
   /// The allocation required per key share.
@@ -86,9 +90,61 @@ pub mod pallet {
   #[pallet::storage]
   pub type InSet<T: Config> =
     StorageMap<_, Identity, (NetworkId, [u8; 16], Public), (), OptionQuery>;
+  impl<T: Config> Pallet<T> {
+    fn in_set_key(
+      network: NetworkId,
+      account: T::AccountId,
+    ) -> (NetworkId, [u8; 16], T::AccountId) {
+      (network, sp_io::hashing::blake2_128(&(network, account).encode()), account)
+    }
+
+    // This exists as InSet, for Serai, is the validators set for the next session, *not* the
+    // 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
+    }
+
+    /// Returns true if the account is included in an active set.
+    pub fn in_active_set(network: NetworkId, account: Public) -> bool {
+      if network == NetworkId::Serai {
+        Self::in_active_serai_set(account)
+      } else {
+        InSet::<T>::contains_key(Self::in_set_key(network, account))
+      }
+    }
+
+    /// Returns true if the account has been definitively included in an active or upcoming set.
+    pub fn in_set(network: NetworkId, account: Public) -> bool {
+      if InSet::<T>::contains_key(Self::in_set_key(network, account)) {
+        return true;
+      }
+
+      if network == NetworkId::Serai {
+        return Self::in_active_serai_set(account);
+      }
+
+      false
+    }
+
+    /// Returns true if the account is present in the latest decided set.
+    ///
+    /// This is useful when working with `allocation` and `total_allocated_stake`, which return the
+    /// latest information.
+    pub fn in_latest_decided_set(network: NetworkId, account: Public) -> bool {
+      InSet::<T>::contains_key(Self::in_set_key(network, account))
+    }
+  }
 
   /// The total stake allocated to this network by the active set of validators.
   #[pallet::storage]
+  #[pallet::getter(fn total_allocated_stake)]
   pub type TotalAllocatedStake<T: Config> = StorageMap<_, Identity, NetworkId, Amount, OptionQuery>;
 
   /// The current amount allocated to a validator set by a validator.
@@ -245,13 +301,6 @@ pub mod pallet {
   }
 
   impl<T: Config> Pallet<T> {
-    fn in_set_key(
-      network: NetworkId,
-      account: T::AccountId,
-    ) -> (NetworkId, [u8; 16], T::AccountId) {
-      (network, sp_io::hashing::blake2_128(&(network, account).encode()), account)
-    }
-
     fn new_set(network: NetworkId) {
       // Update CurrentSession
       let session = if network != NetworkId::Serai {
@@ -261,7 +310,7 @@ pub mod pallet {
         CurrentSession::<T>::set(network, Some(new_session));
         new_session
       } else {
-        Self::session(network)
+        Self::session(network).unwrap_or(Session(0))
       };
 
       // Clear the current InSet
@@ -512,6 +561,9 @@ pub mod pallet {
       }
 
       // Decrease the allocation now
+      // Since we don't also update TotalAllocatedStake here, TotalAllocatedStake may be greater
+      // than the sum of all allocations, according to the Allocations StorageMap
+      // This is intentional as this allocation has only been queued for deallocation at this time
       Self::set_allocation(network, account, Amount(new_allocation));
 
       if let Some(was_bft) = was_bft {
@@ -520,22 +572,8 @@ pub mod pallet {
         }
       }
 
-      // If we're not in-set, allow immediate deallocation
+      // If we're not in-set, or this doesn't decrease our key shares, allow immediate deallocation
-      let mut active = InSet::<T>::contains_key(Self::in_set_key(network, account));
+      let active = Self::in_set(network, account);
-      // If the network is Serai, also check pallet_session's list of active validators, as our
-      // InSet is actually the queued for next session's validators
-      // Only runs if active isn't already true in order to short-circuit
-      if (!active) && (network == NetworkId::Serai) {
-        // 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 {
-            active = true;
-            break;
-          }
-        }
-      }
-      // Also allow immediate deallocation if the key shares remain the same
       if (!active) || (!decreased_key_shares) {
         if active {
           // Since it's being immediately deallocated, decrease TotalAllocatedStake
@@ -548,7 +586,8 @@ pub mod pallet {
       }
 
       // Set it to PendingDeallocations, letting the staking pallet release it on a future session
-      let mut to_unlock_on = Self::session(network);
+      // This unwrap should be fine as this account is active, meaning a session has occurred
+      let mut to_unlock_on = Self::session(network).unwrap();
       if network == NetworkId::Serai {
         // Since the next Serai set will already have been decided, we can only deallocate once the
         // next set ends
@@ -570,7 +609,7 @@ pub mod pallet {
 
     // Checks if this session has completed the handover from the prior session.
     fn handover_completed(network: NetworkId, session: Session) -> bool {
-      let current_session = Self::session(network);
+      let Some(current_session) = Self::session(network) else { return false };
       // No handover occurs on genesis
       if current_session.0 == 0 {
         return true;
@@ -597,7 +636,8 @@ pub mod pallet {
 
     pub fn new_session() {
       for network in serai_primitives::NETWORKS {
-        let current_session = Self::session(network);
+        // 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) {