Remove footgun around session keys/handlers (#3949)

* Remove footgun around session keys/handlers

- `OpaqueKeys` now has an associated type `KeyTypeIdProviders`. This can
be used in the runtime as input for `SessionHandler` from the session
trait.
- `impl_opaque_keys` now works with modules and extracts the `KeyTypeId`
from the module directly.
- Added some checks to the `session` storage initialization that checks
that the `SessionHandler` and `Keys` use the same number of keys and
that the order is equal.

* Update core/sr-primitives/src/traits.rs

substrate/Cargo.lock

@@ -1438,7 +1438,7 @@ dependencies = [
 
 [[package]]
 name = "impl-trait-for-tuples"
-version = "0.1.2"
+version = "0.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 dependencies = [
  "proc-macro2 1.0.6 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -3928,7 +3928,7 @@ dependencies = [
 name = "sr-primitives"
 version = "2.0.0"
 dependencies = [
- "impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-scale-codec 1.0.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "paste 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -4040,7 +4040,7 @@ dependencies = [
 name = "srml-authorship"
 version = "0.1.0"
 dependencies = [
- "impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-scale-codec 1.0.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "sr-io 2.0.0",
  "sr-primitives 2.0.0",
@@ -4248,7 +4248,7 @@ dependencies = [
 name = "srml-finality-tracker"
 version = "2.0.0"
 dependencies = [
- "impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-scale-codec 1.0.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde 1.0.101 (registry+https://github.com/rust-lang/crates.io-index)",
  "sr-io 2.0.0",
@@ -4415,7 +4415,7 @@ dependencies = [
 name = "srml-session"
 version = "2.0.0"
 dependencies = [
- "impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-scale-codec 1.0.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "safe-mix 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -4485,7 +4485,7 @@ name = "srml-support"
 version = "2.0.0"
 dependencies = [
  "bitmask 0.5.0 (registry+https://github.com/rust-lang/crates.io-index)",
- "impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)",
  "once_cell 0.2.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-scale-codec 1.0.6 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -4553,7 +4553,7 @@ name = "srml-system"
 version = "2.0.0"
 dependencies = [
  "criterion 0.2.11 (registry+https://github.com/rust-lang/crates.io-index)",
- "impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-scale-codec 1.0.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "safe-mix 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde 1.0.101 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -4597,7 +4597,7 @@ dependencies = [
 name = "srml-timestamp"
 version = "2.0.0"
 dependencies = [
- "impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-scale-codec 1.0.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde 1.0.101 (registry+https://github.com/rust-lang/crates.io-index)",
  "sr-io 2.0.0",
@@ -4865,7 +4865,7 @@ dependencies = [
 name = "substrate-chain-spec"
 version = "2.0.0"
 dependencies = [
- "impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde 1.0.101 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde_json 1.0.41 (registry+https://github.com/rust-lang/crates.io-index)",
  "sr-primitives 2.0.0",
@@ -7019,7 +7019,7 @@ dependencies = [
 "checksum idna 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "02e2673c30ee86b5b96a9cb52ad15718aa1f966f5ab9ad54a8b95d5ca33120a9"
 "checksum impl-codec 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)" = "3fa0086251524c50fd53b32e7b05eb6d79e2f97221eaf0c53c0ca9c3096f21d3"
 "checksum impl-serde 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "a263dc95daa6c3788c8f7133d86dc2ad89ec5a0c56167f9e3441c5f7f33358c4"
-"checksum impl-trait-for-tuples 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "6947b372790f8948f439bb6aaa6baabdf80be1a207a477ff072f83fb793e428f"
+"checksum impl-trait-for-tuples 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "7ef5550a42e3740a0e71f909d4c861056a284060af885ae7aa6242820f920d9d"
 "checksum indexmap 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "a61202fbe46c4a951e9404a720a0180bcf3212c750d735cb5c4ba4dc551299f3"
 "checksum integer-sqrt 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "ea155abb3ba6f382a75f1418988c05fe82959ed9ce727de427f9cfd425b0c903"
 "checksum interleaved-ordered 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "141340095b15ed7491bd3d4ced9d20cebfb826174b6bb03386381f62b01e3d77"

substrate/core/application-crypto/src/ed25519.rs

@@ -23,6 +23,10 @@ pub use primitives::ed25519::*;
 mod app {
 	use primitives::testing::ED25519;
 	crate::app_crypto!(super, ED25519);
+
+	impl crate::traits::BoundToRuntimeAppPublic for Public {
+		type Public = Self;
+	}
 }
 
 pub use app::Public as AppPublic;

substrate/core/application-crypto/src/sr25519.rs

@@ -23,6 +23,10 @@ pub use primitives::sr25519::*;
 mod app {
 	use primitives::testing::SR25519;
 	crate::app_crypto!(super, SR25519);
+
+	impl crate::traits::BoundToRuntimeAppPublic for Public {
+		type Public = Self;
+	}
 }
 
 pub use app::Public as AppPublic;

substrate/core/application-crypto/src/traits.rs

@@ -127,3 +127,8 @@ pub trait RuntimeAppPublic: Sized  {
 	fn verify<M: AsRef<[u8]>>(&self, msg: &M, signature: &Self::Signature) -> bool;
 }
 
+/// Something that bound to a fixed `RuntimeAppPublic`.
+pub trait BoundToRuntimeAppPublic {
+	/// The `RuntimeAppPublic` this type is bound to.
+	type Public: RuntimeAppPublic;
+}

substrate/core/sr-primitives/src/lib.rs

@@ -60,7 +60,7 @@ pub use generic::{DigestItem, Digest};
 
 /// Re-export this since it's part of the API of this crate.
 pub use primitives::{TypeId, crypto::{key_types, KeyTypeId, CryptoType, AccountId32}};
-pub use app_crypto::RuntimeAppPublic;
+pub use app_crypto::{RuntimeAppPublic, BoundToRuntimeAppPublic};
 
 /// Re-export `RuntimeDebug`, to avoid dependency clutter.
 pub use primitives::RuntimeDebug;

substrate/core/sr-primitives/src/testing.rs

@@ -117,10 +117,10 @@ impl app_crypto::RuntimeAppPublic for UintAuthorityId {
 }
 
 impl OpaqueKeys for UintAuthorityId {
-	type KeyTypeIds = std::iter::Cloned<std::slice::Iter<'static, KeyTypeId>>;
+	type KeyTypeIdProviders = ();
 
-	fn key_ids() -> Self::KeyTypeIds {
-		[key_types::DUMMY].iter().cloned()
+	fn key_ids() -> &'static [KeyTypeId] {
+		&[key_types::DUMMY]
 	}
 
 	fn get_raw(&self, _: KeyTypeId) -> &[u8] {
@@ -132,6 +132,10 @@ impl OpaqueKeys for UintAuthorityId {
 	}
 }
 
+impl crate::BoundToRuntimeAppPublic for UintAuthorityId {
+	type Public = Self;
+}
+
 /// Digest item
 pub type DigestItem = generic::DigestItem<H256>;
 

substrate/core/sr-primitives/src/traits.rs

@@ -982,11 +982,11 @@ pub trait ValidateUnsigned {
 /// Opaque datatype that may be destructured into a series of raw byte slices (which represent
 /// individual keys).
 pub trait OpaqueKeys: Clone {
-	/// An iterator over the type IDs of keys that this holds.
-	type KeyTypeIds: IntoIterator<Item=super::KeyTypeId>;
+	/// Types bound to this opaque keys that provide the key type ids returned.
+	type KeyTypeIdProviders;
 
-	/// Return an iterator over the key-type IDs supported by this set.
-	fn key_ids() -> Self::KeyTypeIds;
+	/// Return the key-type IDs supported by this set.
+	fn key_ids() -> &'static [crate::KeyTypeId];
 	/// Get the raw bytes of key with key-type ID `i`.
 	fn get_raw(&self, i: super::KeyTypeId) -> &[u8];
 	/// Get the decoded key with index `i`.
@@ -1086,22 +1086,25 @@ macro_rules! count {
 }
 
 /// Implement `OpaqueKeys` for a described struct.
-/// Would be much nicer for this to be converted to `derive` code.
 ///
-/// Every field type must be equivalent implement `as_ref()`, which is expected
-/// to hold the standard SCALE-encoded form of that key. This is typically
-/// just the bytes of the key.
+/// Every field type must implement [`BoundToRuntimeAppPublic`](crate::BoundToRuntimeAppPublic).
+/// `KeyTypeIdProviders` is set to the types given as fields.
 ///
 /// ```rust
-/// use sr_primitives::{impl_opaque_keys, KeyTypeId, app_crypto::{sr25519, ed25519}};
-/// use primitives::testing::{SR25519, ED25519};
+/// use sr_primitives::{
+/// 	impl_opaque_keys, KeyTypeId, BoundToRuntimeAppPublic, app_crypto::{sr25519, ed25519}
+/// };
+///
+/// pub struct KeyModule;
+/// impl BoundToRuntimeAppPublic for KeyModule { type Public = ed25519::AppPublic; }
+///
+/// pub struct KeyModule2;
+/// impl BoundToRuntimeAppPublic for KeyModule2 { type Public = sr25519::AppPublic; }
 ///
 /// impl_opaque_keys! {
 /// 	pub struct Keys {
-/// 		#[id(ED25519)]
-/// 		pub ed25519: ed25519::AppPublic,
-/// 		#[id(SR25519)]
-/// 		pub sr25519: sr25519::AppPublic,
+/// 		pub key_module: KeyModule,
+/// 		pub key_module2: KeyModule2,
 /// 	}
 /// }
 /// ```
@@ -1110,7 +1113,6 @@ macro_rules! impl_opaque_keys {
 	(
 		pub struct $name:ident {
 			$(
-				#[id($key_id:expr)]
 				pub $field:ident: $type:ty,
 			)*
 		}
@@ -1119,12 +1121,12 @@ macro_rules! impl_opaque_keys {
 			Default, Clone, PartialEq, Eq,
 			$crate::codec::Encode,
 			$crate::codec::Decode,
-			$crate::RuntimeDebug
+			$crate::RuntimeDebug,
 		)]
 		#[cfg_attr(feature = "std", derive($crate::serde::Serialize, $crate::serde::Deserialize))]
 		pub struct $name {
 			$(
-				pub $field: $type,
+				pub $field: <$type as $crate::BoundToRuntimeAppPublic>::Public,
 			)*
 		}
 
@@ -1137,7 +1139,11 @@ macro_rules! impl_opaque_keys {
 			pub fn generate(seed: Option<&str>) -> $crate::rstd::vec::Vec<u8> {
 				let keys = Self{
 					$(
-						$field: <$type as $crate::app_crypto::RuntimeAppPublic>::generate_pair(seed),
+						$field: <
+							<
+								$type as $crate::BoundToRuntimeAppPublic
+							>::Public as $crate::RuntimeAppPublic
+						>::generate_pair(seed),
 					)*
 				};
 				$crate::codec::Encode::encode(&keys)
@@ -1145,17 +1151,30 @@ macro_rules! impl_opaque_keys {
 		}
 
 		impl $crate::traits::OpaqueKeys for $name {
-			type KeyTypeIds = $crate::rstd::iter::Cloned<
-				$crate::rstd::slice::Iter<'static, $crate::KeyTypeId>
-			>;
+			type KeyTypeIdProviders = ( $( $type, )* );
 
-			fn key_ids() -> Self::KeyTypeIds {
-				[ $($key_id),* ].iter().cloned()
+			fn key_ids() -> &'static [$crate::KeyTypeId] {
+				&[
+					$(
+						<
+							<
+								$type as $crate::BoundToRuntimeAppPublic
+							>::Public as $crate::RuntimeAppPublic
+						>::ID
+					),*
+				]
 			}
 
 			fn get_raw(&self, i: $crate::KeyTypeId) -> &[u8] {
 				match i {
-					$( i if i == $key_id => self.$field.as_ref(), )*
+					$(
+						i if i == <
+							<
+								$type as $crate::BoundToRuntimeAppPublic
+							>::Public as $crate::RuntimeAppPublic
+						>::ID =>
+							self.$field.as_ref(),
+					)*
 					_ => &[],
 				}
 			}

substrate/core/test-runtime/src/lib.rs

@@ -25,15 +25,7 @@ pub mod system;
 use rstd::{prelude::*, marker::PhantomData};
 use codec::{Encode, Decode, Input, Error};
 
-use primitives::{
-	Blake2Hasher,
-	OpaqueMetadata,
-	RuntimeDebug,
-	testing::{
-		ED25519,
-		SR25519,
-	}
-};
+use primitives::{Blake2Hasher, OpaqueMetadata, RuntimeDebug};
 use app_crypto::{ed25519, sr25519, RuntimeAppPublic};
 pub use app_crypto;
 use trie_db::{TrieMut, Trie};
@@ -454,9 +446,7 @@ fn code_using_trie() -> u64 {
 
 impl_opaque_keys! {
 	pub struct SessionKeys {
-		#[id(ED25519)]
 		pub ed25519: ed25519::AppPublic,
-		#[id(SR25519)]
 		pub sr25519: sr25519::AppPublic,
 	}
 }

substrate/node-template/runtime/src/lib.rs

@@ -9,7 +9,7 @@
 include!(concat!(env!("OUT_DIR"), "/wasm_binary.rs"));
 
 use rstd::prelude::*;
-use primitives::{OpaqueMetadata, crypto::key_types};
+use primitives::OpaqueMetadata;
 use sr_primitives::{
 	ApplyResult, transaction_validity::TransactionValidity, generic, create_runtime_str,
 	impl_opaque_keys, MultiSignature
@@ -84,10 +84,8 @@ pub mod opaque {
 
 	impl_opaque_keys! {
 		pub struct SessionKeys {
-			#[id(key_types::AURA)]
-			pub aura: AuraId,
-			#[id(key_types::GRANDPA)]
-			pub grandpa: GrandpaId,
+			pub aura: Aura,
+			pub grandpa: Grandpa,
 		}
 	}
 }

substrate/node/runtime/src/lib.rs

@@ -29,20 +29,18 @@ use node_primitives::{
 	AccountId, AccountIndex, Balance, BlockNumber, Hash, Index,
 	Moment, Signature,
 };
-use babe_primitives::AuthorityId as BabeId;
 use grandpa::fg_primitives;
 use client::{
 	block_builder::api::{self as block_builder_api, InherentData, CheckInherentsResult},
 	runtime_api as client_api, impl_runtime_apis
 };
-use sr_primitives::{
-	Permill, Perbill, ApplyResult, impl_opaque_keys, generic, create_runtime_str, key_types
-};
+use sr_primitives::{Permill, Perbill, ApplyResult, impl_opaque_keys, generic, create_runtime_str};
 use sr_primitives::curve::PiecewiseLinear;
 use sr_primitives::transaction_validity::TransactionValidity;
 use sr_primitives::weights::Weight;
 use sr_primitives::traits::{
 	self, BlakeTwo256, Block as BlockT, NumberFor, StaticLookup, SaturatedConversion,
+	OpaqueKeys,
 };
 use version::RuntimeVersion;
 #[cfg(any(feature = "std", test))]
@@ -211,34 +209,21 @@ impl authorship::Trait for Runtime {
 	type EventHandler = Staking;
 }
 
-// !!!!!!!!!!!!!
-// WARNING!!!!!!  SEE NOTE BELOW BEFORE TOUCHING THIS CODE
-// !!!!!!!!!!!!!
-type SessionHandlers = (Grandpa, Babe, ImOnline);
-
 impl_opaque_keys! {
 	pub struct SessionKeys {
-		#[id(key_types::GRANDPA)]
-		pub grandpa: GrandpaId,
-		#[id(key_types::BABE)]
-		pub babe: BabeId,
-		#[id(key_types::IM_ONLINE)]
-		pub im_online: ImOnlineId,
+		pub grandpa: Grandpa,
+		pub babe: Babe,
+		pub im_online: ImOnline,
 	}
 }
 
-// NOTE: `SessionHandler` and `SessionKeys` are co-dependent: One key will be used for each handler.
-// The number and order of items in `SessionHandler` *MUST* be the same number and order of keys in
-// `SessionKeys`.
-// TODO: Introduce some structure to tie these together to make it a bit less of a footgun. This
-// should be easy, since OneSessionHandler trait provides the `Key` as an associated type. #2858
 parameter_types! {
 	pub const DisabledValidatorsThreshold: Perbill = Perbill::from_percent(17);
 }
 
 impl session::Trait for Runtime {
 	type OnSessionEnding = Staking;
-	type SessionHandler = SessionHandlers;
+	type SessionHandler = <SessionKeys as OpaqueKeys>::KeyTypeIdProviders;
 	type ShouldEndSession = Babe;
 	type Event = Event;
 	type Keys = SessionKeys;

substrate/srml/aura/src/lib.rs

@@ -182,6 +182,10 @@ impl<T: Trait> Module<T> {
 	}
 }
 
+impl<T: Trait> sr_primitives::BoundToRuntimeAppPublic for Module<T> {
+	type Public = T::AuthorityId;
+}
+
 impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
 	type Key = T::AuthorityId;
 

substrate/srml/authority-discovery/src/lib.rs

@@ -107,6 +107,10 @@ impl<T: Trait> Module<T> {
 	}
 }
 
+impl<T: Trait> sr_primitives::BoundToRuntimeAppPublic for Module<T> {
+	type Public = T::AuthorityId;
+}
+
 impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
 	type Key = T::AuthorityId;
 
@@ -139,7 +143,7 @@ mod tests {
 	use runtime_io::TestExternalities;
 	use sr_primitives::{
 		testing::{Header, UintAuthorityId}, traits::{ConvertInto, IdentityLookup, OpaqueKeys},
-		Perbill,
+		Perbill, KeyTypeId,
 	};
 	use support::{impl_outer_origin, parameter_types};
 
@@ -218,6 +222,8 @@ mod tests {
 
 	pub struct TestSessionHandler;
 	impl session::SessionHandler<AuthorityId> for TestSessionHandler {
+		const KEY_TYPE_IDS: &'static [KeyTypeId] = &[key_types::DUMMY];
+
 		fn on_new_session<Ks: OpaqueKeys>(
 			_changed: bool,
 			_validators: &[(AuthorityId, Ks)],

substrate/srml/babe/src/lib.rs

@@ -546,6 +546,10 @@ impl<T: Trait> OnTimestampSet<T::Moment> for Module<T> {
 	fn on_timestamp_set(_moment: T::Moment) { }
 }
 
+impl<T: Trait> sr_primitives::BoundToRuntimeAppPublic for Module<T> {
+	type Public = AuthorityId;
+}
+
 impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
 	type Key = AuthorityId;
 

substrate/srml/babe/src/mock.rs

@@ -20,9 +20,7 @@
 use super::{Trait, Module, GenesisConfig};
 use babe_primitives::AuthorityId;
 use sr_primitives::{
-	traits::IdentityLookup, Perbill,
-	testing::{Header, UintAuthorityId},
-	impl_opaque_keys, key_types::DUMMY,
+	traits::IdentityLookup, Perbill, testing::{Header, UintAuthorityId}, impl_opaque_keys,
 };
 use sr_version::RuntimeVersion;
 use support::{impl_outer_origin, parameter_types};
@@ -71,7 +69,6 @@ impl system::Trait for Test {
 
 impl_opaque_keys! {
 	pub struct MockSessionKeys {
-		#[id(DUMMY)]
 		pub dummy: UintAuthorityId,
 	}
 }

substrate/srml/grandpa/src/lib.rs

@@ -36,8 +36,7 @@ use support::{
 	decl_event, decl_storage, decl_module, dispatch::Result,
 };
 use sr_primitives::{
-	generic::{DigestItem, OpaqueDigestItemId}, traits::Zero,
-	Perbill,
+	generic::{DigestItem, OpaqueDigestItemId}, traits::Zero, Perbill,
 };
 use sr_staking_primitives::{
 	SessionIndex,
@@ -374,6 +373,10 @@ impl<T: Trait> Module<T> {
 	}
 }
 
+impl<T: Trait> sr_primitives::BoundToRuntimeAppPublic for Module<T> {
+	type Public = AuthorityId;
+}
+
 impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T>
 	where T: session::Trait
 {

substrate/srml/im-online/src/lib.rs

@@ -427,8 +427,11 @@ impl<T: Trait> Module<T> {
 	}
 }
 
-impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
+impl<T: Trait> sr_primitives::BoundToRuntimeAppPublic for Module<T> {
+	type Public = T::AuthorityId;
+}
 
+impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
 	type Key = T::AuthorityId;
 
 	fn on_genesis_session<'a, I: 'a>(validators: I)

substrate/srml/session/Cargo.toml

@@ -16,7 +16,7 @@ system = { package = "srml-system", path = "../system", default-features = false
 timestamp = { package = "srml-timestamp", path = "../timestamp", default-features = false }
 substrate-trie = { path = "../../core/trie", default-features = false, optional = true }
 runtime-io ={ package = "sr-io", path = "../../core/sr-io", default-features = false }
-impl-trait-for-tuples = "0.1.2"
+impl-trait-for-tuples = "0.1.3"
 
 [dev-dependencies]
 primitives = { package = "substrate-primitives",  path = "../../core/primitives" }

substrate/srml/session/src/historical.rs

@@ -182,11 +182,9 @@ impl<T: Trait> ProvingTrie<T> {
 
 				// map each key to the owner index.
 				for key_id in T::Keys::key_ids() {
-					let key = keys.get_raw(key_id);
+					let key = keys.get_raw(*key_id);
 					let res = (key_id, key).using_encoded(|k|
-						i.using_encoded(|v|
-							trie.insert(k, v)
-						)
+						i.using_encoded(|v| trie.insert(k, v))
 					);
 
 					let _ = res.map_err(|_| "failed to insert into trie")?;

substrate/srml/session/src/lib.rs

@@ -121,7 +121,7 @@
 
 use rstd::{prelude::*, marker::PhantomData, ops::{Sub, Rem}};
 use codec::Decode;
-use sr_primitives::{KeyTypeId, Perbill, RuntimeAppPublic};
+use sr_primitives::{KeyTypeId, Perbill, RuntimeAppPublic, BoundToRuntimeAppPublic};
 use sr_primitives::weights::SimpleDispatchInfo;
 use sr_primitives::traits::{Convert, Zero, Member, OpaqueKeys};
 use sr_staking_primitives::SessionIndex;
@@ -192,6 +192,12 @@ impl<A> OnSessionEnding<A> for () {
 
 /// Handler for session lifecycle events.
 pub trait SessionHandler<ValidatorId> {
+	/// All the key type ids this session handler can process.
+	///
+	/// The order must be the same as it expects them in
+	/// [`on_new_session`](Self::on_new_session) and [`on_genesis_session`](Self::on_genesis_session).
+	const KEY_TYPE_IDS: &'static [KeyTypeId];
+
 	/// The given validator set will be used for the genesis session.
 	/// It is guaranteed that the given validator set will also be used
 	/// for the second session, therefore the first call to `on_new_session`
@@ -220,7 +226,7 @@ pub trait SessionHandler<ValidatorId> {
 }
 
 /// A session handler for specific key type.
-pub trait OneSessionHandler<ValidatorId> {
+pub trait OneSessionHandler<ValidatorId>: BoundToRuntimeAppPublic {
 	/// The key type expected.
 	type Key: Decode + Default + RuntimeAppPublic;
 
@@ -258,6 +264,10 @@ pub trait OneSessionHandler<ValidatorId> {
 impl<AId> SessionHandler<AId> for Tuple {
 	for_tuples!( where #( Tuple: OneSessionHandler<AId> )* );
 
+	for_tuples!(
+		const KEY_TYPE_IDS: &'static [KeyTypeId] = &[ #( <Tuple::Key as RuntimeAppPublic>::ID ),* ];
+	);
+
 	fn on_genesis_session<Ks: OpaqueKeys>(validators: &[(AId, Ks)]) {
 		for_tuples!(
 			#(
@@ -382,6 +392,20 @@ decl_storage! {
 	add_extra_genesis {
 		config(keys): Vec<(T::ValidatorId, T::Keys)>;
 		build(|config: &GenesisConfig<T>| {
+			if T::SessionHandler::KEY_TYPE_IDS.len() != T::Keys::key_ids().len() {
+				panic!("Number of keys in session handler and session keys does not match");
+			}
+
+			T::SessionHandler::KEY_TYPE_IDS.iter().zip(T::Keys::key_ids()).enumerate()
+				.for_each(|(i, (sk, kk))| {
+					if sk != kk {
+						panic!(
+							"Session handler and session key expect different key type at index: {}",
+							i,
+						);
+					}
+				});
+
 			for (who, keys) in config.keys.iter().cloned() {
 				assert!(
 					<Module<T>>::load_keys(&who).is_none(),
@@ -594,23 +618,23 @@ impl<T: Trait> Module<T> {
 		let old_keys = Self::load_keys(&who);
 
 		for id in T::Keys::key_ids() {
-			let key = keys.get_raw(id);
+			let key = keys.get_raw(*id);
 
 			// ensure keys are without duplication.
 			ensure!(
-				Self::key_owner(id, key).map_or(true, |owner| &owner == who),
+				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 let Some(old) = old_keys.as_ref().map(|k| k.get_raw(*id)) {
 				if key == old {
 					continue;
 				}
 
-				Self::clear_key_owner(id, old);
+				Self::clear_key_owner(*id, old);
 			}
 
-			Self::put_key_owner(id, key, &who);
+			Self::put_key_owner(*id, key, &who);
 		}
 
 		Self::put_keys(&who, &keys);
@@ -621,8 +645,8 @@ impl<T: Trait> Module<T> {
 	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);
+				let key_data = old_keys.get_raw(*id);
+				Self::clear_key_owner(*id, key_data);
 			}
 		}
 	}

substrate/srml/session/src/mock.rs

@@ -28,7 +28,6 @@ use sr_staking_primitives::SessionIndex;
 
 impl_opaque_keys! {
 	pub struct MockSessionKeys {
-		#[id(DUMMY)]
 		pub dummy: UintAuthorityId,
 	}
 }
@@ -67,6 +66,7 @@ impl ShouldEndSession<u64> for TestShouldEndSession {
 
 pub struct TestSessionHandler;
 impl SessionHandler<u64> for TestSessionHandler {
+	const KEY_TYPE_IDS: &'static [sr_primitives::KeyTypeId] = &[UintAuthorityId::ID];
 	fn on_genesis_session<T: OpaqueKeys>(_validators: &[(u64, T)]) {}
 	fn on_new_session<T: OpaqueKeys>(
 		changed: bool,

substrate/srml/staking/src/mock.rs

@@ -17,12 +17,12 @@
 //! Test utilities
 
 use std::{collections::HashSet, cell::RefCell};
-use sr_primitives::Perbill;
+use sr_primitives::{Perbill, KeyTypeId};
 use sr_primitives::curve::PiecewiseLinear;
 use sr_primitives::traits::{IdentityLookup, Convert, OpaqueKeys, OnInitialize, SaturatedConversion};
 use sr_primitives::testing::{Header, UintAuthorityId};
 use sr_staking_primitives::SessionIndex;
-use primitives::H256;
+use primitives::{H256, crypto::key_types};
 use runtime_io;
 use support::{assert_ok, impl_outer_origin, parameter_types, StorageLinkedMap};
 use support::traits::{Currency, Get, FindAuthor};
@@ -52,6 +52,8 @@ thread_local! {
 
 pub struct TestSessionHandler;
 impl session::SessionHandler<AccountId> for TestSessionHandler {
+	const KEY_TYPE_IDS: &'static [KeyTypeId] = &[key_types::DUMMY];
+
 	fn on_genesis_session<Ks: OpaqueKeys>(_validators: &[(AccountId, Ks)]) {}
 
 	fn on_new_session<Ks: OpaqueKeys>(