Merge pull request #1298 from opentensor/feat/validate-serve-axon

Add serve_axon extrinsic validation

pallets/subtensor/src/lib.rs

@@ -1674,6 +1674,10 @@ pub enum CustomTransactionError {
     InsufficientLiquidity,
     SlippageTooHigh,
     TransferDisallowed,
+    HotKeyNotRegisteredInNetwork,
+    InvalidIpAddress,
+    ServingRateLimitExceeded,
+    InvalidPort,
     BadRequest,
 }
 
@@ -1690,6 +1694,10 @@ impl From<CustomTransactionError> for u8 {
             CustomTransactionError::InsufficientLiquidity => 7,
             CustomTransactionError::SlippageTooHigh => 8,
             CustomTransactionError::TransferDisallowed => 9,
+            CustomTransactionError::HotKeyNotRegisteredInNetwork => 10,
+            CustomTransactionError::InvalidIpAddress => 11,
+            CustomTransactionError::ServingRateLimitExceeded => 12,
+            CustomTransactionError::InvalidPort => 13,
             CustomTransactionError::BadRequest => 255,
         }
     }
@@ -1773,6 +1781,22 @@ where
                     CustomTransactionError::TransferDisallowed.into(),
                 )
                 .into()),
+                Error::<T>::HotKeyNotRegisteredInNetwork => Err(InvalidTransaction::Custom(
+                    CustomTransactionError::HotKeyNotRegisteredInNetwork.into(),
+                )
+                .into()),
+                Error::<T>::InvalidIpAddress => Err(InvalidTransaction::Custom(
+                    CustomTransactionError::InvalidIpAddress.into(),
+                )
+                .into()),
+                Error::<T>::ServingRateLimitExceeded => Err(InvalidTransaction::Custom(
+                    CustomTransactionError::ServingRateLimitExceeded.into(),
+                )
+                .into()),
+                Error::<T>::InvalidPort => Err(InvalidTransaction::Custom(
+                    CustomTransactionError::InvalidPort.into(),
+                )
+                .into()),
                 _ => Err(
                     InvalidTransaction::Custom(CustomTransactionError::BadRequest.into()).into(),
                 ),
@@ -2175,6 +2199,32 @@ where
                     })
                 }
             }
+            Some(Call::serve_axon {
+                netuid,
+                version,
+                ip,
+                port,
+                ip_type,
+                protocol,
+                placeholder1,
+                placeholder2,
+            }) => {
+                // Fully validate the user input
+                Self::result_to_validity(
+                    Pallet::<T>::validate_serve_axon(
+                        who,
+                        *netuid,
+                        *version,
+                        *ip,
+                        *port,
+                        *ip_type,
+                        *protocol,
+                        *placeholder1,
+                        *placeholder2,
+                    ),
+                    Self::get_priority_vanilla(),
+                )
+            }
             _ => {
                 if let Some(
                     BalancesCall::transfer_keep_alive { .. }

pallets/subtensor/src/subnets/serving.rs

@@ -69,35 +69,28 @@ impl<T: Config> Pallet<T> {
         // We check the callers (hotkey) signature.
         let hotkey_id = ensure_signed(origin)?;
 
-        // Ensure the hotkey is registered somewhere.
-        ensure!(
-            Self::is_hotkey_registered_on_any_network(&hotkey_id),
-            Error::<T>::HotKeyNotRegisteredInNetwork
-        );
-
-        // Check the ip signature validity.
-        ensure!(Self::is_valid_ip_type(ip_type), Error::<T>::InvalidIpType);
-        ensure!(
-            Self::is_valid_ip_address(ip_type, ip),
-            Error::<T>::InvalidIpAddress
-        );
-
-        // Get the previous axon information.
-        let mut prev_axon = Self::get_axon_info(netuid, &hotkey_id);
-        let current_block: u64 = Self::get_current_block_as_u64();
-        ensure!(
-            Self::axon_passes_rate_limit(netuid, &prev_axon, current_block),
-            Error::<T>::ServingRateLimitExceeded
-        );
+        // Validate user input
+        Self::validate_serve_axon(
+            &hotkey_id,
+            netuid,
+            version,
+            ip,
+            port,
+            ip_type,
+            protocol,
+            placeholder1,
+            placeholder2,
+        )?;
 
-        // Check certificate
+        // Check+insert certificate
         if let Some(certificate) = certificate {
             if let Ok(certificate) = NeuronCertificateOf::try_from(certificate) {
                 NeuronCertificates::<T>::insert(netuid, hotkey_id.clone(), certificate)
             }
         }
 
         // We insert the axon meta.
+        let mut prev_axon = Self::get_axon_info(netuid, &hotkey_id);
         prev_axon.block = Self::get_current_block_as_u64();
         prev_axon.version = version;
         prev_axon.ip = ip;
@@ -176,19 +169,19 @@ impl<T: Config> Pallet<T> {
         // We check the callers (hotkey) signature.
         let hotkey_id = ensure_signed(origin)?;
 
-        // Ensure the hotkey is registered somewhere.
-        ensure!(
-            Self::is_hotkey_registered_on_any_network(&hotkey_id),
-            Error::<T>::HotKeyNotRegisteredInNetwork
-        );
-
         // Check the ip signature validity.
         ensure!(Self::is_valid_ip_type(ip_type), Error::<T>::InvalidIpType);
         ensure!(
             Self::is_valid_ip_address(ip_type, ip),
             Error::<T>::InvalidIpAddress
         );
 
+        // Ensure the hotkey is registered somewhere.
+        ensure!(
+            Self::is_hotkey_registered_on_any_network(&hotkey_id),
+            Error::<T>::HotKeyNotRegisteredInNetwork
+        );
+
         // We get the previous axon info assoicated with this ( netuid, uid )
         let mut prev_prometheus = Self::get_prometheus_info(netuid, &hotkey_id);
         let current_block: u64 = Self::get_current_block_as_u64();
@@ -332,4 +325,55 @@ impl<T: Config> Pallet<T> {
 
         Ok(true)
     }
+
+    pub fn validate_serve_axon(
+        hotkey_id: &T::AccountId,
+        netuid: u16,
+        version: u32,
+        ip: u128,
+        port: u16,
+        ip_type: u8,
+        protocol: u8,
+        placeholder1: u8,
+        placeholder2: u8,
+    ) -> Result<(), Error<T>> {
+        // Ensure the hotkey is registered somewhere.
+        ensure!(
+            Self::is_hotkey_registered_on_any_network(hotkey_id),
+            Error::<T>::HotKeyNotRegisteredInNetwork
+        );
+
+        // Check the ip signature validity.
+        ensure!(Self::is_valid_ip_type(ip_type), Error::<T>::InvalidIpType);
+        ensure!(
+            Self::is_valid_ip_address(ip_type, ip),
+            Error::<T>::InvalidIpAddress
+        );
+
+        // Get the previous axon information.
+        let mut prev_axon = Self::get_axon_info(netuid, hotkey_id);
+        let current_block: u64 = Self::get_current_block_as_u64();
+        ensure!(
+            Self::axon_passes_rate_limit(netuid, &prev_axon, current_block),
+            Error::<T>::ServingRateLimitExceeded
+        );
+
+        // Validate axon data with delegate func
+        prev_axon.block = Self::get_current_block_as_u64();
+        prev_axon.version = version;
+        prev_axon.ip = ip;
+        prev_axon.port = port;
+        prev_axon.ip_type = ip_type;
+        prev_axon.protocol = protocol;
+        prev_axon.placeholder1 = placeholder1;
+        prev_axon.placeholder2 = placeholder2;
+
+        let axon_validated = Self::validate_axon_data(&prev_axon);
+        ensure!(
+            axon_validated.is_ok(),
+            axon_validated.err().unwrap_or(Error::<T>::InvalidPort)
+        );
+
+        Ok(())
+    }
 }

pallets/subtensor/src/tests/serving.rs

@@ -2,8 +2,8 @@ use super::mock::*;
 
 use crate::Error;
 use crate::*;
-use frame_support::assert_noop;
 use frame_support::pallet_prelude::Weight;
+use frame_support::{assert_err, assert_noop};
 use frame_support::{
     assert_ok,
     dispatch::{DispatchClass, DispatchInfo, GetDispatchInfo, Pays},
@@ -1257,3 +1257,60 @@ fn test_set_subnet_identity_dispatch_info_ok() {
         assert_eq!(dispatch_info.pays_fee, Pays::Yes);
     });
 }
+
+// cargo test --package pallet-subtensor --lib -- tests::serving::test_serve_axon_validate --exact --show-output
+#[test]
+fn test_serve_axon_validate() {
+    // Testing the signed extension validate function
+    // correctly filters the `serve_axon` transaction.
+
+    new_test_ext(0).execute_with(|| {
+        let hotkey = U256::from(1);
+        let netuid: u16 = 1;
+        let version: u32 = 2;
+        let ip: u128 = 1676056785;
+        let port: u16 = 128;
+        let ip_type: u8 = 4;
+        let protocol: u8 = 0;
+        let placeholder1: u8 = 0;
+        let placeholder2: u8 = 0;
+
+        // Serve axon bad call
+        let call = RuntimeCall::SubtensorModule(SubtensorCall::serve_axon {
+            netuid,
+            version,
+            ip,
+            port,
+            ip_type,
+            protocol,
+            placeholder1,
+            placeholder2,
+        });
+
+        let info: crate::DispatchInfo =
+            crate::DispatchInfoOf::<<Test as frame_system::Config>::RuntimeCall>::default();
+
+        let extension = crate::SubtensorSignedExtension::<Test>::new();
+        // Submit to the signed extension validate function
+        let result_bad = extension.validate(&hotkey, &call.clone(), &info, 10);
+
+        // Should fail due to insufficient stake
+        assert_err!(
+            result_bad,
+            crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
+                CustomTransactionError::HotKeyNotRegisteredInNetwork.into()
+            ))
+        );
+
+        // Register the hotkey in the subnet and try again
+        let coldkey = U256::from(1);
+        add_network(netuid, 13, 0);
+        register_ok_neuron(netuid, hotkey, coldkey, 0);
+
+        // Submit to the signed extension validate function
+        let result_ok = extension.validate(&hotkey, &call.clone(), &info, 10);
+
+        // Now the call passes
+        assert_ok!(result_ok);
+    });
+}

runtime/src/lib.rs

@@ -229,7 +229,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     //   `spec_version`, and `authoring_version` are the same between Wasm and native.
     // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use
     //   the compatible custom types.
-    spec_version: 235,
+    spec_version: 236,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,