From 9cd78adf535b0230dae7c205b2bb9583ccbc994b Mon Sep 17 00:00:00 2001 From: Luiz Carvalho Date: Sun, 24 Mar 2024 15:01:11 -0300 Subject: [PATCH] feat(module): improve param validation (#43) This commit changes how we validate module params, making it impossible to create invalid addresses, names, or repeated names when updating an existing module. The checks are now in one place. --- pallets/subspace/src/global.rs | 2 +- pallets/subspace/src/lib.rs | 42 ++++---- pallets/subspace/src/module.rs | 132 +++++++++++++++++-------- pallets/subspace/src/registration.rs | 13 +-- pallets/subspace/src/subnet.rs | 33 ++----- pallets/subspace/tests/registration.rs | 94 ++++++++++++++++++ pallets/subspace/tests/weights.rs | 2 +- 7 files changed, 223 insertions(+), 95 deletions(-) diff --git a/pallets/subspace/src/global.rs b/pallets/subspace/src/global.rs index 55728aa20..68c20e0f8 100644 --- a/pallets/subspace/src/global.rs +++ b/pallets/subspace/src/global.rs @@ -321,6 +321,6 @@ impl Pallet { } pub fn rm_from_whitelist(module_key: &T::AccountId) { - LegitWhitelist::::remove(&module_key); + LegitWhitelist::::remove(module_key); } } diff --git a/pallets/subspace/src/lib.rs b/pallets/subspace/src/lib.rs index 8353dd18d..678fe0267 100644 --- a/pallets/subspace/src/lib.rs +++ b/pallets/subspace/src/lib.rs @@ -59,6 +59,7 @@ pub mod pallet { use frame_support::{pallet_prelude::*, traits::Currency}; use frame_system::pallet_prelude::*; + use module::ModuleChangeset; use sp_arithmetic::per_things::Percent; pub use sp_std::{vec, vec::Vec}; @@ -894,10 +895,8 @@ pub mod pallet { // Errors inform users that something went wrong. #[pallet::error] pub enum Error { - ModuleNameAlreadyExists, // --- Thrown when a module name already exists. - NetworkDoesNotExist, // --- Thrown when the network does not exist. + NetworkDoesNotExist, // --- Thrown when the network does not exist. TooFewVotesForNewProposal, - ModuleAddressTooLong, NetworkExist, // --- Thrown when the network already exist. InvalidIpType, /* ---- Thrown when the user tries to serve an module which * is not of type 4 (IPv4) or 6 (IPv6). */ @@ -952,11 +951,7 @@ pub mod pallet { BalanceNotAdded, StakeNotRemoved, SubnetNameNotExists, - ModuleNameTooLong, /* --- Thrown when the user tries to register a module name that is - * too long. */ KeyAlreadyRegistered, // - ModuleNameDoesNotExist, /* --- Thrown when the user tries to remove a module name that - * does not exist. */ EmptyKeys, NotNominator, /* --- Thrown when the user tries to set the nominator and is not the * nominator */ @@ -1002,6 +997,20 @@ pub mod pallet { InvalidMaxBurn, InvalidTargetRegistrationsPerInterval, + // Modules + /// The module name is too long. + ModuleNameTooLong, + /// The module name is invalid. It has to be a UTF-8 encoded string. + InvalidModuleName, + /// The address is too long. + ModuleAddressTooLong, + /// The module address is invalid. + InvalidModuleAddress, + /// The module name does not exist in the subnet. + ModuleNameDoesNotExist, + /// A module with this name already exists in the subnet. + ModuleNameAlreadyExists, + // VOTING ProposalDoesNotExist, VotingPowerIsZero, @@ -1076,7 +1085,9 @@ pub mod pallet { for (uid_usize, (key, name, address, weights)) in self.modules[subnet_idx].iter().enumerate() { - self::Pallet::::append_module(netuid, key, name.clone(), address.clone()); + let changeset = ModuleChangeset::new(name.clone(), address.clone()); + self::Pallet::::append_module(netuid, key, changeset) + .expect("genesis modules are valid"); Weights::::insert(netuid, uid_usize as u16, weights); } } @@ -1192,17 +1203,10 @@ pub mod pallet { let key = ensure_signed(origin.clone())?; ensure!(Self::is_registered(netuid, &key), Error::::NotRegistered); let uid: u16 = Self::get_uid_for_key(netuid, &key); - let mut params = Self::module_params(netuid, uid); - params.name = name; - params.address = address; - if let Some(delegation_fee) = delegation_fee { - ensure!( - delegation_fee >= Self::get_floor_delegation_fee(), - Error::::InvalidMinDelegationFee - ); - params.delegation_fee = delegation_fee; - } - Self::do_update_module(origin, netuid, params) + let params = Self::module_params(netuid, uid); + + let changeset = ModuleChangeset::update(¶ms, name, address, delegation_fee); + Self::do_update_module(origin, netuid, changeset) } #[pallet::weight((Weight::zero(), DispatchClass::Normal, Pays::No))] diff --git a/pallets/subspace/src/module.rs b/pallets/subspace/src/module.rs index 4b65d485f..10b66b35d 100644 --- a/pallets/subspace/src/module.rs +++ b/pallets/subspace/src/module.rs @@ -24,36 +24,88 @@ pub struct ModuleInfo { stats: ModuleStats, } +#[derive(Debug)] +pub struct ModuleChangeset { + pub name: Option>, + pub address: Option>, + pub delegation_fee: Option, +} + +impl ModuleChangeset { + pub fn new(name: Vec, address: Vec) -> Self { + Self { + name: Some(name), + address: Some(address), + delegation_fee: None, + } + } + + pub fn update( + params: &ModuleParams, + name: Vec, + address: Vec, + delegation_fee: Option, + ) -> Self { + Self { + name: (name != params.name).then_some(name), + address: (address != params.address).then_some(address), + delegation_fee, + } + } + + /// Checks whether the module params are valid. Name and address must be non-empty and below the + /// max name length allowed. + pub fn apply( + self, + netuid: u16, + key: T::AccountId, + uid: u16, + ) -> Result<(), sp_runtime::DispatchError> { + let max = MaxNameLength::::get() as usize; + + if let Some(name) = self.name { + ensure!(!name.is_empty(), Error::::InvalidModuleName); + ensure!(name.len() <= max, Error::::ModuleNameTooLong); + core::str::from_utf8(&name).map_err(|_| Error::::InvalidModuleName)?; + ensure!( + !Pallet::::does_module_name_exist(netuid, &name), + Error::::ModuleNameAlreadyExists + ); + + Name::::insert(netuid, uid, name); + } + + if let Some(addr) = self.address { + ensure!(!addr.is_empty(), Error::::InvalidModuleAddress); + ensure!(addr.len() <= max, Error::::ModuleAddressTooLong); + core::str::from_utf8(&addr).map_err(|_| Error::::InvalidModuleAddress)?; + + Address::::insert(netuid, uid, addr); + } + + if let Some(fee) = self.delegation_fee { + let floor = Pallet::::get_floor_delegation_fee(); + ensure!(fee >= floor, Error::::InvalidMinDelegationFee); + + DelegationFee::::insert(netuid, key, fee); + } + + Ok(()) + } +} + impl Pallet { pub fn do_update_module( origin: T::RuntimeOrigin, netuid: u16, - params: ModuleParams, + changeset: ModuleChangeset, ) -> DispatchResult { - // --- 1. We check the callers (key) signature. + // 1. We check the callers (key) signature. let key = ensure_signed(origin)?; let uid: u16 = Self::get_uid_for_key(netuid, &key); - Self::check_module_params(netuid, ¶ms)?; - Self::set_module_params(netuid, uid, params); - // --- 8. Return is successful dispatch. - Ok(()) - } - /// Checks whether the module params are valid. Name and address must be non-empty and below the - /// max name length allowed. - pub fn check_module_params(_netuid: u16, params: &ModuleParams) -> DispatchResult { - let max_name_length = MaxNameLength::::get() as usize; - - assert!(!params.name.is_empty()); - ensure!( - params.name.len() <= max_name_length, - Error::::ModuleNameTooLong - ); - assert!(!params.address.is_empty()); - ensure!( - params.address.len() <= max_name_length, - Error::::ModuleAddressTooLong - ); + // 2. Apply the changeset + changeset.apply::(netuid, key, uid)?; Ok(()) } @@ -67,12 +119,6 @@ impl Pallet { } } - pub fn set_module_params(netuid: u16, uid: u16, module_params: ModuleParams) { - Self::set_module_name(netuid, uid, module_params.name); - Self::set_module_address(netuid, uid, module_params.address); - Self::set_module_delegation_fee(netuid, uid, module_params.delegation_fee); - } - pub fn get_module_address(netuid: u16, uid: u16) -> Vec { Address::::get(netuid, uid) } @@ -190,31 +236,37 @@ impl Pallet { } // Appends the uid to the network (without increasing stake). - pub fn append_module(netuid: u16, key: &T::AccountId, name: Vec, address: Vec) -> u16 { + pub fn append_module( + netuid: u16, + key: &T::AccountId, + changeset: ModuleChangeset, + ) -> Result { // 1. Get the next uid. This is always equal to subnetwork_n. let uid: u16 = Self::get_subnet_n(netuid); let block_number = Self::get_current_block_number(); + log::debug!("append_module( netuid: {netuid:?} | uid: {key:?} | new_key: {uid:?})"); - // 3. Expand with new position. - Emission::::append(netuid, 0); - Incentive::::append(netuid, 0); - Dividends::::append(netuid, 0); - LastUpdate::::append(netuid, block_number); - // 4. Insert new account information. + // 2. Apply the changeset + changeset.apply::(netuid, key.clone(), uid)?; + + // 3. Insert new account information. Keys::::insert(netuid, uid, key); // Make key - uid association. Uids::::insert(netuid, key, uid); // Make uid - key association. RegistrationBlock::::insert(netuid, uid, block_number); // Fill block at registration. - Name::::insert(netuid, uid, name); // Fill module namespace. - Address::::insert(netuid, uid, address); // Fill module info. - DelegationFee::::insert(netuid, key, DelegationFee::::get(netuid, key)); // Make uid - key association. - N::::insert(netuid, N::::get(netuid) + 1); // Increase the number of modules in the network. + N::::mutate(netuid, |n| *n += 1); // Increase the number of modules in the network. + + // 4. Expand with new position. + Emission::::append(netuid, 0); + Incentive::::append(netuid, 0); + Dividends::::append(netuid, 0); + LastUpdate::::append(netuid, block_number); // increase the stake of the new key Self::increase_stake(netuid, key, key, 0); - uid + Ok(uid) } pub fn get_module_stats(netuid: u16, uid: u16) -> ModuleStats { diff --git a/pallets/subspace/src/registration.rs b/pallets/subspace/src/registration.rs index ea819a882..cf068cc9a 100644 --- a/pallets/subspace/src/registration.rs +++ b/pallets/subspace/src/registration.rs @@ -1,3 +1,5 @@ +use crate::module::ModuleChangeset; + use super::*; use frame_support::pallet_prelude::DispatchResult; @@ -104,24 +106,19 @@ impl Pallet { Error::::NotEnoughStakeToRegister ); - // --- 6. Ensure the module key is not already registered, - // and namespace is not already taken. + // --- 6. Ensure the module key is not already registered. ensure!( !Self::key_registered(netuid, &module_key), Error::::KeyAlreadyRegistered ); - ensure!( - !Self::does_module_name_exist(netuid, name.clone()), - Error::::NameAlreadyRegistered - ); - // --- 7. Check if we are exceeding the max allowed modules per network. // If we do deregister slot. Self::check_module_limits(netuid); // --- 8. Register the module. - let uid: u16 = Self::append_module(netuid, &module_key, name, address); + let changeset = ModuleChangeset::new(name, address); + let uid: u16 = Self::append_module(netuid, &module_key, changeset)?; // --- 9. Add the stake to the module, now that it is registered on the network. Self::do_add_stake(origin, netuid, module_key.clone(), stake)?; diff --git a/pallets/subspace/src/subnet.rs b/pallets/subspace/src/subnet.rs index e2201be80..5c119a19e 100644 --- a/pallets/subspace/src/subnet.rs +++ b/pallets/subspace/src/subnet.rs @@ -280,9 +280,9 @@ impl Pallet { } // TODO: see if we can optimize this further - pub fn does_module_name_exist(netuid: u16, name: Vec) -> bool { + pub fn does_module_name_exist(netuid: u16, name: &[u8]) -> bool { as IterableStorageDoubleMap>>::iter_prefix(netuid) - .any(|(_uid, _name)| _name == name) + .any(|(_, existing)| existing == name) } pub fn is_subnet_founder(netuid: u16, key: &T::AccountId) -> bool { @@ -616,37 +616,18 @@ impl Pallet { let uid = Self::get_uid_for_key(netuid, key); Self::get_emission_for_uid(netuid, uid) } + pub fn get_emission_for_uid(netuid: u16, uid: u16) -> u64 { - let vec = Emission::::get(netuid); - if (uid as usize) < vec.len() { - vec[uid as usize] - } else { - 0 - } + Emission::::get(netuid).get(uid as usize).copied().unwrap_or_default() } pub fn get_incentive_for_uid(netuid: u16, uid: u16) -> u16 { - let vec = Incentive::::get(netuid); - if (uid as usize) < vec.len() { - vec[uid as usize] - } else { - 0 - } + Incentive::::get(netuid).get(uid as usize).copied().unwrap_or_default() } pub fn get_dividends_for_uid(netuid: u16, uid: u16) -> u16 { - let vec = Dividends::::get(netuid); - if (uid as usize) < vec.len() { - vec[uid as usize] - } else { - 0 - } + Dividends::::get(netuid).get(uid as usize).copied().unwrap_or_default() } pub fn get_last_update_for_uid(netuid: u16, uid: u16) -> u64 { - let vec = LastUpdate::::get(netuid); - if (uid as usize) < vec.len() { - vec[uid as usize] - } else { - 0 - } + LastUpdate::::get(netuid).get(uid as usize).copied().unwrap_or_default() } pub fn get_global_max_allowed_subnets() -> u16 { diff --git a/pallets/subspace/tests/registration.rs b/pallets/subspace/tests/registration.rs index 65904a4ba..65f7305c6 100644 --- a/pallets/subspace/tests/registration.rs +++ b/pallets/subspace/tests/registration.rs @@ -6,6 +6,7 @@ use sp_core::U256; use log::info; use pallet_subspace::Error; +use sp_runtime::{DispatchResult, Percent}; /******************************************** subscribing::subscribe() tests @@ -195,3 +196,96 @@ fn test_whitelist() { assert!(SubspaceModule::is_in_legit_whitelist(&adding_key)); }); } + +fn register_custom(netuid: u16, key: U256, name: &[u8], addr: &[u8]) -> DispatchResult { + let network: Vec = format!("test{netuid}").as_bytes().to_vec(); + + let origin = get_origin(key); + let is_new_subnet: bool = !SubspaceModule::if_subnet_exist(netuid); + if is_new_subnet { + SubspaceModule::set_max_registrations_per_block(1000) + } + + SubspaceModule::register(origin, network, name.to_vec(), addr.to_vec(), 0, key) +} + +fn test_validation_cases(f: impl Fn(&[u8], &[u8]) -> DispatchResult) { + assert_err!(f(b"", b""), Error::::InvalidModuleName); + assert_err!( + f("o".repeat(100).as_bytes(), b""), + Error::::ModuleNameTooLong + ); + assert_err!(f(b"\xc3\x28", b""), Error::::InvalidModuleName); + + assert_err!(f(b"test", b""), Error::::InvalidModuleAddress); + assert_err!( + f(b"test", "o".repeat(100).as_bytes()), + Error::::ModuleAddressTooLong + ); + assert_err!(f(b"test", b"\xc3\x28"), Error::::InvalidModuleAddress); + + assert_ok!(f(b"test", b"abc")); +} + +#[test] +fn validates_module_on_registration() { + new_test_ext().execute_with(|| { + test_validation_cases(|name, addr| register_custom(0, 0.into(), name, addr)); + + assert_err!( + register_custom(0, 1.into(), b"test", b"0.0.0.0:1"), + Error::::ModuleNameAlreadyExists + ); + }); +} + +#[test] +fn validates_module_on_update() { + new_test_ext().execute_with(|| { + let subnet = 0; + let key_0: U256 = 0.into(); + let origin_0 = get_origin(0.into()); + assert_ok!(register_custom(subnet, key_0, b"test", b"0.0.0.0:1")); + + test_validation_cases(|name, addr| { + SubspaceModule::update_module( + origin_0.clone(), + subnet, + name.to_vec(), + addr.to_vec(), + None, + ) + }); + + let key_1: U256 = 1.into(); + let origin_1 = get_origin(key_1); + assert_ok!(register_custom(0, key_1, b"test2", b"0.0.0.0:2")); + + let update_module = |name: &[u8], addr: &[u8]| { + SubspaceModule::update_module( + origin_1.clone(), + subnet, + name.to_vec(), + addr.to_vec(), + Some(Percent::from_percent(5)), + ) + }; + + assert_err!( + update_module(b"test", b""), + Error::::ModuleNameAlreadyExists + ); + assert_ok!(update_module(b"test2", b"0.0.0.0:2")); + assert_ok!(update_module(b"test3", b"0.0.0.0:3")); + + let params = SubspaceModule::module_params(0, 1); + assert_eq!(params.name, b"test3"); + assert_eq!(params.address, b"0.0.0.0:3"); + + SubspaceModule::set_floor_delegation_fee(Percent::from_percent(10)); + assert_err!( + update_module(b"test3", b"0.0.0.0:3"), + Error::::InvalidMinDelegationFee + ); + }); +} diff --git a/pallets/subspace/tests/weights.rs b/pallets/subspace/tests/weights.rs index f52014534..eef18ac4f 100644 --- a/pallets/subspace/tests/weights.rs +++ b/pallets/subspace/tests/weights.rs @@ -1,7 +1,7 @@ mod mock; use frame_support::{assert_err, assert_ok}; -use pallet_subspace::{Error, GlobalParams}; +use pallet_subspace::Error; use sp_core::U256; use sp_runtime::DispatchError;