Slash and prove membership of prior sessions (#2970)

* skeleton for tracking historical sessions

* refactor OpaqueKeys

* some more skeleton work

* adjust session to new OpaqueKeys API

* further refactoring of key-type-ids

* session gets validator ID parameter

* run up against compiler

* tweak staking to support new session changes

* first run at child storage for deduplication

* Make session use `AccountId` as `ValidatorId`

* run up against child trie issues

* switch to using normal trie but with a fixed prefix

* clear out some println

* add dedup test

* flesh out historical module more

* introduce ExposureOf for staking

* test the historical module

* WASM compiles

* tests all compile

* do some mock change

* fix bulk of tests

* fix staking tests

* test obsolecence mechanic

* Apply suggestions from code review

Co-Authored-By: Bastian Köcher <bkchr@users.noreply.github.com>

* some more style nits

* a couple more nits

* tweak tries

* fix typo thie -> this

substrate/srml/session/src/historical.rs

@@ -0,0 +1,450 @@
+// Copyright 2019 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate 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 General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! An opt-in utility for tracking historical sessions in SRML-session.
+//!
+//! This is generally useful when implementing blockchains that require accountable
+//! safety where validators from some amount f prior sessions must remain slashable.
+//!
+//! Rather than store the full session data for any given session, we instead commit
+//! to the roots of merkle tries containing the session data.
+//!
+//! These roots and proofs of inclusion can be generated at any time during the current session.
+//! Afterwards, the proofs can be fed to a consensus module when reporting misbehavior.
+
+use rstd::prelude::*;
+use parity_codec::{Encode, Decode};
+use primitives::KeyTypeId;
+use primitives::traits::{Convert, OpaqueKeys, Hash as HashT};
+use srml_support::{
+	StorageValue, StorageMap, decl_module, decl_storage,
+};
+use srml_support::{Parameter, print};
+use substrate_trie::{MemoryDB, Trie, TrieMut, TrieDBMut, TrieDB, Recorder};
+
+use super::{SessionIndex, Module as SessionModule};
+
+/// Trait necessary for the historical module.
+pub trait Trait: super::Trait {
+	/// Full identification of the validator.
+	type FullIdentification: Parameter;
+
+	/// A conversion from validator ID to full identification.
+	///
+	/// This should contain any references to economic actors associated with the
+	/// validator, since they may be outdated by the time this is queried from a
+	/// historical trie.
+	///
+	/// This mapping is expected to remain stable in between calls to
+	/// `Self::OnSessionEnding::on_session_ending` which return new validators.
+	type FullIdentificationOf: Convert<Self::ValidatorId, Option<Self::FullIdentification>>;
+}
+
+decl_storage! {
+	trait Store for Module<T: Trait> as Session {
+		/// Mapping from historical session indices to session-data root hash.
+		HistoricalSessions get(historical_root): map SessionIndex => Option<T::Hash>;
+		/// Queued full identifications for queued sessions whose validators have become obsolete.
+		CachedObsolete get(cached_obsolete): map SessionIndex
+			=> Option<Vec<(T::ValidatorId, T::FullIdentification)>>;
+		/// The range of historical sessions we store. [first, last)
+		StoredRange: Option<(SessionIndex, SessionIndex)>;
+	}
+}
+
+decl_module! {
+	pub struct Module<T: Trait> for enum Call where origin: T::Origin { }
+}
+
+impl<T: Trait> Module<T> {
+	/// Prune historical stored session roots up to (but not including)
+	/// `up_to`.
+	pub fn prune_up_to(up_to: SessionIndex) {
+		<Self as Store>::StoredRange::mutate(|range| {
+			let (start, end) = match *range {
+				Some(range) => range,
+				None => return, // nothing to prune.
+			};
+
+			let up_to = rstd::cmp::min(up_to, end);
+
+			if up_to < start {
+				return // out of bounds. harmless.
+			}
+
+			(start..up_to).for_each(<Self as Store>::HistoricalSessions::remove);
+
+			let new_start = up_to;
+			*range = if new_start == end {
+				None // nothing is stored.
+			} else {
+				Some((new_start, end))
+			}
+		})
+	}
+}
+
+/// Specialization of the crate-level `OnSessionEnding` which returns the old
+/// set of full identification when changing the validator set.
+pub trait OnSessionEnding<ValidatorId, FullIdentification>: crate::OnSessionEnding<ValidatorId> {
+	/// Returns the set of new validators, if any, along with the old validators
+	/// and their full identifications.
+	fn on_session_ending(ending: SessionIndex, applied_at: SessionIndex)
+		-> Option<(Vec<ValidatorId>, Vec<(ValidatorId, FullIdentification)>)>;
+}
+
+/// An `OnSessionEnding` implementation that wraps an inner `I` and also
+/// sets the historical trie root of the ending session.
+pub struct NoteHistoricalRoot<T, I>(rstd::marker::PhantomData<(T, I)>);
+
+impl<T: Trait, I> crate::OnSessionEnding<T::ValidatorId> for NoteHistoricalRoot<T, I>
+	where I: OnSessionEnding<T::ValidatorId, T::FullIdentification>
+{
+	fn on_session_ending(ending: SessionIndex, applied_at: SessionIndex) -> Option<Vec<T::ValidatorId>> {
+		StoredRange::mutate(|range| {
+			range.get_or_insert_with(|| (ending, ending)).1 = ending + 1;
+		});
+
+		// do all of this _before_ calling the other `on_session_ending` impl
+		// so that we have e.g. correct exposures from the _current_.
+
+		match ProvingTrie::<T>::generate_for(ending) {
+			Ok(trie) => <HistoricalSessions<T>>::insert(ending, &trie.root),
+			Err(reason) => {
+				print("Failed to generate historical ancestry-inclusion proof.");
+				print(reason);
+			}
+		};
+
+		// trie has been generated for this session, so it's no longer queued.
+		<CachedObsolete<T>>::remove(&ending);
+
+		let (new_validators, old_exposures) = <I as OnSessionEnding<_, _>>::on_session_ending(ending, applied_at)?;
+
+		// every session from `ending+1 .. applied_at` now has obsolete `FullIdentification`
+		// now that a new validator election has occurred.
+		// we cache these in the trie until those sessions themselves end.
+		for obsolete in (ending + 1) .. applied_at {
+			<CachedObsolete<T>>::insert(obsolete, &old_exposures);
+		}
+
+		Some(new_validators)
+	}
+}
+
+type HasherOf<T> = <<T as system::Trait>::Hashing as HashT>::Hasher;
+
+/// A tuple of the validator's ID and their full identification.
+pub type IdentificationTuple<T> = (<T as crate::Trait>::ValidatorId, <T as Trait>::FullIdentification);
+
+/// a trie instance for checking and generating proofs.
+pub struct ProvingTrie<T: Trait> {
+	db: MemoryDB<HasherOf<T>>,
+	root: T::Hash,
+}
+
+impl<T: Trait> ProvingTrie<T> {
+	fn generate_for(now: SessionIndex) -> Result<Self, &'static str> {
+		let mut db = MemoryDB::default();
+		let mut root = Default::default();
+
+		fn build<T: Trait, I>(root: &mut T::Hash, db: &mut MemoryDB<HasherOf<T>>, validators: I)
+			-> Result<(), &'static str>
+			where I: IntoIterator<Item=(T::ValidatorId, Option<T::FullIdentification>)>
+		{
+			let mut trie = TrieDBMut::new(db, root);
+			for (i, (validator, full_id)) in validators.into_iter().enumerate() {
+				let i = i as u32;
+				let keys = match <SessionModule<T>>::load_keys(&validator) {
+					None => continue,
+					Some(k) => k,
+				};
+
+				let full_id = full_id.or_else(|| T::FullIdentificationOf::convert(validator.clone()));
+				let full_id = match full_id {
+					None => return Err("no full identification for a current validator"),
+					Some(full) => (validator, full),
+				};
+
+				// map each key to the owner index.
+				for key_id in T::Keys::key_ids() {
+					let key = keys.get_raw(key_id);
+					let res = (key_id, key).using_encoded(|k|
+						i.using_encoded(|v|
+							trie.insert(k, v)
+						)
+					);
+
+					let _ = res.map_err(|_| "failed to insert into trie")?;
+				}
+
+				// map each owner index to the full identification.
+				let _ = i.using_encoded(|k| full_id.using_encoded(|v| trie.insert(k, v)))
+					.map_err(|_| "failed to insert into trie")?;
+			}
+
+			Ok(())
+		}
+
+		// if the current session's full identifications are obsolete but cached,
+		// use those.
+		if let Some(obsolete) = <CachedObsolete<T>>::get(&now) {
+			build::<T, _>(&mut root, &mut db, obsolete.into_iter().map(|(v, f)| (v, Some(f))))?
+		} else {
+			let validators = <SessionModule<T>>::validators();
+			build::<T, _>(&mut root, &mut db, validators.into_iter().map(|v| (v, None)))?
+		}
+
+		Ok(ProvingTrie {
+			db,
+			root,
+		})
+	}
+
+	fn from_nodes(root: T::Hash, nodes: &[Vec<u8>]) -> Self {
+		use substrate_trie::HashDBT;
+
+		let mut memory_db = MemoryDB::default();
+		for node in nodes {
+			HashDBT::insert(&mut memory_db, &[], &node[..]);
+		}
+
+		ProvingTrie {
+			db: memory_db,
+			root,
+		}
+	}
+
+	/// Prove the full verification data for a given key and key ID.
+	pub fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
+		let trie = TrieDB::new(&self.db, &self.root).ok()?;
+		let mut recorder = Recorder::new();
+		let val_idx = (key_id, key_data).using_encoded(|s| {
+			trie.get_with(s, &mut recorder)
+				.ok()?
+				.and_then(|raw| u32::decode(&mut &*raw))
+		})?;
+
+		val_idx.using_encoded(|s| {
+			trie.get_with(s, &mut recorder)
+				.ok()?
+				.and_then(|raw| <IdentificationTuple<T>>::decode(&mut &*raw))
+		})?;
+
+		Some(recorder.drain().into_iter().map(|r| r.data).collect())
+	}
+
+	/// Access the underlying trie root.
+	pub fn root(&self) -> &T::Hash {
+		&self.root
+	}
+
+	// Check a proof contained within the current memory-db. Returns `None` if the
+	// nodes within the current `MemoryDB` are insufficient to query the item.
+	fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<IdentificationTuple<T>> {
+		let trie = TrieDB::new(&self.db, &self.root).ok()?;
+		let val_idx = (key_id, key_data).using_encoded(|s| trie.get(s))
+			.ok()?
+			.and_then(|raw| u32::decode(&mut &*raw))?;
+
+		val_idx.using_encoded(|s| trie.get(s))
+			.ok()?
+			.and_then(|raw| <IdentificationTuple<T>>::decode(&mut &*raw))
+	}
+
+}
+
+/// Proof of ownership of a specific key.
+#[derive(Encode, Decode, Clone)]
+pub struct Proof {
+	session: SessionIndex,
+	trie_nodes: Vec<Vec<u8>>,
+}
+
+impl<T: Trait, D: AsRef<[u8]>> srml_support::traits::KeyOwnerProofSystem<(KeyTypeId, D)>
+	for Module<T>
+{
+	type Proof = Proof;
+	type FullIdentification = IdentificationTuple<T>;
+
+	fn prove(key: (KeyTypeId, D)) -> Option<Self::Proof> {
+		let session = <SessionModule<T>>::current_index();
+		let trie = ProvingTrie::<T>::generate_for(session).ok()?;
+
+		let (id, data) = key;
+
+		trie.prove(id, data.as_ref()).map(|trie_nodes| Proof {
+			session,
+			trie_nodes,
+		})
+	}
+
+	fn check_proof(key: (KeyTypeId, D), proof: Proof) -> Option<IdentificationTuple<T>> {
+		let (id, data) = key;
+
+		if proof.session == <SessionModule<T>>::current_index() {
+			<SessionModule<T>>::key_owner(id, data.as_ref()).and_then(|owner|
+				T::FullIdentificationOf::convert(owner.clone()).map(move |id| (owner, id))
+			)
+		} else {
+			let root = <HistoricalSessions<T>>::get(&proof.session)?;
+			let trie = ProvingTrie::<T>::from_nodes(root, &proof.trie_nodes);
+
+			trie.query(id, data.as_ref())
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use runtime_io::with_externalities;
+	use substrate_primitives::Blake2Hasher;
+	use primitives::{
+		traits::OnInitialize,
+		testing::{UintAuthorityId, UINT_DUMMY_KEY},
+	};
+	use crate::mock::{
+		NEXT_VALIDATORS, force_new_session,
+		set_next_validators, Test, System, Session,
+	};
+	use srml_support::traits::KeyOwnerProofSystem;
+
+	type Historical = Module<Test>;
+
+	fn new_test_ext() -> runtime_io::TestExternalities<Blake2Hasher> {
+		let mut t = system::GenesisConfig::default().build_storage::<Test>().unwrap().0;
+		t.extend(timestamp::GenesisConfig::<Test> {
+			minimum_period: 5,
+		}.build_storage().unwrap().0);
+		let (storage, _child_storage) = crate::GenesisConfig::<Test> {
+			keys: NEXT_VALIDATORS.with(|l|
+				l.borrow().iter().cloned().map(|i| (i, UintAuthorityId(i))).collect()
+			),
+		}.build_storage().unwrap();
+		t.extend(storage);
+		runtime_io::TestExternalities::new(t)
+	}
+
+	#[test]
+	fn generated_proof_is_good() {
+		with_externalities(&mut new_test_ext(), || {
+			set_next_validators(vec![1, 2]);
+			force_new_session();
+
+			System::set_block_number(1);
+			Session::on_initialize(1);
+
+			let encoded_key_1 = UintAuthorityId(1).encode();
+			let proof = Historical::prove((UINT_DUMMY_KEY, &encoded_key_1[..])).unwrap();
+
+			// proof-checking in the same session is OK.
+			assert!(
+				Historical::check_proof(
+					(UINT_DUMMY_KEY, &encoded_key_1[..]),
+					proof.clone(),
+				).is_some()
+			);
+
+			set_next_validators(vec![1, 2, 4]);
+			force_new_session();
+
+			assert!(Historical::cached_obsolete(&(proof.session + 1)).is_none());
+
+			System::set_block_number(2);
+			Session::on_initialize(2);
+
+			assert!(Historical::cached_obsolete(&(proof.session + 1)).is_some());
+
+			assert!(Historical::historical_root(proof.session).is_some());
+			assert!(Session::current_index() > proof.session);
+
+			// proof-checking in the next session is also OK.
+			assert!(
+				Historical::check_proof(
+					(UINT_DUMMY_KEY, &encoded_key_1[..]),
+					proof.clone(),
+				).is_some()
+			);
+
+			set_next_validators(vec![1, 2, 5]);
+
+			force_new_session();
+			System::set_block_number(3);
+			Session::on_initialize(3);
+
+			assert!(Historical::cached_obsolete(&(proof.session + 1)).is_none());
+		});
+	}
+
+	#[test]
+	fn prune_up_to_works() {
+		with_externalities(&mut new_test_ext(), || {
+			for i in 1..101u64 {
+				set_next_validators(vec![i]);
+				force_new_session();
+
+				System::set_block_number(i);
+				Session::on_initialize(i);
+
+			}
+
+			assert_eq!(StoredRange::get(), Some((0, 100)));
+
+			for i in 1..100 {
+				assert!(Historical::historical_root(i).is_some())
+			}
+
+			Historical::prune_up_to(10);
+			assert_eq!(StoredRange::get(), Some((10, 100)));
+
+			Historical::prune_up_to(9);
+			assert_eq!(StoredRange::get(), Some((10, 100)));
+
+			for i in 10..100 {
+				assert!(Historical::historical_root(i).is_some())
+			}
+
+			Historical::prune_up_to(99);
+			assert_eq!(StoredRange::get(), Some((99, 100)));
+
+			Historical::prune_up_to(100);
+			assert_eq!(StoredRange::get(), None);
+
+			for i in 101..201u64 {
+				set_next_validators(vec![i]);
+				force_new_session();
+
+				System::set_block_number(i);
+				Session::on_initialize(i);
+
+			}
+
+			assert_eq!(StoredRange::get(), Some((100, 200)));
+
+			for i in 101..200 {
+				assert!(Historical::historical_root(i).is_some())
+			}
+
+			Historical::prune_up_to(9999);
+			assert_eq!(StoredRange::get(), None);
+
+			for i in 101..200 {
+				assert!(Historical::historical_root(i).is_none())
+			}
+		});
+	}
+}

substrate/srml/session/src/lib.rs

@@ -29,13 +29,17 @@
 //! <!-- Original author of paragraph: @gavofyork -->
 //!
 //! - **Session:** A session is a period of time that has a constant set of validators. Validators can only join
-//! or exit the validator set at a session change. It is measured in block numbers and set with `set_length`
-//! during a session for use in subsequent sessions.
+//! or exit the validator set at a session change. It is measured in block numbers. The block where a session is
+//! ended is determined by the `ShouldSessionEnd` trait. When the session is ending, a new validator set
+//! can be chosen by `OnSessionEnding` implementations.
 //! - **Session key:** A session key is actually several keys kept together that provide the various signing
 //! functions required by network authorities/validators in pursuit of their duties.
+//! - **Validator ID:** Every account has an associated validator ID. For some simple staking systems, this
+//! may just be the same as the account ID. For staking systems using a stash/controller model,
+//! the validator ID would be the stash account ID of the controller.
 //! - **Session key configuration process:** A session key is set using `set_key` for use in the
-//! next session. It is stored in `NextKeyFor`, a mapping between the caller's `AccountId` and the session
-//! key provided. `set_key` allows users to set their session key prior to becoming a validator.
+//! next session. It is stored in `NextKeyFor`, a mapping between the caller's `ValidatorId` and the session
+//! keys provided. `set_key` allows users to set their session key prior to being selected as validator.
 //! It is a public call since it uses `ensure_signed`, which checks that the origin is a signed account.
 //! As such, the account ID of the origin stored in in `NextKeyFor` may not necessarily be associated with
 //! a block author or a validator. The session keys of accounts are removed once their account balance is zero.
@@ -116,23 +120,30 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 use rstd::{prelude::*, marker::PhantomData, ops::{Sub, Rem}};
-#[cfg(not(feature = "std"))]
-use rstd::alloc::borrow::ToOwned;
-#[cfg(feature = "std")]
-use runtime_io::with_storage;
-use parity_codec::Decode;
-use primitives::traits::{Zero, Member, OpaqueKeys};
+use parity_codec::{Decode, Encode};
+use primitives::KeyTypeId;
+use primitives::traits::{Convert, Zero, Member, OpaqueKeys, TypedKey, Hash};
 use srml_support::{
-	ConsensusEngineId, StorageValue, StorageMap, for_each_tuple, decl_module,
+	dispatch::Result,
+	storage,
+	ConsensusEngineId, StorageValue, for_each_tuple, decl_module,
 	decl_event, decl_storage,
 };
-use srml_support::{ensure, traits::{OnFreeBalanceZero, Get, FindAuthor}, Parameter, print};
-use system::ensure_signed;
+use srml_support::{ensure, traits::{OnFreeBalanceZero, Get, FindAuthor}, Parameter};
+use system::{self, ensure_signed};
+
+#[cfg(test)]
+mod mock;
+
+#[cfg(feature = "historical")]
+pub mod historical;
 
 /// Simple index type with which we can count sessions.
 pub type SessionIndex = u32;
 
+/// Decides whether the session should be ended.
 pub trait ShouldEndSession<BlockNumber> {
+	/// Return `true` if the session should be ended.
 	fn should_end_session(now: BlockNumber) -> bool;
 }
 
@@ -157,28 +168,37 @@ impl<
 	}
 }
 
-pub trait OnSessionEnding<AccountId> {
+/// An event handler for when the session is ending.
+pub trait OnSessionEnding<ValidatorId> {
 	/// Handle the fact that the session is ending, and optionally provide the new validator set.
-	fn on_session_ending(i: SessionIndex) -> Option<Vec<AccountId>>;
+	///
+	/// `ending_index` is the index of the currently ending session.
+	/// The returned validator set, if any, will not be applied until `next_index`.
+	/// `next_index` is guaranteed to be at least `ending_index + 1`, since session indices don't
+	/// repeat.
+	fn on_session_ending(ending_index: SessionIndex, next_index: SessionIndex) -> Option<Vec<ValidatorId>>;
 }
 
 impl<A> OnSessionEnding<A> for () {
-	fn on_session_ending(_: SessionIndex) -> Option<Vec<A>> { None }
+	fn on_session_ending(_: SessionIndex, _: SessionIndex) -> Option<Vec<A>> { None }
 }
 
 /// Handler for when a session keys set changes.
-pub trait SessionHandler<AccountId> {
+pub trait SessionHandler<ValidatorId> {
 	/// Session set has changed; act appropriately.
-	fn on_new_session<Ks: OpaqueKeys>(changed: bool, validators: &[(AccountId, Ks)]);
+	fn on_new_session<Ks: OpaqueKeys>(changed: bool, validators: &[(ValidatorId, Ks)]);
 
 	/// A validator got disabled. Act accordingly until a new session begins.
 	fn on_disabled(validator_index: usize);
 }
 
-pub trait OneSessionHandler<AccountId> {
-	type Key: Decode + Default;
+/// One session-key type handler.
+pub trait OneSessionHandler<ValidatorId> {
+	/// The key type expected.
+	type Key: Decode + Default + TypedKey;
+
 	fn on_new_session<'a, I: 'a>(changed: bool, validators: I)
-		where I: Iterator<Item=(&'a AccountId, Self::Key)>, AccountId: 'a;
+		where I: Iterator<Item=(&'a ValidatorId, Self::Key)>, ValidatorId: 'a;
 	fn on_disabled(i: usize);
 }
 
@@ -193,11 +213,10 @@ macro_rules! impl_session_handlers {
 	( $($t:ident)* ) => {
 		impl<AId, $( $t: OneSessionHandler<AId> ),*> SessionHandler<AId> for ( $( $t , )* ) {
 			fn on_new_session<Ks: OpaqueKeys>(changed: bool, validators: &[(AId, Ks)]) {
-				let mut i: usize = 0;
 				$(
-					i += 1;
 					let our_keys = validators.iter()
-						.map(|k| (&k.0, k.1.get::<$t::Key>(i - 1).unwrap_or_default()));
+						.map(|k| (&k.0, k.1.get::<$t::Key>(<$t::Key as TypedKey>::KEY_TYPE)
+							.unwrap_or_default()));
 					$t::on_new_session(changed, our_keys);
 				)*
 			}
@@ -213,15 +232,17 @@ macro_rules! impl_session_handlers {
 for_each_tuple!(impl_session_handlers);
 
 /// Handler for selecting the genesis validator set.
-pub trait SelectInitialValidators<T: Trait> {
+pub trait SelectInitialValidators<ValidatorId> {
 	/// Returns the initial validator set. If `None` is returned
 	/// all accounts that have session keys set in the genesis block
 	/// will be validators.
-	fn select_initial_validators() -> Option<Vec<T::AccountId>>;
+	fn select_initial_validators() -> Option<Vec<ValidatorId>>;
 }
 
-impl<T: Trait> SelectInitialValidators<T> for () {
-	fn select_initial_validators() -> Option<Vec<T::AccountId>> {
+/// Implementation of `SelectInitialValidators` that does nothing.
+pub struct ConfigValidators;
+impl<V> SelectInitialValidators<V> for ConfigValidators {
+	fn select_initial_validators() -> Option<Vec<V>> {
 		None
 	}
 }
@@ -230,28 +251,35 @@ pub trait Trait: system::Trait {
 	/// The overarching event type.
 	type Event: From<Event> + Into<<Self as system::Trait>::Event>;
 
+	/// A stable ID for a validator.
+	type ValidatorId: Member + Parameter;
+
+	/// A conversion to validator ID to account ID.
+	type ValidatorIdOf: Convert<Self::AccountId, Option<Self::ValidatorId>>;
+
 	/// Indicator for when to end the session.
 	type ShouldEndSession: ShouldEndSession<Self::BlockNumber>;
 
 	/// Handler for when a session is about to end.
-	type OnSessionEnding: OnSessionEnding<Self::AccountId>;
+	type OnSessionEnding: OnSessionEnding<Self::ValidatorId>;
 
 	/// Handler when a session has changed.
-	type SessionHandler: SessionHandler<Self::AccountId>;
+	type SessionHandler: SessionHandler<Self::ValidatorId>;
 
 	/// The keys.
 	type Keys: OpaqueKeys + Member + Parameter + Default;
 
 	/// Select initial validators.
-	type SelectInitialValidators: SelectInitialValidators<Self>;
+	type SelectInitialValidators: SelectInitialValidators<Self::ValidatorId>;
 }
 
-type OpaqueKey = Vec<u8>;
+const DEDUP_KEY_LEN: usize = 13;
+const DEDUP_KEY_PREFIX: &[u8; DEDUP_KEY_LEN] = b":session:keys";
 
 decl_storage! {
 	trait Store for Module<T: Trait> as Session {
 		/// The current set of validators.
-		Validators get(validators): Vec<T::AccountId>;
+		Validators get(validators): Vec<T::ValidatorId>;
 
 		/// Current index of the session.
 		CurrentIndex get(current_index): SessionIndex;
@@ -259,56 +287,48 @@ decl_storage! {
 		/// True if anything has changed in this session.
 		Changed: bool;
 
-		/// The next key to be used for a given validator. At the end of the session, they
-		/// will be moved into the `QueuedKeys` and so changes here will not take effect for
-		/// at least one whole session.
-		NextKeyFor get(next_key_for): map T::AccountId => Option<T::Keys>;
-
 		/// Queued keys changed.
 		QueuedChanged: bool;
 
 		/// The queued keys for the next session. When the next session begins, these keys
 		/// will be used to determine the validator's session keys.
-		QueuedKeys get(queued_keys): Vec<(T::AccountId, T::Keys)>;
+		QueuedKeys get(queued_keys): Vec<(T::ValidatorId, T::Keys)>;
 
-		/// The keys that are currently active.
-		Active: map u32 => Vec<OpaqueKey>;
 	}
 	add_extra_genesis {
-		config(keys): Vec<(T::AccountId, T::Keys)>;
-		build(|storage, _, config: &GenesisConfig<T>| {
-			with_storage(storage, || {
-				let all_validators = config.keys.iter()
-					.map(|(validator, _)| validator.to_owned())
-					.collect::<Vec<_>>();
-				let all_keys = (0..T::Keys::count()).map(|i| (
-					i as u32,
-					config.keys.iter()
-						.map(|x| x.1.get_raw(i).to_vec())
-						.collect::<Vec<OpaqueKey>>(),
-				)).collect::<Vec<(u32, Vec<OpaqueKey>)>>();
-				<Validators<T>>::put(all_validators.clone());
-				for (v, sk) in config.keys.clone() {
-					<NextKeyFor<T>>::insert(v, sk);
-				}
-				for (i, keys) in all_keys {
-					Active::insert(i, keys);
+		config(keys): Vec<(T::ValidatorId, T::Keys)>;
+		build(|
+			storage: &mut primitives::StorageOverlay,
+			_: &mut primitives::ChildrenStorageOverlay,
+			config: &GenesisConfig<T>
+		| {
+			runtime_io::with_storage(storage, || {
+				for (who, keys) in config.keys.iter().cloned() {
+					assert!(
+						<Module<T>>::load_keys(&who).is_none(),
+						"genesis config contained duplicate validator {:?}", who,
+					);
+
+					<Module<T>>::do_set_keys(&who, keys)
+						.expect("genesis config must not contain duplicates; qed");
 				}
 
-				let selected_validators =
-					T::SelectInitialValidators::select_initial_validators()
-					.unwrap_or(all_validators);
-				let selected_keys = selected_validators.iter().map(|validator| {
-					(
-						validator.to_owned(),
-						<NextKeyFor<T>>::get(validator)
-							.unwrap_or_default()
-					)
-				}).collect::<Vec<_>>();
-				<Validators<T>>::put(selected_validators);
-				<QueuedKeys<T>>::put(selected_keys);
-			})
-		})
+				let initial_validators = T::SelectInitialValidators::select_initial_validators()
+					.unwrap_or_else(|| config.keys.iter().map(|(ref v, _)| v.clone()).collect());
+
+				let queued_keys: Vec<_> = initial_validators
+					.iter()
+					.cloned()
+					.map(|v| (
+						v.clone(),
+						<Module<T>>::load_keys(&v).unwrap_or_default(),
+					))
+					.collect();
+
+				<Validators<T>>::put(initial_validators);
+				<QueuedKeys<T>>::put(queued_keys);
+			});
+		});
 	}
 }
 
@@ -331,51 +351,25 @@ decl_module! {
 		/// The dispatch origin of this function must be signed.
 		///
 		/// # <weight>
-		/// - O(1).
+		/// - O(log n) in number of accounts.
 		/// - One extra DB entry.
 		/// # </weight>
-		fn set_keys(origin, keys: T::Keys, proof: Vec<u8>) {
+		fn set_keys(origin, keys: T::Keys, proof: Vec<u8>) -> Result {
 			let who = ensure_signed(origin)?;
 
 			ensure!(keys.ownership_proof_is_valid(&proof), "invalid ownership proof");
 
-			let old_keys = <NextKeyFor<T>>::get(&who);
-			let mut updates = vec![];
+			let who = match T::ValidatorIdOf::convert(who) {
+				Some(val_id) => val_id,
+				None => return Err("no associated validator ID for account."),
+			};
 
-			for i in 0..T::Keys::count() {
-				let new_key = keys.get_raw(i);
-				let maybe_old_key = old_keys.as_ref().map(|o| o.get_raw(i));
-				if maybe_old_key == Some(new_key) {
-					// no change.
-					updates.push(None);
-					continue;
-				}
-				let mut active = Active::get(i as u32);
-				match active.binary_search_by(|k| k[..].cmp(&new_key)) {
-					Ok(_) => return Err("duplicate key provided"),
-					Err(pos) => active.insert(pos, new_key.to_owned()),
-				}
-				if let Some(old_key) = maybe_old_key {
-					match active.binary_search_by(|k| k[..].cmp(&old_key)) {
-						Ok(pos) => { active.remove(pos); }
-						Err(_) => {
-							// unreachable as long as our state is valid. we don't want to panic if
-							// it isn't, though.
-							print("ERROR: active doesn't contain outgoing key");
-						}
-					}
-				}
-				updates.push(Some((i, active)));
-			}
+			Self::do_set_keys(&who, keys)?;
 
-			// Update the active sets.
-			for (i, active) in updates.into_iter().filter_map(|x| x) {
-				Active::insert(i as u32, active);
-			}
-			// Set new keys value for next session.
-			<NextKeyFor<T>>::insert(who, keys);
 			// Something changed.
 			Changed::put(true);
+
+			Ok(())
 		}
 
 		/// Called when a block is finalized. Will rotate session if it is the last
@@ -401,12 +395,14 @@ impl<T: Trait> Module<T> {
 		// Get queued session keys and validators.
 		let session_keys = <QueuedKeys<T>>::get();
 		let validators = session_keys.iter()
-			.map(|(validator, _)| validator.to_owned())
+			.map(|(validator, _)| validator.clone())
 			.collect::<Vec<_>>();
-		<Validators<T>>::put(validators);
+		<Validators<T>>::put(&validators);
+
+		let applied_at = session_index + 2;
 
 		// Get next validator set.
-		let maybe_validators = T::OnSessionEnding::on_session_ending(session_index);
+		let maybe_validators = T::OnSessionEnding::on_session_ending(session_index, applied_at);
 		let next_validators = if let Some(validators) = maybe_validators {
 			next_changed = true;
 			validators
@@ -419,10 +415,11 @@ impl<T: Trait> Module<T> {
 		CurrentIndex::put(session_index);
 
 		// Queue next session keys.
-		let next_session_keys = next_validators.into_iter()
-			.map(|a| { let k = <NextKeyFor<T>>::get(&a).unwrap_or_default(); (a, k) })
+		let queued_amalgamated = next_validators.into_iter()
+			.map(|a| { let k = Self::load_keys(&a).unwrap_or_default(); (a, k) })
 			.collect::<Vec<_>>();
-		<QueuedKeys<T>>::put(next_session_keys);
+
+		<QueuedKeys<T>>::put(queued_amalgamated);
 		QueuedChanged::put(next_changed);
 
 		// Record that this happened.
@@ -439,15 +436,96 @@ impl<T: Trait> Module<T> {
 	}
 
 	/// Disable the validator identified by `c`. (If using with the staking module, this would be
-	/// their *controller* account.)
-	pub fn disable(c: &T::AccountId) -> rstd::result::Result<(), ()> {
+	/// their *stash* account.)
+	pub fn disable(c: &T::ValidatorId) -> rstd::result::Result<(), ()> {
 		Self::validators().iter().position(|i| i == c).map(Self::disable_index).ok_or(())
 	}
+
+	// perform the set_key operation, checking for duplicates.
+	// does not set `Changed`.
+	fn do_set_keys(who: &T::ValidatorId, keys: T::Keys) -> Result {
+		let old_keys = Self::load_keys(&who);
+
+		for id in T::Keys::key_ids() {
+			let key = keys.get_raw(id);
+
+			// ensure keys are without duplication.
+			ensure!(
+				Self::key_owner(id, key).map_or(true, |owner| &owner == who),
+				"registered duplicate key"
+			);
+
+			if let Some(old) = old_keys.as_ref().map(|k| k.get_raw(id)) {
+				if key == old {
+					continue;
+				}
+
+				Self::clear_key_owner(id, old);
+			}
+
+			Self::put_key_owner(id, key, &who);
+		}
+
+		Self::put_keys(&who, &keys);
+
+		Ok(())
+	}
+
+	fn prune_dead_keys(who: &T::ValidatorId) {
+		if let Some(old_keys) = Self::take_keys(who) {
+			for id in T::Keys::key_ids() {
+				let key_data = old_keys.get_raw(id);
+				Self::clear_key_owner(id, key_data);
+			}
+
+			Changed::put(true);
+		}
+	}
+
+	// Child trie storage.
+
+	fn load_keys(v: &T::ValidatorId) -> Option<T::Keys> {
+		storage::unhashed::get(&dedup_trie_key::<T, _>(v))
+	}
+
+	fn take_keys(v: &T::ValidatorId) -> Option<T::Keys> {
+		storage::unhashed::take(&dedup_trie_key::<T, _>(v))
+	}
+
+	fn put_keys(v: &T::ValidatorId, keys: &T::Keys) {
+		storage::unhashed::put(&dedup_trie_key::<T, _>(v), keys)
+	}
+
+	fn key_owner(id: KeyTypeId, key_data: &[u8]) -> Option<T::ValidatorId> {
+		storage::unhashed::get(&dedup_trie_key::<T, _>(&(id, key_data)))
+	}
+
+	fn put_key_owner(id: KeyTypeId, key_data: &[u8], v: &T::ValidatorId) {
+		storage::unhashed::put(&dedup_trie_key::<T, _>(&(id, key_data)), v);
+	}
+
+	fn clear_key_owner(id: KeyTypeId, key_data: &[u8]) {
+		storage::unhashed::kill(&dedup_trie_key::<T, _>(&(id, key_data)));
+	}
 }
 
-impl<T: Trait> OnFreeBalanceZero<T::AccountId> for Module<T> {
-	fn on_free_balance_zero(who: &T::AccountId) {
-		<NextKeyFor<T>>::remove(who);
+fn dedup_trie_key<T: Trait, K: Encode>(key: &K) -> [u8; 32 + DEDUP_KEY_LEN] {
+	key.using_encoded(|s| {
+		// take at most 32 bytes from the hash of the value.
+		let hash = <T as system::Trait>::Hashing::hash(s);
+		let hash: &[u8] = hash.as_ref();
+		let len = rstd::cmp::min(hash.len(), 32);
+
+		let mut data = [0; 32 + DEDUP_KEY_LEN];
+		data[..DEDUP_KEY_LEN].copy_from_slice(DEDUP_KEY_PREFIX);
+		data[DEDUP_KEY_LEN..][..len].copy_from_slice(hash);
+		data
+	})
+}
+
+impl<T: Trait> OnFreeBalanceZero<T::ValidatorId> for Module<T> {
+	fn on_free_balance_zero(who: &T::ValidatorId) {
+		Self::prune_dead_keys(who);
 	}
 }
 
@@ -456,10 +534,10 @@ impl<T: Trait> OnFreeBalanceZero<T::AccountId> for Module<T> {
 /// registering account-ID of that session key index.
 pub struct FindAccountFromAuthorIndex<T, Inner>(rstd::marker::PhantomData<(T, Inner)>);
 
-impl<T: Trait, Inner: FindAuthor<u32>> FindAuthor<T::AccountId>
+impl<T: Trait, Inner: FindAuthor<u32>> FindAuthor<T::ValidatorId>
 	for FindAccountFromAuthorIndex<T, Inner>
 {
-	fn find_author<'a, I>(digests: I) -> Option<T::AccountId>
+	fn find_author<'a, I>(digests: I) -> Option<T::ValidatorId>
 		where I: 'a + IntoIterator<Item=(ConsensusEngineId, &'a [u8])>
 	{
 		let i = Inner::find_author(digests)?;
@@ -472,102 +550,17 @@ impl<T: Trait, Inner: FindAuthor<u32>> FindAuthor<T::AccountId>
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use std::cell::RefCell;
-	use srml_support::{impl_outer_origin, assert_ok};
+	use srml_support::assert_ok;
 	use runtime_io::with_externalities;
-	use substrate_primitives::{H256, Blake2Hasher};
+	use substrate_primitives::Blake2Hasher;
 	use primitives::{
-		traits::{BlakeTwo256, IdentityLookup, OnInitialize}, testing::{Header, UintAuthorityId}
+		traits::OnInitialize,
+		testing::UintAuthorityId,
+	};
+	use mock::{
+		NEXT_VALIDATORS, SESSION_CHANGED, TEST_SESSION_CHANGED, authorities, force_new_session,
+		set_next_validators, set_session_length, session_changed, Test, Origin, System, Session,
 	};
-
-	impl_outer_origin!{
-		pub enum Origin for Test {}
-	}
-
-	thread_local!{
-		static NEXT_VALIDATORS: RefCell<Vec<u64>> = RefCell::new(vec![1, 2, 3]);
-		static AUTHORITIES: RefCell<Vec<UintAuthorityId>> =
-			RefCell::new(vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
-		static FORCE_SESSION_END: RefCell<bool> = RefCell::new(false);
-		static SESSION_LENGTH: RefCell<u64> = RefCell::new(2);
-		static SESSION_CHANGED: RefCell<bool> = RefCell::new(false);
-		static TEST_SESSION_CHANGED: RefCell<bool> = RefCell::new(false);
-	}
-
-	pub struct TestShouldEndSession;
-	impl ShouldEndSession<u64> for TestShouldEndSession {
-		fn should_end_session(now: u64) -> bool {
-			let l = SESSION_LENGTH.with(|l| *l.borrow());
-			now % l == 0 || FORCE_SESSION_END.with(|l| { let r = *l.borrow(); *l.borrow_mut() = false; r })
-		}
-	}
-
-	pub struct TestSessionHandler;
-	impl SessionHandler<u64> for TestSessionHandler {
-		fn on_new_session<T: OpaqueKeys>(changed: bool, validators: &[(u64, T)]) {
-			SESSION_CHANGED.with(|l| *l.borrow_mut() = changed);
-			AUTHORITIES.with(|l|
-				*l.borrow_mut() = validators.iter().map(|(_, id)| id.get::<UintAuthorityId>(0).unwrap_or_default()).collect()
-			);
-		}
-		fn on_disabled(_validator_index: usize) {}
-	}
-
-	pub struct TestOnSessionEnding;
-	impl OnSessionEnding<u64> for TestOnSessionEnding {
-		fn on_session_ending(_: SessionIndex) -> Option<Vec<u64>> {
-			if !TEST_SESSION_CHANGED.with(|l| *l.borrow()) {
-				Some(NEXT_VALIDATORS.with(|l| l.borrow().clone()))
-			} else {
-				None
-			}
-		}
-	}
-
-	fn authorities() -> Vec<UintAuthorityId> {
-		AUTHORITIES.with(|l| l.borrow().to_vec())
-	}
-
-	fn force_new_session() {
-		FORCE_SESSION_END.with(|l| *l.borrow_mut() = true )
-	}
-
-	fn set_session_length(x: u64) {
-		SESSION_LENGTH.with(|l| *l.borrow_mut() = x )
-	}
-
-	fn session_changed() -> bool {
-		SESSION_CHANGED.with(|l| *l.borrow())
-	}
-
-	#[derive(Clone, Eq, PartialEq)]
-	pub struct Test;
-	impl system::Trait for Test {
-		type Origin = Origin;
-		type Index = u64;
-		type BlockNumber = u64;
-		type Hash = H256;
-		type Hashing = BlakeTwo256;
-		type AccountId = u64;
-		type Lookup = IdentityLookup<Self::AccountId>;
-		type Header = Header;
-		type Event = ();
-	}
-	impl timestamp::Trait for Test {
-		type Moment = u64;
-		type OnTimestampSet = ();
-	}
-	impl Trait for Test {
-		type ShouldEndSession = TestShouldEndSession;
-		type OnSessionEnding = TestOnSessionEnding;
-		type SessionHandler = TestSessionHandler;
-		type Keys = UintAuthorityId;
-		type Event = ();
-		type SelectInitialValidators = ();
-	}
-
-	type System = system::Module<Test>;
-	type Session = Module<Test>;
 
 	fn new_test_ext() -> runtime_io::TestExternalities<Blake2Hasher> {
 		TEST_SESSION_CHANGED.with(|l| *l.borrow_mut() = false);
@@ -597,11 +590,36 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn put_get_keys() {
+		with_externalities(&mut new_test_ext(), || {
+			Session::put_keys(&10, &UintAuthorityId(10));
+			assert_eq!(Session::load_keys(&10), Some(UintAuthorityId(10)));
+		})
+	}
+
+	#[test]
+	fn keys_cleared_on_kill() {
+		let mut ext = new_test_ext();
+		with_externalities(&mut ext, || {
+			assert_eq!(Session::validators(), vec![1, 2, 3]);
+			assert_eq!(Session::load_keys(&1), Some(UintAuthorityId(1)));
+
+			let id = <UintAuthorityId as TypedKey>::KEY_TYPE;
+			assert_eq!(Session::key_owner(id, UintAuthorityId(1).get_raw(id)), Some(1));
+
+			Session::on_free_balance_zero(&1);
+			assert_eq!(Session::load_keys(&1), None);
+			assert_eq!(Session::key_owner(id, UintAuthorityId(1).get_raw(id)), None);
+
+			assert!(Changed::get());
+		})
+	}
+
 	#[test]
 	fn authorities_should_track_validators() {
 		with_externalities(&mut new_test_ext(), || {
-			NEXT_VALIDATORS.with(|v| *v.borrow_mut() = vec![1, 2]);
-
+			set_next_validators(vec![1, 2]);
 			force_new_session();
 			initialize_block(1);
 			assert_eq!(Session::queued_keys(), vec![
@@ -620,7 +638,7 @@ mod tests {
 			assert_eq!(Session::validators(), vec![1, 2]);
 			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2)]);
 
-			NEXT_VALIDATORS.with(|v| *v.borrow_mut() = vec![1, 2, 4]);
+			set_next_validators(vec![1, 2, 4]);
 			assert_ok!(Session::set_keys(Origin::signed(4), UintAuthorityId(4), vec![]));
 			force_new_session();
 			initialize_block(3);
@@ -697,6 +715,19 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn duplicates_are_not_allowed() {
+		with_externalities(&mut new_test_ext(), || {
+			System::set_block_number(1);
+			Session::on_initialize(1);
+			assert!(Session::set_keys(Origin::signed(4), UintAuthorityId(1), vec![]).is_err());
+			assert!(Session::set_keys(Origin::signed(1), UintAuthorityId(10), vec![]).is_ok());
+
+			// is fine now that 1 has migrated off.
+			assert!(Session::set_keys(Origin::signed(4), UintAuthorityId(1), vec![]).is_ok());
+		});
+	}
+
 	#[test]
 	fn session_changed_flag_works() {
 		with_externalities(&mut new_test_ext(), || {

substrate/srml/session/src/mock.rs

@@ -0,0 +1,148 @@
+// Copyright 2019 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate 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 General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Mock helpers for Session.
+
+use super::*;
+use std::cell::RefCell;
+use srml_support::impl_outer_origin;
+use substrate_primitives::H256;
+use primitives::{
+	traits::{BlakeTwo256, IdentityLookup, ConvertInto},
+	testing::{Header, UintAuthorityId}
+};
+
+
+impl_outer_origin! {
+	pub enum Origin for Test {}
+}
+
+thread_local! {
+	pub static NEXT_VALIDATORS: RefCell<Vec<u64>> = RefCell::new(vec![1, 2, 3]);
+	pub static AUTHORITIES: RefCell<Vec<UintAuthorityId>> =
+		RefCell::new(vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
+	pub static FORCE_SESSION_END: RefCell<bool> = RefCell::new(false);
+	pub static SESSION_LENGTH: RefCell<u64> = RefCell::new(2);
+	pub static SESSION_CHANGED: RefCell<bool> = RefCell::new(false);
+	pub static TEST_SESSION_CHANGED: RefCell<bool> = RefCell::new(false);
+}
+
+pub struct TestShouldEndSession;
+impl ShouldEndSession<u64> for TestShouldEndSession {
+	fn should_end_session(now: u64) -> bool {
+		let l = SESSION_LENGTH.with(|l| *l.borrow());
+		now % l == 0 || FORCE_SESSION_END.with(|l| { let r = *l.borrow(); *l.borrow_mut() = false; r })
+	}
+}
+
+pub struct TestSessionHandler;
+impl SessionHandler<u64> for TestSessionHandler {
+	fn on_new_session<T: OpaqueKeys>(changed: bool, validators: &[(u64, T)]) {
+		SESSION_CHANGED.with(|l| *l.borrow_mut() = changed);
+		AUTHORITIES.with(|l|
+			*l.borrow_mut() = validators.iter().map(|(_, id)| id.get::<UintAuthorityId>(0).unwrap_or_default()).collect()
+		);
+	}
+	fn on_disabled(_validator_index: usize) {}
+}
+
+pub struct TestOnSessionEnding;
+impl OnSessionEnding<u64> for TestOnSessionEnding {
+	fn on_session_ending(_: SessionIndex, _: SessionIndex) -> Option<Vec<u64>> {
+		if !TEST_SESSION_CHANGED.with(|l| *l.borrow()) {
+			Some(NEXT_VALIDATORS.with(|l| l.borrow().clone()))
+		} else {
+			None
+		}
+	}
+}
+
+#[cfg(feature = "historical")]
+impl crate::historical::OnSessionEnding<u64, u64> for TestOnSessionEnding {
+	fn on_session_ending(_: SessionIndex, _: SessionIndex)
+		-> Option<(Vec<u64>, Vec<(u64, u64)>)>
+	{
+		if !TEST_SESSION_CHANGED.with(|l| *l.borrow()) {
+			let last_validators = Session::validators();
+			let last_identifications = last_validators.into_iter().map(|v| (v, v)).collect();
+			Some((NEXT_VALIDATORS.with(|l| l.borrow().clone()), last_identifications))
+		} else {
+			None
+		}
+	}
+}
+
+pub fn authorities() -> Vec<UintAuthorityId> {
+	AUTHORITIES.with(|l| l.borrow().to_vec())
+}
+
+pub fn force_new_session() {
+	FORCE_SESSION_END.with(|l| *l.borrow_mut() = true )
+}
+
+pub fn set_session_length(x: u64) {
+	SESSION_LENGTH.with(|l| *l.borrow_mut() = x )
+}
+
+pub fn session_changed() -> bool {
+	SESSION_CHANGED.with(|l| *l.borrow())
+}
+
+pub fn set_next_validators(next: Vec<u64>) {
+	NEXT_VALIDATORS.with(|v| *v.borrow_mut() = next);
+}
+
+#[derive(Clone, Eq, PartialEq)]
+pub struct Test;
+impl system::Trait for Test {
+	type Origin = Origin;
+	type Index = u64;
+	type BlockNumber = u64;
+	type Hash = H256;
+	type Hashing = BlakeTwo256;
+	type AccountId = u64;
+	type Lookup = IdentityLookup<Self::AccountId>;
+	type Header = Header;
+	type Event = ();
+}
+impl timestamp::Trait for Test {
+	type Moment = u64;
+	type OnTimestampSet = ();
+}
+
+
+impl Trait for Test {
+	type ShouldEndSession = TestShouldEndSession;
+	#[cfg(feature = "historical")]
+	type OnSessionEnding = crate::historical::NoteHistoricalRoot<Test, TestOnSessionEnding>;
+	#[cfg(not(feature = "historical"))]
+	type OnSessionEnding = TestOnSessionEnding;
+	type SessionHandler = TestSessionHandler;
+	type ValidatorId = u64;
+	type ValidatorIdOf = ConvertInto;
+	type Keys = UintAuthorityId;
+	type Event = ();
+	type SelectInitialValidators = crate::ConfigValidators;
+}
+
+#[cfg(feature = "historical")]
+impl crate::historical::Trait for Test {
+	type FullIdentification = u64;
+	type FullIdentificationOf = primitives::traits::ConvertInto;
+}
+
+pub type System = system::Module<Test>;
+pub type Session = Module<Test>;