MEV Shield V2: Introduce fee refund system with reason tracking

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 = ();
 }
 

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",

pallets/shield/src/lib.rs

@@ -74,6 +74,66 @@ impl<RuntimeCall> ExtrinsicDecryptor<RuntimeCall> 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<T: frame_system::Config> {
+    /// 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<u128>;
+}
+
+/// No-op implementation that skips refunding.
+impl<T: frame_system::Config> EncryptedExtrinsicFees<T> 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<u128> {
+        None
+    }
+}
+
 #[frame_support::pallet]
 pub mod pallet {
     use super::*;
@@ -95,6 +155,9 @@ pub mod pallet {
         /// Decryptor for stored extrinsics.
         type ExtrinsicDecryptor: ExtrinsicDecryptor<<Self as pallet::Config>::RuntimeCall>;
 
+        /// Handles fee-related operations for encrypted extrinsics.
+        type EncryptedExtrinsicFees: EncryptedExtrinsicFees<Self>;
+
         /// 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<T: Config> Pallet<T> {
                 PendingExtrinsics::<T>::remove(index);
                 weight = weight.saturating_add(remove_weight);
 
+                maybe_refund::<T>(
+                    &pending.who,
+                    index,
+                    store_encrypted_weight(),
+                    Weight::zero(),
+                    RefundReason::Expired,
+                );
+
                 Self::deposit_event(Event::ExtrinsicExpired { index });
 
                 continue;
@@ -553,19 +631,38 @@ impl<T: Config> Pallet<T> {
             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::<T>(
+                        &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::<T>(
+                        &who,
+                        index,
+                        charged_weight,
+                        info.call_weight,
+                        RefundReason::Failure,
+                    );
+
                     Self::deposit_event(Event::ExtrinsicDispatchFailed {
                         index,
                         error: e.error,
@@ -574,6 +671,36 @@ impl<T: Config> Pallet<T> {
             }
         }
 
+        fn maybe_refund<T: Config>(
+            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::<T>::deposit_event(Event::ExtrinsicRefunded {
+                    index,
+                    who: who.clone(),
+                    amount,
+                    reason,
+                });
+            }
+        }
+
         weight
     }
 

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<u64>;
 }
 
 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<Option<AuraId>> = const { RefCell::new(None) };
     static MOCK_NEXT_NEXT: RefCell<Option<Option<AuraId>>> = const { RefCell::new(None) };
+    static REFUND_ENABLED: RefCell<bool> = const { RefCell::new(false) };
+    static REFUND_ON_EXPIRATION: RefCell<bool> = const { RefCell::new(false) };
 }
 
 pub struct MockFindAuthors;
@@ -87,6 +99,47 @@ impl pallet_shield::FindAuthors<Test> 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<Test> 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<u128> {
+        use frame_support::traits::{fungible::Balanced, tokens::Precision};
+
+        let diff = charged_weight
+            .ref_time()
+            .saturating_sub(actual_weight.ref_time());
+        if diff > 0 {
+            let _ = <Balances as Balanced<_>>::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::<Test> {
+        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])
 }