diff --git a/pallets/pass/src/lib.rs b/pallets/pass/src/lib.rs index 20ae7dc..d49e116 100644 --- a/pallets/pass/src/lib.rs +++ b/pallets/pass/src/lib.rs @@ -8,15 +8,19 @@ use fc_traits_authn::{ util::AuthorityFromPalletId, Authenticator, DeviceChallengeResponse, DeviceId, HashedUserId, UserAuthenticator, UserChallengeResponse, }; +use frame_support::traits::schedule::DispatchTime; +use frame_support::traits::Bounded; use frame_support::{ pallet_prelude::*, traits::{ fungible::{Inspect, Mutate}, + schedule::v3::{Named, TaskName}, EnsureOriginWithArg, }, PalletId, }; use frame_system::{pallet_prelude::*, RawOrigin}; +use sp_core::blake2_256; use sp_runtime::{ traits::{Dispatchable, TrailingZeroInput}, DispatchResult, @@ -76,6 +80,12 @@ pub mod pallet { Success = Option>, >; + type Scheduler: Named< + BlockNumberFor, + >::RuntimeCall, + Self::PalletsOrigin, + >; + #[cfg(feature = "runtime-benchmarks")] type BenchmarkHelper: BenchmarkHelper; } @@ -156,7 +166,7 @@ pub mod pallet { |_: &OriginFor, device_id: &DeviceId, credential: &CredentialOf, _: &Option>| -> bool { Pallet::::try_authenticate(device_id, credential).is_ok() } - )] + )] #[pallet::call_index(3)] pub fn authenticate( origin: OriginFor, @@ -209,6 +219,16 @@ pub mod pallet { // Turn the result from the `dispatch` into our expected `DispatchResult` type. res.map(|_| ()).map_err(|e| e.error) } + + #[pallet::call_index(6)] + pub fn remove_session_key( + origin: OriginFor, + session_key: T::AccountId, + ) -> DispatchResult { + ensure_root(origin)?; + Self::do_remove_session(&session_key); + Ok(()) + } } } @@ -286,8 +306,11 @@ impl, I: 'static> Pallet { let (account_id, until) = Sessions::::get(&who).ok_or(Error::::SessionNotFound)?; + + // Failsafe: In the extreme case the scheduler needs to defer the cleanup call for a + // certain if frame_system::Pallet::::block_number() > until { - Self::try_remove_session(&who)?; + Self::do_remove_session(&who); return Err(Error::::SessionExpired.into()); } @@ -320,14 +343,15 @@ impl, I: 'static> Pallet { Ok(account_id) } - fn try_remove_session(who: &T::AccountId) -> DispatchResult { + fn do_remove_session(session_key: &T::AccountId) { + Self::cancel_scheduled_session_key_removal(session_key); // Decrements the provider reference of this `Session` key account once it's expired. - // - // NOTE: This might not get called at all. We should explore an alternative (maybe a - // task) to remove the provider references on all expired sessions. - frame_system::Pallet::::dec_providers(who)?; - Sessions::::remove(who); - Ok(()) + let _ = frame_system::Pallet::::dec_providers(session_key).inspect_err(|_| { + log::warn!( + "Failed to remove the last provider for {session_key:?}, which has an active consumer" + ) + }); + Sessions::::remove(session_key); } pub(crate) fn try_add_session( @@ -338,19 +362,13 @@ impl, I: 'static> Pallet { // Let's try to remove an existing session that uses the same session key (if any). This is // so we ensure we decrease the provider counter correctly. if Sessions::::contains_key(session_key) { - Self::try_remove_session(session_key)?; + Self::do_remove_session(session_key); } let block_number = frame_system::Pallet::::block_number(); // Add a consumer reference to this account, since we'll be using // it meanwhile it stays active as a Session. - // - // NOTE: It is possible that this session might not be used at all, and therefore, this - // provider reference never removed. - // - // We should explore an alternative (maybe a task) to remove the provider references on all - // expired sessions. frame_system::Pallet::::inc_providers(session_key); let session_duration = duration @@ -359,6 +377,7 @@ impl, I: 'static> Pallet { let until = block_number + session_duration; Sessions::::insert(session_key.clone(), (account_id.clone(), until)); + Self::schedule_next_removal(session_key, duration)?; Self::deposit_event(Event::::SessionCreated { session_key: session_key.clone(), @@ -367,4 +386,40 @@ impl, I: 'static> Pallet { Ok(()) } + + fn task_name_from_session_key(session_key: &T::AccountId) -> TaskName { + blake2_256(format!("remove_session_key_{session_key}").as_bytes()) + } + + /// Infallibly cancels an already scheduled session key removal + fn cancel_scheduled_session_key_removal(session_key: &T::AccountId) { + T::Scheduler::cancel_named(Self::task_name_from_session_key(session_key)) + .map_or_else(|_| (), |_| ()); + } + + fn schedule_next_removal( + session_key: &T::AccountId, + duration: Option>, + ) -> DispatchResult { + Self::cancel_scheduled_session_key_removal(session_key); + + let duration = duration + .unwrap_or(T::MaxSessionDuration::get()) + .min(T::MaxSessionDuration::get()); + let call: >::RuntimeCall = Call::remove_session_key { + session_key: session_key.clone(), + } + .into(); + + T::Scheduler::schedule_named( + Self::task_name_from_session_key(session_key), + DispatchTime::After(duration), + None, + 0, + frame_system::RawOrigin::Root.into(), + Bounded::Inline(BoundedVec::truncate_from(call.encode())), + )?; + + Ok(()) + } } diff --git a/pallets/pass/src/mock.rs b/pallets/pass/src/mock.rs index 5b3c86d..f9dd34a 100644 --- a/pallets/pass/src/mock.rs +++ b/pallets/pass/src/mock.rs @@ -10,7 +10,7 @@ use frame_support::{ weights::Weight, DebugNoBound, EqNoBound, PalletId, }; -use frame_system::{pallet_prelude::OriginFor, EnsureRoot, EnsureRootWithSuccess}; +use frame_system::{EnsureRoot, EnsureRootWithSuccess}; use scale_info::TypeInfo; use sp_core::{blake2_256, H256}; use sp_io::TestExternalities; @@ -27,15 +27,29 @@ pub type AccountPublic = ::Signer; pub type AccountId = ::AccountId; // Configure a mock runtime to test the pallet. -frame_support::construct_runtime!( - pub enum Test - { - System: frame_system, - Balances: pallet_balances, - Scheduler: pallet_scheduler, - Pass: pallet_pass, - } -); +#[frame_support::runtime] +mod runtime { + #[runtime::runtime] + #[runtime::derive( + RuntimeCall, + RuntimeEvent, + RuntimeError, + RuntimeOrigin, + RuntimeTask, + RuntimeHoldReason, + RuntimeFreezeReason + )] + pub struct Test; + + #[runtime::pallet_index(0)] + pub type System = frame_system; + #[runtime::pallet_index(1)] + pub type Scheduler = pallet_scheduler; + #[runtime::pallet_index(10)] + pub type Balances = pallet_balances; + #[runtime::pallet_index(11)] + pub type Pass = pallet_pass; +} #[derive_impl(frame_system::config_preludes::TestDefaultConfig as frame_system::DefaultConfig)] impl frame_system::Config for Test { @@ -107,11 +121,14 @@ impl Config for Test { type RuntimeCall = RuntimeCall; type PalletId = PassPalletId; type PalletsOrigin = OriginCaller; + type Scheduler = Scheduler; type MaxSessionDuration = ConstU64<10>; #[cfg(feature = "runtime-benchmarks")] type BenchmarkHelper = BenchmarkHelper; } +#[cfg(feature = "runtime-benchmarks")] +use frame_system::pallet_prelude::OriginFor; #[cfg(feature = "runtime-benchmarks")] use pallet_pass::{CredentialOf, DeviceAttestationOf}; diff --git a/pallets/pass/src/tests.rs b/pallets/pass/src/tests.rs index ddf7f18..69396f0 100644 --- a/pallets/pass/src/tests.rs +++ b/pallets/pass/src/tests.rs @@ -443,6 +443,7 @@ mod add_device { mod dispatch { use super::*; + use crate::Sessions; parameter_types! { pub Call: Box = Box::new(RuntimeCall::System(frame_system::Call::remark_with_event { @@ -646,10 +647,7 @@ mod dispatch { run_to(12); - assert_noop!( - Pass::dispatch(RuntimeOrigin::signed(SIGNER), Call::get(), None, None,), - Error::::SessionExpired - ); + assert!(!Sessions::::contains_key(SIGNER)); assert_ok!(Pass::dispatch( RuntimeOrigin::signed(OTHER), @@ -660,10 +658,7 @@ mod dispatch { run_to(20); - assert_noop!( - Pass::dispatch(RuntimeOrigin::signed(OTHER), Call::get(), None, None,), - Error::::SessionExpired - ); + assert!(!Sessions::::contains_key(OTHER)); }); } } diff --git a/pallets/pass/src/weights.rs b/pallets/pass/src/weights.rs index e5d4e29..e809391 100644 --- a/pallets/pass/src/weights.rs +++ b/pallets/pass/src/weights.rs @@ -3,8 +3,8 @@ #![allow(unused_imports)] #![allow(missing_docs)] -use frame_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}}; use core::marker::PhantomData; +use frame_support::weights::Weight; /// Weight functions needed for pallet_remark. pub trait WeightInfo { @@ -14,6 +14,7 @@ pub trait WeightInfo { fn authenticate() -> Weight; fn add_device() -> Weight; fn dispatch() -> Weight; + fn remove_session_key() -> Weight; } /// Weights for pallet_remark using the Substrate node and recommended hardware. @@ -84,6 +85,17 @@ impl WeightInfo for SubstrateWeight { // Standard Error: 0 .saturating_add(Weight::from_parts(1_359, 0)) } + + /// The range of component `l` is `[1, 1048576]`. + fn remove_session_key() -> Weight { + // Proof Size summary in bytes: + // Measured: `0` + // Estimated: `0` + // Minimum execution time: 8_471_000 picoseconds. + Weight::from_parts(8_586_000, 0) + // Standard Error: 0 + .saturating_add(Weight::from_parts(1_359, 0)) + } } // For backwards compatibility and tests @@ -153,4 +165,15 @@ impl WeightInfo for () { // Standard Error: 0 .saturating_add(Weight::from_parts(0, 0)) } + + /// The range of component `l` is `[1, 1048576]`. + fn remove_session_key() -> Weight { + // Proof Size summary in bytes: + // Measured: `0` + // Estimated: `0` + // Minimum execution time: 8_471_000 picoseconds. + Weight::from_parts(0, 0) + // Standard Error: 0 + .saturating_add(Weight::from_parts(0, 0)) + } }