add an upgrade_keys method for pallet-session

frame/session/src/lib.rs

@@ -683,6 +683,55 @@ impl<T: Config> Module<T> {
 		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<Old, F>(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.
+		<NextKeys<T>>::translate::<Old, _>(|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 _ = <QueuedKeys<T>>::translate::<Vec<(T::ValidatorId, Old)>, _>(
+			|k| {
+				k.map(|k| k.into_iter()
+					.map(|(val, old_keys)| (val.clone(), upgrade(val, old_keys)))
+					.collect::<Vec<_>>())
+			}
+		);
+	}
+
 	/// 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

frame/session/src/mock.rs

@@ -40,6 +40,31 @@ impl From<UintAuthorityId> 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 {}
 }

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 = <super::QueuedKeys<Test>>::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 = <super::NextKeys<Test>>::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::<PreUpgradeMockSessionKeys, _>(
+			|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!(<super::NextKeys<Test>>::get(&i), Some(mock_keys_for(i)));
+		}
+	})
+}