diff --git a/substrate/frame/session/src/lib.rs b/substrate/frame/session/src/lib.rs index 883a0cc5ab..dd176219aa 100644 --- a/substrate/frame/session/src/lib.rs +++ b/substrate/frame/session/src/lib.rs @@ -683,6 +683,55 @@ impl Module { Self::validators().iter().position(|i| i == c).map(Self::disable_index).ok_or(()) } + /// Upgrade the key type from some old type to a new type. Supports adding + /// and removing key types. + /// + /// This function should be used with extreme care and only during an + /// `on_runtime_upgrade` block. Misuse of this function can put your blockchain + /// into an unrecoverable state. + /// + /// Care should be taken that the raw versions of the + /// added keys are unique for every `ValidatorId, KeyTypeId` combination. + /// This is an invariant that the session module typically maintains internally. + /// + /// As the actual values of the keys are typically not known at runtime upgrade, + /// it's recommended to initialize the keys to a (unique) dummy value with the expectation + /// that all validators should invoke `set_keys` before those keys are actually + /// required. + pub fn upgrade_keys(upgrade: F) where + Old: OpaqueKeys + Member + Decode, + F: Fn(T::ValidatorId, Old) -> T::Keys, + { + let old_ids = Old::key_ids(); + let new_ids = T::Keys::key_ids(); + + // Translate NextKeys, and key ownership relations at the same time. + >::translate::(|val, old_keys| { + // Clear all key ownership relations. Typically the overlap should + // stay the same, but no guarantees by the upgrade function. + for i in old_ids.iter() { + Self::clear_key_owner(*i, old_keys.get_raw(*i)); + } + + let new_keys = upgrade(val.clone(), old_keys); + + // And now set the new ones. + for i in new_ids.iter() { + Self::put_key_owner(*i, new_keys.get_raw(*i), &val); + } + + Some(new_keys) + }); + + let _ = >::translate::, _>( + |k| { + k.map(|k| k.into_iter() + .map(|(val, old_keys)| (val.clone(), upgrade(val, old_keys))) + .collect::>()) + } + ); + } + /// Perform the set_key operation, checking for duplicates. Does not set `Changed`. /// /// This ensures that the reference counter in system is incremented appropriately and as such diff --git a/substrate/frame/session/src/mock.rs b/substrate/frame/session/src/mock.rs index 0a7f89f5d5..fa71859feb 100644 --- a/substrate/frame/session/src/mock.rs +++ b/substrate/frame/session/src/mock.rs @@ -40,6 +40,31 @@ impl From for MockSessionKeys { } } +pub const KEY_ID_A: KeyTypeId = KeyTypeId([4; 4]); +pub const KEY_ID_B: KeyTypeId = KeyTypeId([9; 4]); + +#[derive(Debug, Clone, codec::Encode, codec::Decode, PartialEq, Eq)] +pub struct PreUpgradeMockSessionKeys { + pub a: [u8; 32], + pub b: [u8; 64], +} + +impl OpaqueKeys for PreUpgradeMockSessionKeys { + type KeyTypeIdProviders = (); + + fn key_ids() -> &'static [KeyTypeId] { + &[KEY_ID_A, KEY_ID_B] + } + + fn get_raw(&self, i: KeyTypeId) -> &[u8] { + match i { + i if i == KEY_ID_A => &self.a[..], + i if i == KEY_ID_B => &self.b[..], + _ => &[], + } + } +} + impl_outer_origin! { pub enum Origin for Test where system = frame_system {} } diff --git a/substrate/frame/session/src/tests.rs b/substrate/frame/session/src/tests.rs index 36857824de..7a33aa5296 100644 --- a/substrate/frame/session/src/tests.rs +++ b/substrate/frame/session/src/tests.rs @@ -25,6 +25,7 @@ use mock::{ SESSION_CHANGED, TEST_SESSION_CHANGED, authorities, force_new_session, set_next_validators, set_session_length, session_changed, Origin, System, Session, reset_before_session_end_called, before_session_end_called, new_test_ext, + PreUpgradeMockSessionKeys, }; fn initialize_block(block: u64) { @@ -308,3 +309,97 @@ fn return_true_if_more_than_third_is_disabled() { assert_eq!(Session::disable_index(3), true); }); } + +#[test] +fn upgrade_keys() { + use frame_support::storage; + use mock::Test; + use sp_core::crypto::key_types::DUMMY; + + // This test assumes certain mocks. + assert_eq!(mock::NEXT_VALIDATORS.with(|l| l.borrow().clone()), vec![1, 2, 3]); + assert_eq!(mock::VALIDATORS.with(|l| l.borrow().clone()), vec![1, 2, 3]); + + new_test_ext().execute_with(|| { + let pre_one = PreUpgradeMockSessionKeys { + a: [1u8; 32], + b: [1u8; 64], + }; + + let pre_two = PreUpgradeMockSessionKeys { + a: [2u8; 32], + b: [2u8; 64], + }; + + let pre_three = PreUpgradeMockSessionKeys { + a: [3u8; 32], + b: [3u8; 64], + }; + + let val_keys = vec![ + (1u64, pre_one), + (2u64, pre_two), + (3u64, pre_three), + ]; + + // Set `QueuedKeys`. + { + let storage_key = >::hashed_key(); + assert!(storage::unhashed::exists(&storage_key)); + storage::unhashed::put(&storage_key, &val_keys); + } + + // Set `NextKeys`. + { + for &(i, ref keys) in val_keys.iter() { + let storage_key = >::hashed_key_for(i); + assert!(storage::unhashed::exists(&storage_key)); + storage::unhashed::put(&storage_key, keys); + } + } + + // Set `KeyOwner`. + { + for &(i, ref keys) in val_keys.iter() { + // clear key owner for `UintAuthorityId` keys set in genesis. + let presumed = UintAuthorityId(i); + let raw_prev = presumed.as_ref(); + + assert_eq!(Session::key_owner(DUMMY, raw_prev), Some(i)); + Session::clear_key_owner(DUMMY, raw_prev); + + Session::put_key_owner(mock::KEY_ID_A, keys.get_raw(mock::KEY_ID_A), &i); + Session::put_key_owner(mock::KEY_ID_B, keys.get_raw(mock::KEY_ID_B), &i); + } + } + + // Do the upgrade and check sanity. + let mock_keys_for = |val| mock::MockSessionKeys { dummy: UintAuthorityId(val) }; + Session::upgrade_keys::( + |val, _old_keys| mock_keys_for(val), + ); + + // Check key ownership. + for (i, ref keys) in val_keys.iter() { + assert!(Session::key_owner(mock::KEY_ID_A, keys.get_raw(mock::KEY_ID_A)).is_none()); + assert!(Session::key_owner(mock::KEY_ID_B, keys.get_raw(mock::KEY_ID_B)).is_none()); + + let migrated_key = UintAuthorityId(*i); + assert_eq!(Session::key_owner(DUMMY, migrated_key.as_ref()), Some(*i)); + } + + // Check queued keys. + assert_eq!( + Session::queued_keys(), + vec![ + (1, mock_keys_for(1)), + (2, mock_keys_for(2)), + (3, mock_keys_for(3)), + ], + ); + + for i in 1u64..4 { + assert_eq!(>::get(&i), Some(mock_keys_for(i))); + } + }) +}