Filter for incorrect sudo calls

Author: grbIzl

Cargo.lock

@@ -55,6 +55,7 @@ sha2.workspace = true
 rand_chacha.workspace = true
 pallet-crowdloan.workspace = true
 pallet-subtensor-proxy.workspace = true
+pallet-sudo.workspace = true
 
 [dev-dependencies]
 pallet-balances = { workspace = true, features = ["std"] }
@@ -117,6 +118,7 @@ std = [
 	"pallet-subtensor-swap/std",
 	"subtensor-swap-interface/std",
 	"pallet-subtensor-utility/std",
+	"pallet-sudo/std",
 	"safe-math/std",
 	"share-pool/std",
 	"ark-serialize/std",

pallets/subtensor/Cargo.toml

@@ -47,6 +47,7 @@ frame_support::construct_runtime!(
         Swap: pallet_subtensor_swap::{Pallet, Call, Storage, Event<T>} = 12,
         Crowdloan: pallet_crowdloan::{Pallet, Call, Storage, Event<T>} = 13,
         Proxy: pallet_subtensor_proxy = 14,
+        Sudo: pallet_sudo::{Pallet, Call, Storage, Event<T>} = 15,
     }
 );
 
@@ -471,6 +472,12 @@ impl InstanceFilter<RuntimeCall> for subtensor_runtime_common::ProxyType {
     }
 }
 
+impl pallet_sudo::Config for Test {
+    type RuntimeEvent = RuntimeEvent;
+    type RuntimeCall = RuntimeCall;
+    type WeightInfo = pallet_sudo::weights::SubstrateWeight<Test>;
+}
+
 mod test_crypto {
     use super::KEY_TYPE;
     use sp_core::{

pallets/subtensor/src/tests/mock.rs

@@ -31,3 +31,4 @@ mod swap_hotkey;
 mod swap_hotkey_with_subnet;
 mod uids;
 mod weights;
+mod tx_extension;

pallets/subtensor/src/tests/mod.rs

@@ -0,0 +1,102 @@
+use frame_support::assert_ok;
+use frame_support::dispatch::GetDispatchInfo;
+use crate::tests::mock::{new_test_ext, RuntimeCall, RuntimeOrigin, SubtensorCall};
+use subtensor_runtime_common::{NetUid, TaoCurrency};
+use sp_core::U256;
+use sp_runtime::traits::{TransactionExtension, TxBaseImplication, ValidateResult};
+use crate::transaction_extension::SubtensorTransactionExtension;
+
+use super::mock::*;
+use crate::*;
+
+fn some_call() -> RuntimeCall {
+    let hotkey = U256::from(0);
+    let amount_staked = TaoCurrency::from(5000);
+    let netuid = NetUid::from(1);
+    RuntimeCall::SubtensorModule(SubtensorCall::add_stake {
+        hotkey,
+        netuid,
+        amount_staked,
+    })
+}
+fn sudo_extrinsic(inner: RuntimeCall) -> RuntimeCall {
+    RuntimeCall::Sudo(pallet_sudo::Call::sudo {
+        call: Box::new(inner),
+    })
+}
+
+fn validate_ext(
+    origin: RuntimeOrigin,
+    call: &RuntimeCall,
+) -> ValidateResult<Option<<Test as frame_system::Config>::AccountId>, RuntimeCall> {
+    let ext = SubtensorTransactionExtension::<Test>::new();
+
+    ext.validate(
+        origin,
+        call,
+        &call.get_dispatch_info(),
+        0,
+        (),
+        &TxBaseImplication(()),
+        TransactionSource::External,
+    )
+}
+#[test]
+fn sudo_signed_by_correct_key_is_valid() {
+    new_test_ext(1).execute_with(|| {
+        let sudo_key: <Test as frame_system::Config>::AccountId = U256::from(42);
+        pallet_sudo::Key::<Test>::put(sudo_key);
+        let sudo_call = sudo_extrinsic(some_call());
+
+        // Signed origin with correct sudo key
+        let origin = RuntimeOrigin::signed(sudo_key);
+        let res = validate_ext(origin, &sudo_call);
+        assert_ok!(res);
+    });
+}
+
+#[test]
+fn sudo_signed_by_wrong_account_is_rejected() {
+    new_test_ext(1).execute_with(|| {
+        let sudo_key: <Test as frame_system::Config>::AccountId = U256::from(42);
+        // Set sudo key in storage
+        pallet_sudo::Key::<Test>::put(sudo_key);
+        let sudo_call = sudo_extrinsic(some_call());
+        // Wrong signer
+        let origin = RuntimeOrigin::signed(U256::from(99));
+        let res = validate_ext(origin, &sudo_call);
+        assert!(
+            matches!(
+                res,
+                Err(TransactionValidityError::Invalid(InvalidTransaction::BadSigner))
+            )
+        );
+    });
+}
+
+#[test]
+fn sudo_when_no_sudo_key_configured_is_rejected() {
+    new_test_ext(1).execute_with(|| {
+        // Remove sudo key
+        pallet_sudo::Key::<Test>::kill();
+        let sudo_call = sudo_extrinsic(some_call());
+        let origin = RuntimeOrigin::signed(U256::from(42));
+        let res = validate_ext(origin, &sudo_call);
+        assert!(
+            matches!(
+                res,
+                Err(TransactionValidityError::Invalid(InvalidTransaction::BadSigner))
+            )
+        );
+    });
+}
+
+#[test]
+fn non_sudo_extrinsic_does_not_trigger_filter() {
+    new_test_ext(1).execute_with(|| {
+        let origin = RuntimeOrigin::signed(U256::from(42));
+        let call = some_call();
+        let res = validate_ext(origin, &call);
+        assert!(res.is_ok());
+    });
+}

pallets/subtensor/src/tests/tx_extension.rs

@@ -13,7 +13,9 @@ use sp_runtime::traits::{
 };
 use sp_runtime::transaction_validity::{
     TransactionSource, TransactionValidity, TransactionValidityError, ValidTransaction,
+    InvalidTransaction
 };
+use pallet_sudo::Call as SudoCall;
 use sp_std::marker::PhantomData;
 use sp_std::vec::Vec;
 use subtensor_macros::freeze_struct;
@@ -85,7 +87,7 @@ where
     }
 }
 
-impl<T: Config + Send + Sync + TypeInfo + pallet_balances::Config>
+impl<T: Config + Send + Sync + TypeInfo + pallet_balances::Config + pallet_sudo::Config>
     TransactionExtension<<T as frame_system::Config>::RuntimeCall>
     for SubtensorTransactionExtension<T>
 where
@@ -94,6 +96,7 @@ where
     <T as frame_system::Config>::RuntimeOrigin: AsSystemOriginSigner<T::AccountId> + Clone,
     <T as frame_system::Config>::RuntimeCall: IsSubType<Call<T>>,
     <T as frame_system::Config>::RuntimeCall: IsSubType<BalancesCall<T>>,
+    <T as frame_system::Config>::RuntimeCall: IsSubType<SudoCall<T>>,
 {
     const IDENTIFIER: &'static str = "SubtensorTransactionExtension";
 
@@ -121,6 +124,21 @@ where
             return Ok((Default::default(), None, origin));
         };
 
+        // Check validity of the signer for sudo call
+        if let Some(_sudo_call) = IsSubType::<pallet_sudo::Call<T>>::is_sub_type(call) {
+            let sudo_key = pallet_sudo::pallet::Key::<T>::get();
+
+            // No sudo key configured → reject
+            let Some(expected_who) = sudo_key else {
+                return Err(InvalidTransaction::BadSigner.into());
+            };
+
+            // Signer does not match the sudo key → reject
+            if *who != expected_who {
+                return Err(InvalidTransaction::BadSigner.into());
+            }
+        }
+
         // Verify ColdkeySwapScheduled map for coldkey
         match call.is_sub_type() {
             // Whitelist