pallet-session: track consumer refs and release deposits for externally set keys (#11197)

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.

---------

Co-authored-by: cmd[bot] <41898282+github-actions[bot]@users.noreply.github.com>
(cherry picked from commit b965285)

Author: sigurpol

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,143 @@ fn existing_validators_without_hold_are_except() {
 	});
 }
 
+#[cfg(feature = "historical")]
+mod externally_set_keys_tracking {
+	use super::*;
+
+	const ACCOUNT: u64 = 1000;
+
+	fn setup_account() {
+		frame_system::Pallet::<Test>::inc_providers(&ACCOUNT);
+		ValidatorAccounts::mutate(|m| {
+			m.insert(ACCOUNT, ACCOUNT);
+		});
+	}
+
+	fn set_local(key: u64) {
+		let keys = UintAuthorityId(key).into();
+		let proof = create_set_keys_proof(ACCOUNT, &UintAuthorityId(key));
+		assert_ok!(Session::set_keys(RuntimeOrigin::signed(ACCOUNT), keys, proof));
+	}
+
+	fn set_remote(key: u64) {
+		<Session as SessionInterface>::set_keys(&ACCOUNT, UintAuthorityId(key).into()).unwrap();
+	}
+
+	fn purge_local() {
+		assert_ok!(Session::purge_keys(RuntimeOrigin::signed(ACCOUNT)));
+	}
+
+	fn purge_remote() {
+		<Session as SessionInterface>::purge_keys(&ACCOUNT).unwrap();
+	}
+
+	fn assert_local_state(consumers_before: u32) {
+		assert!(!ExternallySetKeys::<Test>::contains_key(&ACCOUNT));
+		// +1 from session's inc_consumers, +1 from pallet-balances hold.
+		assert_eq!(System::consumers(&ACCOUNT), consumers_before + 2);
+		assert_eq!(session_hold(ACCOUNT), KeyDeposit::get());
+	}
+
+	fn assert_remote_state(consumers_before: u32) {
+		assert!(ExternallySetKeys::<Test>::contains_key(&ACCOUNT));
+		assert_eq!(System::consumers(&ACCOUNT), consumers_before);
+		assert_eq!(session_hold(ACCOUNT), 0);
+	}
+
+	fn assert_clean_state(consumers_before: u32) {
+		assert!(!ExternallySetKeys::<Test>::contains_key(&ACCOUNT));
+		assert_eq!(System::consumers(&ACCOUNT), consumers_before);
+		assert_eq!(session_hold(ACCOUNT), 0);
+	}
+
+	#[test]
+	fn set_local_purge_local() {
+		new_test_ext().execute_with(|| {
+			setup_account();
+			let consumers_before = System::consumers(&ACCOUNT);
+
+			set_local(ACCOUNT);
+			assert_local_state(consumers_before);
+
+			purge_local();
+			assert_clean_state(consumers_before);
+		});
+	}
+
+	#[test]
+	fn set_local_purge_remote() {
+		new_test_ext().execute_with(|| {
+			setup_account();
+			let consumers_before = System::consumers(&ACCOUNT);
+
+			set_local(ACCOUNT);
+			assert_local_state(consumers_before);
+
+			purge_remote();
+			assert_clean_state(consumers_before);
+		});
+	}
+
+	#[test]
+	fn set_remote_purge_local() {
+		new_test_ext().execute_with(|| {
+			setup_account();
+			let consumers_before = System::consumers(&ACCOUNT);
+
+			set_remote(ACCOUNT);
+			assert_remote_state(consumers_before);
+
+			purge_local();
+			assert_clean_state(consumers_before);
+		});
+	}
+
+	#[test]
+	fn set_remote_purge_remote() {
+		new_test_ext().execute_with(|| {
+			setup_account();
+			let consumers_before = System::consumers(&ACCOUNT);
+
+			set_remote(ACCOUNT);
+			assert_remote_state(consumers_before);
+
+			purge_remote();
+			assert_clean_state(consumers_before);
+		});
+	}
+
+	#[test]
+	fn set_local_to_remote() {
+		new_test_ext().execute_with(|| {
+			setup_account();
+			let consumers_before = System::consumers(&ACCOUNT);
+
+			set_local(ACCOUNT);
+			assert_local_state(consumers_before);
+
+			// Transition to remote: deposit released, consumer decremented.
+			set_remote(70);
+			assert_remote_state(consumers_before);
+		});
+	}
+
+	#[test]
+	fn set_remote_to_local() {
+		new_test_ext().execute_with(|| {
+			setup_account();
+			let consumers_before = System::consumers(&ACCOUNT);
+
+			set_remote(ACCOUNT);
+			assert_remote_state(consumers_before);
+
+			// Transition to local: deposit placed, consumer incremented.
+			set_local(70);
+			assert_local_state(consumers_before);
+		});
+	}
+}
+
 mod disabling_byzantine_threshold {
 	use super::*;
 	use crate::disabling::{DisablingStrategy, UpToLimitDisablingStrategy};