frame/beefy: prune entries in set id session mapping

frame/beefy-mmr/src/mock.rs

@@ -133,6 +133,7 @@ impl pallet_beefy::Config for Test {
 	)>>::IdentificationTuple;
 	type HandleEquivocation = ();
 	type MaxAuthorities = ConstU32<100>;
+	type MaxSetIdSessionEntries = ConstU64<100>;
 	type OnNewValidatorSet = BeefyMmr;
 	type WeightInfo = ();
 }

frame/beefy/src/lib.rs

@@ -99,8 +99,18 @@ pub mod pallet {
 		type HandleEquivocation: HandleEquivocation<Self>;
 
 		/// The maximum number of authorities that can be added.
+		#[pallet::constant]
 		type MaxAuthorities: Get<u32>;
 
+		/// The maximum number of entries to keep in the set id to session index mapping.
+		///
+		/// Since the `SetIdSession` map is only used for validating equivocations this
+		/// value should relate to the bonding duration of whatever staking system is
+		/// being used (if any). If equivocation handling is not enabled then this value
+		/// can be zero.
+		#[pallet::constant]
+		type MaxSetIdSessionEntries: Get<u64>;
+
 		/// A hook to act on the new BEEFY validator set.
 		///
 		/// For some applications it might be beneficial to make the BEEFY validator set available
@@ -136,6 +146,12 @@ pub mod pallet {
 	/// A mapping from BEEFY set ID to the index of the *most recent* session for which its
 	/// members were responsible.
 	///
+	/// This is only used for validating equivocation proofs. An equivocation proof must
+	/// contains a key-ownership proof for a given session, therefore we need a way to tie
+	/// together sessions and BEEFY set ids, i.e. we need to validate that a validator
+	/// was the owner of a given key on a given session, and what the active set ID was
+	/// during that session.
+	///
 	/// TWOX-NOTE: `ValidatorSetId` is not under user control.
 	#[pallet::storage]
 	#[pallet::getter(fn session_for_set)]
@@ -462,9 +478,15 @@ where
 		// We want to have at least one BEEFY mandatory block per session.
 		Self::change_authorities(bounded_next_authorities, bounded_next_queued_authorities);
 
+		let validator_set_id = Self::validator_set_id();
 		// Update the mapping for the new set id that corresponds to the latest session (i.e. now).
 		let session_index = <pallet_session::Pallet<T>>::current_index();
-		SetIdSession::<T>::insert(Self::validator_set_id(), &session_index);
+		SetIdSession::<T>::insert(validator_set_id, &session_index);
+		// Prune old entry if limit reached.
+		let max_set_id_session_entries = T::MaxSetIdSessionEntries::get().max(1);
+		if validator_set_id >= max_set_id_session_entries {
+			SetIdSession::<T>::remove(validator_set_id - max_set_id_session_entries);
+		}
 	}
 
 	fn on_disabled(i: u32) {

frame/beefy/src/mock.rs

@@ -111,6 +111,7 @@ parameter_types! {
 	pub const Period: u64 = 1;
 	pub const ReportLongevity: u64 =
 		BondingDuration::get() as u64 * SessionsPerEra::get() as u64 * Period::get();
+	pub const MaxSetIdSessionEntries: u32 = BondingDuration::get() * SessionsPerEra::get();
 }
 
 impl pallet_beefy::Config for Test {
@@ -125,6 +126,7 @@ impl pallet_beefy::Config for Test {
 	type HandleEquivocation =
 		super::EquivocationHandler<u64, Self::KeyOwnerIdentification, Offences, ReportLongevity>;
 	type MaxAuthorities = ConstU32<100>;
+	type MaxSetIdSessionEntries = MaxSetIdSessionEntries;
 	type OnNewValidatorSet = ();
 	type WeightInfo = ();
 }

frame/beefy/src/tests.rs

@@ -162,6 +162,39 @@ fn validator_set_updates_work() {
 	});
 }
 
+#[test]
+fn cleans_up_old_set_id_session_mappings() {
+	new_test_ext(vec![1, 2, 3, 4]).execute_with(|| {
+		let max_set_id_session_entries = MaxSetIdSessionEntries::get();
+
+		// we have 3 sessions per era
+		let era_limit = max_set_id_session_entries / 3;
+		// sanity check against division precision loss
+		assert_eq!(0, max_set_id_session_entries % 3);
+		// go through `max_set_id_session_entries` sessions
+		start_era(era_limit);
+
+		// we should have a session id mapping for all the set ids from
+		// `max_set_id_session_entries` eras we have observed
+		for i in 1..=max_set_id_session_entries {
+			assert!(Beefy::session_for_set(i as u64).is_some());
+		}
+
+		// go through another `max_set_id_session_entries` sessions
+		start_era(era_limit * 2);
+
+		// we should keep tracking the new mappings for new sessions
+		for i in max_set_id_session_entries + 1..=max_set_id_session_entries * 2 {
+			assert!(Beefy::session_for_set(i as u64).is_some());
+		}
+
+		// but the old ones should have been pruned by now
+		for i in 1..=max_set_id_session_entries {
+			assert!(Beefy::session_for_set(i as u64).is_none());
+		}
+	});
+}
+
 /// Returns a list with 3 authorities with known keys:
 /// Alice, Bob and Charlie.
 pub fn test_authorities() -> Vec<BeefyId> {