Make author_mapping::set_keys take 1 input like pallet_session::set_keys

pallets/author-mapping/src/benchmarks.rs

@@ -83,7 +83,7 @@ benchmarks! {
 		let caller = create_funded_user::<T>();
 		let id = nimbus_id(1u8);
 		let key: T::Keys = nimbus_id(2u8).into();
-	}: _(RawOrigin::Signed(caller.clone()), id.clone(), key.clone())
+	}: _(RawOrigin::Signed(caller.clone()), (id.clone(), key.clone()))
 	verify {
 		assert_eq!(Pallet::<T>::account_id_of(&id), Some(caller));
 		assert_eq!(Pallet::<T>::keys_of(&id), Some(key));
@@ -97,14 +97,13 @@ benchmarks! {
 		let second_keys: T::Keys = nimbus_id(3u8).into();
 		assert_ok!(Pallet::<T>::register_keys(
 				RawOrigin::Signed(caller.clone()).into(),
-				first_id.clone(),
-				first_keys.clone(),
+				(first_id.clone(),
+				first_keys.clone()),
 			)
 		);
 	}: _(RawOrigin::Signed(caller.clone()),
-		first_id.clone(),
-		second_id.clone(),
-		second_keys.clone()) verify {
+		(second_id.clone(),
+		second_keys.clone())) verify {
 		assert_eq!(Pallet::<T>::account_id_of(&first_id), None);
 		assert_eq!(Pallet::<T>::keys_of(&first_id), None);
 		assert_eq!(Pallet::<T>::account_id_of(&second_id), Some(caller));

pallets/author-mapping/src/lib.rs

@@ -91,6 +91,8 @@ pub mod pallet {
 		CannotAffordSecurityDeposit,
 		/// The NimbusId in question is already associated and cannot be overwritten
 		AlreadyAssociated,
+		/// No existing NimbusId can be found for the account
+		OldAuthorIdNotFound,
 	}
 
 	#[pallet::event]
@@ -173,6 +175,7 @@ pub mod pallet {
 				..stored_info
 			};
 			MappingWithDeposit::<T>::insert(&new_author_id, &new_stored_info);
+			NimbusLookup::<T>::insert(&account_id, &new_author_id);
 
 			<Pallet<T>>::deposit_event(Event::AuthorRotated {
 				new_author_id: new_author_id,
@@ -203,6 +206,7 @@ pub mod pallet {
 			);
 
 			MappingWithDeposit::<T>::remove(&author_id);
+			NimbusLookup::<T>::remove(&account_id);
 
 			T::DepositCurrency::unreserve(&account_id, stored_info.deposit);
 
@@ -217,13 +221,9 @@ pub mod pallet {
 
 		/// Add association and set session keys
 		#[pallet::weight(<T as Config>::WeightInfo::register_keys())]
-		pub fn register_keys(
-			origin: OriginFor<T>,
-			author_id: NimbusId,
-			keys: T::Keys,
-		) -> DispatchResult {
+		pub fn register_keys(origin: OriginFor<T>, keys: (NimbusId, T::Keys)) -> DispatchResult {
 			let account_id = ensure_signed(origin)?;
-
+			let (author_id, keys) = keys;
 			ensure!(
 				MappingWithDeposit::<T>::get(&author_id).is_none(),
 				Error::<T>::AlreadyAssociated
@@ -246,21 +246,18 @@ pub mod pallet {
 		/// No new security deposit is required. Will replace `update_association` which is kept
 		/// now for backwards compatibility reasons.
 		#[pallet::weight(<T as Config>::WeightInfo::set_keys())]
-		pub fn set_keys(
-			origin: OriginFor<T>,
-			old_author_id: NimbusId,
-			new_author_id: NimbusId,
-			new_keys: T::Keys,
-		) -> DispatchResult {
+		pub fn set_keys(origin: OriginFor<T>, keys: (NimbusId, T::Keys)) -> DispatchResult {
 			let account_id = ensure_signed(origin)?;
-
+			let old_author_id =
+				Self::nimbus_id_of(&account_id).ok_or(Error::<T>::OldAuthorIdNotFound)?;
 			let stored_info = MappingWithDeposit::<T>::try_get(&old_author_id)
 				.map_err(|_| Error::<T>::AssociationNotFound)?;
 
 			ensure!(
 				account_id == stored_info.account,
 				Error::<T>::NotYourAssociation
 			);
+			let (new_author_id, new_keys) = keys;
 			ensure!(
 				MappingWithDeposit::<T>::get(&new_author_id).is_none(),
 				Error::<T>::AlreadyAssociated
@@ -274,6 +271,7 @@ pub mod pallet {
 					..stored_info
 				},
 			);
+			NimbusLookup::<T>::insert(&account_id, new_author_id.clone());
 
 			<Pallet<T>>::deposit_event(Event::AuthorRotated {
 				new_author_id,
@@ -293,7 +291,7 @@ pub mod pallet {
 		) -> DispatchResult {
 			let deposit = T::DepositAmount::get();
 
-			T::DepositCurrency::reserve(&account_id, deposit)
+			T::DepositCurrency::reserve(account_id, deposit)
 				.map_err(|_| Error::<T>::CannotAffordSecurityDeposit)?;
 
 			let info = RegistrationInfo {
@@ -302,7 +300,8 @@ pub mod pallet {
 				keys,
 			};
 
-			MappingWithDeposit::<T>::insert(&author_id, &info);
+			MappingWithDeposit::<T>::insert(author_id, info);
+			NimbusLookup::<T>::insert(account_id, author_id);
 
 			Ok(())
 		}
@@ -315,6 +314,12 @@ pub mod pallet {
 	pub type MappingWithDeposit<T: Config> =
 		StorageMap<_, Blake2_128Concat, NimbusId, RegistrationInfo<T>, OptionQuery>;
 
+	#[pallet::storage]
+	#[pallet::getter(fn nimbus_lookup)]
+	/// We maintain a reverse mapping from AccountIds to NimbusIDS
+	pub type NimbusLookup<T: Config> =
+		StorageMap<_, Twox64Concat, T::AccountId, NimbusId, OptionQuery>;
+
 	#[pallet::genesis_config]
 	/// Genesis config for author mapping pallet
 	pub struct GenesisConfig<T: Config> {
@@ -366,5 +371,9 @@ pub mod pallet {
 		pub fn keys_of(author_id: &NimbusId) -> Option<T::Keys> {
 			Self::account_and_deposit_of(author_id).map(|info| info.keys)
 		}
+		/// A helper function to lookup NimbusId associated with a given AccountId
+		pub fn nimbus_id_of(account_id: &T::AccountId) -> Option<NimbusId> {
+			NimbusLookup::<T>::get(account_id)
+		}
 	}
 }

pallets/author-mapping/src/migrations.rs

@@ -14,17 +14,86 @@
 // You should have received a copy of the GNU General Public License
 // along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
 
-use crate::{BalanceOf, Config, MappingWithDeposit, RegistrationInfo};
+use crate::{BalanceOf, Config, Event, MappingWithDeposit, NimbusLookup, Pallet, RegistrationInfo};
 use frame_support::{
 	pallet_prelude::PhantomData,
-	traits::{Get, OnRuntimeUpgrade},
+	traits::{Get, OnRuntimeUpgrade, ReservableCurrency},
 	weights::Weight,
 };
 use nimbus_primitives::NimbusId;
 use parity_scale_codec::{Decode, Encode};
 #[cfg(feature = "try-runtime")]
 use scale_info::prelude::format;
 
+/// Migrates MappingWithDeposit map value from RegistrationInfo to RegistrationInformation,
+/// thereby adding a keys: T::Keys field to the value to support VRF keys that can be looked up
+/// via NimbusId.
+pub struct AddAccountIdToNimbusLookup<T>(PhantomData<T>);
+impl<T: Config> OnRuntimeUpgrade for AddAccountIdToNimbusLookup<T> {
+	fn on_runtime_upgrade() -> Weight {
+		log::info!(target: "AddAccountIdToNimbusLookup", "running migration");
+
+		let mut read_count = 0u64;
+		let mut write_count = 0u64;
+		<MappingWithDeposit<T>>::translate(|nimbus_id, registration_info: RegistrationInfo<T>| {
+			read_count += 2u64;
+			if NimbusLookup::<T>::get(&registration_info.account).is_none() {
+				<NimbusLookup<T>>::insert(&registration_info.account, nimbus_id);
+				write_count += 2u64;
+				Some(registration_info)
+			} else {
+				// revoke the additional association and return the funds
+				T::DepositCurrency::unreserve(
+					&registration_info.account,
+					registration_info.deposit,
+				);
+
+				<Pallet<T>>::deposit_event(Event::AuthorDeRegistered {
+					author_id: nimbus_id,
+					account_id: registration_info.account,
+					keys: registration_info.keys,
+				});
+				write_count += 1u64;
+				None
+			}
+		});
+		// return weight
+		read_count.saturating_mul(T::DbWeight::get().read)
+			+ write_count.saturating_mul(T::DbWeight::get().write)
+	}
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<(), &'static str> {
+		use frame_support::traits::OnRuntimeUpgradeHelpersExt;
+		let mut nimbus_set: Vec<NimbusId> = Vec::new();
+		for (nimbus_id, info) in <MappingWithDeposit<T>>::iter() {
+			if !nimbus_set.contains(&nimbus_id) {
+				Self::set_temp_storage(
+					info.account,
+					&format!("MappingWithDeposit{:?}Account", nimbus_id)[..],
+				);
+				nimbus_set.push(nimbus_id);
+			}
+		}
+		Ok(())
+	}
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade() -> Result<(), &'static str> {
+		use frame_support::traits::OnRuntimeUpgradeHelpersExt;
+		for (nimbus_id, _) in <MappingWithDeposit<T>>::iter() {
+			let old_account: T::AccountId =
+				Self::get_temp_storage(&format!("MappingWithDeposit{:?}Account", nimbus_id)[..])
+					.expect("qed");
+			let maybe_account_of_nimbus = <NimbusLookup<T>>::get(old_account);
+			assert_eq!(
+				Some(nimbus_id),
+				maybe_account_of_nimbus,
+				"New NimbusLookup dne expected NimbusID"
+			);
+		}
+		Ok(())
+	}
+}
+
 /// Migrates MappingWithDeposit map value from RegistrationInfo to RegistrationInformation,
 /// thereby adding a keys: T::Keys field to the value to support VRF keys that can be looked up
 /// via NimbusId.

pallets/author-mapping/src/tests.rs

@@ -297,8 +297,7 @@ fn eligible_account_can_full_register() {
 		.execute_with(|| {
 			assert_ok!(AuthorMapping::register_keys(
 				Origin::signed(2),
-				TestAuthor::Bob.into(),
-				TestAuthor::Alice.into(),
+				(TestAuthor::Bob.into(), TestAuthor::Alice.into()),
 			));
 
 			assert_eq!(Balances::free_balance(&2), 900);
@@ -328,8 +327,7 @@ fn cannot_register_keys_without_deposit() {
 			assert_noop!(
 				AuthorMapping::register_keys(
 					Origin::signed(2),
-					TestAuthor::Alice.into(),
-					TestAuthor::Bob.into(),
+					(TestAuthor::Alice.into(), TestAuthor::Bob.into()),
 				),
 				Error::<Runtime>::CannotAffordSecurityDeposit
 			);
@@ -348,8 +346,7 @@ fn double_full_registration_counts_twice_as_much() {
 			// Register once as Bob
 			assert_ok!(AuthorMapping::register_keys(
 				Origin::signed(2),
-				TestAuthor::Bob.into(),
-				TestAuthor::Charlie.into(),
+				(TestAuthor::Bob.into(), TestAuthor::Charlie.into()),
 			));
 
 			assert_eq!(Balances::free_balance(&2), 900);
@@ -371,8 +368,7 @@ fn double_full_registration_counts_twice_as_much() {
 			// Register again as Alice
 			assert_ok!(AuthorMapping::register_keys(
 				Origin::signed(2),
-				TestAuthor::Alice.into(),
-				TestAuthor::Bob.into(),
+				(TestAuthor::Alice.into(), TestAuthor::Bob.into()),
 			));
 
 			assert_eq!(Balances::free_balance(&2), 800);
@@ -409,8 +405,7 @@ fn full_registered_author_cannot_be_overwritten() {
 			assert_noop!(
 				AuthorMapping::register_keys(
 					Origin::signed(2),
-					TestAuthor::Alice.into(),
-					TestAuthor::Bob.into()
+					(TestAuthor::Alice.into(), TestAuthor::Bob.into()),
 				),
 				Error::<Runtime>::AlreadyAssociated
 			);
@@ -426,9 +421,7 @@ fn registered_can_full_rotate() {
 		.execute_with(|| {
 			assert_ok!(AuthorMapping::set_keys(
 				Origin::signed(2),
-				TestAuthor::Bob.into(),
-				TestAuthor::Charlie.into(),
-				TestAuthor::Charlie.into(),
+				(TestAuthor::Charlie.into(), TestAuthor::Charlie.into())
 			));
 
 			assert_eq!(AuthorMapping::account_id_of(&TestAuthor::Bob.into()), None);
@@ -453,11 +446,9 @@ fn unregistered_author_cannot_be_full_rotated() {
 		assert_noop!(
 			AuthorMapping::set_keys(
 				Origin::signed(2),
-				TestAuthor::Alice.into(),
-				TestAuthor::Bob.into(),
-				TestAuthor::Bob.into(),
+				(TestAuthor::Bob.into(), TestAuthor::Bob.into()),
 			),
-			Error::<Runtime>::AssociationNotFound
+			Error::<Runtime>::OldAuthorIdNotFound
 		);
 	})
 }
@@ -472,11 +463,9 @@ fn registered_author_cannot_be_full_rotated_by_non_owner() {
 			assert_noop!(
 				AuthorMapping::set_keys(
 					Origin::signed(2),
-					TestAuthor::Alice.into(),
-					TestAuthor::Bob.into(),
-					TestAuthor::Bob.into(),
+					(TestAuthor::Bob.into(), TestAuthor::Bob.into())
 				),
-				Error::<Runtime>::NotYourAssociation
+				Error::<Runtime>::OldAuthorIdNotFound
 			);
 		})
 }
@@ -491,9 +480,7 @@ fn full_rotating_to_the_same_author_id_leaves_registration_in_tact() {
 			assert_noop!(
 				AuthorMapping::set_keys(
 					Origin::signed(1),
-					TestAuthor::Alice.into(),
-					TestAuthor::Alice.into(),
-					TestAuthor::Alice.into(),
+					(TestAuthor::Alice.into(), TestAuthor::Alice.into())
 				),
 				Error::<Runtime>::AlreadyAssociated
 			);

precompiles/author-mapping/AuthorMappingInterface.sol

@@ -46,15 +46,10 @@ interface AuthorMapping {
 
     /**
      * Set keys
-     * Selector: a8259c85
+     * Selector: 47f92fc4
      *
-     * @param old_author_id The old nimbusId to be replaced
      * @param new_author_id The new nimbusId to be associated
      * @param new_keys The new session keys
      */
-    function set_keys(
-        bytes32 old_author_id,
-        bytes32 new_author_id,
-        bytes32 new_keys
-    ) external;
+    function set_keys(bytes32 new_author_id, bytes32 new_keys) external;
 }