From 6bdbc078f359387276f0b5ef2f17d3984b6da093 Mon Sep 17 00:00:00 2001 From: Shamil Gadelshin Date: Wed, 1 Apr 2026 13:40:23 +0300 Subject: [PATCH 1/2] Introduce refund to mev-shield V2 --- Cargo.lock | 1 + pallets/shield/Cargo.toml | 2 + pallets/shield/src/lib.rs | 129 ++++++++- pallets/shield/src/mock.rs | 72 +++++ pallets/shield/src/tests.rs | 416 ++++++++++++++++++++++++++++ pallets/subtensor/src/tests/mock.rs | 1 + runtime/src/lib.rs | 32 +++ 7 files changed, 652 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index e7cf36d22e..a0edb9d2b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10610,6 +10610,7 @@ dependencies = [ "log", "ml-kem", "pallet-aura", + "pallet-balances", "pallet-subtensor-utility", "pallet-timestamp", "parity-scale-codec", diff --git a/pallets/shield/Cargo.toml b/pallets/shield/Cargo.toml index 8888b249b9..575ff6a9e8 100644 --- a/pallets/shield/Cargo.toml +++ b/pallets/shield/Cargo.toml @@ -48,6 +48,7 @@ rand_chacha = { workspace = true, optional = true } [dev-dependencies] stc-shield.workspace = true pallet-subtensor-utility.workspace = true +pallet-balances.workspace = true rand_chacha.workspace = true pallet-timestamp.workspace = true pallet-aura.workspace = true @@ -77,6 +78,7 @@ std = [ "sp-consensus-aura?/std", "rand_chacha?/std", "pallet-subtensor-utility/std", + "pallet-balances/std", "stc-shield/std", "subtensor-runtime-common/std", "sp-keystore/std", diff --git a/pallets/shield/src/lib.rs b/pallets/shield/src/lib.rs index f96374590e..f0c6fed8dc 100644 --- a/pallets/shield/src/lib.rs +++ b/pallets/shield/src/lib.rs @@ -74,6 +74,66 @@ impl ExtrinsicDecryptor for () { } } +/// Reason a fee refund is being issued for an encrypted extrinsic. +#[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Encode, + Decode, + TypeInfo, + MaxEncodedLen, + codec::DecodeWithMemTracking, +)] +pub enum RefundReason { + /// The extrinsic was dispatched successfully. + Success, + /// The extrinsic dispatch failed. + Failure, + /// The extrinsic expired without being dispatched. + Expired, +} + +/// Handles fee-related operations for encrypted extrinsics, including +/// refunding the fee difference between what was charged at `store_encrypted` +/// time and the actual weight consumed during `on_initialize` dispatch. +pub trait EncryptedExtrinsicFees { + /// Whether refunding is enabled. + fn refund_enabled() -> bool; + + /// Whether to refund fees when an encrypted extrinsic expires without dispatch. + fn refund_on_expiration() -> bool; + + fn refund( + who: &T::AccountId, + charged_weight: Weight, + actual_weight: Weight, + reason: RefundReason, + ) -> Option; +} + +/// No-op implementation that skips refunding. +impl EncryptedExtrinsicFees for () { + fn refund_enabled() -> bool { + false + } + + fn refund_on_expiration() -> bool { + false + } + + fn refund( + _who: &T::AccountId, + _charged_weight: Weight, + _actual_weight: Weight, + _reason: RefundReason, + ) -> Option { + None + } +} + #[frame_support::pallet] pub mod pallet { use super::*; @@ -95,6 +155,9 @@ pub mod pallet { /// Decryptor for stored extrinsics. type ExtrinsicDecryptor: ExtrinsicDecryptor<::RuntimeCall>; + /// Handles fee-related operations for encrypted extrinsics. + type EncryptedExtrinsicFees: EncryptedExtrinsicFees; + /// Weight information for extrinsics in this pallet. type WeightInfo: WeightInfo; } @@ -226,6 +289,13 @@ pub mod pallet { MaxExtrinsicWeightSet { value: u64 }, /// Extrinsic exceeded the per-extrinsic weight limit and was removed. ExtrinsicWeightExceeded { index: u32 }, + /// A fee refund was issued for an encrypted extrinsic. + ExtrinsicRefunded { + index: u32, + who: T::AccountId, + amount: u128, + reason: RefundReason, + }, } #[pallet::error] @@ -510,6 +580,14 @@ impl Pallet { PendingExtrinsics::::remove(index); weight = weight.saturating_add(remove_weight); + maybe_refund::( + &pending.who, + index, + store_encrypted_weight(), + Weight::zero(), + RefundReason::Expired, + ); + Self::deposit_event(Event::ExtrinsicExpired { index }); continue; @@ -553,19 +631,38 @@ impl Pallet { weight = weight.saturating_add(remove_weight); // Dispatch the extrinsic - let origin: T::RuntimeOrigin = frame_system::RawOrigin::Signed(pending.who).into(); + let who = pending.who; + let origin: T::RuntimeOrigin = frame_system::RawOrigin::Signed(who.clone()).into(); let result = call.dispatch(origin); + let charged_weight = store_encrypted_weight(); + match result { Ok(post_info) => { let actual_weight = post_info.actual_weight.unwrap_or(info.call_weight); weight = weight.saturating_add(actual_weight); + maybe_refund::( + &who, + index, + charged_weight, + actual_weight, + RefundReason::Success, + ); + Self::deposit_event(Event::ExtrinsicDispatched { index }); } Err(e) => { weight = weight.saturating_add(info.call_weight); + maybe_refund::( + &who, + index, + charged_weight, + info.call_weight, + RefundReason::Failure, + ); + Self::deposit_event(Event::ExtrinsicDispatchFailed { index, error: e.error, @@ -574,6 +671,36 @@ impl Pallet { } } + fn maybe_refund( + who: &T::AccountId, + index: u32, + charged_weight: Weight, + actual_weight: Weight, + reason: RefundReason, + ) { + let enabled = match reason { + RefundReason::Expired => T::EncryptedExtrinsicFees::refund_on_expiration(), + RefundReason::Success | RefundReason::Failure => { + T::EncryptedExtrinsicFees::refund_enabled() + } + }; + if !enabled { + return; + } + + if let Some(amount) = + T::EncryptedExtrinsicFees::refund(who, charged_weight, actual_weight, reason) + && amount > 0 + { + Pallet::::deposit_event(Event::ExtrinsicRefunded { + index, + who: who.clone(), + amount, + reason, + }); + } + } + weight } diff --git a/pallets/shield/src/mock.rs b/pallets/shield/src/mock.rs index 1bb6fa018a..074e69ced7 100644 --- a/pallets/shield/src/mock.rs +++ b/pallets/shield/src/mock.rs @@ -24,6 +24,7 @@ construct_runtime!( Aura: pallet_aura = 2, MevShield: pallet_shield = 3, Utility: pallet_subtensor_utility = 4, + Balances: pallet_balances = 5, } ); @@ -32,11 +33,13 @@ const SLOT_DURATION: u64 = 6000; parameter_types! { pub const SlotDuration: u64 = SLOT_DURATION; pub const MaxAuthorities: u32 = 32; + pub const ExistentialDeposit: u64 = 1; } #[derive_impl(frame_system::config_preludes::TestDefaultConfig)] impl frame_system::Config for Test { type Block = Block; + type AccountData = pallet_balances::AccountData; } impl pallet_timestamp::Config for Test { @@ -60,9 +63,18 @@ impl pallet_subtensor_utility::Config for Test { type WeightInfo = (); } +#[derive_impl(pallet_balances::config_preludes::TestDefaultConfig)] +impl pallet_balances::Config for Test { + type AccountStore = System; + type Balance = u64; + type ExistentialDeposit = ExistentialDeposit; +} + thread_local! { static MOCK_CURRENT: RefCell> = const { RefCell::new(None) }; static MOCK_NEXT_NEXT: RefCell>> = const { RefCell::new(None) }; + static REFUND_ENABLED: RefCell = const { RefCell::new(false) }; + static REFUND_ON_EXPIRATION: RefCell = const { RefCell::new(false) }; } pub struct MockFindAuthors; @@ -87,6 +99,47 @@ impl pallet_shield::FindAuthors for MockFindAuthors { } } +/// Mock fee handler that deposits actual balance refunds. +/// Uses a simple 1:1 weight-to-fee mapping (ref_time = balance units). +pub struct MockEncryptedExtrinsicFees; + +impl pallet_shield::EncryptedExtrinsicFees for MockEncryptedExtrinsicFees { + fn refund_enabled() -> bool { + REFUND_ENABLED.with(|e| *e.borrow()) + } + + fn refund_on_expiration() -> bool { + REFUND_ON_EXPIRATION.with(|e| *e.borrow()) + } + + fn refund( + who: &u64, + charged_weight: sp_weights::Weight, + actual_weight: sp_weights::Weight, + _reason: pallet_shield::RefundReason, + ) -> Option { + use frame_support::traits::{fungible::Balanced, tokens::Precision}; + + let diff = charged_weight + .ref_time() + .saturating_sub(actual_weight.ref_time()); + if diff > 0 { + let _ = >::deposit(who, diff, Precision::BestEffort); + Some(diff as u128) + } else { + None + } + } +} + +pub fn enable_refund(enabled: bool) { + REFUND_ENABLED.with(|e| *e.borrow_mut() = enabled); +} + +pub fn enable_refund_on_expiration(enabled: bool) { + REFUND_ON_EXPIRATION.with(|e| *e.borrow_mut() = enabled); +} + /// Mock decryptor that just decodes the bytes without decryption. pub struct MockDecryptor; @@ -101,6 +154,7 @@ impl pallet_shield::Config for Test { type FindAuthors = MockFindAuthors; type RuntimeCall = RuntimeCall; type ExtrinsicDecryptor = MockDecryptor; + type EncryptedExtrinsicFees = MockEncryptedExtrinsicFees; type WeightInfo = (); } @@ -115,6 +169,24 @@ pub fn new_test_ext() -> sp_io::TestExternalities { ext } +/// Create test externalities with funded accounts. +pub fn new_test_ext_with_balances(balances: Vec<(u64, u64)>) -> sp_io::TestExternalities { + let mut t = RuntimeGenesisConfig::default() + .build_storage() + .expect("valid genesis"); + pallet_balances::GenesisConfig:: { + balances, + dev_accounts: None, + } + .assimilate_storage(&mut t) + .expect("balances storage should build ok"); + let mut ext = sp_io::TestExternalities::new(t); + ext.register_extension(sp_keystore::KeystoreExt::new( + sp_keystore::testing::MemoryKeystore::new(), + )); + ext +} + pub fn valid_pk() -> ShieldEncKey { BoundedVec::truncate_from(vec![0x42; MLKEM768_ENC_KEY_LEN]) } diff --git a/pallets/shield/src/tests.rs b/pallets/shield/src/tests.rs index 622b5be74a..a57ff0445c 100644 --- a/pallets/shield/src/tests.rs +++ b/pallets/shield/src/tests.rs @@ -1162,3 +1162,419 @@ mod encrypted_extrinsics_tests { }); } } + +// --------------------------------------------------------------------------- +// Refund tests — verify actual balance changes via MockEncryptedExtrinsicFees +// --------------------------------------------------------------------------- + +mod refund_tests { + use super::*; + use crate::mock::{ + Balances, enable_refund, enable_refund_on_expiration, new_test_ext_with_balances, + }; + use crate::{ExtrinsicLifetime, PendingExtrinsics, RefundReason, STORE_ENCRYPTED_WEIGHT}; + use frame_support::dispatch::GetDispatchInfo; + use frame_support::traits::Hooks; + + const INITIAL_BALANCE: u64 = 100_000_000_000_000; + + // Test 1: refund disabled — balance unchanged after successful dispatch, no refund event + #[test] + fn refund_disabled_no_balance_change_on_success() { + new_test_ext_with_balances(vec![(1, INITIAL_BALANCE)]).execute_with(|| { + enable_refund(false); + System::set_block_number(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()), + )); + + let balance_before = Balances::free_balance(1); + MevShield::on_initialize(2); + + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 0 }.into()); + assert_eq!(Balances::free_balance(1), balance_before); + + // No refund event when refund is disabled + assert!(System::events().iter().all(|e| !matches!( + e.event, + RuntimeEvent::MevShield(crate::Event::::ExtrinsicRefunded { .. }) + ))); + }); + } + + // Test 2: refund disabled — balance unchanged after failed dispatch + #[test] + fn refund_disabled_no_balance_change_on_failure() { + new_test_ext_with_balances(vec![(1, INITIAL_BALANCE)]).execute_with(|| { + enable_refund(false); + System::set_block_number(1); + + let call = RuntimeCall::System(frame_system::Call::set_heap_pages { pages: 64 }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call.encode()), + )); + + let balance_before = Balances::free_balance(1); + MevShield::on_initialize(2); + + System::assert_has_event( + crate::Event::::ExtrinsicDispatchFailed { + index: 0, + error: sp_runtime::DispatchError::BadOrigin, + } + .into(), + ); + assert_eq!(Balances::free_balance(1), balance_before); + + assert!(System::events().iter().all(|e| !matches!( + e.event, + RuntimeEvent::MevShield(crate::Event::::ExtrinsicRefunded { .. }) + ))); + }); + } + + // Test 3: refund on successful dispatch — partial refund (charged - actual) + refund event + #[test] + fn refund_deposits_partial_balance_on_successful_dispatch() { + new_test_ext_with_balances(vec![(42, INITIAL_BALANCE)]).execute_with(|| { + enable_refund(true); + System::set_block_number(1); + + let call = RuntimeCall::System(frame_system::Call::remark { + remark: vec![1, 2, 3], + }); + let actual_weight = call.get_dispatch_info().call_weight.ref_time(); + + // The call has nonzero weight, so refund must be partial + assert!(actual_weight > 0, "call must have nonzero weight"); + let expected_refund = STORE_ENCRYPTED_WEIGHT - actual_weight; + assert!( + expected_refund < STORE_ENCRYPTED_WEIGHT, + "refund must be less than the charged fee" + ); + + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(42), + BoundedVec::truncate_from(call.encode()), + )); + + let balance_before = Balances::free_balance(42); + MevShield::on_initialize(2); + + System::assert_has_event(crate::Event::::ExtrinsicDispatched { index: 0 }.into()); + System::assert_has_event( + crate::Event::::ExtrinsicRefunded { + index: 0, + who: 42, + amount: expected_refund as u128, + reason: RefundReason::Success, + } + .into(), + ); + // User receives exactly the overpayment, not the full charged fee + assert_eq!(Balances::free_balance(42), balance_before + expected_refund,); + }); + } + + // Test 4: refund on failed dispatch — partial refund using call_weight + refund event + #[test] + fn refund_deposits_partial_balance_on_failed_dispatch() { + new_test_ext_with_balances(vec![(99, INITIAL_BALANCE)]).execute_with(|| { + enable_refund(true); + System::set_block_number(1); + + // set_heap_pages has significant weight (~103M), making the partial + // refund clearly visible compared to the 20B charged fee. + let call = RuntimeCall::System(frame_system::Call::set_heap_pages { pages: 64 }); + let actual_weight = call.get_dispatch_info().call_weight.ref_time(); + + assert!(actual_weight > 0, "call must have nonzero weight"); + let expected_refund = STORE_ENCRYPTED_WEIGHT - actual_weight; + assert!( + expected_refund < STORE_ENCRYPTED_WEIGHT, + "refund must be less than the charged fee" + ); + + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(99), + BoundedVec::truncate_from(call.encode()), + )); + + let balance_before = Balances::free_balance(99); + MevShield::on_initialize(2); + + System::assert_has_event( + crate::Event::::ExtrinsicDispatchFailed { + index: 0, + error: sp_runtime::DispatchError::BadOrigin, + } + .into(), + ); + System::assert_has_event( + crate::Event::::ExtrinsicRefunded { + index: 0, + who: 99, + amount: expected_refund as u128, + reason: RefundReason::Failure, + } + .into(), + ); + assert_eq!(Balances::free_balance(99), balance_before + expected_refund,); + }); + } + + // Test 5a: refund_on_expiration enabled — full refund (actual_weight = 0) + refund event + #[test] + fn refund_on_expiration_deposits_full_fee() { + new_test_ext_with_balances(vec![(7, INITIAL_BALANCE)]).execute_with(|| { + enable_refund_on_expiration(true); + System::set_block_number(1); + + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(7), + BoundedVec::truncate_from(call.encode()), + )); + + let balance_before = Balances::free_balance(7); + + let lifetime = ExtrinsicLifetime::::get(); + let expired_block = 1 + lifetime as u64 + 1; + System::set_block_number(expired_block); + MevShield::on_initialize(expired_block); + + System::assert_has_event(crate::Event::::ExtrinsicExpired { index: 0 }.into()); + System::assert_has_event( + crate::Event::::ExtrinsicRefunded { + index: 0, + who: 7, + amount: STORE_ENCRYPTED_WEIGHT as u128, + reason: RefundReason::Expired, + } + .into(), + ); + // Full refund: charged_weight - 0 = STORE_ENCRYPTED_WEIGHT + assert_eq!( + Balances::free_balance(7), + balance_before + STORE_ENCRYPTED_WEIGHT, + ); + }); + } + + // Test 5b: refund_on_expiration disabled — no balance change on expiration + #[test] + fn no_refund_on_expiration_when_disabled() { + new_test_ext_with_balances(vec![(7, INITIAL_BALANCE)]).execute_with(|| { + enable_refund_on_expiration(false); + System::set_block_number(1); + + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(7), + BoundedVec::truncate_from(call.encode()), + )); + + let balance_before = Balances::free_balance(7); + + let lifetime = ExtrinsicLifetime::::get(); + let expired_block = 1 + lifetime as u64 + 1; + System::set_block_number(expired_block); + MevShield::on_initialize(expired_block); + + System::assert_has_event(crate::Event::::ExtrinsicExpired { index: 0 }.into()); + assert_eq!(Balances::free_balance(7), balance_before); + + assert!(System::events().iter().all(|e| !matches!( + e.event, + RuntimeEvent::MevShield(crate::Event::::ExtrinsicRefunded { .. }) + ))); + }); + } + + // Test 6: no balance change on decode failure + #[test] + fn no_refund_on_decode_failure() { + new_test_ext_with_balances(vec![(1, INITIAL_BALANCE)]).execute_with(|| { + enable_refund(true); + System::set_block_number(1); + + let invalid_bytes = BoundedVec::truncate_from(vec![0xFF, 0xFF]); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + invalid_bytes, + )); + + let balance_before = Balances::free_balance(1); + MevShield::on_initialize(2); + + System::assert_has_event( + crate::Event::::ExtrinsicDecodeFailed { index: 0 }.into(), + ); + assert_eq!(Balances::free_balance(1), balance_before); + + assert!(System::events().iter().all(|e| !matches!( + e.event, + RuntimeEvent::MevShield(crate::Event::::ExtrinsicRefunded { .. }) + ))); + }); + } + + // Test 7: no balance change when per-extrinsic weight exceeded + #[test] + fn no_refund_on_weight_exceeded() { + new_test_ext_with_balances(vec![(1, INITIAL_BALANCE)]).execute_with(|| { + enable_refund(true); + System::set_block_number(1); + + assert_ok!(MevShield::set_max_extrinsic_weight( + RuntimeOrigin::root(), + 0, + )); + + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call.encode()), + )); + + let balance_before = Balances::free_balance(1); + MevShield::on_initialize(2); + + System::assert_has_event( + crate::Event::::ExtrinsicWeightExceeded { index: 0 }.into(), + ); + assert_eq!(Balances::free_balance(1), balance_before); + + assert!(System::events().iter().all(|e| !matches!( + e.event, + RuntimeEvent::MevShield(crate::Event::::ExtrinsicRefunded { .. }) + ))); + }); + } + + // Test 8: no balance change when extrinsic is postponed + #[test] + fn no_refund_on_postponement() { + new_test_ext_with_balances(vec![(1, INITIAL_BALANCE)]).execute_with(|| { + enable_refund(true); + System::set_block_number(1); + + assert_ok!(MevShield::set_on_initialize_weight( + RuntimeOrigin::root(), + 0, + )); + + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(1), + BoundedVec::truncate_from(call.encode()), + )); + + let balance_before = Balances::free_balance(1); + MevShield::on_initialize(2); + + System::assert_has_event(crate::Event::::ExtrinsicPostponed { index: 0 }.into()); + assert_eq!(PendingExtrinsics::::count(), 1); + assert_eq!(Balances::free_balance(1), balance_before); + + assert!(System::events().iter().all(|e| !matches!( + e.event, + RuntimeEvent::MevShield(crate::Event::::ExtrinsicRefunded { .. }) + ))); + }); + } + + // Test 9: multiple users — each gets correct partial refund + #[test] + fn refund_multiple_users_correct_balances() { + new_test_ext_with_balances(vec![ + (10, INITIAL_BALANCE), + (20, INITIAL_BALANCE), + (30, INITIAL_BALANCE), + ]) + .execute_with(|| { + enable_refund(true); + System::set_block_number(1); + + // Index 0: lightweight remark from account 10 (will succeed) + let call_ok = RuntimeCall::System(frame_system::Call::remark { remark: vec![0xAA] }); + let weight_ok = call_ok.get_dispatch_info().call_weight.ref_time(); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(10), + BoundedVec::truncate_from(call_ok.encode()), + )); + + // Index 1: heavier set_heap_pages from account 20 (will fail with BadOrigin) + let call_fail = RuntimeCall::System(frame_system::Call::set_heap_pages { pages: 64 }); + let weight_fail = call_fail.get_dispatch_info().call_weight.ref_time(); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(20), + BoundedVec::truncate_from(call_fail.encode()), + )); + + // Index 2: another remark from account 30 (will succeed) + let call_ok2 = RuntimeCall::System(frame_system::Call::remark { remark: vec![0xBB] }); + let weight_ok2 = call_ok2.get_dispatch_info().call_weight.ref_time(); + assert_ok!(MevShield::store_encrypted( + RuntimeOrigin::signed(30), + BoundedVec::truncate_from(call_ok2.encode()), + )); + + let bal10 = Balances::free_balance(10); + let bal20 = Balances::free_balance(20); + let bal30 = Balances::free_balance(30); + + MevShield::on_initialize(2); + + // Heavier call gets smaller refund + assert!( + weight_fail > weight_ok, + "set_heap_pages should be heavier than remark" + ); + let refund_ok = STORE_ENCRYPTED_WEIGHT - weight_ok; + let refund_fail = STORE_ENCRYPTED_WEIGHT - weight_fail; + assert!(refund_fail < refund_ok, "heavier call → smaller refund"); + + assert_eq!(Balances::free_balance(10), bal10 + refund_ok); + assert_eq!(Balances::free_balance(20), bal20 + refund_fail); + assert_eq!( + Balances::free_balance(30), + bal30 + (STORE_ENCRYPTED_WEIGHT - weight_ok2) + ); + + System::assert_has_event( + crate::Event::::ExtrinsicRefunded { + index: 0, + who: 10, + amount: refund_ok as u128, + reason: RefundReason::Success, + } + .into(), + ); + System::assert_has_event( + crate::Event::::ExtrinsicRefunded { + index: 1, + who: 20, + amount: refund_fail as u128, + reason: RefundReason::Failure, + } + .into(), + ); + System::assert_has_event( + crate::Event::::ExtrinsicRefunded { + index: 2, + who: 30, + amount: (STORE_ENCRYPTED_WEIGHT - weight_ok2) as u128, + reason: RefundReason::Success, + } + .into(), + ); + }); + } +} diff --git a/pallets/subtensor/src/tests/mock.rs b/pallets/subtensor/src/tests/mock.rs index fa16b3d0f2..225bf46aaf 100644 --- a/pallets/subtensor/src/tests/mock.rs +++ b/pallets/subtensor/src/tests/mock.rs @@ -103,6 +103,7 @@ impl pallet_shield::Config for Test { type FindAuthors = (); type RuntimeCall = RuntimeCall; type ExtrinsicDecryptor = (); + type EncryptedExtrinsicFees = (); type WeightInfo = (); } diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index a077185341..224bd7754a 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -147,11 +147,43 @@ impl pallet_shield::FindAuthors for FindAuraAuthors { } } +pub struct ShieldWeightRefund; +impl pallet_shield::EncryptedExtrinsicFees for ShieldWeightRefund { + fn refund_enabled() -> bool { + true + } + + fn refund_on_expiration() -> bool { + true + } + + fn refund( + who: &AccountId, + charged_weight: Weight, + actual_weight: Weight, + _reason: pallet_shield::RefundReason, + ) -> Option { + use frame_support::traits::{fungible::Balanced, tokens::Precision}; + use frame_support::weights::WeightToFee; + + let diff = charged_weight.saturating_sub(actual_weight); + let refund_fee = + ::weight_to_fee(&diff); + if refund_fee > 0u128.into() { + let _ = >::deposit(who, refund_fee, Precision::BestEffort); + Some(refund_fee.into()) + } else { + None + } + } +} + impl pallet_shield::Config for Runtime { type AuthorityId = AuraId; type FindAuthors = FindAuraAuthors; type RuntimeCall = RuntimeCall; type ExtrinsicDecryptor = (); + type EncryptedExtrinsicFees = ShieldWeightRefund; type WeightInfo = pallet_shield::weights::SubstrateWeight; } From 41bb98336cb70a3942edcd6df4190996d6d7cfac Mon Sep 17 00:00:00 2001 From: Shamil Gadelshin Date: Thu, 9 Apr 2026 13:48:49 +0300 Subject: [PATCH 2/2] Update eco tests mock --- eco-tests/src/mock.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/eco-tests/src/mock.rs b/eco-tests/src/mock.rs index 65cb6eac03..f3d8264039 100644 --- a/eco-tests/src/mock.rs +++ b/eco-tests/src/mock.rs @@ -90,6 +90,7 @@ impl pallet_shield::Config for Test { type FindAuthors = (); type RuntimeCall = RuntimeCall; type ExtrinsicDecryptor = (); + type EncryptedExtrinsicFees = (); type WeightInfo = (); }