diff --git a/pallets/shield/src/lib.rs b/pallets/shield/src/lib.rs index 6b742687f5..44e35e9597 100644 --- a/pallets/shield/src/lib.rs +++ b/pallets/shield/src/lib.rs @@ -3,24 +3,29 @@ extern crate alloc; +use alloc::vec; use chacha20poly1305::{ KeyInit, XChaCha20Poly1305, XNonce, aead::{Aead, Payload}, }; -use frame_support::{pallet_prelude::*, traits::IsSubType}; -use frame_system::{ensure_none, ensure_signed, pallet_prelude::*}; +use frame_support::{ + dispatch::{GetDispatchInfo, PostDispatchInfo}, + pallet_prelude::*, + traits::{ConstU64, IsSubType}, +}; +use frame_system::{ensure_none, ensure_root, ensure_signed, pallet_prelude::*}; use ml_kem::{ Ciphertext, EncodedSizeUser, MlKem768, MlKem768Params, kem::{Decapsulate, DecapsulationKey}, }; use sp_io::hashing::twox_128; use sp_runtime::traits::{Applyable, Block as BlockT, Checkable, Hash}; +use sp_runtime::traits::{Dispatchable, Saturating}; use stp_shield::{ INHERENT_IDENTIFIER, InherentType, LOG_TARGET, MLKEM768_ENC_KEY_LEN, ShieldEncKey, ShieldedTransaction, }; - -use alloc::vec; +use subtensor_macros::freeze_struct; pub use pallet::*; @@ -45,6 +50,19 @@ type ApplyableCallOf = ::Call; const MAX_EXTRINSIC_DEPTH: u32 = 8; +/// Trait for decrypting stored extrinsics before dispatch. +pub trait ExtrinsicDecryptor { + /// Decrypt the stored bytes and return the decoded RuntimeCall. + fn decrypt(data: &[u8]) -> Result; +} + +/// Default implementation that always returns an error. +impl ExtrinsicDecryptor for () { + fn decrypt(_data: &[u8]) -> Result { + Err(DispatchError::Other("ExtrinsicDecryptor not implemented")) + } +} + #[frame_support::pallet] pub mod pallet { use super::*; @@ -56,6 +74,14 @@ pub mod pallet { /// A way to find the current and next block author. type FindAuthors: FindAuthors; + + /// The overarching call type for dispatching stored extrinsics. + type RuntimeCall: Parameter + + Dispatchable + + GetDispatchInfo; + + /// Decryptor for stored extrinsics. + type ExtrinsicDecryptor: ExtrinsicDecryptor<::RuntimeCall>; } #[pallet::pallet] @@ -93,11 +119,89 @@ pub mod pallet { pub type HasMigrationRun = StorageMap<_, Identity, BoundedVec, bool, ValueQuery>; + /// Maximum size of a single encoded call. + pub type MaxCallSize = ConstU32<8192>; + + /// Default maximum number of pending extrinsics. + pub type DefaultMaxPendingExtrinsics = ConstU32<100>; + + /// Configurable maximum number of pending extrinsics. + /// Defaults to 100 if not explicitly set via `set_max_pending_extrinsics`. + #[pallet::storage] + pub type MaxPendingExtrinsicsLimit = + StorageValue<_, u32, ValueQuery, DefaultMaxPendingExtrinsics>; + + /// Default extrinsic lifetime in blocks. + pub const DEFAULT_EXTRINSIC_LIFETIME: u32 = 10; + + /// Configurable extrinsic lifetime (max block difference between submission and execution). + /// Defaults to 10 blocks if not explicitly set. + #[pallet::storage] + pub type ExtrinsicLifetime = + StorageValue<_, u32, ValueQuery, ConstU32>; + + /// Default maximum weight allowed for on_initialize processing. + pub const DEFAULT_ON_INITIALIZE_WEIGHT: u64 = 500_000_000_000; + + /// Absolute maximum weight for on_initialize: half the total block weight (2s of 4s). + pub const MAX_ON_INITIALIZE_WEIGHT: u64 = 2_000_000_000_000; + + /// Configurable maximum weight for on_initialize processing. + /// Defaults to 500_000_000_000 ref_time if not explicitly set. + #[pallet::storage] + pub type OnInitializeWeight = + StorageValue<_, u64, ValueQuery, ConstU64>; + + /// A pending extrinsic stored for later execution. + #[freeze_struct("c5749ec89253be61")] + #[derive(Clone, Encode, Decode, TypeInfo, MaxEncodedLen, PartialEq, Debug)] + #[scale_info(skip_type_params(T))] + pub struct PendingExtrinsic { + /// The account that submitted the extrinsic. + pub who: T::AccountId, + /// The encoded call data. + pub call: BoundedVec, + /// The block number when the extrinsic was submitted. + pub submitted_at: BlockNumberFor, + } + + /// Storage map for encrypted extrinsics to be executed in on_initialize. + /// Uses u32 index for O(1) insertion and removal. + #[pallet::storage] + pub type PendingExtrinsics = + StorageMap<_, Identity, u32, PendingExtrinsic, OptionQuery>; + + /// Next index to use when inserting a pending extrinsic (unique auto-increment). + #[pallet::storage] + pub type NextPendingExtrinsicIndex = StorageValue<_, u32, ValueQuery>; + + /// Number of pending extrinsics currently stored (for limit checking). + #[pallet::storage] + pub type PendingExtrinsicCount = StorageValue<_, u32, ValueQuery>; + #[pallet::event] #[pallet::generate_deposit(pub(super) fn deposit_event)] pub enum Event { /// Encrypted wrapper accepted. EncryptedSubmitted { id: T::Hash, who: T::AccountId }, + /// Encrypted extrinsic was stored for later execution. + ExtrinsicStored { index: u32, who: T::AccountId }, + /// Extrinsic decode failed during on_initialize. + ExtrinsicDecodeFailed { index: u32 }, + /// Extrinsic dispatch failed during on_initialize. + ExtrinsicDispatchFailed { index: u32, error: DispatchError }, + /// Extrinsic was successfully dispatched during on_initialize. + ExtrinsicDispatched { index: u32 }, + /// Extrinsic expired (exceeded max block lifetime). + ExtrinsicExpired { index: u32 }, + /// Extrinsic postponed due to weight limit. + ExtrinsicPostponed { index: u32 }, + /// Maximum pending extrinsics limit was updated. + MaxPendingExtrinsicsNumberSet { value: u32 }, + /// Maximum on_initialize weight was updated. + OnInitializeWeightSet { value: u64 }, + /// Extrinsic lifetime was updated. + ExtrinsicLifetimeSet { value: u32 }, } #[pallet::error] @@ -106,10 +210,18 @@ pub mod pallet { BadEncKeyLen, /// Unreachable. Unreachable, + /// Too many pending extrinsics in storage. + TooManyPendingExtrinsics, + /// Weight exceeds the absolute maximum (half of total block weight). + WeightExceedsAbsoluteMax, } #[pallet::hooks] impl Hooks> for Pallet { + fn on_initialize(_block_number: BlockNumberFor) -> Weight { + Self::process_pending_extrinsics() + } + fn on_runtime_upgrade() -> frame_support::weights::Weight { let mut weight = frame_support::weights::Weight::from_parts(0, 0); @@ -229,6 +341,84 @@ pub mod pallet { Self::deposit_event(Event::EncryptedSubmitted { id, who }); Ok(()) } + + /// Store an encrypted extrinsic for later execution in on_initialize. + #[pallet::call_index(2)] + #[pallet::weight(Weight::from_parts(10_000_000, 0) + .saturating_add(T::DbWeight::get().reads(2_u64)) + .saturating_add(T::DbWeight::get().writes(3_u64)))] + pub fn store_encrypted( + origin: OriginFor, + call: BoundedVec, + ) -> DispatchResult { + let who = ensure_signed(origin)?; + + let count = PendingExtrinsicCount::::get(); + + ensure!( + count < MaxPendingExtrinsicsLimit::::get(), + Error::::TooManyPendingExtrinsics + ); + + let index = NextPendingExtrinsicIndex::::get(); + let pending = PendingExtrinsic { + who: who.clone(), + call, + submitted_at: frame_system::Pallet::::block_number(), + }; + PendingExtrinsics::::insert(index, pending); + + NextPendingExtrinsicIndex::::put(index.saturating_add(1)); + PendingExtrinsicCount::::put(count.saturating_add(1)); + + Self::deposit_event(Event::ExtrinsicStored { index, who }); + Ok(()) + } + + /// Set the maximum number of pending extrinsics allowed in the queue. + #[pallet::call_index(3)] + #[pallet::weight(T::DbWeight::get().writes(1_u64))] + pub fn set_max_pending_extrinsics_number( + origin: OriginFor, + value: u32, + ) -> DispatchResult { + ensure_root(origin)?; + + MaxPendingExtrinsicsLimit::::put(value); + + Self::deposit_event(Event::MaxPendingExtrinsicsNumberSet { value }); + Ok(()) + } + + /// Set the maximum weight allowed for on_initialize processing. + /// Rejects values exceeding the absolute limit (half of total block weight). + #[pallet::call_index(4)] + #[pallet::weight(T::DbWeight::get().writes(1_u64))] + pub fn set_on_initialize_weight(origin: OriginFor, value: u64) -> DispatchResult { + ensure_root(origin)?; + + ensure!( + value <= MAX_ON_INITIALIZE_WEIGHT, + Error::::WeightExceedsAbsoluteMax + ); + + OnInitializeWeight::::put(value); + + Self::deposit_event(Event::OnInitializeWeightSet { value }); + Ok(()) + } + + /// Set the extrinsic lifetime (max blocks between submission and execution). + #[pallet::call_index(5)] + #[pallet::weight(T::DbWeight::get().writes(1_u64))] + pub fn set_stored_extrinsic_lifetime(origin: OriginFor, value: u32) -> DispatchResult { + ensure_root(origin)?; + + ExtrinsicLifetime::::put(value); + + Self::deposit_event(Event::ExtrinsicLifetimeSet { value }); + Ok(()) + } } #[pallet::inherent] @@ -255,6 +445,98 @@ pub mod pallet { } impl Pallet { + /// Process pending encrypted extrinsics up to the weight limit. + /// Returns the total weight consumed. + pub fn process_pending_extrinsics() -> Weight { + let next_index = NextPendingExtrinsicIndex::::get(); + let count = PendingExtrinsicCount::::get(); + + let mut weight = T::DbWeight::get().reads(2); + + if count == 0 { + return weight; + } + + let start_index = next_index.saturating_sub(count); + let current_block = frame_system::Pallet::::block_number(); + + // Process extrinsics + for index in start_index..next_index { + let Some(pending) = PendingExtrinsics::::get(index) else { + weight = weight.saturating_add(T::DbWeight::get().reads(1)); + + continue; + }; + + // Check if the extrinsic has expired + let age = current_block.saturating_sub(pending.submitted_at); + if age > ExtrinsicLifetime::::get().into() { + remove_pending_extrinsic::(index, &mut weight); + + Self::deposit_event(Event::ExtrinsicExpired { index }); + + continue; + } + + let call = match T::ExtrinsicDecryptor::decrypt(&pending.call) { + Ok(call) => call, + Err(_) => { + remove_pending_extrinsic::(index, &mut weight); + + Self::deposit_event(Event::ExtrinsicDecodeFailed { index }); + + continue; + } + }; + + // Check if dispatching would exceed weight limit + let info = call.get_dispatch_info(); + let dispatch_weight = T::DbWeight::get() + .writes(2) + .saturating_add(info.call_weight); + + let max_weight = Weight::from_parts(OnInitializeWeight::::get(), 0); + + if weight.saturating_add(dispatch_weight).any_gt(max_weight) { + Self::deposit_event(Event::ExtrinsicPostponed { index }); + break; + } + + // We're going to execute it - remove the item from storage + remove_pending_extrinsic::(index, &mut weight); + + // Dispatch the extrinsic + let origin: T::RuntimeOrigin = frame_system::RawOrigin::Signed(pending.who).into(); + let result = call.dispatch(origin); + + match result { + Ok(post_info) => { + let actual_weight = post_info.actual_weight.unwrap_or(info.call_weight); + weight = weight.saturating_add(actual_weight); + + Self::deposit_event(Event::ExtrinsicDispatched { index }); + } + Err(e) => { + weight = weight.saturating_add(info.call_weight); + + Self::deposit_event(Event::ExtrinsicDispatchFailed { + index, + error: e.error, + }); + } + } + } + + /// Remove a pending extrinsic from storage and decrement count. + fn remove_pending_extrinsic(index: u32, weight: &mut Weight) { + PendingExtrinsics::::remove(index); + PendingExtrinsicCount::::mutate(|c| *c = c.saturating_sub(1)); + *weight = weight.saturating_add(T::DbWeight::get().writes(2)); + } + + weight + } + pub fn try_decode_shielded_tx( uxt: ExtrinsicOf, ) -> Option diff --git a/pallets/shield/src/mock.rs b/pallets/shield/src/mock.rs index 5a2aef7d80..530a0b913d 100644 --- a/pallets/shield/src/mock.rs +++ b/pallets/shield/src/mock.rs @@ -1,6 +1,8 @@ use crate as pallet_shield; use stp_shield::MLKEM768_ENC_KEY_LEN; +use codec::Decode; +use frame_support::pallet_prelude::DispatchError; use frame_support::traits::{ConstBool, ConstU64}; use frame_support::{BoundedVec, construct_runtime, derive_impl, parameter_types}; use sp_consensus_aura::sr25519::AuthorityId as AuraId; @@ -85,9 +87,20 @@ impl pallet_shield::FindAuthors for MockFindAuthors { } } +/// Mock decryptor that just decodes the bytes without decryption. +pub struct MockDecryptor; + +impl pallet_shield::ExtrinsicDecryptor for MockDecryptor { + fn decrypt(data: &[u8]) -> Result { + RuntimeCall::decode(&mut &data[..]).map_err(|_| DispatchError::Other("decode failed")) + } +} + impl pallet_shield::Config for Test { type AuthorityId = AuraId; type FindAuthors = MockFindAuthors; + type RuntimeCall = RuntimeCall; + type ExtrinsicDecryptor = MockDecryptor; } pub fn new_test_ext() -> sp_io::TestExternalities { diff --git a/pallets/shield/src/tests.rs b/pallets/shield/src/tests.rs index 04eb29126b..14e906ac2f 100644 --- a/pallets/shield/src/tests.rs +++ b/pallets/shield/src/tests.rs @@ -1,9 +1,9 @@ use crate::mock::*; use crate::{ - AuthorKeys, CurrentKey, Error, HasMigrationRun, NextKey, NextKeyExpiresAt, PendingKey, - PendingKeyExpiresAt, + AuthorKeys, CurrentKey, Error, ExtrinsicLifetime, HasMigrationRun, MaxPendingExtrinsicsLimit, + NextKey, NextKeyExpiresAt, NextPendingExtrinsicIndex, OnInitializeWeight, PendingExtrinsic, + PendingExtrinsicCount, PendingExtrinsics, PendingKey, PendingKeyExpiresAt, }; - use codec::Encode; use frame_support::{BoundedVec, assert_noop, assert_ok}; use sp_runtime::testing::TestSignature; @@ -460,3 +460,627 @@ mod migration_tests { count } } + +// --------------------------------------------------------------------------- +// Encrypted extrinsics storage tests +// --------------------------------------------------------------------------- + +mod encrypted_extrinsics_tests { + use super::*; + use frame_support::traits::Hooks; + + #[test] + fn store_encrypted_works() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + let call = RuntimeCall::System(frame_system::Call::remark { + remark: vec![1, 2, 3], + }); + let encoded_call = BoundedVec::truncate_from(call.encode()); + let who: u64 = 1; + + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(who), + encoded_call.clone(), + )); + + // Verify the extrinsic was stored at index 0 with account ID + let expected = PendingExtrinsic:: { + who, + call: encoded_call, + submitted_at: 1, + }; + assert_eq!(PendingExtrinsics::::get(0), Some(expected)); + assert_eq!(NextPendingExtrinsicIndex::::get(), 1); + assert_eq!(PendingExtrinsicCount::::get(), 1); + + // Verify event was emitted with index + System::assert_last_event( + crate::Event::::ExtrinsicStored { index: 0, who }.into(), + ); + }); + } + + #[test] + fn on_initialize_decodes_and_dispatches() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Store an encoded remark call + let call = RuntimeCall::System(frame_system::Call::remark { + remark: vec![1, 2, 3], + }); + let encoded_call = BoundedVec::truncate_from(call.encode()); + + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + encoded_call, + )); + + // Verify there's a pending extrinsic + assert_eq!(NextPendingExtrinsicIndex::::get(), 1); + assert_eq!(PendingExtrinsicCount::::get(), 1); + assert!(PendingExtrinsics::::get(0).is_some()); + + // Run on_initialize + MevShield::on_initialize(2); + + // Verify storage was cleared but NextPendingExtrinsicIndex stays (unique auto-increment) + assert!(PendingExtrinsics::::get(0).is_none()); + assert_eq!(NextPendingExtrinsicIndex::::get(), 1); + assert_eq!(PendingExtrinsicCount::::get(), 0); + + // Verify ExtrinsicDispatched event was emitted + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 0 }.into()); + }); + } + + #[test] + fn on_initialize_handles_decode_failure() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Store invalid bytes that can't be decoded as a call + let invalid_bytes = BoundedVec::truncate_from(vec![0xFF, 0xFF, 0xFF, 0xFF]); + + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + invalid_bytes, + )); + + // Run on_initialize + MevShield::on_initialize(2); + + // Verify storage was cleared + assert!(PendingExtrinsics::::get(0).is_none()); + + // Verify ExtrinsicDecodeFailed event was emitted + System::assert_has_event( + crate::Event::::ExtrinsicDecodeFailed { index: 0 }.into(), + ); + }); + } + + #[test] + fn on_initialize_handles_dispatch_failure() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Test with multiple calls to ensure the iteration works correctly. + + let call1 = RuntimeCall::System(frame_system::Call::remark { + remark: vec![1, 2, 3], + }); + let call2 = RuntimeCall::System(frame_system::Call::remark { + remark: vec![4, 5, 6], + }); + + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call1.encode()), + )); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call2.encode()), + )); + + // Verify there are 2 pending extrinsics + assert_eq!(NextPendingExtrinsicIndex::::get(), 2); + assert_eq!(PendingExtrinsicCount::::get(), 2); + assert!(PendingExtrinsics::::get(0).is_some()); + assert!(PendingExtrinsics::::get(1).is_some()); + + // Run on_initialize + MevShield::on_initialize(2); + + // Verify storage was cleared + assert!(PendingExtrinsics::::get(0).is_none()); + assert!(PendingExtrinsics::::get(1).is_none()); + + // Verify both calls were dispatched + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 0 }.into()); + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 1 }.into()); + }); + } + + #[test] + fn store_encrypted_rejects_when_full() { + new_test_ext().execute_with(|| { + let max = MaxPendingExtrinsicsLimit::::get(); + + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + let encoded_call = BoundedVec::truncate_from(call.encode()); + + // Fill up the pending extrinsics storage to max + for _ in 0..max { + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + encoded_call.clone(), + )); + } + + // The next one should fail + assert_noop!( + MevShield::store_encrypted(RuntimeOrigin::signed(1), encoded_call), + Error::::TooManyPendingExtrinsics + ); + }); + } + + #[test] + fn on_initialize_processes_mixed_success_and_failure() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Store a valid call + let valid_call = RuntimeCall::System(frame_system::Call::remark { + remark: vec![1, 2, 3], + }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(valid_call.encode()), + )); + + // Store invalid bytes + let invalid_bytes = BoundedVec::truncate_from(vec![0xFF, 0xFF]); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + invalid_bytes, + )); + + // Store another valid call + let valid_call2 = RuntimeCall::System(frame_system::Call::remark { + remark: vec![4, 5, 6], + }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(valid_call2.encode()), + )); + + // Run on_initialize + MevShield::on_initialize(2); + + // Verify storage was cleared + assert!(PendingExtrinsics::::get(0).is_none()); + assert!(PendingExtrinsics::::get(1).is_none()); + assert!(PendingExtrinsics::::get(2).is_none()); + + // Verify correct events were emitted + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 0 }.into()); + System::assert_has_event( + crate::Event::::ExtrinsicDecodeFailed { index: 1 }.into(), + ); + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 2 }.into()); + }); + } + + #[test] + fn on_initialize_expires_old_extrinsics() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Store an extrinsic at block 1 + let call = RuntimeCall::System(frame_system::Call::remark { + remark: vec![1, 2, 3], + }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call.encode()), + )); + + // Verify the extrinsic was stored with submitted_at = 1 + let pending = PendingExtrinsics::::get(0).unwrap(); + assert_eq!(pending.submitted_at, 1); + + // Run on_initialize at block 12 (1 + 10 + 1 = 12, which is > MAX_EXTRINSIC_LIFETIME) + // MAX_EXTRINSIC_LIFETIME is 10, so at block 12, age is 11 which exceeds the limit + System::set_block_number(12); + MevShield::on_initialize(12); + + // Verify storage was cleared + assert!(PendingExtrinsics::::get(0).is_none()); + assert_eq!(PendingExtrinsicCount::::get(), 0); + + // Verify ExtrinsicExpired event was emitted (not ExtrinsicDispatched) + System::assert_has_event(crate::Event::::ExtrinsicExpired { index: 0 }.into()); + }); + } + + #[test] + fn on_initialize_does_not_expire_recent_extrinsics() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Store an extrinsic at block 1 + let call = RuntimeCall::System(frame_system::Call::remark { + remark: vec![1, 2, 3], + }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call.encode()), + )); + + // Run on_initialize at block 11 (age is 10, which equals MAX_EXTRINSIC_LIFETIME) + // Should NOT expire since we check age > MAX, not age >= + System::set_block_number(11); + MevShield::on_initialize(11); + + // Verify storage was cleared (extrinsic was dispatched, not expired) + assert!(PendingExtrinsics::::get(0).is_none()); + + // Verify ExtrinsicDispatched event was emitted (not ExtrinsicExpired) + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 0 }.into()); + }); + } + + #[test] + fn on_initialize_emits_dispatch_failed_on_bad_origin() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // set_heap_pages requires Root origin, so dispatching with Signed will fail + let call = RuntimeCall::System(frame_system::Call::set_heap_pages { pages: 10 }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call.encode()), + )); + + // Run on_initialize + MevShield::on_initialize(2); + + // Verify storage was cleared + assert!(PendingExtrinsics::::get(0).is_none()); + assert_eq!(PendingExtrinsicCount::::get(), 0); + + // Verify ExtrinsicDispatchFailed event was emitted + System::assert_has_event( + crate::Event::::ExtrinsicDispatchFailed { + index: 0, + error: sp_runtime::DispatchError::BadOrigin, + } + .into(), + ); + }); + } + + #[test] + fn on_initialize_handles_missing_slots() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Manually create a gap in indices by directly manipulating storage + let call = RuntimeCall::System(frame_system::Call::remark { + remark: vec![1, 2, 3], + }); + let pending = PendingExtrinsic:: { + who: 1, + call: BoundedVec::truncate_from(call.encode()), + submitted_at: 1, + }; + + // Insert at index 5, leaving 0-4 empty + PendingExtrinsics::::insert(5, pending); + NextPendingExtrinsicIndex::::put(6); + PendingExtrinsicCount::::put(1); + + // Run on_initialize - should handle the gap and process index 5 + MevShield::on_initialize(2); + + // Verify the extrinsic at index 5 was processed + assert!(PendingExtrinsics::::get(5).is_none()); + assert_eq!(PendingExtrinsicCount::::get(), 0); + + // Verify ExtrinsicDispatched event for index 5 + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 5 }.into()); + }); + } + + #[test] + fn multiple_accounts_dispatch_with_correct_origins() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + let user_a: u64 = 100; + let user_b: u64 = 200; + + // User A submits a remark_with_event + let call_a = + RuntimeCall::System(frame_system::Call::remark_with_event { remark: vec![0xAA] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(user_a), + BoundedVec::truncate_from(call_a.encode()), + )); + + // User B submits a remark_with_event + let call_b = + RuntimeCall::System(frame_system::Call::remark_with_event { remark: vec![0xBB] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(user_b), + BoundedVec::truncate_from(call_b.encode()), + )); + + // Run on_initialize + MevShield::on_initialize(2); + + // Verify both events have correct senders + let hash_a = ::Hashing::hash(&[0xAAu8]); + let hash_b = ::Hashing::hash(&[0xBBu8]); + + System::assert_has_event( + frame_system::Event::::Remarked { + sender: user_a, + hash: hash_a, + } + .into(), + ); + System::assert_has_event( + frame_system::Event::::Remarked { + sender: user_b, + hash: hash_b, + } + .into(), + ); + }); + } + + #[test] + fn expiration_mixed_with_valid_extrinsics() { + new_test_ext().execute_with(|| { + // Submit first extrinsic at block 1 + System::set_block_number(1); + let old_call = RuntimeCall::System(frame_system::Call::remark { remark: vec![0x01] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(old_call.encode()), + )); + + // Submit second extrinsic at block 10 + System::set_block_number(10); + let new_call = RuntimeCall::System(frame_system::Call::remark { remark: vec![0x02] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(2), + BoundedVec::truncate_from(new_call.encode()), + )); + + // Run on_initialize at block 12 + // First extrinsic: age = 12 - 1 = 11 > 10, should expire + // Second extrinsic: age = 12 - 10 = 2 <= 10, should dispatch + System::set_block_number(12); + MevShield::on_initialize(12); + + // Verify both were removed from storage + assert!(PendingExtrinsics::::get(0).is_none()); + assert!(PendingExtrinsics::::get(1).is_none()); + assert_eq!(PendingExtrinsicCount::::get(), 0); + + // Verify first expired, second dispatched + System::assert_has_event(crate::Event::::ExtrinsicExpired { index: 0 }.into()); + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 1 }.into()); + }); + } + + #[test] + fn set_max_pending_extrinsics_number_works() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Default is 100 + assert_eq!(MaxPendingExtrinsicsLimit::::get(), 100); + + assert_ok!(MevShield::set_max_pending_extrinsics_number( + RuntimeOrigin::root(), + 50, + )); + + assert_eq!(MaxPendingExtrinsicsLimit::::get(), 50); + + System::assert_last_event( + crate::Event::::MaxPendingExtrinsicsNumberSet { value: 50 }.into(), + ); + }); + } + + #[test] + fn set_max_pending_extrinsics_number_rejects_signed_origin() { + new_test_ext().execute_with(|| { + assert_noop!( + MevShield::set_max_pending_extrinsics_number(RuntimeOrigin::signed(1), 50), + sp_runtime::DispatchError::BadOrigin + ); + }); + } + + #[test] + fn set_max_pending_extrinsics_number_enforced_on_store() { + new_test_ext().execute_with(|| { + // Set limit to 2 + assert_ok!(MevShield::set_max_pending_extrinsics_number( + RuntimeOrigin::root(), + 2, + )); + + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + let encoded_call = BoundedVec::truncate_from(call.encode()); + + // First two should succeed + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + encoded_call.clone(), + )); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + encoded_call.clone(), + )); + + // Third should fail + assert_noop!( + MevShield::store_encrypted(RuntimeOrigin::signed(1), encoded_call), + Error::::TooManyPendingExtrinsics + ); + }); + } + + #[test] + fn set_on_initialize_weight_works() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + assert_eq!( + OnInitializeWeight::::get(), + crate::DEFAULT_ON_INITIALIZE_WEIGHT + ); + + assert_ok!(MevShield::set_on_initialize_weight( + RuntimeOrigin::root(), + 1_000_000, + )); + + assert_eq!(OnInitializeWeight::::get(), 1_000_000); + + System::assert_last_event( + crate::Event::::OnInitializeWeightSet { value: 1_000_000 }.into(), + ); + }); + } + + #[test] + fn set_on_initialize_weight_rejects_signed_origin() { + new_test_ext().execute_with(|| { + assert_noop!( + MevShield::set_on_initialize_weight(RuntimeOrigin::signed(1), 1_000_000), + sp_runtime::DispatchError::BadOrigin + ); + }); + } + + #[test] + fn set_on_initialize_weight_rejects_above_absolute_max() { + new_test_ext().execute_with(|| { + // Exactly at absolute max should succeed + assert_ok!(MevShield::set_on_initialize_weight( + RuntimeOrigin::root(), + crate::MAX_ON_INITIALIZE_WEIGHT, + )); + + // Above absolute max should fail + assert_noop!( + MevShield::set_on_initialize_weight( + RuntimeOrigin::root(), + crate::MAX_ON_INITIALIZE_WEIGHT + 1, + ), + Error::::WeightExceedsAbsoluteMax + ); + }); + } + + #[test] + fn set_on_initialize_weight_enforced_on_processing() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Set weight to 0 so nothing can be processed + assert_ok!(MevShield::set_on_initialize_weight( + RuntimeOrigin::root(), + 0, + )); + + // Store an extrinsic + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call.encode()), + )); + + assert_eq!(PendingExtrinsicCount::::get(), 1); + + // Run on_initialize — should postpone due to weight limit + MevShield::on_initialize(2); + + // Extrinsic should still be pending (postponed) + assert_eq!(PendingExtrinsicCount::::get(), 1); + System::assert_has_event(crate::Event::::ExtrinsicPostponed { index: 0 }.into()); + }); + } + + #[test] + fn set_stored_extrinsic_lifetime_works() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + assert_eq!( + ExtrinsicLifetime::::get(), + crate::DEFAULT_EXTRINSIC_LIFETIME + ); + + assert_ok!(MevShield::set_stored_extrinsic_lifetime( + RuntimeOrigin::root(), + 20 + )); + + assert_eq!(ExtrinsicLifetime::::get(), 20); + + System::assert_last_event( + crate::Event::::ExtrinsicLifetimeSet { value: 20 }.into(), + ); + }); + } + + #[test] + fn set_stored_extrinsic_lifetime_rejects_signed_origin() { + new_test_ext().execute_with(|| { + assert_noop!( + MevShield::set_stored_extrinsic_lifetime(RuntimeOrigin::signed(1), 20), + sp_runtime::DispatchError::BadOrigin + ); + }); + } + + #[test] + fn set_stored_extrinsic_lifetime_enforced_on_expiration() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // Set lifetime to 2 blocks + assert_ok!(MevShield::set_stored_extrinsic_lifetime( + RuntimeOrigin::root(), + 2 + )); + + // Store an extrinsic at block 1 + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call.encode()), + )); + + // At block 4: age = 4 - 1 = 3 > 2, should expire + System::set_block_number(4); + MevShield::on_initialize(4); + + assert!(PendingExtrinsics::::get(0).is_none()); + assert_eq!(PendingExtrinsicCount::::get(), 0); + System::assert_has_event(crate::Event::::ExtrinsicExpired { index: 0 }.into()); + }); + } +} diff --git a/pallets/subtensor/src/tests/mock.rs b/pallets/subtensor/src/tests/mock.rs index 772909638b..38c2330b48 100644 --- a/pallets/subtensor/src/tests/mock.rs +++ b/pallets/subtensor/src/tests/mock.rs @@ -100,6 +100,8 @@ impl pallet_balances::Config for Test { impl pallet_shield::Config for Test { type AuthorityId = sp_core::sr25519::Public; type FindAuthors = (); + type RuntimeCall = RuntimeCall; + type ExtrinsicDecryptor = (); } pub struct NoNestingCallFilter; diff --git a/runtime/src/check_mortality.rs b/runtime/src/check_mortality.rs index 6d8316ba01..2ec5233ed1 100644 --- a/runtime/src/check_mortality.rs +++ b/runtime/src/check_mortality.rs @@ -52,36 +52,42 @@ impl core::fmt::Debug for CheckMortality< } } -impl TransactionExtension - for CheckMortality +impl + TransactionExtension<::RuntimeCall> for CheckMortality where - T::RuntimeCall: Dispatchable + IsSubType>, + ::RuntimeCall: Dispatchable + IsSubType>, T: pallet_shield::Config, { const IDENTIFIER: &'static str = "CheckMortality"; - type Implicit = as TransactionExtension>::Implicit; - type Val = as TransactionExtension>::Val; - type Pre = as TransactionExtension>::Pre; + type Implicit = as TransactionExtension< + ::RuntimeCall, + >>::Implicit; + type Val = as TransactionExtension< + ::RuntimeCall, + >>::Val; + type Pre = as TransactionExtension< + ::RuntimeCall, + >>::Pre; fn implicit(&self) -> Result { CheckMortalitySubstrate::::from(self.0).implicit() } - fn weight(&self, call: &T::RuntimeCall) -> sp_weights::Weight { + fn weight(&self, call: &::RuntimeCall) -> sp_weights::Weight { CheckMortalitySubstrate::::from(self.0).weight(call) } fn validate( &self, origin: T::RuntimeOrigin, - call: &T::RuntimeCall, - info: &DispatchInfoOf, + call: &::RuntimeCall, + info: &DispatchInfoOf<::RuntimeCall>, len: usize, self_implicit: Self::Implicit, inherited_implication: &impl Implication, source: TransactionSource, - ) -> ValidateResult { + ) -> ValidateResult::RuntimeCall> { if let Some(ShieldCall::submit_encrypted { .. }) = IsSubType::>::is_sub_type(call) { @@ -109,8 +115,8 @@ where self, val: Self::Val, origin: &T::RuntimeOrigin, - call: &T::RuntimeCall, - info: &DispatchInfoOf, + call: &::RuntimeCall, + info: &DispatchInfoOf<::RuntimeCall>, len: usize, ) -> Result { CheckMortalitySubstrate::::from(self.0).prepare(val, origin, call, info, len) diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index 5fd4e5b401..6c8da6f98d 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -149,6 +149,8 @@ impl pallet_shield::FindAuthors for FindAuraAuthors { impl pallet_shield::Config for Runtime { type AuthorityId = AuraId; type FindAuthors = FindAuraAuthors; + type RuntimeCall = RuntimeCall; + type ExtrinsicDecryptor = (); } parameter_types! {