Merge pull request #2250 from opentensor/mevshield-fix-spam-attack-vector

MevShield Charge & Refund for `announce_next_key`

Author: sam0x17

pallets/shield/src/lib.rs

@@ -249,12 +249,13 @@ pub mod pallet {
                 .saturating_add(T::DbWeight::get().reads(1_u64))
                 .saturating_add(T::DbWeight::get().writes(1_u64)),
             DispatchClass::Operational,
-            Pays::No
+            Pays::Yes
         ))]
+        #[allow(clippy::useless_conversion)]
         pub fn announce_next_key(
             origin: OriginFor<T>,
             public_key: BoundedVec<u8, ConstU32<2048>>,
-        ) -> DispatchResult {
+        ) -> DispatchResultWithPostInfo {
             // Only a current Aura validator may call this (signed account ∈ Aura authorities)
             T::AuthorityOrigin::ensure_validator(origin)?;
 
@@ -264,9 +265,13 @@ pub mod pallet {
                 Error::<T>::BadPublicKeyLen
             );
 
-            NextKey::<T>::put(public_key.clone());
+            NextKey::<T>::put(public_key);
 
-            Ok(())
+            // Refund the fee on success by setting pays_fee = Pays::No
+            Ok(PostDispatchInfo {
+                actual_weight: None,
+                pays_fee: Pays::No,
+            })
         }
 
         /// Users submit an encrypted wrapper.
@@ -340,7 +345,7 @@ pub mod pallet {
             Weight::from_parts(77_280_000, 0)
                 .saturating_add(T::DbWeight::get().reads(4_u64))
                 .saturating_add(T::DbWeight::get().writes(1_u64)),
-            DispatchClass::Operational,
+            DispatchClass::Normal,
             Pays::No
         ))]
         #[allow(clippy::useless_conversion)]
@@ -425,7 +430,7 @@ pub mod pallet {
             Weight::from_parts(13_260_000, 0)
                 .saturating_add(T::DbWeight::get().reads(1_u64))
                 .saturating_add(T::DbWeight::get().writes(1_u64)),
-            DispatchClass::Operational,
+            DispatchClass::Normal,
             Pays::No
         ))]
         pub fn mark_decryption_failed(
@@ -481,7 +486,7 @@ pub mod pallet {
                         // Only allow locally-submitted / already-in-block txs.
                         TransactionSource::Local | TransactionSource::InBlock => {
                             ValidTransaction::with_tag_prefix("mev-shield-exec")
-                                .priority(u64::MAX)
+                                .priority(1u64)
                                 .longevity(64) // long because propagate(false)
                                 .and_provides(id) // dedupe by wrapper id
                                 .propagate(false) // CRITICAL: no gossip, stays on author node
@@ -494,7 +499,7 @@ pub mod pallet {
                     match source {
                         TransactionSource::Local | TransactionSource::InBlock => {
                             ValidTransaction::with_tag_prefix("mev-shield-failed")
-                                .priority(u64::MAX)
+                                .priority(1u64)
                                 .longevity(64) // long because propagate(false)
                                 .and_provides(id) // dedupe by wrapper id
                                 .propagate(false) // CRITICAL: no gossip, stays on author node

pallets/shield/src/tests.rs

@@ -660,3 +660,72 @@ fn mark_decryption_failed_removes_submission_and_emits_event() {
         assert_noop!(res, pallet_mev_shield::Error::<Test>::MissingSubmission);
     });
 }
+
+#[test]
+fn announce_next_key_charges_then_refunds_fee() {
+    new_test_ext().execute_with(|| {
+        const KYBER_PK_LEN: usize = 1184;
+
+        // ---------------------------------------------------------------------
+        // 1. Seed Aura authorities with a single validator and derive account.
+        // ---------------------------------------------------------------------
+        let validator_pair = test_sr25519_pair();
+        let validator_account: AccountId32 = validator_pair.public().into();
+        let validator_aura_id: <Test as pallet_aura::Config>::AuthorityId =
+            validator_pair.public().into();
+
+        let authorities: BoundedVec<
+            <Test as pallet_aura::Config>::AuthorityId,
+            <Test as pallet_aura::Config>::MaxAuthorities,
+        > = BoundedVec::truncate_from(vec![validator_aura_id]);
+        pallet_aura::Authorities::<Test>::put(authorities);
+
+        // ---------------------------------------------------------------------
+        // 2. Build a valid Kyber public key and the corresponding RuntimeCall.
+        // ---------------------------------------------------------------------
+        let pk_bytes = vec![42u8; KYBER_PK_LEN];
+        let bounded_pk: BoundedVec<u8, FrameConstU32<2048>> =
+            BoundedVec::truncate_from(pk_bytes.clone());
+
+        let runtime_call = RuntimeCall::MevShield(MevShieldCall::<Test>::announce_next_key {
+            public_key: bounded_pk.clone(),
+        });
+
+        // ---------------------------------------------------------------------
+        // 3. Pre-dispatch: DispatchInfo must say Pays::Yes.
+        // ---------------------------------------------------------------------
+        let pre_info = <RuntimeCall as frame_support::dispatch::GetDispatchInfo>::get_dispatch_info(
+            &runtime_call,
+        );
+
+        assert_eq!(
+            pre_info.pays_fee,
+            frame_support::dispatch::Pays::Yes,
+            "announce_next_key must be declared as fee-paying at pre-dispatch"
+        );
+
+        // ---------------------------------------------------------------------
+        // 4. Dispatch via the pallet function.
+        // ---------------------------------------------------------------------
+        let post = MevShield::announce_next_key(
+            RuntimeOrigin::signed(validator_account.clone()),
+            bounded_pk.clone(),
+        )
+        .expect("announce_next_key should succeed for an Aura validator");
+
+        // Post-dispatch info should switch pays_fee from Yes -> No (refund).
+        assert_eq!(
+            post.pays_fee,
+            frame_support::dispatch::Pays::No,
+            "announce_next_key must refund the previously chargeable fee"
+        );
+
+        // And we don't override the actual weight (None => use pre-dispatch weight).
+        assert!(
+            post.actual_weight.is_none(),
+            "announce_next_key should not override actual_weight in PostDispatchInfo"
+        );
+        let next = NextKey::<Test>::get().expect("NextKey should be set by announce_next_key");
+        assert_eq!(next, pk_bytes);
+    });
+}