pallet-session: track consumer refs and release deposits for externally set keys

prdoc/pr_11197.prdoc

@@ -0,0 +1,8 @@
+title: 'pallet-session: track consumer refs and release deposits for externally set keys'
+doc:
+- audience: Runtime Dev
+  description: |-
+    Tracks whether session keys were set externally (via `SessionInterface`, e.g. from AH) or locally. Transitions between the two paths correctly manage the key deposit and consumer ref in both directions.
+crates:
+- name: pallet-session
+  bump: minor

substrate/frame/session/src/lib.rs

@@ -626,6 +626,14 @@ pub mod pallet {
 	pub type KeyOwner<T: Config> =
 		StorageMap<_, Twox64Concat, (KeyTypeId, Vec<u8>), T::ValidatorId, OptionQuery>;
 
+	/// Accounts whose keys were set via `SessionInterface` (external path) without
+	/// incrementing the consumer reference or placing a key deposit. `do_purge_keys`
+	/// only decrements consumers for accounts that were registered through the local
+	/// session pallet.
+	#[pallet::storage]
+	pub type ExternallySetKeys<T: Config> =
+		StorageMap<_, Twox64Concat, T::AccountId, (), OptionQuery>;
+
 	#[pallet::event]
 	#[pallet::generate_deposit(pub(super) fn deposit_event)]
 	pub enum Event<T: Config> {
@@ -920,9 +928,12 @@ impl<T: Config> Pallet<T> {
 
 		let old_keys = Self::inner_set_keys(&who, keys)?;
 
-		// Place deposit on hold if this is a new registration (i.e. old_keys is None).
-		// The hold call itself will return an error if funds are insufficient.
-		if old_keys.is_none() {
+		// Place deposit and increment consumer if this is a new local registration,
+		// or if transitioning from external to local management.
+		// We also clear `ExternallySetKeys` if set.
+		let needs_local_setup =
+			old_keys.is_none() || ExternallySetKeys::<T>::take(account).is_some();
+		if needs_local_setup {
 			let deposit = T::KeyDeposit::get();
 			if !deposit.is_zero() {
 				T::Currency::hold(&HoldReason::Keys.into(), account, deposit)?;
@@ -996,7 +1007,10 @@ impl<T: Config> Pallet<T> {
 			frame_support::traits::tokens::Precision::BestEffort,
 		);
 
-		frame_system::Pallet::<T>::dec_consumers(account);
+		if ExternallySetKeys::<T>::take(account).is_none() {
+			// Consumer was incremented locally via `do_set_keys`, so decrement it.
+			frame_system::Pallet::<T>::dec_consumers(account);
+		}
 
 		Ok(())
 	}
@@ -1196,7 +1210,17 @@ impl<T: Config + historical::Config> SessionInterface for Pallet<T> {
 	fn set_keys(account: &Self::AccountId, keys: Self::Keys) -> DispatchResult {
 		let who = T::ValidatorIdOf::convert(account.clone())
 			.ok_or(Error::<T>::NoAssociatedValidatorId)?;
-		Self::inner_set_keys(&who, keys)?;
+		let old_keys = Self::inner_set_keys(&who, keys)?;
+		// Transitioning from local to external: clean up deposit and consumer ref.
+		if old_keys.is_some() && !ExternallySetKeys::<T>::contains_key(account) {
+			let _ = T::Currency::release_all(
+				&HoldReason::Keys.into(),
+				account,
+				frame_support::traits::tokens::Precision::BestEffort,
+			);
+			frame_system::Pallet::<T>::dec_consumers(account);
+		}
+		ExternallySetKeys::<T>::insert(account, ());
 		Ok(())
 	}
 
@@ -1209,6 +1233,14 @@ impl<T: Config + historical::Config> SessionInterface for Pallet<T> {
 			let key_data = old_keys.get_raw(*id);
 			Self::clear_key_owner(*id, key_data);
 		}
+		let _ = T::Currency::release_all(
+			&HoldReason::Keys.into(),
+			account,
+			frame_support::traits::tokens::Precision::BestEffort,
+		);
+		if ExternallySetKeys::<T>::take(account).is_none() {
+			frame_system::Pallet::<T>::dec_consumers(account);
+		}
 		Ok(())
 	}
 

substrate/frame/session/src/tests.rs

@@ -627,6 +627,142 @@ fn existing_validators_without_hold_are_except() {
 	});
 }
 
+#[cfg(feature = "historical")]
+mod externally_set_keys_tracking {
+	use super::*;
+
+	fn setup_account(account_id: u64) {
+		frame_system::Pallet::<Test>::inc_providers(&account_id);
+		ValidatorAccounts::mutate(|m| {
+			m.insert(account_id, account_id);
+		});
+	}
+
+	#[test]
+	fn purge_skips_dec_consumers_for_externally_set_keys() {
+		// GIVEN: keys set externally (no inc_consumers, no deposit)
+		// WHEN: purge (locally or externally)
+		// THEN: consumer count unchanged, no deposit
+		let account_id = 1000;
+		new_test_ext().execute_with(|| {
+			setup_account(account_id);
+			let consumers_before = System::consumers(&account_id);
+
+			<Session as SessionInterface>::set_keys(
+				&account_id,
+				UintAuthorityId(account_id).into(),
+			)
+			.unwrap();
+			assert!(ExternallySetKeys::<Test>::contains_key(&account_id));
+			assert_eq!(System::consumers(&account_id), consumers_before);
+			assert_eq!(session_hold(account_id), 0);
+
+			// External-to-external update: no change.
+			<Session as SessionInterface>::set_keys(&account_id, UintAuthorityId(70).into())
+				.unwrap();
+			assert!(ExternallySetKeys::<Test>::contains_key(&account_id));
+			assert_eq!(System::consumers(&account_id), consumers_before);
+			assert_eq!(session_hold(account_id), 0);
+
+			// Purge locally
+			frame_support::hypothetically!({
+				assert_ok!(Session::purge_keys(RuntimeOrigin::signed(account_id)));
+				assert_eq!(System::consumers(&account_id), consumers_before);
+				assert_eq!(session_hold(account_id), 0);
+				assert!(!ExternallySetKeys::<Test>::contains_key(&account_id));
+			});
+
+			// Purge externally
+			<Session as SessionInterface>::purge_keys(&account_id).unwrap();
+			assert_eq!(System::consumers(&account_id), consumers_before);
+			assert_eq!(session_hold(account_id), 0);
+			assert!(!ExternallySetKeys::<Test>::contains_key(&account_id));
+		});
+	}
+
+	#[test]
+	fn external_purge_releases_deposit_from_prior_local_registration() {
+		// GIVEN: keys set locally (deposit held)
+		// WHEN: external purge_keys
+		// THEN: deposit released
+		let account_id = 1000;
+		new_test_ext().execute_with(|| {
+			setup_account(account_id);
+
+			let keys = UintAuthorityId(account_id).into();
+			let proof = create_set_keys_proof(account_id, &UintAuthorityId(account_id));
+			assert_ok!(Session::set_keys(RuntimeOrigin::signed(account_id), keys, proof));
+			assert_eq!(session_hold(account_id), KeyDeposit::get());
+
+			// External purge releases the deposit.
+			<Session as SessionInterface>::purge_keys(&account_id).unwrap();
+			assert_eq!(session_hold(account_id), 0);
+		});
+	}
+
+	#[test]
+	fn purge_decrements_consumers_for_locally_set_keys() {
+		// GIVEN: keys set locally (inc_consumers + hold)
+		// WHEN: purge locally
+		// THEN: consumer count returns to baseline, deposit released
+		let account_id = 1000;
+		new_test_ext().execute_with(|| {
+			setup_account(account_id);
+			let consumers_before = System::consumers(&account_id);
+
+			let keys = UintAuthorityId(account_id).into();
+			let proof = create_set_keys_proof(account_id, &UintAuthorityId(account_id));
+			assert_ok!(Session::set_keys(RuntimeOrigin::signed(account_id), keys, proof));
+			assert!(!ExternallySetKeys::<Test>::contains_key(&account_id));
+			assert!(System::consumers(&account_id) > consumers_before);
+			assert_eq!(session_hold(account_id), KeyDeposit::get());
+
+			assert_ok!(Session::purge_keys(RuntimeOrigin::signed(account_id)));
+			assert_eq!(System::consumers(&account_id), consumers_before);
+			assert_eq!(session_hold(account_id), 0);
+		});
+	}
+
+	#[test]
+	fn full_round_trip_local_external_local() {
+		// local → external → local → purge
+		let account_id = 1000;
+		new_test_ext().execute_with(|| {
+			setup_account(account_id);
+			let consumers_before = System::consumers(&account_id);
+
+			// 1. Local set: deposit + consumer.
+			let keys = UintAuthorityId(account_id).into();
+			let proof = create_set_keys_proof(account_id, &UintAuthorityId(account_id));
+			assert_ok!(Session::set_keys(RuntimeOrigin::signed(account_id), keys, proof));
+			assert!(!ExternallySetKeys::<Test>::contains_key(&account_id));
+			let consumers_after_local = System::consumers(&account_id);
+			assert!(consumers_after_local > consumers_before);
+			assert_eq!(session_hold(account_id), KeyDeposit::get());
+
+			// 2. External set: cleans up deposit + consumer.
+			<Session as SessionInterface>::set_keys(&account_id, UintAuthorityId(70).into())
+				.unwrap();
+			assert!(ExternallySetKeys::<Test>::contains_key(&account_id));
+			assert_eq!(System::consumers(&account_id), consumers_before);
+			assert_eq!(session_hold(account_id), 0);
+
+			// 3. Local set again: deposit + consumer restored.
+			let keys = UintAuthorityId(80).into();
+			let proof = create_set_keys_proof(account_id, &UintAuthorityId(80));
+			assert_ok!(Session::set_keys(RuntimeOrigin::signed(account_id), keys, proof));
+			assert!(!ExternallySetKeys::<Test>::contains_key(&account_id));
+			assert_eq!(System::consumers(&account_id), consumers_after_local);
+			assert_eq!(session_hold(account_id), KeyDeposit::get());
+
+			// 4. Purge: back to baseline.
+			assert_ok!(Session::purge_keys(RuntimeOrigin::signed(account_id)));
+			assert_eq!(System::consumers(&account_id), consumers_before);
+			assert_eq!(session_hold(account_id), 0);
+		});
+	}
+}
+
 mod disabling_byzantine_threshold {
 	use super::*;
 	use crate::disabling::{DisablingStrategy, UpToLimitDisablingStrategy};