diff --git a/Cargo.lock b/Cargo.lock index eeffdabd0f..8bea2f5c3d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8222,6 +8222,7 @@ dependencies = [ "num-traits", "pallet-commitments", "pallet-drand", + "pallet-rate-limiting", "pallet-shield", "pallet-subtensor", "pallet-subtensor-swap-rpc", @@ -8342,6 +8343,8 @@ dependencies = [ "pallet-nomination-pools-runtime-api", "pallet-offences", "pallet-preimage", + "pallet-rate-limiting", + "pallet-rate-limiting-runtime-api", "pallet-registry", "pallet-safe-mode", "pallet-scheduler", @@ -8365,6 +8368,7 @@ dependencies = [ "precompile-utils", "rand_chacha 0.3.1", "scale-info", + "serde", "serde_json", "sha2 0.10.9", "smallvec", @@ -8771,6 +8775,7 @@ dependencies = [ "pallet-subtensor", "pallet-subtensor-swap", "parity-scale-codec", + "rate-limiting-interface", "scale-info", "sp-consensus-aura", "sp-consensus-grandpa", @@ -10303,6 +10308,49 @@ dependencies = [ "sp-runtime", ] +[[package]] +name = "pallet-rate-limiting" +version = "0.1.0" +dependencies = [ + "frame-benchmarking", + "frame-support", + "frame-system", + "parity-scale-codec", + "rate-limiting-interface", + "scale-info", + "serde", + "sp-core", + "sp-io", + "sp-runtime", + "sp-std", + "subtensor-runtime-common", +] + +[[package]] +name = "pallet-rate-limiting-rpc" +version = "0.1.0" +dependencies = [ + "jsonrpsee", + "pallet-rate-limiting-runtime-api", + "sp-api", + "sp-blockchain", + "sp-runtime", + "subtensor-runtime-common", +] + +[[package]] +name = "pallet-rate-limiting-runtime-api" +version = "0.1.0" +dependencies = [ + "pallet-rate-limiting", + "parity-scale-codec", + "scale-info", + "serde", + "sp-api", + "sp-std", + "subtensor-runtime-common", +] + [[package]] name = "pallet-recovery" version = "41.0.0" @@ -10780,6 +10828,7 @@ dependencies = [ "pallet-crowdloan", "pallet-drand", "pallet-preimage", + "pallet-rate-limiting", "pallet-scheduler", "pallet-subtensor-proxy", "pallet-subtensor-swap", @@ -10789,6 +10838,7 @@ dependencies = [ "polkadot-runtime-common", "rand 0.8.5", "rand_chacha 0.3.1", + "rate-limiting-interface", "safe-math", "scale-info", "serde", @@ -13681,6 +13731,17 @@ dependencies = [ "rand_core 0.9.3", ] +[[package]] +name = "rate-limiting-interface" +version = "0.1.0" +dependencies = [ + "frame-support", + "parity-scale-codec", + "scale-info", + "serde", + "sp-std", +] + [[package]] name = "raw-cpuid" version = "11.6.0" @@ -18033,6 +18094,7 @@ dependencies = [ "pallet-subtensor-utility", "pallet-timestamp", "parity-scale-codec", + "rate-limiting-interface", "scale-info", "sp-core", "sp-io", @@ -18105,6 +18167,7 @@ dependencies = [ "pallet-evm-precompile-modexp", "pallet-evm-precompile-sha3fips", "pallet-evm-precompile-simple", + "pallet-rate-limiting", "pallet-subtensor", "pallet-subtensor-proxy", "pallet-subtensor-swap", @@ -18176,6 +18239,7 @@ dependencies = [ "pallet-subtensor-swap", "pallet-transaction-payment", "parity-scale-codec", + "rate-limiting-interface", "scale-info", "smallvec", "sp-consensus-aura", diff --git a/Cargo.toml b/Cargo.toml index 1d65a3cd5e..48630c5f1d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,8 @@ members = [ "common", "node", "pallets/*", + "pallets/rate-limiting/runtime-api", + "pallets/rate-limiting/rpc", "precompiles", "primitives/*", "runtime", @@ -59,6 +61,9 @@ pallet-subtensor = { path = "pallets/subtensor", default-features = false } pallet-subtensor-swap = { path = "pallets/swap", default-features = false } pallet-subtensor-swap-runtime-api = { path = "pallets/swap/runtime-api", default-features = false } pallet-subtensor-swap-rpc = { path = "pallets/swap/rpc", default-features = false } +pallet-rate-limiting = { path = "pallets/rate-limiting", default-features = false } +pallet-rate-limiting-runtime-api = { path = "pallets/rate-limiting/runtime-api", default-features = false } +pallet-rate-limiting-rpc = { path = "pallets/rate-limiting/rpc", default-features = false } procedural-fork = { path = "support/procedural-fork", default-features = false } safe-math = { path = "primitives/safe-math", default-features = false } share-pool = { path = "primitives/share-pool", default-features = false } @@ -70,6 +75,7 @@ subtensor-runtime-common = { default-features = false, path = "common" } subtensor-swap-interface = { default-features = false, path = "pallets/swap-interface" } subtensor-transaction-fee = { default-features = false, path = "pallets/transaction-fee" } subtensor-chain-extensions = { default-features = false, path = "chain-extensions" } +rate-limiting-interface = { default-features = false, path = "pallets/rate-limiting-interface" } ed25519-dalek = { version = "2.1.0", default-features = false } async-trait = "0.1" diff --git a/chain-extensions/Cargo.toml b/chain-extensions/Cargo.toml index 9727439e7a..61ec5b12e6 100644 --- a/chain-extensions/Cargo.toml +++ b/chain-extensions/Cargo.toml @@ -36,6 +36,9 @@ subtensor-swap-interface.workspace = true num_enum.workspace = true substrate-fixed.workspace = true +[dev-dependencies] +rate-limiting-interface.workspace = true + [lints] workspace = true diff --git a/chain-extensions/src/mock.rs b/chain-extensions/src/mock.rs index 98ea096199..ae99c4031e 100644 --- a/chain-extensions/src/mock.rs +++ b/chain-extensions/src/mock.rs @@ -18,6 +18,7 @@ use pallet_contracts::HoldReason as ContractsHoldReason; use pallet_subtensor::*; use pallet_subtensor_proxy as pallet_proxy; use pallet_subtensor_utility as pallet_utility; +use rate_limiting_interface::{RateLimitingInfo, TryIntoRateLimitTarget}; use sp_core::{ConstU64, H256, U256, offchain::KeyTypeId}; use sp_runtime::Perbill; use sp_runtime::{ @@ -25,7 +26,9 @@ use sp_runtime::{ traits::{BlakeTwo256, Convert, IdentityLookup}, }; use sp_std::{cell::RefCell, cmp::Ordering, sync::OnceLock}; -use subtensor_runtime_common::{AlphaCurrency, NetUid, TaoCurrency}; +use subtensor_runtime_common::{ + AlphaCurrency, NetUid, TaoCurrency, rate_limiting::RateLimitUsageKey, +}; type Block = frame_system::mocking::MockBlock; @@ -291,7 +294,6 @@ parameter_types! { pub const InitialMinChildKeyTake: u16 = 0; // 0 %; pub const InitialMaxChildKeyTake: u16 = 11_796; // 18 %; pub const InitialWeightsVersionKey: u16 = 0; - pub const InitialServingRateLimit: u64 = 0; // No limit. pub const InitialTxRateLimit: u64 = 0; // Disable rate limit for testing pub const InitialTxDelegateTakeRateLimit: u64 = 1; // 1 block take rate limit for testing pub const InitialTxChildKeyTakeRateLimit: u64 = 1; // 1 block take rate limit for testing @@ -377,7 +379,6 @@ impl pallet_subtensor::Config for Test { type InitialWeightsVersionKey = InitialWeightsVersionKey; type InitialMaxDifficulty = InitialMaxDifficulty; type InitialMinDifficulty = InitialMinDifficulty; - type InitialServingRateLimit = InitialServingRateLimit; type InitialTxRateLimit = InitialTxRateLimit; type InitialTxDelegateTakeRateLimit = InitialTxDelegateTakeRateLimit; type InitialBurn = InitialBurn; @@ -411,6 +412,7 @@ impl pallet_subtensor::Config for Test { type GetCommitments = (); type MaxImmuneUidsPercentage = MaxImmuneUidsPercentage; type CommitmentsInterface = CommitmentsI; + type RateLimiting = NoRateLimiting; type EvmKeyAssociateRateLimit = EvmKeyAssociateRateLimit; } @@ -449,6 +451,33 @@ impl CommitmentsInterface for CommitmentsI { fn purge_netuid(_netuid: NetUid) {} } +pub struct NoRateLimiting; + +impl RateLimitingInfo for NoRateLimiting { + type GroupId = subtensor_runtime_common::rate_limiting::GroupId; + type CallMetadata = RuntimeCall; + type Limit = BlockNumber; + type Scope = subtensor_runtime_common::NetUid; + type UsageKey = RateLimitUsageKey; + + fn rate_limit(_target: TargetArg, _scope: Option) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + None + } + + fn last_seen( + _target: TargetArg, + _usage_key: Option, + ) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + None + } +} + parameter_types! { pub MaximumSchedulerWeight: Weight = Perbill::from_percent(80) * BlockWeights::get().max_block; diff --git a/common/src/lib.rs b/common/src/lib.rs index 658f8b2e01..eac4d49db5 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -17,6 +17,7 @@ use subtensor_macros::freeze_struct; pub use currency::*; mod currency; +pub mod rate_limiting; /// Balance of an account. pub type Balance = u64; diff --git a/common/src/rate_limiting.rs b/common/src/rate_limiting.rs new file mode 100644 index 0000000000..20c9e2d629 --- /dev/null +++ b/common/src/rate_limiting.rs @@ -0,0 +1,110 @@ +//! Shared rate-limiting types. +//! +//! Note: `pallet-rate-limiting` supports multiple independent instances, and is intended to be used +//! as “one instance per pallet” with pallet-specific scope/usage-key types and resolvers. +//! +//! The scope/usage-key types in this module are centralized today due to the current state of +//! `pallet-subtensor` (a large, centralized pallet) and its coupling with `pallet-admin-utils`, +//! which share a single `pallet-rate-limiting` instance and resolver implementation in the runtime. +//! +//! For new pallets, it is strongly recommended to: +//! - define their own `LimitScope` and `UsageKey` types (do not extend `RateLimitUsageKey` here), +//! - provide pallet-local scope/usage resolvers, +//! - and use a dedicated `pallet-rate-limiting` instance. +//! +//! Long-term, we should move away from these shared types by refactoring `pallet-subtensor` into +//! smaller pallets with dedicated `pallet-rate-limiting` instances. + +use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; +use frame_support::pallet_prelude::Parameter; +use scale_info::TypeInfo; +use serde::{Deserialize, Serialize}; + +use crate::{MechId, NetUid}; + +/// Identifier type for rate-limiting groups. +pub type GroupId = u32; + +/// Group id for serving-related calls. +pub const GROUP_SERVE: GroupId = 0; +/// Group id for delegate-take related calls. +pub const GROUP_DELEGATE_TAKE: GroupId = 1; +/// Group id for subnet weight-setting calls. +pub const GROUP_WEIGHTS_SUBNET: GroupId = 2; +/// Group id for network registration calls. +pub const GROUP_REGISTER_NETWORK: GroupId = 3; +/// Group id for owner hyperparameter calls. +pub const GROUP_OWNER_HPARAMS: GroupId = 4; +/// Group id for staking operations. +pub const GROUP_STAKING_OPS: GroupId = 5; +/// Group id for key swap calls. +pub const GROUP_SWAP_KEYS: GroupId = 6; + +/// Usage-key type currently shared by the centralized `pallet-subtensor` rate-limiting instance. +/// +/// Do not add new variants for new pallets. Prefer defining pallet-specific types and using a +/// dedicated `pallet-rate-limiting` instance per pallet. +#[derive( + Serialize, + Deserialize, + Encode, + Decode, + DecodeWithMemTracking, + Clone, + PartialEq, + Eq, + PartialOrd, + Ord, + Debug, + TypeInfo, + MaxEncodedLen, +)] +#[scale_info(skip_type_params(AccountId))] +pub enum RateLimitUsageKey { + Account(AccountId), + Subnet(NetUid), + AccountSubnet { + account: AccountId, + netuid: NetUid, + }, + ColdkeyHotkeySubnet { + coldkey: AccountId, + hotkey: AccountId, + netuid: NetUid, + }, + SubnetNeuron { + netuid: NetUid, + uid: u16, + }, + SubnetMechanismNeuron { + netuid: NetUid, + mecid: MechId, + uid: u16, + }, + AccountSubnetServing { + account: AccountId, + netuid: NetUid, + endpoint: ServingEndpoint, + }, +} + +#[derive( + Serialize, + Deserialize, + Encode, + Decode, + DecodeWithMemTracking, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Debug, + TypeInfo, + MaxEncodedLen, +)] +pub enum ServingEndpoint { + Axon, + Prometheus, +} diff --git a/contract-tests/README.md b/contract-tests/README.md index 78294603d3..90068e5553 100644 --- a/contract-tests/README.md +++ b/contract-tests/README.md @@ -36,6 +36,9 @@ npx papi add devnet -w ws://localhost:9944 If the runtime is upgrade, need to get the metadata again. ```bash +cd contract-tests/bittensor +cargo contract build --release +cd .. sh get-metadata.sh ``` diff --git a/contract-tests/package.json b/contract-tests/package.json index 26136346bb..6af55b1d9e 100644 --- a/contract-tests/package.json +++ b/contract-tests/package.json @@ -1,6 +1,6 @@ { "scripts": { - "test": "mocha --timeout 999999 --retries 3 --file src/setup.ts --require ts-node/register test/*test.ts" + "test": "TS_NODE_PREFER_TS_EXTS=1 TS_NODE_TRANSPILE_ONLY=1 mocha --timeout 999999 --retries 3 --file src/setup.ts --require ts-node/register --extension ts \"test/**/*.ts\"" }, "keywords": [], "author": "", diff --git a/contract-tests/test/subnet.precompile.hyperparameter.test.ts b/contract-tests/test/subnet.precompile.hyperparameter.test.ts index 87968b6e9f..aee28708b0 100644 --- a/contract-tests/test/subnet.precompile.hyperparameter.test.ts +++ b/contract-tests/test/subnet.precompile.hyperparameter.test.ts @@ -91,7 +91,65 @@ describe("Test the Subnet precompile contract", () => { const tx = await contract.setServingRateLimit(netuid, newValue); await tx.wait(); - let onchainValue = await api.query.SubtensorModule.ServingRateLimit.getValue(netuid) + const unwrapEnum = (value: unknown): { tag: string; value: unknown } | null => { + if (!value || typeof value !== "object") { + return null; + } + if ("type" in value && "value" in value) { + return { tag: (value as { type: string }).type, value: (value as { value: unknown }).value }; + } + const keys = Object.keys(value); + if (keys.length === 1) { + const key = keys[0]; + return { tag: key, value: (value as Record)[key] }; + } + return null; + }; + + const toNumber = (value: unknown): number | undefined => { + if (typeof value === "number") { + return value; + } + if (typeof value === "bigint") { + return Number(value); + } + if (typeof value === "string") { + const parsed = Number(value); + return Number.isNaN(parsed) ? undefined : parsed; + } + return undefined; + }; + + const extractRateLimit = (limits: unknown, scope: number): number | undefined => { + const decoded = unwrapEnum(limits); + if (!decoded) { + return undefined; + } + if (decoded.tag === "Global") { + const kind = unwrapEnum(decoded.value); + return kind?.tag === "Exact" ? toNumber(kind.value) : undefined; + } + if (decoded.tag !== "Scoped") { + return undefined; + } + const scopedEntries = decoded.value instanceof Map + ? Array.from(decoded.value.entries()) + : Array.isArray(decoded.value) + ? decoded.value + : []; + const entry = scopedEntries.find( + (item: unknown) => + Array.isArray(item) && item.length === 2 && toNumber(item[0]) === scope, + ) as [unknown, unknown] | undefined; + if (!entry) { + return undefined; + } + const kind = unwrapEnum(entry[1]); + return kind?.tag === "Exact" ? toNumber(kind.value) : undefined; + }; + + const limits = await api.query.RateLimiting.Limits.getValue({ Group: 0 } as any); + const onchainValue = extractRateLimit(limits, netuid); let valueFromContract = Number( diff --git a/contract-tests/test/wasm.contract.test.ts b/contract-tests/test/wasm.contract.test.ts index 26d5c87924..c42dbf5555 100644 --- a/contract-tests/test/wasm.contract.test.ts +++ b/contract-tests/test/wasm.contract.test.ts @@ -6,12 +6,23 @@ import { contracts } from "../.papi/descriptors"; import { getInkClient, InkClient, } from "@polkadot-api/ink-contracts" import { forceSetBalanceToSs58Address, startCall, burnedRegister } from "../src/subtensor"; import fs from "fs" +import path from "path"; import { convertPublicKeyToSs58 } from "../src/address-utils"; import { addNewSubnetwork, sendWasmContractExtrinsic } from "../src/subtensor"; import { tao } from "../src/balance-math"; -const bittensorWasmPath = "./bittensor/target/ink/bittensor.wasm" -const bittensorBytecode = fs.readFileSync(bittensorWasmPath) +const bittensorWasmPath = path.resolve(__dirname, "../bittensor/target/ink/bittensor.wasm") +const loadBittensorBytecode = () => { + if (!fs.existsSync(bittensorWasmPath)) { + throw new Error( + `Missing Ink wasm at ${bittensorWasmPath}. Run ` + + "`cd contract-tests/bittensor && cargo contract build --release` to generate it." + ) + } + + return fs.readFileSync(bittensorWasmPath) +} +let bittensorBytecode: Buffer; describe("Test wasm contract", () => { @@ -60,6 +71,7 @@ describe("Test wasm contract", () => { before(async () => { + bittensorBytecode = loadBittensorBytecode() // init variables got from await and async api = await getDevnetApi() @@ -584,4 +596,4 @@ describe("Test wasm contract", () => { assert.ok(result !== undefined) }) -}); \ No newline at end of file +}); diff --git a/node/Cargo.toml b/node/Cargo.toml index 1d2351c265..06b96c0eba 100644 --- a/node/Cargo.toml +++ b/node/Cargo.toml @@ -70,6 +70,7 @@ frame-system.workspace = true pallet-transaction-payment.workspace = true pallet-commitments.workspace = true pallet-drand.workspace = true +pallet-rate-limiting.workspace = true sp-crypto-ec-utils = { workspace = true, default-features = true, features = [ "bls12-381", ] } @@ -173,6 +174,7 @@ runtime-benchmarks = [ "polkadot-sdk/runtime-benchmarks", "pallet-subtensor/runtime-benchmarks", "pallet-shield/runtime-benchmarks", + "pallet-rate-limiting/runtime-benchmarks", ] pow-faucet = [] diff --git a/node/src/benchmarking.rs b/node/src/benchmarking.rs index 5430a75d9e..2f401e3e95 100644 --- a/node/src/benchmarking.rs +++ b/node/src/benchmarking.rs @@ -143,8 +143,12 @@ pub fn create_benchmark_extrinsic( pallet_subtensor::transaction_extension::SubtensorTransactionExtension::< runtime::Runtime, >::new(), - pallet_drand::drand_priority::DrandPriority::::new(), - frame_metadata_hash_extension::CheckMetadataHash::::new(true), + // Keep the same order while staying under the 12-item tuple limit. + ( + pallet_drand::drand_priority::DrandPriority::::new(), + frame_metadata_hash_extension::CheckMetadataHash::::new(true), + ), + pallet_rate_limiting::RateLimitTransactionExtension::::new(), ); let raw_payload = runtime::SignedPayload::from_raw( @@ -161,8 +165,8 @@ pub fn create_benchmark_extrinsic( (), (), (), + ((), None), (), - None, ), ); let signature = raw_payload.using_encoded(|e| sender.sign(e)); diff --git a/node/src/mev_shield/author.rs b/node/src/mev_shield/author.rs index 99000d4ac6..2501bcec20 100644 --- a/node/src/mev_shield/author.rs +++ b/node/src/mev_shield/author.rs @@ -396,8 +396,12 @@ where pallet_subtensor::transaction_extension::SubtensorTransactionExtension::< runtime::Runtime, >::new(), - pallet_drand::drand_priority::DrandPriority::::new(), - frame_metadata_hash_extension::CheckMetadataHash::::new(false), + // Keep the same order while staying under the 12-item tuple limit. + ( + pallet_drand::drand_priority::DrandPriority::::new(), + frame_metadata_hash_extension::CheckMetadataHash::::new(false), + ), + pallet_rate_limiting::RateLimitTransactionExtension::::new(), ); // 3) Manually construct the `Implicit` tuple that the runtime will also derive. @@ -431,8 +435,12 @@ where (), // ChargeTransactionPaymentWrapper::Implicit = () (), // SudoTransactionExtension::Implicit = () (), // SubtensorTransactionExtension::Implicit = () - (), // DrandPriority::Implicit = () - None, // CheckMetadataHash::Implicit = Option<[u8; 32]> + // Match the nested tuple shape used by TransactionExtensions. + ( + (), // DrandPriority::Implicit = () + None, // CheckMetadataHash::Implicit = Option<[u8; 32]> + ), + (), // RateLimitTransactionExtension::Implicit = () ); // 4) Build the exact signable payload from call + extra + implicit. diff --git a/pallets/admin-utils/Cargo.toml b/pallets/admin-utils/Cargo.toml index 61cdba4cbf..aa6d7e593c 100644 --- a/pallets/admin-utils/Cargo.toml +++ b/pallets/admin-utils/Cargo.toml @@ -38,6 +38,7 @@ sp-core.workspace = true sp-io.workspace = true sp-tracing.workspace = true sp-consensus-aura.workspace = true +rate-limiting-interface.workspace = true pallet-balances = { workspace = true, features = ["std"] } pallet-scheduler.workspace = true pallet-grandpa.workspace = true @@ -75,6 +76,7 @@ std = [ "substrate-fixed/std", "subtensor-swap-interface/std", "subtensor-runtime-common/std", + "rate-limiting-interface/std", ] runtime-benchmarks = [ "frame-benchmarking/runtime-benchmarks", diff --git a/pallets/admin-utils/src/benchmarking.rs b/pallets/admin-utils/src/benchmarking.rs index 7b8124144d..105a8b8ddb 100644 --- a/pallets/admin-utils/src/benchmarking.rs +++ b/pallets/admin-utils/src/benchmarking.rs @@ -64,11 +64,19 @@ mod benchmarks { } #[benchmark] + #[allow(deprecated)] fn sudo_set_serving_rate_limit() { // disable admin freeze window pallet_subtensor::Pallet::::set_admin_freeze_window(0); - #[extrinsic_call] - _(RawOrigin::Root, 1u16.into()/*netuid*/, 100u64/*serving_rate_limit*/)/*sudo_set_serving_rate_limit*/; + #[block] + { + #[allow(deprecated)] + let _ = AdminUtils::::sudo_set_serving_rate_limit( + RawOrigin::Root.into(), + 1u16.into(), /*netuid*/ + 100u64, /*serving_rate_limit*/ + ); + } } #[benchmark] diff --git a/pallets/admin-utils/src/lib.rs b/pallets/admin-utils/src/lib.rs index c3999312b3..18cddb0d42 100644 --- a/pallets/admin-utils/src/lib.rs +++ b/pallets/admin-utils/src/lib.rs @@ -117,6 +117,8 @@ pub mod pallet { MaxAllowedUidsGreaterThanDefaultMaxAllowedUids, /// Bad parameter value InvalidValue, + /// The called extrinsic has been deprecated. + Deprecated, } /// Enum for specifying the type of precompile operation. #[derive( @@ -220,31 +222,22 @@ pub mod pallet { } /// The extrinsic sets the serving rate limit for a subnet. - /// It is only callable by the root account or subnet owner. - /// The extrinsic will call the Subtensor pallet to set the serving rate limit. + /// + /// Deprecated: serving rate limits are now configured via `pallet-rate-limiting` on the + /// serving group target (`GROUP_SERVE`) with `scope = Some(netuid)`. #[pallet::call_index(3)] #[pallet::weight(Weight::from_parts(22_980_000, 0) .saturating_add(::DbWeight::get().reads(2_u64)) .saturating_add(::DbWeight::get().writes(1_u64)))] + #[deprecated( + note = "deprecated: configure via pallet-rate-limiting::set_rate_limit(target=Group(GROUP_SERVE), scope=Some(netuid), ...)" + )] pub fn sudo_set_serving_rate_limit( - origin: OriginFor, - netuid: NetUid, - serving_rate_limit: u64, + _origin: OriginFor, + _netuid: NetUid, + _serving_rate_limit: u64, ) -> DispatchResult { - let maybe_owner = pallet_subtensor::Pallet::::ensure_sn_owner_or_root_with_limits( - origin, - netuid, - &[Hyperparameter::ServingRateLimit.into()], - )?; - pallet_subtensor::Pallet::::ensure_admin_window_open(netuid)?; - pallet_subtensor::Pallet::::set_serving_rate_limit(netuid, serving_rate_limit); - log::debug!("ServingRateLimitSet( serving_rate_limit: {serving_rate_limit:?} ) "); - pallet_subtensor::Pallet::::record_owner_rl( - maybe_owner, - netuid, - &[Hyperparameter::ServingRateLimit.into()], - ); - Ok(()) + Err(Error::::Deprecated.into()) } /// The extrinsic sets the minimum difficulty for a subnet. diff --git a/pallets/admin-utils/src/tests/mock.rs b/pallets/admin-utils/src/tests/mock.rs index 0140808baa..211a72ef08 100644 --- a/pallets/admin-utils/src/tests/mock.rs +++ b/pallets/admin-utils/src/tests/mock.rs @@ -8,6 +8,7 @@ use frame_support::{ }; use frame_system::{self as system, offchain::CreateTransactionBase}; use frame_system::{EnsureRoot, limits}; +use rate_limiting_interface::{RateLimitingInfo, TryIntoRateLimitTarget}; use sp_consensus_aura::sr25519::AuthorityId as AuraId; use sp_consensus_grandpa::AuthorityList as GrandpaAuthorityList; use sp_core::U256; @@ -19,7 +20,7 @@ use sp_runtime::{ }; use sp_std::cmp::Ordering; use sp_weights::Weight; -use subtensor_runtime_common::{NetUid, TaoCurrency}; +use subtensor_runtime_common::{NetUid, TaoCurrency, rate_limiting::RateLimitUsageKey}; type Block = frame_system::mocking::MockBlock; // Configure a mock runtime to test the pallet. @@ -103,7 +104,6 @@ parameter_types! { pub const InitialMinChildKeyTake: u16 = 0; // Allow 0 % pub const InitialMaxChildKeyTake: u16 = 11_796; // 18 %; pub const InitialWeightsVersionKey: u16 = 0; - pub const InitialServingRateLimit: u64 = 0; // No limit. pub const InitialTxRateLimit: u64 = 0; // Disable rate limit for testing pub const InitialTxDelegateTakeRateLimit: u64 = 0; // Disable rate limit for testing pub const InitialTxChildKeyTakeRateLimit: u64 = 0; // Disable rate limit for testing @@ -189,7 +189,6 @@ impl pallet_subtensor::Config for Test { type InitialWeightsVersionKey = InitialWeightsVersionKey; type InitialMaxDifficulty = InitialMaxDifficulty; type InitialMinDifficulty = InitialMinDifficulty; - type InitialServingRateLimit = InitialServingRateLimit; type InitialTxRateLimit = InitialTxRateLimit; type InitialTxDelegateTakeRateLimit = InitialTxDelegateTakeRateLimit; type InitialTxChildKeyTakeRateLimit = InitialTxChildKeyTakeRateLimit; @@ -224,6 +223,7 @@ impl pallet_subtensor::Config for Test { type GetCommitments = (); type MaxImmuneUidsPercentage = MaxImmuneUidsPercentage; type CommitmentsInterface = CommitmentsI; + type RateLimiting = NoRateLimiting; type EvmKeyAssociateRateLimit = EvmKeyAssociateRateLimit; } @@ -356,6 +356,33 @@ impl pallet_subtensor::CommitmentsInterface for CommitmentsI { fn purge_netuid(_netuid: NetUid) {} } +pub struct NoRateLimiting; + +impl RateLimitingInfo for NoRateLimiting { + type GroupId = subtensor_runtime_common::rate_limiting::GroupId; + type CallMetadata = RuntimeCall; + type Limit = u64; + type Scope = subtensor_runtime_common::NetUid; + type UsageKey = RateLimitUsageKey; + + fn rate_limit(_target: TargetArg, _scope: Option) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + None + } + + fn last_seen( + _target: TargetArg, + _usage_key: Option, + ) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + None + } +} + pub struct GrandpaInterfaceImpl; impl crate::GrandpaInterface for GrandpaInterfaceImpl { fn schedule_change( diff --git a/pallets/admin-utils/src/tests/mod.rs b/pallets/admin-utils/src/tests/mod.rs index 024871e60f..a5fb053806 100644 --- a/pallets/admin-utils/src/tests/mod.rs +++ b/pallets/admin-utils/src/tests/mod.rs @@ -44,26 +44,30 @@ fn test_sudo_set_default_take() { } #[test] +#[allow(deprecated)] fn test_sudo_set_serving_rate_limit() { new_test_ext().execute_with(|| { let netuid = NetUid::from(3); let to_be_set: u64 = 10; let init_value: u64 = SubtensorModule::get_serving_rate_limit(netuid); - assert_eq!( + assert_noop!( AdminUtils::sudo_set_serving_rate_limit( <::RuntimeOrigin>::signed(U256::from(1)), netuid, to_be_set ), - Err(DispatchError::BadOrigin) + Error::::Deprecated + ); + assert_eq!(SubtensorModule::get_serving_rate_limit(netuid), init_value); + assert_noop!( + AdminUtils::sudo_set_serving_rate_limit( + <::RuntimeOrigin>::root(), + netuid, + to_be_set + ), + Error::::Deprecated ); assert_eq!(SubtensorModule::get_serving_rate_limit(netuid), init_value); - assert_ok!(AdminUtils::sudo_set_serving_rate_limit( - <::RuntimeOrigin>::root(), - netuid, - to_be_set - )); - assert_eq!(SubtensorModule::get_serving_rate_limit(netuid), to_be_set); }); } diff --git a/pallets/rate-limiting-interface/Cargo.toml b/pallets/rate-limiting-interface/Cargo.toml new file mode 100644 index 0000000000..8f352d0c58 --- /dev/null +++ b/pallets/rate-limiting-interface/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "rate-limiting-interface" +version = "0.1.0" +edition.workspace = true + +[lints] +workspace = true + +[dependencies] +codec = { workspace = true, features = ["derive"], default-features = false } +frame-support = { workspace = true, default-features = false } +scale-info = { workspace = true, features = ["derive"], default-features = false } +serde = { workspace = true, features = ["derive"], default-features = false } +sp-std = { workspace = true, default-features = false } + +[features] +default = ["std"] +std = [ + "codec/std", + "frame-support/std", + "scale-info/std", + "serde/std", + "sp-std/std", +] diff --git a/pallets/rate-limiting-interface/README.md b/pallets/rate-limiting-interface/README.md new file mode 100644 index 0000000000..d8671fc0ba --- /dev/null +++ b/pallets/rate-limiting-interface/README.md @@ -0,0 +1,3 @@ +# `rate-limiting-interface` + +Small, `no_std`-friendly interface crate that defines [`RateLimitingInfo`](src/lib.rs). diff --git a/pallets/rate-limiting-interface/src/lib.rs b/pallets/rate-limiting-interface/src/lib.rs new file mode 100644 index 0000000000..4bf0dab22f --- /dev/null +++ b/pallets/rate-limiting-interface/src/lib.rs @@ -0,0 +1,267 @@ +#![cfg_attr(not(feature = "std"), no_std)] + +//! Read-only interface for querying rate limits and last-seen usage. + +use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; +use frame_support::traits::GetCallMetadata; +use scale_info::TypeInfo; +use serde::{Deserialize, Serialize}; +use sp_std::vec::Vec; + +/// Read-only queries for rate-limiting configuration and usage tracking. +pub trait RateLimitingInfo { + /// Group id type used by rate-limiting targets. + type GroupId; + /// Call type used for name/index resolution. + type CallMetadata: GetCallMetadata; + /// Numeric type used for returned values (commonly a block number / block span type). + type Limit; + /// Optional configuration scope (for example per-network `netuid`). + type Scope; + /// Optional usage key used to refine "last seen" tracking. + type UsageKey; + + /// Returns the configured limit for `target` and optional `scope`. + fn rate_limit(target: TargetArg, scope: Option) -> Option + where + TargetArg: TryIntoRateLimitTarget; + + /// Returns when `target` was last observed for the optional `usage_key`. + fn last_seen( + target: TargetArg, + usage_key: Option, + ) -> Option + where + TargetArg: TryIntoRateLimitTarget; +} + +/// Target identifier for rate limit and usage configuration. +#[derive( + Serialize, + Deserialize, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + MaxEncodedLen, + Debug, +)] +pub enum RateLimitTarget { + /// Per-transaction configuration keyed by pallet/extrinsic indices. + Transaction(TransactionIdentifier), + /// Shared configuration for a named group. + Group(GroupId), +} + +impl RateLimitTarget { + /// Returns the transaction identifier when the target represents a single extrinsic. + pub fn as_transaction(&self) -> Option<&TransactionIdentifier> { + match self { + RateLimitTarget::Transaction(identifier) => Some(identifier), + RateLimitTarget::Group(_) => None, + } + } + + /// Returns the group identifier when the target represents a group configuration. + pub fn as_group(&self) -> Option<&GroupId> { + match self { + RateLimitTarget::Transaction(_) => None, + RateLimitTarget::Group(id) => Some(id), + } + } +} + +impl From for RateLimitTarget { + fn from(identifier: TransactionIdentifier) -> Self { + Self::Transaction(identifier) + } +} + +/// Identifies a runtime call by pallet and extrinsic indices. +#[derive( + Serialize, + Deserialize, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + MaxEncodedLen, + Debug, +)] +pub struct TransactionIdentifier { + /// Pallet variant index. + pub pallet_index: u8, + /// Call variant index within the pallet. + pub extrinsic_index: u8, +} + +impl TransactionIdentifier { + /// Builds a new identifier from pallet/extrinsic indices. + pub const fn new(pallet_index: u8, extrinsic_index: u8) -> Self { + Self { + pallet_index, + extrinsic_index, + } + } + + /// Attempts to build an identifier from a SCALE-encoded call by reading the first two bytes. + pub fn from_call(call: &Call) -> Option { + call.using_encoded(|encoded| { + let pallet_index = *encoded.get(0)?; + let extrinsic_index = *encoded.get(1)?; + Some(Self::new(pallet_index, extrinsic_index)) + }) + } + + /// Resolves pallet/extrinsic names for this identifier using call metadata. + pub fn names(&self) -> Option<(&'static str, &'static str)> { + let modules = Call::get_module_names(); + let pallet_name = *modules.get(self.pallet_index as usize)?; + let call_names = Call::get_call_names(pallet_name); + let extrinsic_name = *call_names.get(self.extrinsic_index as usize)?; + Some((pallet_name, extrinsic_name)) + } + + /// Resolves a pallet/extrinsic name pair into a transaction identifier. + pub fn for_call_names( + pallet_name: &str, + extrinsic_name: &str, + ) -> Option { + let modules = Call::get_module_names(); + let pallet_pos = modules.iter().position(|name| *name == pallet_name)?; + let call_names = Call::get_call_names(pallet_name); + let extrinsic_pos = call_names.iter().position(|name| *name == extrinsic_name)?; + let pallet_index = u8::try_from(pallet_pos).ok()?; + let extrinsic_index = u8::try_from(extrinsic_pos).ok()?; + Some(Self::new(pallet_index, extrinsic_index)) + } +} + +/// Conversion into a concrete [`RateLimitTarget`]. +pub trait TryIntoRateLimitTarget { + type Error; + + fn try_into_rate_limit_target( + self, + ) -> Result, Self::Error>; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RateLimitTargetConversionError { + InvalidUtf8, + UnknownCall, +} + +impl TryIntoRateLimitTarget for RateLimitTarget { + type Error = core::convert::Infallible; + + fn try_into_rate_limit_target( + self, + ) -> Result, Self::Error> { + Ok(self) + } +} + +impl TryIntoRateLimitTarget for GroupId { + type Error = core::convert::Infallible; + + fn try_into_rate_limit_target( + self, + ) -> Result, Self::Error> { + Ok(RateLimitTarget::Group(self)) + } +} + +impl TryIntoRateLimitTarget for (Vec, Vec) { + type Error = RateLimitTargetConversionError; + + fn try_into_rate_limit_target( + self, + ) -> Result, Self::Error> { + let (pallet, extrinsic) = self; + let pallet_name = sp_std::str::from_utf8(&pallet) + .map_err(|_| RateLimitTargetConversionError::InvalidUtf8)?; + let extrinsic_name = sp_std::str::from_utf8(&extrinsic) + .map_err(|_| RateLimitTargetConversionError::InvalidUtf8)?; + + let identifier = TransactionIdentifier::for_call_names::(pallet_name, extrinsic_name) + .ok_or(RateLimitTargetConversionError::UnknownCall)?; + + Ok(RateLimitTarget::Transaction(identifier)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use codec::Encode; + use frame_support::traits::CallMetadata; + + #[derive(Clone, Copy, Debug, Encode)] + struct DummyCall(u8, u8); + + impl GetCallMetadata for DummyCall { + fn get_module_names() -> &'static [&'static str] { + &["P0", "P1"] + } + + fn get_call_names(module: &str) -> &'static [&'static str] { + match module { + "P0" => &["C0"], + "P1" => &["C0", "C1", "C2", "C3", "C4"], + _ => &[], + } + } + + fn get_call_metadata(&self) -> CallMetadata { + CallMetadata { + function_name: "unused", + pallet_name: "unused", + } + } + } + + #[test] + fn transaction_identifier_from_call_reads_first_two_bytes() { + let id = TransactionIdentifier::from_call(&DummyCall(1, 4)).expect("identifier"); + assert_eq!(id, TransactionIdentifier::new(1, 4)); + } + + #[test] + fn transaction_identifier_names_resolves_metadata() { + let id = TransactionIdentifier::new(1, 4); + assert_eq!(id.names::(), Some(("P1", "C4"))); + } + + #[test] + fn transaction_identifier_for_call_names_resolves_indices() { + let id = TransactionIdentifier::for_call_names::("P1", "C4").expect("id"); + assert_eq!(id, TransactionIdentifier::new(1, 4)); + } + + #[test] + fn rate_limit_target_accessors_work() { + let tx = RateLimitTarget::::Transaction(TransactionIdentifier::new(1, 4)); + assert!(tx.as_group().is_none()); + assert_eq!( + tx.as_transaction().copied(), + Some(TransactionIdentifier::new(1, 4)) + ); + + let group = RateLimitTarget::::Group(7); + assert!(group.as_transaction().is_none()); + assert_eq!(group.as_group().copied(), Some(7)); + } +} diff --git a/pallets/rate-limiting/Cargo.toml b/pallets/rate-limiting/Cargo.toml new file mode 100644 index 0000000000..00bae918dd --- /dev/null +++ b/pallets/rate-limiting/Cargo.toml @@ -0,0 +1,49 @@ +[package] +name = "pallet-rate-limiting" +version = "0.1.0" +edition.workspace = true + +[lints] +workspace = true + +[package.metadata.docs.rs] +targets = ["x86_64-unknown-linux-gnu"] + +[dependencies] +codec = { workspace = true, features = ["derive"] } +frame-benchmarking = { workspace = true, optional = true } +frame-support.workspace = true +frame-system.workspace = true +scale-info = { workspace = true, features = ["derive"] } +serde = { workspace = true, features = ["derive"] } +sp-std.workspace = true +sp-runtime.workspace = true +rate-limiting-interface.workspace = true +subtensor-runtime-common.workspace = true + +[dev-dependencies] +sp-core.workspace = true +sp-io.workspace = true +sp-runtime.workspace = true + +[features] +default = ["std"] +std = [ + "codec/std", + "frame-benchmarking?/std", + "frame-support/std", + "frame-system/std", + "rate-limiting-interface/std", + "scale-info/std", + "serde/std", + "sp-std/std", + "sp-runtime/std", + "subtensor-runtime-common/std", +] +runtime-benchmarks = [ + "frame-benchmarking", +] +try-runtime = [ + "frame-support/try-runtime", + "frame-system/try-runtime", +] diff --git a/pallets/rate-limiting/rpc/Cargo.toml b/pallets/rate-limiting/rpc/Cargo.toml new file mode 100644 index 0000000000..d5bf689e8b --- /dev/null +++ b/pallets/rate-limiting/rpc/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "pallet-rate-limiting-rpc" +version = "0.1.0" +description = "RPC interface for the rate limiting pallet" +edition.workspace = true + +[dependencies] +jsonrpsee = { workspace = true, features = ["client-core", "server", "macros"] } +sp-api.workspace = true +sp-blockchain.workspace = true +sp-runtime.workspace = true +pallet-rate-limiting-runtime-api.workspace = true +subtensor-runtime-common = { workspace = true, default-features = false } + +[features] +default = ["std"] +std = [ + "sp-api/std", + "sp-runtime/std", + "pallet-rate-limiting-runtime-api/std", + "subtensor-runtime-common/std", +] diff --git a/pallets/rate-limiting/rpc/src/lib.rs b/pallets/rate-limiting/rpc/src/lib.rs new file mode 100644 index 0000000000..ca7452a7a0 --- /dev/null +++ b/pallets/rate-limiting/rpc/src/lib.rs @@ -0,0 +1,82 @@ +//! RPC interface for the rate limiting pallet. + +use jsonrpsee::{ + core::RpcResult, + proc_macros::rpc, + types::{ErrorObjectOwned, error::ErrorObject}, +}; +use sp_api::ProvideRuntimeApi; +use sp_blockchain::HeaderBackend; +use sp_runtime::traits::Block as BlockT; +use std::sync::Arc; + +pub use pallet_rate_limiting_runtime_api::{RateLimitRpcResponse, RateLimitingRuntimeApi}; + +#[rpc(client, server)] +pub trait RateLimitingRpcApi { + #[method(name = "rateLimiting_getRateLimit")] + fn get_rate_limit( + &self, + pallet: Vec, + extrinsic: Vec, + at: Option, + ) -> RpcResult>; +} + +/// Error type of this RPC api. +pub enum Error { + /// The call to runtime failed. + RuntimeError(String), +} + +impl From for ErrorObjectOwned { + fn from(e: Error) -> Self { + match e { + Error::RuntimeError(e) => ErrorObject::owned(1, e, None::<()>), + } + } +} + +impl From for i32 { + fn from(e: Error) -> i32 { + match e { + Error::RuntimeError(_) => 1, + } + } +} + +/// RPC implementation for the rate limiting pallet. +pub struct RateLimiting { + client: Arc, + _marker: std::marker::PhantomData, +} + +impl RateLimiting { + /// Creates a new instance of the rate limiting RPC helper. + pub fn new(client: Arc) -> Self { + Self { + client, + _marker: Default::default(), + } + } +} + +impl RateLimitingRpcApiServer<::Hash> for RateLimiting +where + Block: BlockT, + C: ProvideRuntimeApi + HeaderBackend + Send + Sync + 'static, + C::Api: RateLimitingRuntimeApi, +{ + fn get_rate_limit( + &self, + pallet: Vec, + extrinsic: Vec, + at: Option<::Hash>, + ) -> RpcResult> { + let api = self.client.runtime_api(); + let at = at.unwrap_or_else(|| self.client.info().best_hash); + + api.get_rate_limit(at, pallet, extrinsic) + .map_err(|e| Error::RuntimeError(format!("Unable to fetch rate limit: {e:?}")).into()) + } +} diff --git a/pallets/rate-limiting/runtime-api/Cargo.toml b/pallets/rate-limiting/runtime-api/Cargo.toml new file mode 100644 index 0000000000..2847d865dd --- /dev/null +++ b/pallets/rate-limiting/runtime-api/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "pallet-rate-limiting-runtime-api" +version = "0.1.0" +description = "Runtime API for the rate limiting pallet" +edition.workspace = true + +[dependencies] +codec = { workspace = true, features = ["derive"] } +scale-info = { workspace = true, features = ["derive"] } +sp-api.workspace = true +sp-std.workspace = true +pallet-rate-limiting.workspace = true +subtensor-runtime-common = { workspace = true, default-features = false } +serde = { workspace = true, features = ["derive"], optional = true } + +[features] +default = ["std"] +std = [ + "codec/std", + "scale-info/std", + "sp-api/std", + "sp-std/std", + "pallet-rate-limiting/std", + "subtensor-runtime-common/std", + "serde", +] diff --git a/pallets/rate-limiting/runtime-api/src/lib.rs b/pallets/rate-limiting/runtime-api/src/lib.rs new file mode 100644 index 0000000000..98b55e9a26 --- /dev/null +++ b/pallets/rate-limiting/runtime-api/src/lib.rs @@ -0,0 +1,25 @@ +#![cfg_attr(not(feature = "std"), no_std)] + +use codec::{Decode, Encode}; +use pallet_rate_limiting::RateLimitKind; +use scale_info::TypeInfo; +use sp_std::vec::Vec; +use subtensor_runtime_common::BlockNumber; + +#[cfg(feature = "std")] +use serde::{Deserialize, Serialize}; + +#[cfg_attr(feature = "std", derive(Serialize, Deserialize))] +#[derive(Clone, Debug, Decode, Encode, Eq, PartialEq, TypeInfo)] +pub struct RateLimitRpcResponse { + pub global: Option>, + pub contextual: Vec<(Vec, RateLimitKind)>, + pub default_limit: BlockNumber, + pub resolved: Option, +} + +sp_api::decl_runtime_apis! { + pub trait RateLimitingRuntimeApi { + fn get_rate_limit(pallet: Vec, extrinsic: Vec) -> Option; + } +} diff --git a/pallets/rate-limiting/src/benchmarking.rs b/pallets/rate-limiting/src/benchmarking.rs new file mode 100644 index 0000000000..265733e113 --- /dev/null +++ b/pallets/rate-limiting/src/benchmarking.rs @@ -0,0 +1,185 @@ +//! Benchmarking setup for pallet-rate-limiting +#![cfg(feature = "runtime-benchmarks")] +#![allow(clippy::arithmetic_side_effects)] + +use codec::Decode; +use frame_benchmarking::v2::*; +use frame_system::{RawOrigin, pallet_prelude::BlockNumberFor}; +use sp_runtime::traits::{One, Saturating}; + +use super::*; +use crate::CallReadOnly; + +pub trait BenchmarkHelper { + fn sample_call() -> Call; +} + +impl BenchmarkHelper for () +where + Call: Decode, +{ + fn sample_call() -> Call { + Decode::decode(&mut &[][..]).expect("Provide a call via BenchmarkHelper::sample_call") + } +} + +fn sample_call() -> Box<::RuntimeCall> +where + T::BenchmarkHelper: BenchmarkHelper<::RuntimeCall>, +{ + Box::new(T::BenchmarkHelper::sample_call()) +} + +fn seed_group(name: &[u8], sharing: GroupSharing) -> ::GroupId { + Pallet::::create_group(RawOrigin::Root.into(), name.to_vec(), sharing) + .expect("group created"); + Pallet::::next_group_id().saturating_sub(::GroupId::one()) +} + +fn register_call_with_group( + group: Option<::GroupId>, +) -> TransactionIdentifier { + let call = sample_call::(); + let identifier = TransactionIdentifier::from_call(call.as_ref()).expect("id"); + Pallet::::register_call(RawOrigin::Root.into(), call, group).expect("registered"); + identifier +} + +#[benchmarks] +mod benchmarks { + use super::*; + use sp_std::vec::Vec; + + #[benchmark] + fn register_call() { + let call = sample_call::(); + let identifier = TransactionIdentifier::from_call(call.as_ref()).expect("id"); + let target = RateLimitTarget::Transaction(identifier); + + #[extrinsic_call] + _(RawOrigin::Root, call, None); + + assert!(Limits::::contains_key(target)); + } + + #[benchmark] + fn set_rate_limit() { + let call = sample_call::(); + let identifier = TransactionIdentifier::from_call(call.as_ref()).expect("id"); + let target = RateLimitTarget::Transaction(identifier); + Limits::::insert(target, RateLimit::global(RateLimitKind::Default)); + + let limit = RateLimitKind::>::Exact(BlockNumberFor::::from(10u32)); + + #[extrinsic_call] + _(RawOrigin::Root, target, None, limit); + + let stored = Limits::::get(target).expect("limit stored"); + assert!( + matches!(stored, RateLimit::Global(RateLimitKind::Exact(span)) if span == BlockNumberFor::::from(10u32)) + ); + } + + #[benchmark] + fn assign_call_to_group() { + let group = seed_group::(b"grp", GroupSharing::UsageOnly); + let identifier = register_call_with_group::(None); + + #[extrinsic_call] + _(RawOrigin::Root, identifier, group, false); + + assert_eq!(CallGroups::::get(identifier), Some(group)); + assert_eq!(CallReadOnly::::get(identifier), Some(false)); + assert!(GroupMembers::::get(group).contains(&identifier)); + } + + #[benchmark] + fn remove_call_from_group() { + let group = seed_group::(b"team", GroupSharing::ConfigOnly); + let identifier = register_call_with_group::(Some(group)); + + #[extrinsic_call] + _(RawOrigin::Root, identifier); + + assert!(CallGroups::::get(identifier).is_none()); + assert!(!GroupMembers::::get(group).contains(&identifier)); + } + + #[benchmark] + fn create_group() { + let name = b"bench".to_vec(); + let sharing = GroupSharing::ConfigAndUsage; + + #[extrinsic_call] + _(RawOrigin::Root, name.clone(), sharing); + + let group = Pallet::::next_group_id().saturating_sub(::GroupId::one()); + let details = Groups::::get(group).expect("group stored"); + let stored: Vec = details.name.into(); + assert_eq!(stored, name); + assert_eq!(details.sharing, sharing); + } + + #[benchmark] + fn update_group() { + let group = seed_group::(b"old", GroupSharing::UsageOnly); + let new_name = b"new".to_vec(); + let new_sharing = GroupSharing::ConfigAndUsage; + + #[extrinsic_call] + _( + RawOrigin::Root, + group, + Some(new_name.clone()), + Some(new_sharing), + ); + + let details = Groups::::get(group).expect("group exists"); + let stored: Vec = details.name.into(); + assert_eq!(stored, new_name); + assert_eq!(details.sharing, new_sharing); + } + + #[benchmark] + fn delete_group() { + let group = seed_group::(b"delete", GroupSharing::UsageOnly); + + #[extrinsic_call] + _(RawOrigin::Root, group); + + assert!(Groups::::get(group).is_none()); + } + + #[benchmark] + fn deregister_call() { + let group = seed_group::(b"dreg", GroupSharing::ConfigAndUsage); + let identifier = register_call_with_group::(Some(group)); + let target = RateLimitTarget::Transaction(identifier); + let usage_target = Pallet::::usage_target(&identifier).expect("usage target"); + LastSeen::::insert( + usage_target, + None::, + BlockNumberFor::::from(1u32), + ); + + #[extrinsic_call] + _(RawOrigin::Root, identifier, None, true); + + assert!(Limits::::get(target).is_none()); + assert!(LastSeen::::get(usage_target, None::).is_none()); + assert!(CallGroups::::get(identifier).is_none()); + assert!(!GroupMembers::::get(group).contains(&identifier)); + } + + #[benchmark] + fn set_default_rate_limit() { + let block_span = BlockNumberFor::::from(10u32); + + #[extrinsic_call] + _(RawOrigin::Root, block_span); + + assert_eq!(DefaultLimit::::get(), block_span); + } + + impl_benchmark_test_suite!(Pallet, crate::mock::new_test_ext(), crate::mock::Test); +} diff --git a/pallets/rate-limiting/src/lib.rs b/pallets/rate-limiting/src/lib.rs new file mode 100644 index 0000000000..217611c002 --- /dev/null +++ b/pallets/rate-limiting/src/lib.rs @@ -0,0 +1,1395 @@ +#![cfg_attr(not(feature = "std"), no_std)] + +//! Rate limiting for runtime calls with optional contextual restrictions. +//! +//! # Overview +//! +//! `pallet-rate-limiting` lets a runtime restrict how frequently particular calls can execute. +//! Limits are stored on-chain, keyed by explicit [`RateLimitTarget`] values. A target is either a +//! single [`TransactionIdentifier`] (the pallet/extrinsic indices) or a named *group* managed by +//! the admin APIs. Groups provide a way to give multiple calls the same configuration and/or usage +//! tracking without duplicating storage. Each target entry stores either a global span or a set of +//! scoped spans resolved at runtime. The pallet exposes a handful of extrinsics, restricted by +//! [`Config::AdminOrigin`], to manage this data: +//! +//! - [`register_call`](pallet::Pallet::register_call): register a call for rate limiting, seed its +//! initial configuration using [`Config::LimitScopeResolver`], and optionally place it into a +//! group. +//! - [`set_rate_limit`](pallet::Pallet::set_rate_limit): assign or override the limit at a specific +//! target/scope by supplying a [`RateLimitKind`] span. +//! - [`assign_call_to_group`](pallet::Pallet::assign_call_to_group) and +//! [`remove_call_from_group`](pallet::Pallet::remove_call_from_group): manage group membership for +//! registered calls. +//! - [`set_call_read_only`](pallet::Pallet::set_call_read_only): for grouped calls, choose whether +//! successful dispatches should update the shared usage row (`false` by default). +//! - [`deregister_call`](pallet::Pallet::deregister_call): remove scoped configuration or wipe the +//! registration entirely. +//! - [`set_default_rate_limit`](pallet::Pallet::set_default_rate_limit): set the global default +//! block span used by `RateLimitKind::Default` entries. +//! +//! The pallet also tracks the last block in which a target was observed, per optional *usage key*. +//! A usage key may refine tracking beyond the limit scope (for example combining a `netuid` with a +//! hyperparameter), so the two concepts are explicitly separated in the configuration. When the +//! admin puts several calls into a group and marks usage as shared, each dispatch still runs the +//! resolver: the group only chooses the storage target, while the resolver output (or `None`) picks +//! the row under that target. Calls that resolve to the same usage key update the same timestamp; +//! calls that resolve to different keys keep isolated timers even when they share a group. The same +//! rule applies to limit scopes—grouping funnels configuration into the same target, but the scope +//! resolver decides whether that entry is global or per-context. +//! +//! Each storage map is namespaced by pallet instance; runtimes can deploy multiple independent +//! instances to manage distinct rate-limiting scopes (in the global sense). +//! +//! # Transaction extension +//! +//! Enforcement happens via [`RateLimitTransactionExtension`], which implements +//! `sp_runtime::traits::TransactionExtension`. The extension consults `Limits`, fetches the current +//! block, and decides whether the call is eligible. If successful, it returns metadata that causes +//! [`LastSeen`](pallet::LastSeen) to update after dispatch. A rejected call yields +//! `InvalidTransaction::Custom(1)`. +//! +//! To enable the extension, add it to your runtime's transaction extension tuple. For example: +//! +//! ```ignore +//! pub type TransactionExtensions = ( +//! // ... other extensions ... +//! pallet_rate_limiting::RateLimitTransactionExtension, +//! ); +//! ``` +//! +//! # Context resolvers +//! +//! The pallet relies on two resolvers: +//! +//! - [`Config::LimitScopeResolver`], which determines how limits are stored (for example by +//! returning a `netuid`). The resolver can also signal that a call should bypass rate limiting or +//! adjust the effective span at validation time. When it returns `None`, the configuration is +//! stored as a global fallback. +//! - [`Config::UsageResolver`], which decides how executions are tracked in +//! [`LastSeen`](pallet::LastSeen). This can refine the limit scope (for example by returning a +//! tuple of `(netuid, hyperparameter)`). +//! +//! Each resolver receives the origin and call and may return `Some(identifier)` when scoping is +//! required, or `None` to use the global entry. Extrinsics such as +//! [`set_rate_limit`](pallet::Pallet::set_rate_limit) automatically consult these resolvers. When a +//! call belongs to a group the pallet still runs the resolver—instead of indexing storage at the +//! transaction-level target, it indexes at the group target. Resolving to different contexts keeps +//! independent limit/usage rows even though the calls share a group; resolving to the same context +//! causes them to share enforcement state. +//! +//! ```ignore +//! pub struct WeightsContextResolver; +//! +//! // Limits are scoped per netuid. +//! pub struct ScopeResolver; +//! impl pallet_rate_limiting::RateLimitScopeResolver< +//! RuntimeOrigin, +//! RuntimeCall, +//! NetUid, +//! BlockNumber, +//! > for ScopeResolver { +//! fn context(origin: &RuntimeOrigin, call: &RuntimeCall) -> Option { +//! match call { +//! RuntimeCall::Subtensor(pallet_subtensor::Call::set_weights { netuid, .. }) => { +//! Some(*netuid) +//! } +//! _ => None, +//! } +//! } +//! +//! fn should_bypass(origin: &RuntimeOrigin, _call: &RuntimeCall) -> BypassDecision { +//! if matches!(origin, RuntimeOrigin::Root) { +//! BypassDecision::bypass_and_skip() +//! } else { +//! BypassDecision::enforce_and_record() +//! } +//! } +//! +//! fn adjust_span(_origin: &RuntimeOrigin, _call: &RuntimeCall, span: BlockNumber) -> BlockNumber { +//! span +//! } +//! } +//! +//! // Usage tracking distinguishes hyperparameter + netuid. +//! pub struct UsageResolver; +//! impl pallet_rate_limiting::RateLimitUsageResolver< +//! RuntimeOrigin, +//! RuntimeCall, +//! (NetUid, HyperParam), +//! > for UsageResolver { +//! fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option<(NetUid, HyperParam)> { +//! match call { +//! RuntimeCall::Subtensor(pallet_subtensor::Call::set_hyperparam { +//! netuid, +//! hyper, +//! .. +//! }) => Some((*netuid, *hyper)), +//! _ => None, +//! } +//! } +//! } +//! +//! impl pallet_rate_limiting::Config for Runtime { +//! type RuntimeCall = RuntimeCall; +//! type LimitScope = NetUid; +//! type LimitScopeResolver = ScopeResolver; +//! type UsageKey = (NetUid, HyperParam); +//! type UsageResolver = UsageResolver; +//! type AdminOrigin = frame_system::EnsureRoot; +//! } +//! ``` + +#[cfg(feature = "runtime-benchmarks")] +pub use benchmarking::BenchmarkHelper; +pub use pallet::*; +pub use rate_limiting_interface::{RateLimitTarget, TransactionIdentifier}; +pub use rate_limiting_interface::{RateLimitingInfo, TryIntoRateLimitTarget}; +pub use tx_extension::RateLimitTransactionExtension; +pub use types::{ + BypassDecision, EnsureLimitSettingRule, GroupSharing, RateLimit, RateLimitGroup, RateLimitKind, + RateLimitScopeResolver, RateLimitUsageResolver, +}; + +#[cfg(feature = "runtime-benchmarks")] +mod benchmarking; +mod tx_extension; +mod types; + +#[cfg(test)] +mod mock; + +#[cfg(test)] +mod tests; + +#[frame_support::pallet] +pub mod pallet { + use codec::Codec; + use frame_support::{ + BoundedBTreeSet, BoundedVec, + pallet_prelude::*, + traits::{BuildGenesisConfig, EnsureOrigin, GetCallMetadata}, + }; + use frame_system::pallet_prelude::*; + use sp_runtime::traits::{ + AtLeast32BitUnsigned, DispatchOriginOf, Dispatchable, Member, One, Saturating, Zero, + }; + use sp_std::{boxed::Box, convert::TryFrom, marker::PhantomData, vec::Vec}; + + #[cfg(feature = "runtime-benchmarks")] + use crate::benchmarking::BenchmarkHelper as BenchmarkHelperTrait; + use crate::types::{ + BypassDecision, EnsureLimitSettingRule, GroupSharing, RateLimit, RateLimitGroup, + RateLimitKind, RateLimitScopeResolver, RateLimitTarget, RateLimitUsageResolver, + TransactionIdentifier, + }; + + type GroupNameOf = BoundedVec>::MaxGroupNameLength>; + type GroupMembersOf = + BoundedBTreeSet>::MaxGroupMembers>; + type GroupDetailsOf = RateLimitGroup<>::GroupId, GroupNameOf>; + + /// Configuration trait for the rate limiting pallet. + #[pallet::config] + pub trait Config: frame_system::Config + where + BlockNumberFor: MaybeSerializeDeserialize, + <>::RuntimeCall as Dispatchable>::RuntimeOrigin: + From<::RuntimeOrigin>, + { + /// The overarching runtime call type. + type RuntimeCall: Parameter + + Codec + + GetCallMetadata + + Dispatchable + + IsType<::RuntimeCall>; + + /// Origin permitted to configure rate limits. + type AdminOrigin: EnsureOrigin>; + + /// Rule type that decides which origins may call [`Pallet::set_rate_limit`]. + type LimitSettingRule: Parameter + Member + MaxEncodedLen + MaybeSerializeDeserialize; + + /// Default rule applied when a target does not have an explicit entry in + /// [`LimitSettingRules`]. + type DefaultLimitSettingRule: Get; + + /// Origin checker invoked when setting a rate limit, parameterized by the stored rule. + type LimitSettingOrigin: EnsureLimitSettingRule, Self::LimitSettingRule, Self::LimitScope>; + + /// Scope identifier used to namespace stored rate limits. + type LimitScope: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; + + /// Resolves the scope for the given runtime call when configuring limits. + type LimitScopeResolver: RateLimitScopeResolver< + DispatchOriginOf<>::RuntimeCall>, + >::RuntimeCall, + Self::LimitScope, + BlockNumberFor, + >; + + /// Usage key tracked in [`LastSeen`] for rate-limited calls. + type UsageKey: Parameter + Clone + PartialEq + Eq + Ord + MaybeSerializeDeserialize; + + /// Resolves the usage key for the given runtime call when enforcing limits. + type UsageResolver: RateLimitUsageResolver< + DispatchOriginOf<>::RuntimeCall>, + >::RuntimeCall, + Self::UsageKey, + >; + + /// Identifier assigned to managed groups. + type GroupId: Parameter + + Member + + Copy + + MaybeSerializeDeserialize + + MaxEncodedLen + + AtLeast32BitUnsigned + + Default; + + /// Maximum number of extrinsics that may belong to a single group. + #[pallet::constant] + type MaxGroupMembers: Get; + + /// Maximum length (in bytes) of a group name. + #[pallet::constant] + type MaxGroupNameLength: Get; + + /// Helper used to construct runtime calls for benchmarking. + #[cfg(feature = "runtime-benchmarks")] + type BenchmarkHelper: BenchmarkHelperTrait<>::RuntimeCall>; + } + + /// Storage mapping from rate limit target to its configured rate limit. + #[pallet::storage] + #[pallet::getter(fn limits)] + pub type Limits, I: 'static = ()> = StorageMap< + _, + Blake2_128Concat, + RateLimitTarget<>::GroupId>, + RateLimit<>::LimitScope, BlockNumberFor>, + OptionQuery, + >; + + #[pallet::type_value] + pub fn DefaultLimitSettingRuleFor, I: 'static>() -> T::LimitSettingRule { + T::DefaultLimitSettingRule::get() + } + + /// Stores the rule used to authorize [`Pallet::set_rate_limit`] per call/group target. + #[pallet::storage] + #[pallet::getter(fn limit_setting_rule)] + pub type LimitSettingRules, I: 'static = ()> = StorageMap< + _, + Blake2_128Concat, + RateLimitTarget<>::GroupId>, + >::LimitSettingRule, + ValueQuery, + DefaultLimitSettingRuleFor, + >; + + /// Tracks when a rate-limited target was last observed per usage key. + #[pallet::storage] + pub type LastSeen, I: 'static = ()> = StorageDoubleMap< + _, + Blake2_128Concat, + RateLimitTarget<>::GroupId>, + Blake2_128Concat, + Option<>::UsageKey>, + BlockNumberFor, + OptionQuery, + >; + + /// Default block span applied when an extrinsic uses the default rate limit. + #[pallet::storage] + #[pallet::getter(fn default_limit)] + pub type DefaultLimit, I: 'static = ()> = + StorageValue<_, BlockNumberFor, ValueQuery>; + + /// Maps a transaction identifier to its assigned group. + #[pallet::storage] + #[pallet::getter(fn call_group)] + pub type CallGroups, I: 'static = ()> = StorageMap< + _, + Blake2_128Concat, + TransactionIdentifier, + >::GroupId, + OptionQuery, + >; + + /// Tracks whether a grouped call should skip writing usage metadata on success. + #[pallet::storage] + #[pallet::getter(fn call_read_only)] + pub type CallReadOnly, I: 'static = ()> = + StorageMap<_, Blake2_128Concat, TransactionIdentifier, bool, OptionQuery>; + + /// Metadata for each configured group. + #[pallet::storage] + #[pallet::getter(fn groups)] + pub type Groups, I: 'static = ()> = StorageMap< + _, + Blake2_128Concat, + >::GroupId, + GroupDetailsOf, + OptionQuery, + >; + + /// Tracks membership for each group. + #[pallet::storage] + #[pallet::getter(fn group_members)] + pub type GroupMembers, I: 'static = ()> = StorageMap< + _, + Blake2_128Concat, + >::GroupId, + GroupMembersOf, + ValueQuery, + >; + + /// Enforces unique group names. + #[pallet::storage] + #[pallet::getter(fn group_id_by_name)] + pub type GroupNameIndex, I: 'static = ()> = + StorageMap<_, Blake2_128Concat, GroupNameOf, >::GroupId, OptionQuery>; + + /// Identifier used for the next group creation. + #[pallet::storage] + #[pallet::getter(fn next_group_id)] + pub type NextGroupId, I: 'static = ()> = + StorageValue<_, >::GroupId, ValueQuery>; + + /// Events emitted by the rate limiting pallet. + #[pallet::event] + #[pallet::generate_deposit(pub(super) fn deposit_event)] + pub enum Event, I: 'static = ()> { + /// A call was registered for rate limiting. + CallRegistered { + /// Identifier of the registered transaction. + transaction: TransactionIdentifier, + /// Scope seeded during registration (if any). + scope: Option<>::LimitScope>, + /// Optional group assignment applied at registration time. + group: Option<>::GroupId>, + /// Pallet name associated with the transaction. + pallet: Vec, + /// Extrinsic name associated with the transaction. + extrinsic: Vec, + }, + /// A rate limit was set or updated for the specified target. + RateLimitSet { + /// Target whose configuration changed. + target: RateLimitTarget<>::GroupId>, + /// Identifier of the transaction when the target represents a call. + transaction: Option, + /// Limit scope to which the configuration applies, if any. + scope: Option<>::LimitScope>, + /// The rate limit policy applied to the target. + limit: RateLimitKind>, + /// Pallet name associated with the transaction, when available. + pallet: Option>, + /// Extrinsic name associated with the transaction, when available. + extrinsic: Option>, + }, + /// The rule that authorizes [`Pallet::set_rate_limit`] was updated for a target. + LimitSettingRuleUpdated { + /// Target whose limit-setting rule changed. + target: RateLimitTarget<>::GroupId>, + /// Updated rule. + rule: >::LimitSettingRule, + }, + /// A rate-limited call was deregistered or had a scoped entry cleared. + CallDeregistered { + /// Target whose configuration changed. + target: RateLimitTarget<>::GroupId>, + /// Identifier of the transaction when the target represents a call. + transaction: Option, + /// Limit scope from which the configuration was cleared, if any. + scope: Option<>::LimitScope>, + /// Pallet name associated with the transaction, when available. + pallet: Option>, + /// Extrinsic name associated with the transaction, when available. + extrinsic: Option>, + }, + /// The default rate limit was set or updated. + DefaultRateLimitSet { + /// The new default limit expressed in blocks. + block_span: BlockNumberFor, + }, + /// A group was created. + GroupCreated { + /// Identifier of the new group. + group: >::GroupId, + /// Human readable group name. + name: Vec, + /// Sharing policy configured for the group. + sharing: GroupSharing, + }, + /// A group's metadata or policy changed. + GroupUpdated { + /// Identifier of the group. + group: >::GroupId, + /// Human readable name. + name: Vec, + /// Updated sharing configuration. + sharing: GroupSharing, + }, + /// A group was deleted. + GroupDeleted { + /// Identifier of the removed group. + group: >::GroupId, + }, + /// A transaction was assigned to or removed from a group. + CallGroupUpdated { + /// Identifier of the transaction. + transaction: TransactionIdentifier, + /// Updated group assignment (None when cleared). + group: Option<>::GroupId>, + }, + /// A grouped call toggled whether it writes usage after enforcement. + CallReadOnlyUpdated { + /// Identifier of the transaction. + transaction: TransactionIdentifier, + /// Group to which the call belongs. + group: >::GroupId, + /// Current read-only flag. + read_only: bool, + }, + } + + /// Errors that can occur while configuring rate limits. + #[pallet::error] + pub enum Error { + /// Failed to extract the pallet and extrinsic indices from the call. + InvalidRuntimeCall, + /// Attempted to remove a limit that is not present. + MissingRateLimit, + /// Group metadata was not found. + UnknownGroup, + /// Attempted to create or rename a group to an existing name. + DuplicateGroupName, + /// Group name exceeds the configured maximum length. + GroupNameTooLong, + /// Operation requires the group to have no members. + GroupHasMembers, + /// Adding a member would exceed the configured limit. + GroupMemberLimitExceeded, + /// Call already belongs to the requested group. + CallAlreadyInGroup, + /// Call is not assigned to a group. + CallNotInGroup, + /// Operation requires the call to be registered first. + CallNotRegistered, + /// Attempted to register a call that already exists. + CallAlreadyRegistered, + /// Rate limit for this call must be configured via its group target. + MustTargetGroup, + /// Resolver failed to supply a required context value. + MissingScope, + /// Group cannot be removed because configuration or usage entries remain. + GroupInUse, + } + + #[pallet::genesis_config] + pub struct GenesisConfig, I: 'static = ()> { + pub default_limit: BlockNumberFor, + pub limits: Vec<( + RateLimitTarget<>::GroupId>, + Option<>::LimitScope>, + RateLimitKind>, + )>, + pub groups: Vec<(>::GroupId, Vec, GroupSharing)>, + } + + impl, I: 'static> Default for GenesisConfig { + fn default() -> Self { + Self { + default_limit: Zero::zero(), + limits: Vec::new(), + groups: Vec::new(), + } + } + } + + #[pallet::genesis_build] + impl, I: 'static> BuildGenesisConfig for GenesisConfig { + fn build(&self) { + DefaultLimit::::put(self.default_limit); + + // Seed groups first so limit targets can reference them. + let mut max_group: >::GroupId = Zero::zero(); + for (group_id, name, sharing) in &self.groups { + let bounded = GroupNameOf::::try_from(name.clone()) + .expect("Genesis group name exceeds MaxGroupNameLength"); + + assert!( + !Groups::::contains_key(group_id), + "Duplicate group id in genesis config" + ); + assert!( + !GroupNameIndex::::contains_key(&bounded), + "Duplicate group name in genesis config" + ); + + Groups::::insert( + group_id, + RateLimitGroup { + id: *group_id, + name: bounded.clone(), + sharing: *sharing, + }, + ); + GroupNameIndex::::insert(&bounded, *group_id); + GroupMembers::::insert(*group_id, GroupMembersOf::::new()); + if *group_id > max_group { + max_group = *group_id; + } + } + let next = max_group.saturating_add(One::one()); + NextGroupId::::put(next); + + for (identifier, scope, kind) in &self.limits { + if let RateLimitTarget::Group(group) = identifier { + assert!( + Groups::::contains_key(group), + "Genesis limit references unknown group" + ); + } + let target = *identifier; + Limits::::mutate(target, |entry| match scope { + None => { + *entry = Some(RateLimit::global(*kind)); + } + Some(sc) => { + if let Some(config) = entry { + config.upsert_scope(sc.clone(), *kind); + } else { + *entry = Some(RateLimit::scoped_single(sc.clone(), *kind)); + } + } + }); + } + } + } + + #[pallet::pallet] + #[pallet::without_storage_info] + pub struct Pallet(PhantomData<(T, I)>); + + impl, I: 'static> Pallet { + /// Returns `true` when the given transaction identifier passes its configured rate limit + /// within the provided usage scope. + pub fn is_within_limit( + origin: &DispatchOriginOf<>::RuntimeCall>, + call: &>::RuntimeCall, + identifier: &TransactionIdentifier, + scope: &Option<>::LimitScope>, + usage_key: &Option<>::UsageKey>, + ) -> Result { + let bypass: BypassDecision = + >::LimitScopeResolver::should_bypass(origin, call); + if bypass.bypass_enforcement { + return Ok(true); + } + + let target = Self::config_target(identifier)?; + Self::ensure_scope_available(&target, scope)?; + + let Some(block_span) = Self::effective_span(origin, call, &target, scope) else { + return Ok(true); + }; + + let usage_target = Self::usage_target(identifier)?; + Ok(Self::within_span(&usage_target, usage_key, block_span)) + } + + /// Resolves the configured span for the provided target/scope, applying the pallet default + /// when the stored value uses [`RateLimitKind::Default`]. + pub fn resolved_limit( + target: &RateLimitTarget<>::GroupId>, + scope: &Option<>::LimitScope>, + ) -> Option> { + let config = Limits::::get(target)?; + let kind = config.kind_for(scope.as_ref())?; + Some(match *kind { + RateLimitKind::Default => DefaultLimit::::get(), + RateLimitKind::Exact(block_span) => block_span, + }) + } + + /// Resolves the span for a target/scope and applies the configured span adjustment (e.g., + /// tempo scaling) using the pallet's scope resolver. + pub fn effective_span( + origin: &DispatchOriginOf<>::RuntimeCall>, + call: &>::RuntimeCall, + target: &RateLimitTarget<>::GroupId>, + scope: &Option<>::LimitScope>, + ) -> Option> { + let span = Self::resolved_limit(target, scope)?; + Some(>::LimitScopeResolver::adjust_span( + origin, call, span, + )) + } + + pub(crate) fn within_span( + target: &RateLimitTarget<>::GroupId>, + usage_key: &Option<>::UsageKey>, + block_span: BlockNumberFor, + ) -> bool { + if block_span.is_zero() { + return true; + } + + if let Some(last) = LastSeen::::get(target, usage_key) { + let current = frame_system::Pallet::::block_number(); + let delta = current.saturating_sub(last); + if delta < block_span { + return false; + } + } + + true + } + + pub(crate) fn should_record_usage( + identifier: &TransactionIdentifier, + usage_target: &RateLimitTarget<>::GroupId>, + ) -> bool { + match usage_target { + RateLimitTarget::Group(_) => { + !CallReadOnly::::get(identifier).unwrap_or(false) + } + RateLimitTarget::Transaction(_) => true, + } + } + + /// Inserts or updates the cached usage timestamp for a rate-limited call. + /// + /// This is primarily intended for migrations that need to hydrate the new tracking storage + /// from legacy pallets. + pub fn record_last_seen( + target: RateLimitTarget<>::GroupId>, + usage_key: Option<>::UsageKey>, + block_number: BlockNumberFor, + ) { + LastSeen::::insert(target, usage_key, block_number); + } + + /// Migrates a stored rate limit configuration from one scope to another. + /// + /// Returns `true` when an entry was moved. Passing identical `from`/`to` scopes simply + /// checks that a configuration exists. + pub fn migrate_limit_scope( + target: RateLimitTarget<>::GroupId>, + from: Option<>::LimitScope>, + to: Option<>::LimitScope>, + ) -> bool { + if from == to { + return Limits::::contains_key(target); + } + + let mut migrated = false; + Limits::::mutate(target, |maybe_config| { + if let Some(config) = maybe_config { + match (from.as_ref(), to.as_ref()) { + (None, Some(target)) => { + if let RateLimit::Global(kind) = config { + *config = RateLimit::scoped_single(target.clone(), *kind); + migrated = true; + } + } + (Some(source), Some(target)) => { + if let RateLimit::Scoped(map) = config { + if let Some(kind) = map.remove(source) { + map.insert(target.clone(), kind); + migrated = true; + } + } + } + (Some(source), None) => { + if let RateLimit::Scoped(map) = config { + if map.len() == 1 && map.contains_key(source) { + if let Some(kind) = map.remove(source) { + *config = RateLimit::global(kind); + migrated = true; + } + } + } + } + _ => {} + } + } + }); + + migrated + } + + /// Migrates the cached usage information for a rate-limited call to a new key. + /// + /// Returns `true` when an entry was moved. Passing identical keys simply checks that an + /// entry exists. + pub fn migrate_usage_key( + target: RateLimitTarget<>::GroupId>, + from: Option<>::UsageKey>, + to: Option<>::UsageKey>, + ) -> bool { + if from == to { + return LastSeen::::contains_key(target, to); + } + + let Some(block) = LastSeen::::take(target, from) else { + return false; + }; + + LastSeen::::insert(target, to, block); + true + } + + /// Returns the configured limit for the specified pallet/extrinsic names, if any. + pub fn limit_for_call_names( + pallet_name: &str, + extrinsic_name: &str, + scope: Option<>::LimitScope>, + ) -> Option>> { + let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; + let target = Self::config_target(&identifier).ok()?; + Limits::::get(target).and_then(|config| config.kind_for(scope.as_ref()).copied()) + } + + /// Returns the resolved block span for the specified pallet/extrinsic names, if any. + pub fn resolved_limit_for_call_names( + pallet_name: &str, + extrinsic_name: &str, + scope: Option<>::LimitScope>, + ) -> Option> { + let identifier = Self::identifier_for_call_names(pallet_name, extrinsic_name)?; + let target = Self::config_target(&identifier).ok()?; + Self::resolved_limit(&target, &scope) + } + + /// Looks up the transaction identifier for a pallet/extrinsic name pair. + pub fn identifier_for_call_names( + pallet_name: &str, + extrinsic_name: &str, + ) -> Option { + let modules = >::RuntimeCall::get_module_names(); + let pallet_pos = modules.iter().position(|name| *name == pallet_name)?; + let call_names = >::RuntimeCall::get_call_names(pallet_name); + let extrinsic_pos = call_names.iter().position(|name| *name == extrinsic_name)?; + let pallet_index = u8::try_from(pallet_pos).ok()?; + let extrinsic_index = u8::try_from(extrinsic_pos).ok()?; + Some(TransactionIdentifier::new(pallet_index, extrinsic_index)) + } + + fn ensure_call_registered(identifier: &TransactionIdentifier) -> DispatchResult { + let target = RateLimitTarget::Transaction(*identifier); + ensure!( + Limits::::contains_key(target), + Error::::CallNotRegistered + ); + Ok(()) + } + + fn ensure_call_unregistered(identifier: &TransactionIdentifier) -> DispatchResult { + let target = RateLimitTarget::Transaction(*identifier); + ensure!( + !Limits::::contains_key(target), + Error::::CallAlreadyRegistered + ); + Ok(()) + } + + /// Returns true when the call has been registered (either directly or via a group). + pub fn is_registered(identifier: &TransactionIdentifier) -> bool { + let tx_target = RateLimitTarget::Transaction(*identifier); + Limits::::contains_key(tx_target) || CallGroups::::contains_key(identifier) + } + + fn call_metadata( + identifier: &TransactionIdentifier, + ) -> Result<(Vec, Vec), DispatchError> { + let (pallet_name, extrinsic_name) = identifier + .names::<>::RuntimeCall>() + .ok_or(Error::::InvalidRuntimeCall)?; + Ok(( + Vec::from(pallet_name.as_bytes()), + Vec::from(extrinsic_name.as_bytes()), + )) + } + + /// Returns the storage target used to store configuration for the provided identifier, + /// respecting any configured group assignment. + pub fn config_target( + identifier: &TransactionIdentifier, + ) -> Result>::GroupId>, DispatchError> { + Self::target_for(identifier, GroupSharing::config_uses_group) + } + + pub(crate) fn usage_target( + identifier: &TransactionIdentifier, + ) -> Result>::GroupId>, DispatchError> { + Self::target_for(identifier, GroupSharing::usage_uses_group) + } + + fn target_for( + identifier: &TransactionIdentifier, + predicate: impl Fn(GroupSharing) -> bool, + ) -> Result>::GroupId>, DispatchError> { + let group = Self::group_assignment(identifier)?; + Ok(Self::target_from_details( + identifier, + group.as_ref(), + predicate, + )) + } + + fn group_assignment( + identifier: &TransactionIdentifier, + ) -> Result>, DispatchError> { + let Some(group) = CallGroups::::get(identifier) else { + return Ok(None); + }; + let details = Self::ensure_group_details(group)?; + Ok(Some(details)) + } + + fn target_from_details( + identifier: &TransactionIdentifier, + details: Option<&GroupDetailsOf>, + predicate: impl Fn(GroupSharing) -> bool, + ) -> RateLimitTarget<>::GroupId> { + if let Some(details) = details { + if predicate(details.sharing) { + return RateLimitTarget::Group(details.id); + } + } + RateLimitTarget::Transaction(*identifier) + } + + fn ensure_group_details( + group: >::GroupId, + ) -> Result, DispatchError> { + Groups::::get(group).ok_or(Error::::UnknownGroup.into()) + } + + fn ensure_scope_available( + target: &RateLimitTarget<>::GroupId>, + scope: &Option<>::LimitScope>, + ) -> Result<(), DispatchError> { + if scope.is_some() { + return Ok(()); + } + + if let Some(RateLimit::Scoped(map)) = Limits::::get(target) { + if !map.is_empty() { + return Err(Error::::MissingScope.into()); + } + } + + Ok(()) + } + + fn bounded_group_name(name: Vec) -> Result, DispatchError> { + GroupNameOf::::try_from(name).map_err(|_| Error::::GroupNameTooLong.into()) + } + + fn ensure_group_name_available( + name: &GroupNameOf, + current: Option<>::GroupId>, + ) -> DispatchResult { + if let Some(existing) = GroupNameIndex::::get(name) { + ensure!(Some(existing) == current, Error::::DuplicateGroupName); + } + Ok(()) + } + + fn ensure_group_deletable(group: >::GroupId) -> DispatchResult { + ensure!( + GroupMembers::::get(group).is_empty(), + Error::::GroupHasMembers + ); + let target = RateLimitTarget::Group(group); + ensure!( + !Limits::::contains_key(target), + Error::::GroupInUse + ); + ensure!( + LastSeen::::iter_prefix(target).next().is_none(), + Error::::GroupInUse + ); + Ok(()) + } + + fn insert_call_into_group( + identifier: &TransactionIdentifier, + group: >::GroupId, + ) -> DispatchResult { + GroupMembers::::try_mutate(group, |members| -> DispatchResult { + match members.try_insert(*identifier) { + Ok(true) => Ok(()), + Ok(false) => Err(Error::::CallAlreadyInGroup.into()), + Err(_) => Err(Error::::GroupMemberLimitExceeded.into()), + } + })?; + Ok(()) + } + + fn detach_call_from_group( + identifier: &TransactionIdentifier, + group: >::GroupId, + ) -> bool { + GroupMembers::::mutate(group, |members| members.remove(identifier)) + } + } + + #[pallet::call] + impl, I: 'static> Pallet { + /// Registers a call for rate limiting and seeds its initial configuration. + #[pallet::call_index(0)] + #[pallet::weight(T::DbWeight::get().reads_writes(3, 3))] + pub fn register_call( + origin: OriginFor, + call: Box<>::RuntimeCall>, + group: Option<>::GroupId>, + ) -> DispatchResult { + let resolver_origin: DispatchOriginOf<>::RuntimeCall> = + Into::>::RuntimeCall>>::into(origin.clone()); + let scope = + >::LimitScopeResolver::context(&resolver_origin, call.as_ref()); + + T::AdminOrigin::ensure_origin(origin)?; + + let identifier = TransactionIdentifier::from_call(call.as_ref()) + .ok_or(Error::::InvalidRuntimeCall)?; + Self::ensure_call_unregistered(&identifier)?; + + let target = RateLimitTarget::Transaction(identifier); + + if let Some(ref sc) = scope { + Limits::::insert( + target, + RateLimit::scoped_single(sc.clone(), RateLimitKind::Default), + ); + } else { + Limits::::insert(target, RateLimit::global(RateLimitKind::Default)); + } + + let mut assigned_group = None; + if let Some(group_id) = group { + Self::ensure_group_details(group_id)?; + Self::insert_call_into_group(&identifier, group_id)?; + CallGroups::::insert(&identifier, group_id); + CallReadOnly::::insert(&identifier, false); + assigned_group = Some(group_id); + } + + let (pallet, extrinsic) = Self::call_metadata(&identifier)?; + Self::deposit_event(Event::CallRegistered { + transaction: identifier, + scope: scope.clone(), + group: assigned_group, + pallet: pallet.clone(), + extrinsic: extrinsic.clone(), + }); + + if let Some(group_id) = assigned_group { + Self::deposit_event(Event::CallGroupUpdated { + transaction: identifier, + group: Some(group_id), + }); + Self::deposit_event(Event::CallReadOnlyUpdated { + transaction: identifier, + group: group_id, + read_only: false, + }); + } + + Ok(()) + } + + /// Configures a rate limit for either a transaction or group target. + #[pallet::call_index(1)] + #[pallet::weight(T::DbWeight::get().reads_writes(2, 2))] + pub fn set_rate_limit( + origin: OriginFor, + target: RateLimitTarget<>::GroupId>, + scope: Option<>::LimitScope>, + limit: RateLimitKind>, + ) -> DispatchResult { + let rule = LimitSettingRules::::get(&target); + T::LimitSettingOrigin::ensure_origin(origin, &rule, &scope)?; + + let (transaction, pallet, extrinsic) = match target { + RateLimitTarget::Transaction(identifier) => { + Self::ensure_call_registered(&identifier)?; + if let Some(group) = CallGroups::::get(&identifier) { + let details = Self::ensure_group_details(group)?; + ensure!( + !details.sharing.config_uses_group(), + Error::::MustTargetGroup + ); + } + let (pallet, extrinsic) = Self::call_metadata(&identifier)?; + (Some(identifier), Some(pallet), Some(extrinsic)) + } + RateLimitTarget::Group(group) => { + Self::ensure_group_details(group)?; + (None, None, None) + } + }; + + if let Some(ref scoped) = scope { + Limits::::mutate(target, |slot| match slot { + Some(config) => config.upsert_scope(scoped.clone(), limit), + None => *slot = Some(RateLimit::scoped_single(scoped.clone(), limit)), + }); + } else { + Limits::::insert(target, RateLimit::global(limit)); + } + + Self::deposit_event(Event::RateLimitSet { + target, + transaction, + scope, + limit, + pallet, + extrinsic, + }); + Ok(()) + } + + /// Sets the rule used to authorize [`Pallet::set_rate_limit`] for the provided target. + #[pallet::call_index(10)] + #[pallet::weight(T::DbWeight::get().reads_writes(2, 1))] + pub fn set_limit_setting_rule( + origin: OriginFor, + target: RateLimitTarget<>::GroupId>, + rule: >::LimitSettingRule, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + match target { + RateLimitTarget::Transaction(identifier) => { + Self::ensure_call_registered(&identifier)?; + } + RateLimitTarget::Group(group) => { + Self::ensure_group_details(group)?; + } + } + + LimitSettingRules::::insert(target, rule.clone()); + Self::deposit_event(Event::LimitSettingRuleUpdated { target, rule }); + + Ok(()) + } + + /// Assigns a registered call to the specified group and optionally marks it as read-only + /// for usage tracking. + #[pallet::call_index(2)] + #[pallet::weight(T::DbWeight::get().reads_writes(3, 3))] + pub fn assign_call_to_group( + origin: OriginFor, + transaction: TransactionIdentifier, + group: >::GroupId, + read_only: bool, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + Self::ensure_call_registered(&transaction)?; + Self::ensure_group_details(group)?; + + let current = CallGroups::::get(&transaction); + ensure!(current.is_none(), Error::::CallAlreadyInGroup); + Self::insert_call_into_group(&transaction, group)?; + CallGroups::::insert(&transaction, group); + CallReadOnly::::insert(&transaction, read_only); + + Self::deposit_event(Event::CallGroupUpdated { + transaction, + group: Some(group), + }); + Self::deposit_event(Event::CallReadOnlyUpdated { + transaction, + group, + read_only, + }); + + Ok(()) + } + + /// Removes a registered call from its current group assignment. + #[pallet::call_index(3)] + #[pallet::weight(T::DbWeight::get().reads_writes(2, 2))] + pub fn remove_call_from_group( + origin: OriginFor, + transaction: TransactionIdentifier, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + Self::ensure_call_registered(&transaction)?; + let Some(group) = CallGroups::::take(&transaction) else { + return Err(Error::::CallNotInGroup.into()); + }; + CallReadOnly::::remove(&transaction); + Self::detach_call_from_group(&transaction, group); + + Self::deposit_event(Event::CallGroupUpdated { + transaction, + group: None, + }); + + Ok(()) + } + + /// Sets the default rate limit that applies when an extrinsic uses [`RateLimitKind::Default`]. + #[pallet::call_index(4)] + #[pallet::weight(T::DbWeight::get().writes(1))] + pub fn set_default_rate_limit( + origin: OriginFor, + block_span: BlockNumberFor, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + DefaultLimit::::put(block_span); + Self::deposit_event(Event::DefaultRateLimitSet { block_span }); + Ok(()) + } + + /// Creates a new rate-limiting group with the provided name and sharing configuration. + #[pallet::call_index(5)] + #[pallet::weight(T::DbWeight::get().reads_writes(1, 3))] + pub fn create_group( + origin: OriginFor, + name: Vec, + sharing: GroupSharing, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + let bounded = Self::bounded_group_name(name)?; + Self::ensure_group_name_available(&bounded, None)?; + + let group = NextGroupId::::mutate(|current| { + let next = current.saturating_add(One::one()); + sp_std::mem::replace(current, next) + }); + + Groups::::insert( + group, + RateLimitGroup { + id: group, + name: bounded.clone(), + sharing, + }, + ); + GroupNameIndex::::insert(&bounded, group); + GroupMembers::::insert(group, GroupMembersOf::::new()); + + let name_bytes: Vec = bounded.into(); + Self::deposit_event(Event::GroupCreated { + group, + name: name_bytes, + sharing, + }); + Ok(()) + } + + /// Updates the metadata or sharing configuration of an existing group. + #[pallet::call_index(6)] + #[pallet::weight(T::DbWeight::get().reads_writes(3, 3))] + pub fn update_group( + origin: OriginFor, + group: >::GroupId, + name: Option>, + sharing: Option, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + Groups::::try_mutate(group, |maybe_details| -> DispatchResult { + let details = maybe_details.as_mut().ok_or(Error::::UnknownGroup)?; + + if let Some(new_name) = name { + let bounded = Self::bounded_group_name(new_name)?; + Self::ensure_group_name_available(&bounded, Some(group))?; + GroupNameIndex::::remove(&details.name); + GroupNameIndex::::insert(&bounded, group); + details.name = bounded; + } + + if let Some(new_sharing) = sharing { + details.sharing = new_sharing; + } + + Ok(()) + })?; + + let updated = Self::ensure_group_details(group)?; + let name_bytes: Vec = updated.name.clone().into(); + Self::deposit_event(Event::GroupUpdated { + group, + name: name_bytes, + sharing: updated.sharing, + }); + + Ok(()) + } + + /// Deletes an existing group. The group must be empty and unused. + #[pallet::call_index(7)] + #[pallet::weight(T::DbWeight::get().reads_writes(3, 3))] + pub fn delete_group( + origin: OriginFor, + group: >::GroupId, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + Self::ensure_group_deletable(group)?; + + let details = Groups::::take(group).ok_or(Error::::UnknownGroup)?; + GroupNameIndex::::remove(&details.name); + GroupMembers::::remove(group); + + Self::deposit_event(Event::GroupDeleted { group }); + + Ok(()) + } + + /// Deregisters a call or removes a scoped entry from its configuration. + #[pallet::call_index(8)] + #[pallet::weight(T::DbWeight::get().reads_writes(4, 4))] + pub fn deregister_call( + origin: OriginFor, + transaction: TransactionIdentifier, + scope: Option<>::LimitScope>, + clear_usage: bool, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + Self::ensure_call_registered(&transaction)?; + let target = Self::config_target(&transaction)?; + let tx_target = RateLimitTarget::Transaction(transaction); + let usage_target = Self::usage_target(&transaction)?; + + match &scope { + Some(sc) => { + let mut removed = false; + Limits::::mutate_exists(target, |maybe_config| { + if let Some(RateLimit::Scoped(map)) = maybe_config { + if map.remove(sc).is_some() { + removed = true; + if map.is_empty() { + *maybe_config = None; + } + } + } + }); + ensure!(removed, Error::::MissingRateLimit); + + if let Some(group) = CallGroups::::take(&transaction) { + CallReadOnly::::remove(&transaction); + Self::detach_call_from_group(&transaction, group); + Self::deposit_event(Event::CallGroupUpdated { + transaction, + group: None, + }); + } + } + None => { + Limits::::remove(target); + if target != tx_target { + Limits::::remove(tx_target); + } + + if let Some(group) = CallGroups::::take(&transaction) { + CallReadOnly::::remove(&transaction); + Self::detach_call_from_group(&transaction, group); + Self::deposit_event(Event::CallGroupUpdated { + transaction, + group: None, + }); + } + } + } + + if clear_usage { + let _ = LastSeen::::clear_prefix(&usage_target, u32::MAX, None); + } + + let (pallet, extrinsic) = Self::call_metadata(&transaction)?; + Self::deposit_event(Event::CallDeregistered { + target, + transaction: Some(transaction), + scope, + pallet: Some(pallet), + extrinsic: Some(extrinsic), + }); + + Ok(()) + } + + /// Updates whether a grouped call should skip writing usage metadata after enforcement. + /// + /// The call must already be assigned to a group. + #[pallet::call_index(9)] + #[pallet::weight(T::DbWeight::get().reads_writes(2, 1))] + pub fn set_call_read_only( + origin: OriginFor, + transaction: TransactionIdentifier, + read_only: bool, + ) -> DispatchResult { + T::AdminOrigin::ensure_origin(origin)?; + + Self::ensure_call_registered(&transaction)?; + let group = + CallGroups::::get(&transaction).ok_or(Error::::CallNotInGroup)?; + CallReadOnly::::insert(&transaction, read_only); + + Self::deposit_event(Event::CallReadOnlyUpdated { + transaction, + group, + read_only, + }); + + Ok(()) + } + } +} + +impl, I: 'static> RateLimitingInfo for pallet::Pallet { + type GroupId = >::GroupId; + type CallMetadata = >::RuntimeCall; + type Limit = frame_system::pallet_prelude::BlockNumberFor; + type Scope = >::LimitScope; + type UsageKey = >::UsageKey; + + fn rate_limit(target: TargetArg, scope: Option) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + let raw_target = target + .try_into_rate_limit_target::() + .ok()?; + let config_target = match raw_target { + // A transaction identifier may be assigned to a group; resolve the effective storage + // target. + RateLimitTarget::Transaction(identifier) => Self::config_target(&identifier).ok()?, + _ => raw_target, + }; + Self::resolved_limit(&config_target, &scope) + } + + fn last_seen( + target: TargetArg, + usage_key: Option, + ) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + let raw_target = target + .try_into_rate_limit_target::() + .ok()?; + let usage_target = match raw_target { + // A transaction identifier may be assigned to a group; resolve the effective storage + // target. + RateLimitTarget::Transaction(identifier) => Self::usage_target(&identifier).ok()?, + _ => raw_target, + }; + pallet::LastSeen::::get(usage_target, usage_key) + } +} diff --git a/pallets/rate-limiting/src/mock.rs b/pallets/rate-limiting/src/mock.rs new file mode 100644 index 0000000000..16e470be3e --- /dev/null +++ b/pallets/rate-limiting/src/mock.rs @@ -0,0 +1,211 @@ +#![allow(dead_code)] + +use core::convert::TryInto; + +use frame_support::{ + derive_impl, + dispatch::DispatchResult, + sp_runtime::{ + BuildStorage, + traits::{BlakeTwo256, IdentityLookup}, + }, + traits::{ConstU16, ConstU32, ConstU64, EnsureOrigin, Everything}, +}; +use frame_system::{EnsureRoot, ensure_signed}; +use serde::{Deserialize, Serialize}; +use sp_core::H256; +use sp_io::TestExternalities; +use sp_std::vec::Vec; + +use crate as pallet_rate_limiting; +use crate::TransactionIdentifier; + +pub type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic; +pub type Block = frame_system::mocking::MockBlock; + +frame_support::construct_runtime!( + pub enum Test { + System: frame_system, + RateLimiting: pallet_rate_limiting, + } +); + +#[derive_impl(frame_system::config_preludes::TestDefaultConfig)] +impl frame_system::Config for Test { + type BaseCallFilter = Everything; + type BlockWeights = (); + type BlockLength = (); + type DbWeight = (); + type RuntimeOrigin = RuntimeOrigin; + type RuntimeCall = RuntimeCall; + type Nonce = u64; + type Hash = H256; + type Hashing = BlakeTwo256; + type AccountId = u64; + type Lookup = IdentityLookup; + type RuntimeEvent = RuntimeEvent; + type BlockHashCount = ConstU64<250>; + type Version = (); + type PalletInfo = PalletInfo; + type AccountData = (); + type OnNewAccount = (); + type OnKilledAccount = (); + type SystemWeightInfo = (); + type SS58Prefix = ConstU16<42>; + type OnSetCode = (); + type MaxConsumers = ConstU32<16>; + type Block = Block; +} + +pub type LimitScope = u16; +pub type UsageKey = u16; +pub type GroupId = u32; + +#[derive( + codec::Encode, + codec::Decode, + codec::DecodeWithMemTracking, + Serialize, + Deserialize, + Clone, + Copy, + PartialEq, + Eq, + scale_info::TypeInfo, + codec::MaxEncodedLen, + Debug, +)] +pub enum LimitSettingRule { + RootOnly, + AnySigned, +} + +frame_support::parameter_types! { + pub const DefaultLimitSettingRule: LimitSettingRule = LimitSettingRule::RootOnly; +} + +pub struct LimitSettingOrigin; + +impl pallet_rate_limiting::EnsureLimitSettingRule + for LimitSettingOrigin +{ + fn ensure_origin( + origin: RuntimeOrigin, + rule: &LimitSettingRule, + _scope: &Option, + ) -> DispatchResult { + match rule { + LimitSettingRule::RootOnly => EnsureRoot::::ensure_origin(origin) + .map(|_| ()) + .map_err(Into::into), + LimitSettingRule::AnySigned => { + let _ = ensure_signed(origin)?; + Ok(()) + } + } + } +} + +pub struct TestScopeResolver; +pub struct TestUsageResolver; + +impl pallet_rate_limiting::RateLimitScopeResolver + for TestScopeResolver +{ + fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option { + match call { + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { + (*block_span).try_into().ok() + } + RuntimeCall::RateLimiting(_) => Some(1), + _ => None, + } + } + + fn should_bypass( + _origin: &RuntimeOrigin, + call: &RuntimeCall, + ) -> pallet_rate_limiting::types::BypassDecision { + match call { + RuntimeCall::RateLimiting(RateLimitingCall::remove_call_from_group { .. }) => { + pallet_rate_limiting::types::BypassDecision::bypass_and_skip() + } + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { .. }) => { + pallet_rate_limiting::types::BypassDecision::bypass_and_record() + } + _ => pallet_rate_limiting::types::BypassDecision::enforce_and_record(), + } + } + + fn adjust_span(_origin: &RuntimeOrigin, call: &RuntimeCall, span: u64) -> u64 { + if matches!( + call, + RuntimeCall::RateLimiting(RateLimitingCall::deregister_call { .. }) + ) { + span.saturating_mul(2) + } else { + span + } + } +} + +impl pallet_rate_limiting::RateLimitUsageResolver + for TestUsageResolver +{ + fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option> { + match call { + RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span }) => { + (*block_span).try_into().ok().map(|key| vec![key]) + } + RuntimeCall::RateLimiting(_) => Some(vec![1]), + _ => None, + } + } +} + +impl pallet_rate_limiting::Config for Test { + type RuntimeCall = RuntimeCall; + type LimitScope = LimitScope; + type LimitScopeResolver = TestScopeResolver; + type UsageKey = UsageKey; + type UsageResolver = TestUsageResolver; + type AdminOrigin = EnsureRoot; + type LimitSettingRule = LimitSettingRule; + type DefaultLimitSettingRule = DefaultLimitSettingRule; + type LimitSettingOrigin = LimitSettingOrigin; + type GroupId = GroupId; + type MaxGroupMembers = ConstU32<32>; + type MaxGroupNameLength = ConstU32<64>; + #[cfg(feature = "runtime-benchmarks")] + type BenchmarkHelper = BenchHelper; +} + +#[cfg(feature = "runtime-benchmarks")] +pub struct BenchHelper; + +#[cfg(feature = "runtime-benchmarks")] +impl crate::BenchmarkHelper for BenchHelper { + fn sample_call() -> RuntimeCall { + RuntimeCall::System(frame_system::Call::remark { remark: Vec::new() }) + } +} + +pub type RateLimitingCall = crate::Call; + +pub fn new_test_ext() -> TestExternalities { + let storage = frame_system::GenesisConfig::::default() + .build_storage() + .expect("genesis build succeeds"); + + let mut ext = TestExternalities::new(storage); + ext.execute_with(|| System::set_block_number(1)); + ext +} + +pub(crate) fn identifier_for(call: &RuntimeCall) -> TransactionIdentifier { + TransactionIdentifier::from_call(call).expect("identifier for call") +} + +pub(crate) fn pop_last_event() -> RuntimeEvent { + System::events().pop().expect("event expected").event +} diff --git a/pallets/rate-limiting/src/tests.rs b/pallets/rate-limiting/src/tests.rs new file mode 100644 index 0000000000..5fc79a9362 --- /dev/null +++ b/pallets/rate-limiting/src/tests.rs @@ -0,0 +1,766 @@ +use frame_support::{assert_noop, assert_ok}; +use sp_std::vec::Vec; + +use crate::{ + CallGroups, CallReadOnly, Config, GroupMembers, GroupSharing, LastSeen, LimitSettingRules, + Limits, RateLimit, RateLimitKind, RateLimitTarget, TransactionIdentifier, mock::*, + pallet::Error, +}; +use frame_support::traits::Get; + +fn target(identifier: TransactionIdentifier) -> RateLimitTarget { + RateLimitTarget::Transaction(identifier) +} + +fn remark_call() -> RuntimeCall { + RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }) +} + +fn scoped_call() -> RuntimeCall { + RuntimeCall::RateLimiting(RateLimitingCall::set_rate_limit { + target: RateLimitTarget::Transaction(TransactionIdentifier::new(0, 0)), + scope: Some(1), + limit: RateLimitKind::Default, + }) +} + +fn register(call: RuntimeCall, group: Option) -> TransactionIdentifier { + let identifier = identifier_for(&call); + assert_ok!(RateLimiting::register_call( + RuntimeOrigin::root(), + Box::new(call), + group + )); + identifier +} + +fn create_group(name: &[u8], sharing: GroupSharing) -> GroupId { + assert_ok!(RateLimiting::create_group( + RuntimeOrigin::root(), + name.to_vec(), + sharing, + )); + RateLimiting::next_group_id().saturating_sub(1) +} + +fn last_event() -> RuntimeEvent { + pop_last_event() +} + +#[test] +fn set_rate_limit_respects_limit_setting_rule() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + let tx_target = target(identifier); + + // Default rule is root-only. + assert_noop!( + RateLimiting::set_rate_limit( + RuntimeOrigin::signed(1), + tx_target, + None, + RateLimitKind::Exact(1), + ), + sp_runtime::DispatchError::BadOrigin + ); + + // Root updates the limit-setting rule for this transaction target. + assert_ok!(RateLimiting::set_limit_setting_rule( + RuntimeOrigin::root(), + tx_target, + LimitSettingRule::AnySigned, + )); + + assert_eq!( + LimitSettingRules::::get(tx_target), + LimitSettingRule::AnySigned + ); + + // Now any signed origin may set the limit for this target. + assert_ok!(RateLimiting::set_rate_limit( + RuntimeOrigin::signed(1), + tx_target, + None, + RateLimitKind::Exact(7), + )); + }); +} + +#[test] +fn register_call_seeds_global_limit() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + let tx_target = target(identifier); + let stored = Limits::::get(tx_target).expect("limit"); + assert!(matches!(stored, RateLimit::Global(RateLimitKind::Default))); + + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::CallRegistered { transaction, .. }) + if transaction == identifier + )); + }); +} + +#[test] +fn register_call_seeds_scoped_limit() { + new_test_ext().execute_with(|| { + let identifier = register(scoped_call(), None); + let tx_target = target(identifier); + let stored = Limits::::get(tx_target).expect("limit"); + match stored { + RateLimit::Scoped(map) => { + assert_eq!(map.get(&1u16), Some(&RateLimitKind::Default)); + } + _ => panic!("expected scoped entry"), + } + + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::CallRegistered { transaction, scope, .. }) + if transaction == identifier && scope == Some(1u16) + )); + }); +} + +#[test] +fn set_rate_limit_updates_transaction_target() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + let tx_target = target(identifier); + let limit = RateLimitKind::Exact(9); + assert_ok!(RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + tx_target, + None, + limit, + )); + let stored = Limits::::get(tx_target).expect("limit"); + assert!(matches!(stored, RateLimit::Global(RateLimitKind::Exact(9)))); + + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::RateLimitSet { + target: RateLimitTarget::Transaction(t), + limit: RateLimitKind::Exact(9), + .. + }) if t == identifier + )); + }); +} + +#[test] +fn set_rate_limit_requires_registration_and_group_targeting() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + let target = target(identifier); + + // Unregistered call. + let unknown = TransactionIdentifier::new(99, 0); + assert_noop!( + RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + RateLimitTarget::Transaction(unknown), + None, + RateLimitKind::Exact(1), + ), + Error::::CallNotRegistered + ); + + // Group requires targeting the group. + let group = create_group(b"cfg", GroupSharing::ConfigAndUsage); + assert_ok!(RateLimiting::assign_call_to_group( + RuntimeOrigin::root(), + identifier, + group, + false, + )); + assert_noop!( + RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + target, + None, + RateLimitKind::Exact(2), + ), + Error::::MustTargetGroup + ); + }); +} + +#[test] +fn set_rate_limit_respects_group_config_sharing() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + let group = create_group(b"test", GroupSharing::ConfigAndUsage); + // Consume group creation event to keep ordering predictable. + let created = last_event(); + assert!(matches!( + created, + RuntimeEvent::RateLimiting(crate::Event::GroupCreated { group: g, .. }) if g == group + )); + assert_ok!(RateLimiting::assign_call_to_group( + RuntimeOrigin::root(), + identifier, + group, + false, + )); + let events: Vec<_> = System::events() + .into_iter() + .map(|e| e.event) + .filter(|evt| matches!(evt, RuntimeEvent::RateLimiting(_))) + .collect(); + assert!(events.iter().any(|evt| { + matches!( + evt, + RuntimeEvent::RateLimiting(crate::Event::CallReadOnlyUpdated { + transaction, + group: g, + read_only: false, + }) if *transaction == identifier && *g == group + ) + })); + assert!(events.iter().any(|evt| { + matches!( + evt, + RuntimeEvent::RateLimiting(crate::Event::CallGroupUpdated { + transaction, + group: Some(g), + }) if *transaction == identifier && *g == group + ) + })); + assert_noop!( + RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + RateLimitTarget::Transaction(identifier), + None, + RateLimitKind::Exact(5), + ), + Error::::MustTargetGroup + ); + }); +} + +#[test] +fn assign_and_remove_group_membership() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + let group = create_group(b"team", GroupSharing::UsageOnly); + assert_ok!(RateLimiting::assign_call_to_group( + RuntimeOrigin::root(), + identifier, + group, + false, + )); + assert_eq!(CallGroups::::get(identifier), Some(group)); + assert_eq!(CallReadOnly::::get(identifier), Some(false)); + assert!(GroupMembers::::get(group).contains(&identifier)); + assert_ok!(RateLimiting::remove_call_from_group( + RuntimeOrigin::root(), + identifier, + )); + assert!(CallGroups::::get(identifier).is_none()); + + // Last event should signal removal. + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::CallGroupUpdated { transaction, group: None }) + if transaction == identifier + )); + }); +} + +#[test] +fn set_rate_limit_on_group_updates_storage() { + new_test_ext().execute_with(|| { + let group = create_group(b"grp", GroupSharing::ConfigOnly); + let target = RateLimitTarget::Group(group); + assert_ok!(RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + target, + None, + RateLimitKind::Exact(3), + )); + assert!(matches!( + Limits::::get(target), + Some(RateLimit::Global(RateLimitKind::Exact(3))) + )); + + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::RateLimitSet { + target: RateLimitTarget::Group(g), + limit: RateLimitKind::Exact(3), + .. + }) if g == group + )); + }); +} + +#[test] +fn create_and_delete_group_emit_events() { + new_test_ext().execute_with(|| { + assert_ok!(RateLimiting::create_group( + RuntimeOrigin::root(), + b"ev".to_vec(), + GroupSharing::UsageOnly, + )); + let group = RateLimiting::next_group_id().saturating_sub(1); + let created = last_event(); + assert!(matches!( + created, + RuntimeEvent::RateLimiting(crate::Event::GroupCreated { group: g, .. }) if g == group + )); + + assert_ok!(RateLimiting::delete_group(RuntimeOrigin::root(), group)); + let deleted = last_event(); + assert!(matches!( + deleted, + RuntimeEvent::RateLimiting(crate::Event::GroupDeleted { group: g }) if g == group + )); + }); +} + +#[test] +fn deregister_call_scope_removes_entry() { + new_test_ext().execute_with(|| { + let identifier = register(scoped_call(), None); + let tx_target = target(identifier); + assert_ok!(RateLimiting::set_rate_limit( + RuntimeOrigin::root(), + tx_target, + Some(2u16), + RateLimitKind::Exact(4), + )); + LastSeen::::insert(tx_target, Some(9u16), 10); + assert_ok!(RateLimiting::deregister_call( + RuntimeOrigin::root(), + identifier, + Some(2u16), + false, + )); + match Limits::::get(tx_target) { + Some(RateLimit::Scoped(map)) => { + assert!(map.contains_key(&1u16)); + assert!(!map.contains_key(&2u16)); + } + other => panic!("unexpected config: {:?}", other), + } + // usage remains intact when clear_usage is false + assert_eq!(LastSeen::::get(tx_target, Some(9u16)), Some(10)); + + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::CallDeregistered { + target, + transaction: Some(t), + scope: Some(sc), + .. + }) if target == tx_target && t == identifier && sc == 2u16 + )); + + // No group assigned in this test. + assert!(CallGroups::::get(identifier).is_none()); + }); +} + +#[test] +fn register_call_rejects_duplicates_and_unknown_group() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + // Duplicate should fail. + assert_noop!( + RateLimiting::register_call(RuntimeOrigin::root(), Box::new(remark_call()), None), + Error::::CallAlreadyRegistered + ); + + // Unknown group should fail. + assert_noop!( + RateLimiting::register_call(RuntimeOrigin::root(), Box::new(scoped_call()), Some(99)), + Error::::UnknownGroup + ); + + assert!(Limits::::contains_key(target(identifier))); + }); +} + +#[test] +fn group_name_limits_and_uniqueness_enforced() { + new_test_ext().execute_with(|| { + // Overlong name. + let max_name = <::MaxGroupNameLength as Get>::get() as usize; + let long_name = vec![0u8; max_name + 1]; + assert_noop!( + RateLimiting::create_group(RuntimeOrigin::root(), long_name, GroupSharing::UsageOnly), + Error::::GroupNameTooLong + ); + + // Duplicate names rejected on create and update. + let first = create_group(b"alpha", GroupSharing::UsageOnly); + let second = create_group(b"beta", GroupSharing::UsageOnly); + + assert_noop!( + RateLimiting::create_group( + RuntimeOrigin::root(), + b"alpha".to_vec(), + GroupSharing::UsageOnly + ), + Error::::DuplicateGroupName + ); + + assert_noop!( + RateLimiting::update_group( + RuntimeOrigin::root(), + second, + Some(b"alpha".to_vec()), + None + ), + Error::::DuplicateGroupName + ); + + // Unknown group update. + assert_noop!( + RateLimiting::update_group(RuntimeOrigin::root(), 99, None, None), + Error::::UnknownGroup + ); + + assert_eq!( + RateLimiting::groups(first).unwrap().name.into_inner(), + b"alpha".to_vec() + ); + + // Updating first group emits event. + assert_ok!(RateLimiting::update_group( + RuntimeOrigin::root(), + first, + Some(b"gamma".to_vec()), + None, + )); + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::GroupUpdated { group, .. }) if group == first + )); + }); +} + +#[test] +fn group_member_limit_and_removal_errors() { + new_test_ext().execute_with(|| { + let group = create_group(b"cap", GroupSharing::UsageOnly); + + let max_members = <::MaxGroupMembers as Get>::get(); + GroupMembers::::mutate(group, |members| { + for i in 0..max_members { + let _ = members.try_insert(TransactionIdentifier::new(0, (i + 1) as u8)); + } + }); + + // Next insert should fail. + let extra = register(remark_call(), None); + assert_noop!( + RateLimiting::assign_call_to_group(RuntimeOrigin::root(), extra, group, false), + Error::::GroupMemberLimitExceeded + ); + + // Removing a call not in a group errors. + assert_noop!( + RateLimiting::remove_call_from_group(RuntimeOrigin::root(), extra), + Error::::CallNotInGroup + ); + }); +} + +#[test] +fn set_call_read_only_requires_group() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + assert_noop!( + RateLimiting::set_call_read_only(RuntimeOrigin::root(), identifier, true), + Error::::CallNotInGroup + ); + }); +} + +#[test] +fn set_call_read_only_updates_assignment_and_emits_event() { + new_test_ext().execute_with(|| { + let group = create_group(b"ro", GroupSharing::UsageOnly); + let identifier = register(remark_call(), None); + assert_ok!(RateLimiting::assign_call_to_group( + RuntimeOrigin::root(), + identifier, + group, + false, + )); + + assert_ok!(RateLimiting::set_call_read_only( + RuntimeOrigin::root(), + identifier, + true + )); + + assert_eq!(CallGroups::::get(identifier), Some(group)); + assert_eq!(CallReadOnly::::get(identifier), Some(true)); + + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::CallReadOnlyUpdated { + transaction, + group: g, + read_only: true, + }) if transaction == identifier && g == group + )); + }); +} + +#[test] +fn cannot_delete_group_in_use_or_unknown() { + new_test_ext().execute_with(|| { + let group = create_group(b"busy", GroupSharing::ConfigOnly); + let identifier = register(remark_call(), Some(group)); + let target = RateLimitTarget::Group(group); + Limits::::insert(target, RateLimit::global(RateLimitKind::Exact(1))); + LastSeen::::insert(target, None::, 10); + + // Remove member so only config/usage keep the group in-use. + assert_ok!(RateLimiting::remove_call_from_group( + RuntimeOrigin::root(), + identifier + )); + + // Cannot delete when in use. + assert_noop!( + RateLimiting::delete_group(RuntimeOrigin::root(), group), + Error::::GroupInUse + ); + + // Clear state then delete. + Limits::::remove(target); + let _ = LastSeen::::clear_prefix(&target, u32::MAX, None); + assert_ok!(RateLimiting::delete_group(RuntimeOrigin::root(), group)); + + // Unknown group. + assert_noop!( + RateLimiting::delete_group(RuntimeOrigin::root(), 999), + Error::::UnknownGroup + ); + }); +} + +#[test] +fn deregister_call_clears_registration() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + let tx_target = target(identifier); + LastSeen::::insert(tx_target, None::, 5); + assert_ok!(RateLimiting::deregister_call( + RuntimeOrigin::root(), + identifier, + None, + true, + )); + assert!(Limits::::get(tx_target).is_none()); + assert!(LastSeen::::get(tx_target, None::).is_none()); + assert!(CallGroups::::get(identifier).is_none()); + + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::CallDeregistered { + target, + transaction: Some(t), + scope: None, + .. + }) if target == tx_target && t == identifier + )); + }); +} + +#[test] +fn deregister_errors_for_unknown_or_missing_scope() { + new_test_ext().execute_with(|| { + let unknown = TransactionIdentifier::new(10, 1); + assert_noop!( + RateLimiting::deregister_call(RuntimeOrigin::root(), unknown, None, true), + Error::::CallNotRegistered + ); + + let identifier = register(scoped_call(), None); + let tx_target = target(identifier); + // Removing a non-existent scoped entry fails. + assert_noop!( + RateLimiting::deregister_call(RuntimeOrigin::root(), identifier, Some(99u16), false), + Error::::MissingRateLimit + ); + + // Removing the last scoped entry clears Limits and LastSeen. + LastSeen::::insert(tx_target, Some(1u16), 5); + assert_ok!(RateLimiting::deregister_call( + RuntimeOrigin::root(), + identifier, + Some(1u16), + true, + )); + assert!(Limits::::get(tx_target).is_none()); + assert!(LastSeen::::get(tx_target, Some(1u16)).is_none()); + }); +} + +#[test] +fn is_within_limit_detects_rate_limited_scope() { + new_test_ext().execute_with(|| { + let call = scoped_call(); + let identifier = identifier_for(&call); + let tx_target = target(identifier); + Limits::::insert( + tx_target, + RateLimit::scoped_single(1u16, RateLimitKind::Exact(3)), + ); + LastSeen::::insert(tx_target, Some(1u16), 9); + System::set_block_number(11); + let result = RateLimiting::is_within_limit( + &RuntimeOrigin::signed(1), + &call, + &identifier, + &Some(1u16), + &Some(1u16), + ) + .expect("ok"); + assert!(!result); + }); +} + +#[test] +fn migrate_usage_key_tracks_scope() { + new_test_ext().execute_with(|| { + let call = scoped_call(); + let identifier = identifier_for(&call); + let tx_target = target(identifier); + LastSeen::::insert(tx_target, Some(6u16), 10); + assert!(RateLimiting::migrate_usage_key( + tx_target, + Some(6u16), + Some(7u16) + )); + assert_eq!(LastSeen::::get(tx_target, Some(7u16)), Some(10)); + }); +} + +#[test] +fn migrate_limit_scope_covers_transitions() { + new_test_ext().execute_with(|| { + let identifier = register(remark_call(), None); + let tx_target = target(identifier); + + // global -> scoped + assert!(RateLimiting::migrate_limit_scope( + tx_target, + None, + Some(42u16) + )); + match Limits::::get(tx_target) { + Some(RateLimit::Scoped(map)) => { + assert_eq!(map.get(&42u16), Some(&RateLimitKind::Default)) + } + other => panic!("unexpected config: {:?}", other), + } + + // scoped -> scoped + assert!(RateLimiting::migrate_limit_scope( + tx_target, + Some(42u16), + Some(43u16) + )); + match Limits::::get(tx_target) { + Some(RateLimit::Scoped(map)) => { + assert_eq!(map.get(&43u16), Some(&RateLimitKind::Default)) + } + other => panic!("unexpected config: {:?}", other), + } + + // scoped -> global (only entry) + assert!(RateLimiting::migrate_limit_scope( + tx_target, + Some(43u16), + None + )); + assert!(matches!( + Limits::::get(tx_target), + Some(RateLimit::Global(RateLimitKind::Default)) + )); + + // no-op when scopes identical + assert!(RateLimiting::migrate_limit_scope(tx_target, None, None)); + }); +} + +#[test] +fn set_default_limit_updates_span_and_resolves_in_enforcement() { + new_test_ext().execute_with(|| { + assert_eq!(RateLimiting::default_limit(), 0); + assert_ok!(RateLimiting::set_default_rate_limit( + RuntimeOrigin::root(), + 5 + )); + let event = last_event(); + assert!(matches!( + event, + RuntimeEvent::RateLimiting(crate::Event::DefaultRateLimitSet { block_span: 5 }) + )); + assert_eq!(RateLimiting::default_limit(), 5); + + let call = remark_call(); + let identifier = register(call.clone(), None); + let tx_target = target(identifier); + + System::set_block_number(10); + // No last-seen yet, first call passes. + assert!( + RateLimiting::is_within_limit( + &RuntimeOrigin::signed(1), + &call, + &identifier, + &None, + &None, + ) + .unwrap() + ); + + LastSeen::::insert(tx_target, None::, 12); + System::set_block_number(15); + // Span 5 should block when delta < 5. + assert!( + !RateLimiting::is_within_limit( + &RuntimeOrigin::signed(1), + &call, + &identifier, + &None, + &None, + ) + .unwrap() + ); + }); +} + +#[test] +fn limit_for_call_names_prefers_scoped_value() { + new_test_ext().execute_with(|| { + let call = scoped_call(); + let identifier = identifier_for(&call); + Limits::::insert( + target(identifier), + RateLimit::scoped_single(9u16, RateLimitKind::Exact(8)), + ); + let fetched = + RateLimiting::limit_for_call_names("RateLimiting", "set_rate_limit", Some(9u16)) + .expect("limit"); + assert_eq!(fetched, RateLimitKind::Exact(8)); + }); +} diff --git a/pallets/rate-limiting/src/tx_extension.rs b/pallets/rate-limiting/src/tx_extension.rs new file mode 100644 index 0000000000..303649c9c9 --- /dev/null +++ b/pallets/rate-limiting/src/tx_extension.rs @@ -0,0 +1,641 @@ +use codec::{Decode, DecodeWithMemTracking, Encode}; +use frame_support::{ + dispatch::{DispatchInfo, DispatchResult, PostDispatchInfo}, + pallet_prelude::Weight, + sp_runtime::{ + traits::{ + DispatchInfoOf, DispatchOriginOf, Dispatchable, Implication, TransactionExtension, + ValidateResult, Zero, + }, + transaction_validity::{ + InvalidTransaction, TransactionSource, TransactionValidityError, ValidTransaction, + }, + }, +}; +use scale_info::TypeInfo; +use sp_std::{marker::PhantomData, result::Result, vec, vec::Vec}; + +use crate::{ + Config, LastSeen, Pallet, + types::{ + RateLimitScopeResolver, RateLimitTarget, RateLimitUsageResolver, TransactionIdentifier, + }, +}; + +/// Identifier returned in the transaction metadata for the rate limiting extension. +const IDENTIFIER: &str = "RateLimitTransactionExtension"; + +/// Custom error code used to signal a rate limit violation. +const RATE_LIMIT_DENIED: u8 = 1; + +/// Transaction extension that enforces pallet rate limiting rules. +#[derive(Default, Encode, Decode, DecodeWithMemTracking, TypeInfo)] +pub struct RateLimitTransactionExtension(PhantomData<(T, I)>) +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo; + +impl RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ + pub fn new() -> Self { + Self(PhantomData) + } +} + +impl Clone for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ + fn clone(&self) -> Self { + Self(PhantomData) + } +} + +impl PartialEq for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ + fn eq(&self, _other: &Self) -> bool { + true + } +} + +impl Eq for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ +} + +impl core::fmt::Debug for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo, +{ + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(IDENTIFIER) + } +} + +impl TransactionExtension<>::RuntimeCall> + for RateLimitTransactionExtension +where + T: Config + Send + Sync + TypeInfo, + I: 'static + TypeInfo + Send + Sync, + >::RuntimeCall: Dispatchable, +{ + const IDENTIFIER: &'static str = IDENTIFIER; + + type Implicit = (); + type Val = Option<( + RateLimitTarget<>::GroupId>, + Option>::UsageKey>>, + bool, + )>; + type Pre = Option<( + RateLimitTarget<>::GroupId>, + Option>::UsageKey>>, + bool, + )>; + + fn weight(&self, _call: &>::RuntimeCall) -> Weight { + Weight::zero() + } + + fn validate( + &self, + origin: DispatchOriginOf<>::RuntimeCall>, + call: &>::RuntimeCall, + _info: &DispatchInfoOf<>::RuntimeCall>, + _len: usize, + _self_implicit: Self::Implicit, + _inherited_implication: &impl Implication, + _source: TransactionSource, + ) -> ValidateResult>::RuntimeCall> { + let Some(identifier) = TransactionIdentifier::from_call(call) else { + return Err(TransactionValidityError::Invalid(InvalidTransaction::Call)); + }; + + if !Pallet::::is_registered(&identifier) { + return Ok((ValidTransaction::default(), None, origin)); + } + + let scope = >::LimitScopeResolver::context(&origin, call); + let usage = >::UsageResolver::context(&origin, call); + + let config_target = Pallet::::config_target(&identifier) + .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; + let usage_target = Pallet::::usage_target(&identifier) + .map_err(|_| TransactionValidityError::Invalid(InvalidTransaction::Call))?; + let bypass = >::LimitScopeResolver::should_bypass(&origin, call); + let should_record = + bypass.record_usage && Pallet::::should_record_usage(&identifier, &usage_target); + + let Some(block_span) = + Pallet::::effective_span(&origin, call, &config_target, &scope) + else { + return Ok((ValidTransaction::default(), None, origin)); + }; + + if bypass.bypass_enforcement { + return Ok(( + ValidTransaction::default(), + should_record.then_some((usage_target, usage, true)), + origin, + )); + } + + if block_span.is_zero() { + return Ok((ValidTransaction::default(), None, origin)); + } + + let usage_keys: Vec>::UsageKey>> = match usage.clone() { + None => vec![None], + Some(keys) => keys.into_iter().map(Some).collect(), + }; + + let within_limit = usage_keys + .iter() + .all(|key| Pallet::::within_span(&usage_target, key, block_span)); + + if !within_limit { + return Err(TransactionValidityError::Invalid( + InvalidTransaction::Custom(RATE_LIMIT_DENIED), + )); + } + + Ok(( + ValidTransaction::default(), + Some((usage_target, usage, should_record)), + origin, + )) + } + + fn prepare( + self, + val: Self::Val, + _origin: &DispatchOriginOf<>::RuntimeCall>, + _call: &>::RuntimeCall, + _info: &DispatchInfoOf<>::RuntimeCall>, + _len: usize, + ) -> Result { + Ok(val) + } + + fn post_dispatch( + pre: Self::Pre, + _info: &DispatchInfoOf<>::RuntimeCall>, + _post_info: &mut PostDispatchInfo, + _len: usize, + result: &DispatchResult, + ) -> Result<(), TransactionValidityError> { + if result.is_ok() { + if let Some((target, usage, should_record)) = pre { + if !should_record { + return Ok(()); + } + let block_number = frame_system::Pallet::::block_number(); + match usage { + None => LastSeen::::insert( + target, + None::<>::UsageKey>, + block_number, + ), + Some(keys) => { + for key in keys { + LastSeen::::insert(target, Some(key), block_number); + } + } + } + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use codec::Encode; + use frame_support::{ + assert_ok, + dispatch::{GetDispatchInfo, PostDispatchInfo}, + }; + use sp_runtime::{ + traits::{TransactionExtension, TxBaseImplication}, + transaction_validity::{InvalidTransaction, TransactionSource, TransactionValidityError}, + }; + + use crate::{ + GroupSharing, LastSeen, Limits, + types::{RateLimit, RateLimitKind}, + }; + + use super::*; + use crate::mock::*; + + fn remark_call() -> RuntimeCall { + RuntimeCall::System(frame_system::Call::::remark { remark: Vec::new() }) + } + + fn bypass_call() -> RuntimeCall { + RuntimeCall::RateLimiting(RateLimitingCall::remove_call_from_group { + transaction: TransactionIdentifier::new(0, 0), + }) + } + + fn adjustable_call() -> RuntimeCall { + RuntimeCall::RateLimiting(RateLimitingCall::deregister_call { + transaction: TransactionIdentifier::new(0, 0), + scope: None, + clear_usage: false, + }) + } + + fn new_tx_extension() -> RateLimitTransactionExtension { + RateLimitTransactionExtension(Default::default()) + } + + fn target_for_call(call: &RuntimeCall) -> RateLimitTarget { + RateLimitTarget::Transaction(identifier_for(call)) + } + + fn validate_with_tx_extension( + extension: &RateLimitTransactionExtension, + call: &RuntimeCall, + ) -> Result< + ( + sp_runtime::transaction_validity::ValidTransaction, + Option<(RateLimitTarget, Option>, bool)>, + RuntimeOrigin, + ), + TransactionValidityError, + > { + let info = call.get_dispatch_info(); + let len = call.encode().len(); + extension.validate( + RuntimeOrigin::signed(42), + call, + &info, + len, + (), + &TxBaseImplication(()), + TransactionSource::External, + ) + } + + #[test] + fn tx_extension_allows_calls_without_limit() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = remark_call(); + + let (_valid, val, _origin) = + validate_with_tx_extension(&extension, &call).expect("valid"); + assert!(val.is_none()); + + let info = call.get_dispatch_info(); + let len = call.encode().len(); + let origin_for_prepare = RuntimeOrigin::signed(42); + let pre = extension + .clone() + .prepare(val.clone(), &origin_for_prepare, &call, &info, len) + .expect("prepare succeeds"); + + let mut post = PostDispatchInfo::default(); + RateLimitTransactionExtension::::post_dispatch( + pre, + &info, + &mut post, + len, + &Ok(()), + ) + .expect("post_dispatch succeeds"); + + let target = target_for_call(&call); + assert_eq!(LastSeen::::get(target, None::), None); + }); + } + + #[test] + fn tx_extension_honors_bypass_signal() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = bypass_call(); + + let (valid, val, _) = + validate_with_tx_extension(&extension, &call).expect("bypass should succeed"); + assert_eq!(valid.priority, 0); + assert!(val.is_none()); + + let identifier = identifier_for(&call); + let target = RateLimitTarget::Transaction(identifier); + Limits::::insert(target, RateLimit::global(RateLimitKind::Exact(3))); + LastSeen::::insert(target, None::, 1); + + let (_valid, post_val, _) = + validate_with_tx_extension(&extension, &call).expect("still bypassed"); + assert!(post_val.is_none()); + }); + } + + #[test] + fn tx_extension_applies_adjusted_span() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = adjustable_call(); + let identifier = identifier_for(&call); + let target = RateLimitTarget::Transaction(identifier); + Limits::::insert(target, RateLimit::global(RateLimitKind::Exact(4))); + LastSeen::::insert(target, Some(1u16), 10); + + System::set_block_number(14); + + // Stored span (4) would allow the call, but adjusted span (8) should block it. + let err = validate_with_tx_extension(&extension, &call) + .expect_err("adjusted span should apply"); + match err { + TransactionValidityError::Invalid(InvalidTransaction::Custom(code)) => { + assert_eq!(code, RATE_LIMIT_DENIED); + } + other => panic!("unexpected error: {:?}", other), + } + }); + } + + #[test] + fn tx_extension_records_usage_on_bypass() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { + block_span: 2, + }); + let identifier = identifier_for(&call); + let target = RateLimitTarget::Transaction(identifier); + + assert_ok!(RateLimiting::register_call( + RuntimeOrigin::root(), + Box::new(call.clone()), + None, + )); + + System::set_block_number(5); + + let (_valid, val, origin) = + validate_with_tx_extension(&extension, &call).expect("bypass should succeed"); + assert!(val.is_some(), "bypass decision should still record usage"); + + let info = call.get_dispatch_info(); + let len = call.encode().len(); + let pre = extension + .clone() + .prepare(val.clone(), &origin, &call, &info, len) + .expect("prepare succeeds"); + + let mut post = PostDispatchInfo::default(); + RateLimitTransactionExtension::::post_dispatch( + pre, + &info, + &mut post, + len, + &Ok(()), + ) + .expect("post_dispatch succeeds"); + + assert_eq!( + LastSeen::::get(target, Some(2u16)), + Some(5u64.into()) + ); + }); + } + + #[test] + fn tx_extension_records_last_seen_for_successful_call() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = remark_call(); + let identifier = identifier_for(&call); + let target = RateLimitTarget::Transaction(identifier); + Limits::::insert(target, RateLimit::global(RateLimitKind::Exact(5))); + + System::set_block_number(10); + + let (_valid, val, _) = validate_with_tx_extension(&extension, &call).expect("valid"); + assert!(val.is_some()); + + let info = call.get_dispatch_info(); + let len = call.encode().len(); + let origin_for_prepare = RuntimeOrigin::signed(42); + let pre = extension + .clone() + .prepare(val.clone(), &origin_for_prepare, &call, &info, len) + .expect("prepare succeeds"); + + let mut post = PostDispatchInfo::default(); + RateLimitTransactionExtension::::post_dispatch( + pre, + &info, + &mut post, + len, + &Ok(()), + ) + .expect("post_dispatch succeeds"); + + assert_eq!( + LastSeen::::get(target, None::), + Some(10) + ); + }); + } + + #[test] + fn tx_extension_rejects_when_call_occurs_too_soon() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = remark_call(); + let identifier = identifier_for(&call); + let target = RateLimitTarget::Transaction(identifier); + Limits::::insert(target, RateLimit::global(RateLimitKind::Exact(5))); + LastSeen::::insert(target, None::, 20); + + System::set_block_number(22); + + let err = + validate_with_tx_extension(&extension, &call).expect_err("should be rate limited"); + match err { + TransactionValidityError::Invalid(InvalidTransaction::Custom(code)) => { + assert_eq!(code, 1); + } + other => panic!("unexpected error: {:?}", other), + } + }); + } + + #[test] + fn tx_extension_skips_last_seen_when_span_zero() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + let call = remark_call(); + let identifier = identifier_for(&call); + let target = RateLimitTarget::Transaction(identifier); + Limits::::insert(target, RateLimit::global(RateLimitKind::Exact(0))); + + System::set_block_number(30); + + let (_valid, val, _) = validate_with_tx_extension(&extension, &call).expect("valid"); + assert!(val.is_none()); + + let info = call.get_dispatch_info(); + let len = call.encode().len(); + let origin_for_prepare = RuntimeOrigin::signed(42); + let pre = extension + .clone() + .prepare(val.clone(), &origin_for_prepare, &call, &info, len) + .expect("prepare succeeds"); + + let mut post = PostDispatchInfo::default(); + RateLimitTransactionExtension::::post_dispatch( + pre, + &info, + &mut post, + len, + &Ok(()), + ) + .expect("post_dispatch succeeds"); + + assert_eq!(LastSeen::::get(target, None::), None); + }); + } + + #[test] + fn tx_extension_skips_write_for_read_only_group_member() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + assert_ok!(RateLimiting::create_group( + RuntimeOrigin::root(), + b"use-ro".to_vec(), + GroupSharing::UsageOnly, + )); + let group = RateLimiting::next_group_id().saturating_sub(1); + + let call = remark_call(); + let identifier = identifier_for(&call); + assert_ok!(RateLimiting::register_call( + RuntimeOrigin::root(), + Box::new(call.clone()), + Some(group), + )); + assert_ok!(RateLimiting::set_call_read_only( + RuntimeOrigin::root(), + identifier, + true + )); + + let tx_target = RateLimitTarget::Transaction(identifier); + let usage_target = RateLimitTarget::Group(group); + Limits::::insert(tx_target, RateLimit::global(RateLimitKind::Exact(2))); + LastSeen::::insert(usage_target, Some(1u16), 2); + + System::set_block_number(5); + + let (_valid, val, _) = validate_with_tx_extension(&extension, &call).expect("valid"); + let info = call.get_dispatch_info(); + let len = call.encode().len(); + let origin_for_prepare = RuntimeOrigin::signed(42); + let pre = extension + .clone() + .prepare(val.clone(), &origin_for_prepare, &call, &info, len) + .expect("prepare succeeds"); + + let mut post = PostDispatchInfo::default(); + RateLimitTransactionExtension::::post_dispatch( + pre, + &info, + &mut post, + len, + &Ok(()), + ) + .expect("post_dispatch succeeds"); + + // Usage key should remain untouched because the call is read-only. + assert_eq!(LastSeen::::get(usage_target, Some(1u16)), Some(2)); + }); + } + + #[test] + fn tx_extension_respects_usage_group_sharing() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + assert_ok!(RateLimiting::create_group( + RuntimeOrigin::root(), + b"use".to_vec(), + GroupSharing::UsageOnly, + )); + let group = RateLimiting::next_group_id().saturating_sub(1); + + let call = remark_call(); + let identifier = identifier_for(&call); + assert_ok!(RateLimiting::register_call( + RuntimeOrigin::root(), + Box::new(call.clone()), + Some(group), + )); + + let tx_target = RateLimitTarget::Transaction(identifier); + let usage_target = RateLimitTarget::Group(group); + Limits::::insert(tx_target, RateLimit::global(RateLimitKind::Exact(5))); + LastSeen::::insert(usage_target, None::, 10); + System::set_block_number(12); + + let err = validate_with_tx_extension(&extension, &call) + .expect_err("usage grouping should rate limit"); + match err { + TransactionValidityError::Invalid(InvalidTransaction::Custom(code)) => { + assert_eq!(code, RATE_LIMIT_DENIED); + } + other => panic!("unexpected error: {:?}", other), + } + }); + } + + #[test] + fn tx_extension_respects_config_group_sharing() { + new_test_ext().execute_with(|| { + let extension = new_tx_extension(); + assert_ok!(RateLimiting::create_group( + RuntimeOrigin::root(), + b"cfg".to_vec(), + GroupSharing::ConfigOnly, + )); + let group = RateLimiting::next_group_id().saturating_sub(1); + + let call = remark_call(); + let identifier = identifier_for(&call); + assert_ok!(RateLimiting::register_call( + RuntimeOrigin::root(), + Box::new(call.clone()), + Some(group), + )); + + let tx_target = RateLimitTarget::Transaction(identifier); + let group_target = RateLimitTarget::Group(group); + Limits::::remove(tx_target); + Limits::::insert(group_target, RateLimit::global(RateLimitKind::Exact(5))); + LastSeen::::insert(tx_target, None::, 10); + System::set_block_number(12); + + let err = validate_with_tx_extension(&extension, &call) + .expect_err("config grouping should rate limit"); + match err { + TransactionValidityError::Invalid(InvalidTransaction::Custom(code)) => { + assert_eq!(code, RATE_LIMIT_DENIED); + } + other => panic!("unexpected error: {:?}", other), + } + }); + } +} diff --git a/pallets/rate-limiting/src/types.rs b/pallets/rate-limiting/src/types.rs new file mode 100644 index 0000000000..f6f54b472f --- /dev/null +++ b/pallets/rate-limiting/src/types.rs @@ -0,0 +1,236 @@ +use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; +use frame_support::dispatch::DispatchResult; +pub use rate_limiting_interface::{RateLimitTarget, TransactionIdentifier}; +use scale_info::TypeInfo; +use serde::{Deserialize, Serialize}; +use sp_std::{collections::btree_map::BTreeMap, vec::Vec}; + +/// Resolves the optional identifier within which a rate limit applies and can optionally adjust +/// enforcement behaviour. +pub trait RateLimitScopeResolver { + /// Returns `Some(scope)` when the limit should be applied per-scope, or `None` for global + /// limits. + fn context(origin: &Origin, call: &Call) -> Option; + + /// Returns how the call should interact with enforcement and usage tracking. + fn should_bypass(_origin: &Origin, _call: &Call) -> BypassDecision { + BypassDecision::enforce_and_record() + } + + /// Optionally adjusts the effective span used during enforcement. Defaults to the original + /// `span`. + fn adjust_span(_origin: &Origin, _call: &Call, span: Span) -> Span { + span + } +} + +/// Controls whether enforcement should run and whether usage should be recorded for a call. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BypassDecision { + pub bypass_enforcement: bool, + pub record_usage: bool, +} + +impl BypassDecision { + pub const fn new(bypass_enforcement: bool, record_usage: bool) -> Self { + Self { + bypass_enforcement, + record_usage, + } + } + + pub const fn enforce_and_record() -> Self { + Self::new(false, true) + } + + pub const fn bypass_and_record() -> Self { + Self::new(true, true) + } + + pub const fn bypass_and_skip() -> Self { + Self::new(true, false) + } +} + +/// Resolves the optional usage tracking key applied when enforcing limits. +pub trait RateLimitUsageResolver { + /// Returns `Some(keys)` to track usage per key, or `None` for global usage tracking. + /// + /// When multiple keys are returned, the rate limit is enforced against each key and all are + /// recorded on success. + fn context(origin: &Origin, call: &Call) -> Option>; +} + +/// Origin check performed when configuring a rate limit. +/// +/// `pallet-rate-limiting` supports configuring a distinct "who may set limits" rule per call/group +/// target. This trait is invoked by [`pallet::Pallet::set_rate_limit`] after loading the rule from +/// storage, allowing runtimes to implement arbitrary permissioning logic. +/// +/// Note: the hook receives the provided `scope` (if any). Some policies (for example "subnet owner") +/// require a scope value (such as `netuid`) in order to validate the caller. +pub trait EnsureLimitSettingRule { + fn ensure_origin(origin: Origin, rule: &Rule, scope: &Option) -> DispatchResult; +} + +/// Sharing mode configured for a group. +#[derive( + Serialize, + Deserialize, + Clone, + Copy, + PartialEq, + Eq, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + MaxEncodedLen, + Debug, +)] +pub enum GroupSharing { + /// Limits remain per transaction; usage is shared by the group. + UsageOnly, + /// Limits are shared by the group; usage remains per transaction. + ConfigOnly, + /// Both limits and usage are shared by the group. + ConfigAndUsage, +} + +impl GroupSharing { + /// Returns `true` when configuration for this group should use the group target key. + pub fn config_uses_group(self) -> bool { + matches!( + self, + GroupSharing::ConfigOnly | GroupSharing::ConfigAndUsage + ) + } + + /// Returns `true` when usage tracking for this group should use the group target key. + pub fn usage_uses_group(self) -> bool { + matches!(self, GroupSharing::UsageOnly | GroupSharing::ConfigAndUsage) + } +} + +/// Metadata describing a configured group. +#[derive( + Serialize, + Deserialize, + Clone, + PartialEq, + Eq, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + MaxEncodedLen, + Debug, +)] +pub struct RateLimitGroup { + /// Stable identifier assigned to the group. + pub id: GroupId, + /// Human readable group name. + pub name: Name, + /// Sharing configuration enforced for the group. + pub sharing: GroupSharing, +} + +/// Policy describing the block span enforced by a rate limit. +#[derive( + Serialize, + Deserialize, + Clone, + Copy, + PartialEq, + Eq, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + MaxEncodedLen, + Debug, +)] +pub enum RateLimitKind { + /// Use the pallet-level default rate limit. + Default, + /// Apply an exact rate limit measured in blocks. + Exact(BlockNumber), +} + +/// Stored rate limit configuration for a transaction identifier. +/// +/// The configuration is mutually exclusive: either the call is globally limited or it stores a set +/// of per-scope spans. +#[derive( + Serialize, + Deserialize, + Clone, + PartialEq, + Eq, + Encode, + Decode, + DecodeWithMemTracking, + TypeInfo, + Debug, +)] +#[serde( + bound = "Scope: Ord + serde::Serialize + serde::de::DeserializeOwned, BlockNumber: serde::Serialize + serde::de::DeserializeOwned" +)] +pub enum RateLimit { + /// Global span applied to every invocation. + Global(RateLimitKind), + /// Per-scope spans keyed by `Scope`. + Scoped(BTreeMap>), +} + +impl RateLimit +where + Scope: Ord, +{ + /// Convenience helper to build a global configuration. + pub fn global(kind: RateLimitKind) -> Self { + Self::Global(kind) + } + + /// Convenience helper to build a scoped configuration containing a single entry. + pub fn scoped_single(scope: Scope, kind: RateLimitKind) -> Self { + let mut map = BTreeMap::new(); + map.insert(scope, kind); + Self::Scoped(map) + } + + /// Returns the span configured for the provided scope, if any. + pub fn kind_for(&self, scope: Option<&Scope>) -> Option<&RateLimitKind> { + match self { + RateLimit::Global(kind) => Some(kind), + RateLimit::Scoped(map) => scope.and_then(|key| map.get(key)), + } + } + + /// Inserts or updates a scoped entry, converting from a global configuration if needed. + pub fn upsert_scope(&mut self, scope: Scope, kind: RateLimitKind) { + match self { + RateLimit::Global(_) => { + let mut map = BTreeMap::new(); + map.insert(scope, kind); + *self = RateLimit::Scoped(map); + } + RateLimit::Scoped(map) => { + map.insert(scope, kind); + } + } + } + + /// Removes a scoped entry, returning whether one existed. + pub fn remove_scope(&mut self, scope: &Scope) -> bool { + match self { + RateLimit::Global(_) => false, + RateLimit::Scoped(map) => map.remove(scope).is_some(), + } + } + + /// Returns true when the scoped configuration contains no entries. + pub fn is_scoped_empty(&self) -> bool { + matches!(self, RateLimit::Scoped(map) if map.is_empty()) + } +} diff --git a/pallets/subtensor/Cargo.toml b/pallets/subtensor/Cargo.toml index 2e35a89d19..cd9410de8b 100644 --- a/pallets/subtensor/Cargo.toml +++ b/pallets/subtensor/Cargo.toml @@ -55,6 +55,8 @@ sha2.workspace = true rand_chacha.workspace = true pallet-crowdloan.workspace = true pallet-subtensor-proxy.workspace = true +pallet-rate-limiting.workspace = true +rate-limiting-interface.workspace = true [dev-dependencies] pallet-balances = { workspace = true, features = ["std"] } @@ -114,6 +116,8 @@ std = [ "pallet-crowdloan/std", "pallet-drand/std", "pallet-subtensor-proxy/std", + "pallet-rate-limiting/std", + "rate-limiting-interface/std", "pallet-subtensor-swap/std", "subtensor-swap-interface/std", "pallet-subtensor-utility/std", diff --git a/pallets/subtensor/src/benchmarks.rs b/pallets/subtensor/src/benchmarks.rs index 476be905e9..67b7f473b0 100644 --- a/pallets/subtensor/src/benchmarks.rs +++ b/pallets/subtensor/src/benchmarks.rs @@ -157,7 +157,6 @@ mod pallet_benchmarks { netuid, caller.clone() )); - Subtensor::::set_serving_rate_limit(netuid, 0); #[extrinsic_call] _( @@ -195,7 +194,6 @@ mod pallet_benchmarks { netuid, caller.clone() )); - Subtensor::::set_serving_rate_limit(netuid, 0); #[extrinsic_call] _( diff --git a/pallets/subtensor/src/coinbase/root.rs b/pallets/subtensor/src/coinbase/root.rs index 328ce3805c..83ff49abbf 100644 --- a/pallets/subtensor/src/coinbase/root.rs +++ b/pallets/subtensor/src/coinbase/root.rs @@ -325,7 +325,6 @@ impl Pallet { LastAdjustmentBlock::::remove(netuid); // --- 16. Serving / rho / curves, and other per-net controls. - ServingRateLimit::::remove(netuid); Rho::::remove(netuid); AlphaSigmoidSteepness::::remove(netuid); diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index ef2d44e68b..c44381dc2b 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -908,12 +908,6 @@ pub mod pallet { 0 } - /// Default value for serving rate limit. - #[pallet::type_value] - pub fn DefaultServingRateLimit() -> u64 { - T::InitialServingRateLimit::get() - } - /// Default value for weight commit/reveal enabled. #[pallet::type_value] pub fn DefaultCommitRevealWeightsEnabled() -> bool { @@ -1241,7 +1235,7 @@ pub mod pallet { /// ================== /// ==== Coinbase ==== /// ================== - /// --- ITEM ( global_block_emission ) + /// --- ITEM ( global_block_emission ) #[pallet::storage] pub type BlockEmission = StorageValue<_, u64, ValueQuery, DefaultBlockEmission>; @@ -1281,7 +1275,7 @@ pub mod pallet { #[pallet::storage] pub type TotalStake = StorageValue<_, TaoCurrency, ValueQuery, DefaultZeroTao>; - /// --- ITEM ( moving_alpha ) -- subnet moving alpha. + /// --- ITEM ( moving_alpha ) -- subnet moving alpha. #[pallet::storage] pub type SubnetMovingAlpha = StorageValue<_, I96F32, ValueQuery, DefaultMovingAlpha>; @@ -1674,11 +1668,6 @@ pub mod pallet { pub type RecycleOrBurn = StorageMap<_, Identity, NetUid, RecycleOrBurnEnum, ValueQuery, DefaultRecycleOrBurn>; - /// --- MAP ( netuid ) --> serving_rate_limit - #[pallet::storage] - pub type ServingRateLimit = - StorageMap<_, Identity, NetUid, u64, ValueQuery, DefaultServingRateLimit>; - /// --- MAP ( netuid ) --> Rho #[pallet::storage] pub type Rho = StorageMap<_, Identity, NetUid, u16, ValueQuery, DefaultRho>; @@ -2445,7 +2434,6 @@ pub enum CustomTransactionError { TransferDisallowed, HotKeyNotRegisteredInNetwork, InvalidIpAddress, - ServingRateLimitExceeded, InvalidPort, BadRequest, ZeroMaxAmount, @@ -2472,7 +2460,6 @@ impl From for u8 { CustomTransactionError::TransferDisallowed => 9, CustomTransactionError::HotKeyNotRegisteredInNetwork => 10, CustomTransactionError::InvalidIpAddress => 11, - CustomTransactionError::ServingRateLimitExceeded => 12, CustomTransactionError::InvalidPort => 13, CustomTransactionError::BadRequest => 255, CustomTransactionError::ZeroMaxAmount => 14, diff --git a/pallets/subtensor/src/macros/config.rs b/pallets/subtensor/src/macros/config.rs index a735bde1e1..b6dc13059d 100644 --- a/pallets/subtensor/src/macros/config.rs +++ b/pallets/subtensor/src/macros/config.rs @@ -8,6 +8,7 @@ mod config { use crate::{CommitmentsInterface, GetAlphaForTao, GetTaoForAlpha}; use pallet_commitments::GetCommitments; + use rate_limiting_interface::RateLimitingInfo; use subtensor_swap_interface::{SwapEngine, SwapHandler}; /// Configure the pallet by specifying the parameters and types on which it depends. @@ -56,6 +57,17 @@ mod config { /// Interface to clean commitments on network dissolution. type CommitmentsInterface: CommitmentsInterface; + /// Read-only interface for querying rate limiting configuration and usage. + type RateLimiting: RateLimitingInfo< + GroupId = subtensor_runtime_common::rate_limiting::GroupId, + CallMetadata = ::RuntimeCall, + Limit = BlockNumberFor, + Scope = subtensor_runtime_common::NetUid, + UsageKey = subtensor_runtime_common::rate_limiting::RateLimitUsageKey< + Self::AccountId, + >, + >; + /// Rate limit for associating an EVM key. type EvmKeyAssociateRateLimit: Get; @@ -171,9 +183,6 @@ mod config { /// Initial weights version key. #[pallet::constant] type InitialWeightsVersionKey: Get; - /// Initial serving rate limit. - #[pallet::constant] - type InitialServingRateLimit: Get; /// Initial transaction rate limit. #[pallet::constant] type InitialTxRateLimit: Get; diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index 8c0b2210ec..06f6e6375e 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -814,10 +814,6 @@ mod dispatches { /// /// * 'InvalidIpAddress': /// - The numerically encoded ip address does not resolve to a proper ip. - /// - /// * 'ServingRateLimitExceeded': - /// - Attempting to set prometheus information withing the rate limit min. - /// #[pallet::call_index(4)] #[pallet::weight((Weight::from_parts(33_010_000, 0) .saturating_add(T::DbWeight::get().reads(4)) @@ -898,10 +894,6 @@ mod dispatches { /// /// * 'InvalidIpAddress': /// - The numerically encoded ip address does not resolve to a proper ip. - /// - /// * 'ServingRateLimitExceeded': - /// - Attempting to set prometheus information withing the rate limit min. - /// #[pallet::call_index(40)] #[pallet::weight((Weight::from_parts(32_510_000, 0) .saturating_add(T::DbWeight::get().reads(4)) diff --git a/pallets/subtensor/src/macros/errors.rs b/pallets/subtensor/src/macros/errors.rs index 6c3d7a35df..f87388d49b 100644 --- a/pallets/subtensor/src/macros/errors.rs +++ b/pallets/subtensor/src/macros/errors.rs @@ -79,6 +79,7 @@ mod errors { SettingWeightsTooFast, /// A validator is attempting to set weights from a validator with incorrect weight version. IncorrectWeightVersionKey, + /// DEPRECATED /// An axon or prometheus serving exceeded the rate limit for a registered neuron. ServingRateLimitExceeded, /// The caller is attempting to set weights with more UIDs than allowed. diff --git a/pallets/subtensor/src/macros/events.rs b/pallets/subtensor/src/macros/events.rs index a06e035d86..baf39b2994 100644 --- a/pallets/subtensor/src/macros/events.rs +++ b/pallets/subtensor/src/macros/events.rs @@ -101,7 +101,7 @@ mod events { MinDifficultySet(NetUid, u64), /// setting max difficulty on a network. MaxDifficultySet(NetUid, u64), - /// setting the prometheus serving rate limit. + /// [DEPRECATED] setting the prometheus serving rate limit. ServingRateLimitSet(NetUid, u64), /// setting burn on a network. BurnSet(NetUid, TaoCurrency), diff --git a/pallets/subtensor/src/subnets/serving.rs b/pallets/subtensor/src/subnets/serving.rs index d11eb479d3..0791484d88 100644 --- a/pallets/subtensor/src/subnets/serving.rs +++ b/pallets/subtensor/src/subnets/serving.rs @@ -51,10 +51,6 @@ impl Pallet { /// /// * 'InvalidIpAddress': /// - The numerically encoded ip address does not resolve to a proper ip. - /// - /// * 'ServingRateLimitExceeded': - /// - Attempting to set prometheus information withing the rate limit min. - /// pub fn do_serve_axon( origin: T::RuntimeOrigin, netuid: NetUid, @@ -155,10 +151,6 @@ impl Pallet { /// /// * 'InvalidIpAddress': /// - The numerically encoded ip address does not resolve to a proper ip. - /// - /// * 'ServingRateLimitExceeded': - /// - Attempting to set prometheus information withing the rate limit min. - /// pub fn do_serve_prometheus( origin: T::RuntimeOrigin, netuid: NetUid, @@ -185,11 +177,6 @@ impl Pallet { // 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(); - ensure!( - Self::prometheus_passes_rate_limit(netuid, &prev_prometheus, current_block), - Error::::ServingRateLimitExceeded - ); // We insert the prometheus meta. prev_prometheus.block = Self::get_current_block_as_u64(); @@ -220,26 +207,6 @@ impl Pallet { --==[[ Helper functions ]]==-- *********************************/ - pub fn axon_passes_rate_limit( - netuid: NetUid, - prev_axon_info: &AxonInfoOf, - current_block: u64, - ) -> bool { - let rate_limit: u64 = Self::get_serving_rate_limit(netuid); - let last_serve = prev_axon_info.block; - rate_limit == 0 || last_serve == 0 || current_block.saturating_sub(last_serve) >= rate_limit - } - - pub fn prometheus_passes_rate_limit( - netuid: NetUid, - prev_prometheus_info: &PrometheusInfoOf, - current_block: u64, - ) -> bool { - let rate_limit: u64 = Self::get_serving_rate_limit(netuid); - let last_serve = prev_prometheus_info.block; - rate_limit == 0 || last_serve == 0 || current_block.saturating_sub(last_serve) >= rate_limit - } - pub fn get_axon_info(netuid: NetUid, hotkey: &T::AccountId) -> AxonInfoOf { if let Some(axons) = Axons::::get(netuid, hotkey) { axons @@ -345,11 +312,6 @@ impl Pallet { // 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::::ServingRateLimitExceeded - ); // Validate axon data with delegate func prev_axon.block = Self::get_current_block_as_u64(); diff --git a/pallets/subtensor/src/tests/mock.rs b/pallets/subtensor/src/tests/mock.rs index 090fcf8f75..8cb3478b34 100644 --- a/pallets/subtensor/src/tests/mock.rs +++ b/pallets/subtensor/src/tests/mock.rs @@ -20,6 +20,7 @@ use frame_system as system; use frame_system::{EnsureRoot, RawOrigin, limits, offchain::CreateTransactionBase}; use pallet_subtensor_proxy as pallet_proxy; use pallet_subtensor_utility as pallet_utility; +use rate_limiting_interface::{RateLimitingInfo, TryIntoRateLimitTarget}; use sp_core::{ConstU64, Get, H256, U256, offchain::KeyTypeId}; use sp_runtime::Perbill; use sp_runtime::{ @@ -28,7 +29,7 @@ use sp_runtime::{ }; use sp_std::{cell::RefCell, cmp::Ordering, sync::OnceLock}; use sp_tracing::tracing_subscriber; -use subtensor_runtime_common::{NetUid, TaoCurrency}; +use subtensor_runtime_common::{NetUid, TaoCurrency, rate_limiting::RateLimitUsageKey}; use subtensor_swap_interface::{Order, SwapHandler}; use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; type Block = frame_system::mocking::MockBlock; @@ -178,7 +179,6 @@ parameter_types! { pub const InitialMinChildKeyTake: u16 = 0; // 0 %; pub const InitialMaxChildKeyTake: u16 = 11_796; // 18 %; pub const InitialWeightsVersionKey: u16 = 0; - pub const InitialServingRateLimit: u64 = 0; // No limit. pub const InitialTxRateLimit: u64 = 0; // Disable rate limit for testing pub const InitialTxDelegateTakeRateLimit: u64 = 1; // 1 block take rate limit for testing pub const InitialTxChildKeyTakeRateLimit: u64 = 1; // 1 block take rate limit for testing @@ -264,7 +264,6 @@ impl crate::Config for Test { type InitialWeightsVersionKey = InitialWeightsVersionKey; type InitialMaxDifficulty = InitialMaxDifficulty; type InitialMinDifficulty = InitialMinDifficulty; - type InitialServingRateLimit = InitialServingRateLimit; type InitialTxRateLimit = InitialTxRateLimit; type InitialTxDelegateTakeRateLimit = InitialTxDelegateTakeRateLimit; type InitialBurn = InitialBurn; @@ -298,6 +297,7 @@ impl crate::Config for Test { type GetCommitments = (); type MaxImmuneUidsPercentage = MaxImmuneUidsPercentage; type CommitmentsInterface = CommitmentsI; + type RateLimiting = NoRateLimiting; type EvmKeyAssociateRateLimit = EvmKeyAssociateRateLimit; } @@ -336,6 +336,33 @@ impl CommitmentsInterface for CommitmentsI { fn purge_netuid(_netuid: NetUid) {} } +pub struct NoRateLimiting; + +impl RateLimitingInfo for NoRateLimiting { + type GroupId = subtensor_runtime_common::rate_limiting::GroupId; + type CallMetadata = RuntimeCall; + type Limit = BlockNumber; + type Scope = subtensor_runtime_common::NetUid; + type UsageKey = RateLimitUsageKey; + + fn rate_limit(_target: TargetArg, _scope: Option) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + None + } + + fn last_seen( + _target: TargetArg, + _usage_key: Option, + ) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + None + } +} + parameter_types! { pub MaximumSchedulerWeight: Weight = Perbill::from_percent(80) * BlockWeights::get().max_block; diff --git a/pallets/subtensor/src/tests/networks.rs b/pallets/subtensor/src/tests/networks.rs index 8214d58be0..8a8850b178 100644 --- a/pallets/subtensor/src/tests/networks.rs +++ b/pallets/subtensor/src/tests/networks.rs @@ -388,7 +388,6 @@ fn dissolve_clears_all_per_subnet_storages() { PendingOwnerCut::::insert(net, AlphaCurrency::from(1)); BlocksSinceLastStep::::insert(net, 1u64); LastMechansimStepBlock::::insert(net, 1u64); - ServingRateLimit::::insert(net, 1u64); Rho::::insert(net, 1u16); AlphaSigmoidSteepness::::insert(net, 1i16); @@ -548,7 +547,6 @@ fn dissolve_clears_all_per_subnet_storages() { assert!(!PendingOwnerCut::::contains_key(net)); assert!(!BlocksSinceLastStep::::contains_key(net)); assert!(!LastMechansimStepBlock::::contains_key(net)); - assert!(!ServingRateLimit::::contains_key(net)); assert!(!Rho::::contains_key(net)); assert!(!AlphaSigmoidSteepness::::contains_key(net)); diff --git a/pallets/subtensor/src/tests/serving.rs b/pallets/subtensor/src/tests/serving.rs index b52666bf26..d56787cccd 100644 --- a/pallets/subtensor/src/tests/serving.rs +++ b/pallets/subtensor/src/tests/serving.rs @@ -281,23 +281,19 @@ fn test_axon_serving_rate_limit_exceeded() { placeholder1, placeholder2 )); - SubtensorModule::set_serving_rate_limit(netuid, 2); run_to_block(2); // Go to block 2 - // Needs to be 2 blocks apart, we are only 1 block apart - assert_eq!( - SubtensorModule::serve_axon( - <::RuntimeOrigin>::signed(hotkey_account_id), - netuid, - version, - ip, - port, - ip_type, - protocol, - placeholder1, - placeholder2 - ), - Err(Error::::ServingRateLimitExceeded.into()) - ); + // Rate limiting is enforced by the transaction extension, not the pallet call. + assert_ok!(SubtensorModule::serve_axon( + <::RuntimeOrigin>::signed(hotkey_account_id), + netuid, + version, + ip, + port, + ip_type, + protocol, + placeholder1, + placeholder2 + )); }); } @@ -479,19 +475,15 @@ fn test_prometheus_serving_rate_limit_exceeded() { port, ip_type )); - SubtensorModule::set_serving_rate_limit(netuid, 1); - // Same block, need 1 block to pass - assert_eq!( - SubtensorModule::serve_prometheus( - <::RuntimeOrigin>::signed(hotkey_account_id), - netuid, - version, - ip, - port, - ip_type - ), - Err(Error::::ServingRateLimitExceeded.into()) - ); + // Rate limiting is enforced by the transaction extension, not the pallet call. + assert_ok!(SubtensorModule::serve_prometheus( + <::RuntimeOrigin>::signed(hotkey_account_id), + netuid, + version, + ip, + port, + ip_type + )); }); } diff --git a/pallets/subtensor/src/transaction_extension.rs b/pallets/subtensor/src/transaction_extension.rs index cf1d410ea9..e227e2d483 100644 --- a/pallets/subtensor/src/transaction_extension.rs +++ b/pallets/subtensor/src/transaction_extension.rs @@ -70,9 +70,6 @@ where CustomTransactionError::HotKeyNotRegisteredInNetwork.into() } Error::::InvalidIpAddress => CustomTransactionError::InvalidIpAddress.into(), - Error::::ServingRateLimitExceeded => { - CustomTransactionError::ServingRateLimitExceeded.into() - } Error::::InvalidPort => CustomTransactionError::InvalidPort.into(), _ => CustomTransactionError::BadRequest.into(), }) diff --git a/pallets/subtensor/src/utils/misc.rs b/pallets/subtensor/src/utils/misc.rs index 10fc0535f0..05c1af7d6f 100644 --- a/pallets/subtensor/src/utils/misc.rs +++ b/pallets/subtensor/src/utils/misc.rs @@ -1,12 +1,15 @@ use super::*; use crate::Error; use crate::system::{ensure_signed, ensure_signed_or_root, pallet_prelude::BlockNumberFor}; +use rate_limiting_interface::RateLimitingInfo; use safe_math::*; use sp_core::Get; use sp_core::U256; -use sp_runtime::Saturating; +use sp_runtime::{SaturatedConversion, Saturating}; use substrate_fixed::types::{I32F32, I64F64, U64F64, U96F32}; -use subtensor_runtime_common::{AlphaCurrency, NetUid, NetUidStorageIndex, TaoCurrency}; +use subtensor_runtime_common::{ + AlphaCurrency, NetUid, NetUidStorageIndex, TaoCurrency, rate_limiting, +}; impl Pallet { pub fn ensure_subnet_owner_or_root( @@ -428,11 +431,9 @@ impl Pallet { } pub fn get_serving_rate_limit(netuid: NetUid) -> u64 { - ServingRateLimit::::get(netuid) - } - pub fn set_serving_rate_limit(netuid: NetUid, serving_rate_limit: u64) { - ServingRateLimit::::insert(netuid, serving_rate_limit); - Self::deposit_event(Event::ServingRateLimitSet(netuid, serving_rate_limit)); + T::RateLimiting::rate_limit(rate_limiting::GROUP_SERVE, Some(netuid)) + .unwrap_or_default() + .saturated_into() } pub fn get_min_difficulty(netuid: NetUid) -> u64 { diff --git a/pallets/subtensor/src/utils/rate_limiting.rs b/pallets/subtensor/src/utils/rate_limiting.rs index 85f58cfc64..468aecd1c1 100644 --- a/pallets/subtensor/src/utils/rate_limiting.rs +++ b/pallets/subtensor/src/utils/rate_limiting.rs @@ -1,3 +1,5 @@ +use codec::{Decode, Encode}; +use scale_info::TypeInfo; use subtensor_runtime_common::NetUid; use super::*; diff --git a/pallets/transaction-fee/Cargo.toml b/pallets/transaction-fee/Cargo.toml index b5f08ac7c3..86bea7c843 100644 --- a/pallets/transaction-fee/Cargo.toml +++ b/pallets/transaction-fee/Cargo.toml @@ -27,6 +27,7 @@ pallet-scheduler = { workspace = true, default-features = false, optional = true [dev-dependencies] frame-executive.workspace = true pallet-evm-chain-id.workspace = true +rate-limiting-interface.workspace = true scale-info.workspace = true sp-consensus-aura.workspace = true sp-consensus-grandpa.workspace = true @@ -56,6 +57,7 @@ std = [ "pallet-subtensor/std", "pallet-subtensor-swap/std", "pallet-transaction-payment/std", + "rate-limiting-interface/std", "scale-info/std", "sp-runtime/std", "sp-std/std", diff --git a/pallets/transaction-fee/src/tests/mock.rs b/pallets/transaction-fee/src/tests/mock.rs index ee5b1693ba..c5f25b54e8 100644 --- a/pallets/transaction-fee/src/tests/mock.rs +++ b/pallets/transaction-fee/src/tests/mock.rs @@ -12,6 +12,7 @@ use frame_system::{ self as system, EnsureRoot, RawOrigin, limits, offchain::CreateTransactionBase, }; pub use pallet_subtensor::*; +use rate_limiting_interface::{RateLimitingInfo, TryIntoRateLimitTarget}; pub use sp_core::U256; use sp_core::{ConstU64, H256}; use sp_runtime::{ @@ -21,7 +22,9 @@ use sp_runtime::{ }; use sp_std::cmp::Ordering; use sp_weights::Weight; -pub use subtensor_runtime_common::{AlphaCurrency, Currency, NetUid, TaoCurrency}; +pub use subtensor_runtime_common::{ + AlphaCurrency, Currency, NetUid, TaoCurrency, rate_limiting::RateLimitUsageKey, +}; use subtensor_swap_interface::{Order, SwapHandler}; use crate::SubtensorTxFeeHandler; @@ -168,7 +171,6 @@ parameter_types! { pub const InitialMinChildKeyTake: u16 = 0; // Allow 0 % pub const InitialMaxChildKeyTake: u16 = 11_796; // 18 %; pub const InitialWeightsVersionKey: u16 = 0; - pub const InitialServingRateLimit: u64 = 0; // No limit. pub const InitialTxRateLimit: u64 = 0; // Disable rate limit for testing pub const InitialTxDelegateTakeRateLimit: u64 = 0; // Disable rate limit for testing pub const InitialTxChildKeyTakeRateLimit: u64 = 0; // Disable rate limit for testing @@ -254,7 +256,6 @@ impl pallet_subtensor::Config for Test { type InitialWeightsVersionKey = InitialWeightsVersionKey; type InitialMaxDifficulty = InitialMaxDifficulty; type InitialMinDifficulty = InitialMinDifficulty; - type InitialServingRateLimit = InitialServingRateLimit; type InitialTxRateLimit = InitialTxRateLimit; type InitialTxDelegateTakeRateLimit = InitialTxDelegateTakeRateLimit; type InitialTxChildKeyTakeRateLimit = InitialTxChildKeyTakeRateLimit; @@ -289,6 +290,7 @@ impl pallet_subtensor::Config for Test { type GetCommitments = (); type MaxImmuneUidsPercentage = MaxImmuneUidsPercentage; type CommitmentsInterface = CommitmentsI; + type RateLimiting = NoRateLimiting; type EvmKeyAssociateRateLimit = EvmKeyAssociateRateLimit; } @@ -421,6 +423,33 @@ impl pallet_subtensor::CommitmentsInterface for CommitmentsI { fn purge_netuid(_netuid: NetUid) {} } +pub struct NoRateLimiting; + +impl RateLimitingInfo for NoRateLimiting { + type GroupId = subtensor_runtime_common::rate_limiting::GroupId; + type CallMetadata = RuntimeCall; + type Limit = u64; + type Scope = subtensor_runtime_common::NetUid; + type UsageKey = RateLimitUsageKey; + + fn rate_limit(_target: TargetArg, _scope: Option) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + None + } + + fn last_seen( + _target: TargetArg, + _usage_key: Option, + ) -> Option + where + TargetArg: TryIntoRateLimitTarget, + { + None + } +} + parameter_types! { pub MaximumSchedulerWeight: Weight = Perbill::from_percent(80) * BlockWeights::get().max_block; diff --git a/precompiles/Cargo.toml b/precompiles/Cargo.toml index 12460dcfbb..599b404350 100644 --- a/precompiles/Cargo.toml +++ b/precompiles/Cargo.toml @@ -34,6 +34,7 @@ substrate-fixed.workspace = true pallet-subtensor.workspace = true pallet-subtensor-swap.workspace = true pallet-admin-utils.workspace = true +pallet-rate-limiting.workspace = true subtensor-swap-interface.workspace = true pallet-crowdloan.workspace = true @@ -50,6 +51,7 @@ std = [ "log/std", "pallet-admin-utils/std", "pallet-balances/std", + "pallet-rate-limiting/std", "pallet-evm-precompile-dispatch/std", "pallet-evm-precompile-modexp/std", "pallet-evm-precompile-sha3fips/std", diff --git a/precompiles/src/lib.rs b/precompiles/src/lib.rs index 8069a1eb92..fe08fbec5a 100644 --- a/precompiles/src/lib.rs +++ b/precompiles/src/lib.rs @@ -66,12 +66,17 @@ where + pallet_subtensor::Config + pallet_subtensor_swap::Config + pallet_proxy::Config - + pallet_crowdloan::Config, + + pallet_crowdloan::Config + + pallet_rate_limiting::Config< + LimitScope = subtensor_runtime_common::NetUid, + GroupId = subtensor_runtime_common::rate_limiting::GroupId, + >, R::AccountId: From<[u8; 32]> + ByteArray + Into<[u8; 32]>, ::RuntimeCall: From> + From> + From> + From> + + From> + From> + GetDispatchInfo + Dispatchable, @@ -93,12 +98,17 @@ where + pallet_subtensor::Config + pallet_subtensor_swap::Config + pallet_proxy::Config - + pallet_crowdloan::Config, + + pallet_crowdloan::Config + + pallet_rate_limiting::Config< + LimitScope = subtensor_runtime_common::NetUid, + GroupId = subtensor_runtime_common::rate_limiting::GroupId, + >, R::AccountId: From<[u8; 32]> + ByteArray + Into<[u8; 32]>, ::RuntimeCall: From> + From> + From> + From> + + From> + From> + GetDispatchInfo + Dispatchable, @@ -149,12 +159,17 @@ where + pallet_subtensor::Config + pallet_subtensor_swap::Config + pallet_proxy::Config - + pallet_crowdloan::Config, + + pallet_crowdloan::Config + + pallet_rate_limiting::Config< + LimitScope = subtensor_runtime_common::NetUid, + GroupId = subtensor_runtime_common::rate_limiting::GroupId, + >, R::AccountId: From<[u8; 32]> + ByteArray + Into<[u8; 32]>, ::RuntimeCall: From> + From> + From> + From> + + From> + From> + GetDispatchInfo + Dispatchable diff --git a/precompiles/src/subnet.rs b/precompiles/src/subnet.rs index b7f5cdb098..bdda32a6ce 100644 --- a/precompiles/src/subnet.rs +++ b/precompiles/src/subnet.rs @@ -4,9 +4,10 @@ use frame_support::dispatch::{GetDispatchInfo, PostDispatchInfo}; use frame_support::traits::ConstU32; use frame_system::RawOrigin; use pallet_evm::{AddressMapping, PrecompileHandle}; +use pallet_rate_limiting::{RateLimitKind, RateLimitTarget}; use precompile_utils::{EvmResult, prelude::BoundedString}; use sp_core::H256; -use sp_runtime::traits::Dispatchable; +use sp_runtime::traits::{Dispatchable, SaturatedConversion}; use sp_std::vec; use subtensor_runtime_common::{Currency, NetUid}; @@ -19,10 +20,15 @@ where R: frame_system::Config + pallet_evm::Config + pallet_subtensor::Config - + pallet_admin_utils::Config, + + pallet_admin_utils::Config + + pallet_rate_limiting::Config< + LimitScope = NetUid, + GroupId = subtensor_runtime_common::rate_limiting::GroupId, + >, R::AccountId: From<[u8; 32]>, ::RuntimeCall: From> + From> + + From> + GetDispatchInfo + Dispatchable, ::AddressMapping: AddressMapping, @@ -36,10 +42,15 @@ where R: frame_system::Config + pallet_evm::Config + pallet_subtensor::Config - + pallet_admin_utils::Config, + + pallet_admin_utils::Config + + pallet_rate_limiting::Config< + LimitScope = NetUid, + GroupId = subtensor_runtime_common::rate_limiting::GroupId, + >, R::AccountId: From<[u8; 32]>, ::RuntimeCall: From> + From> + + From> + GetDispatchInfo + Dispatchable, ::AddressMapping: AddressMapping, @@ -141,9 +152,9 @@ where #[precompile::public("getServingRateLimit(uint16)")] #[precompile::view] fn get_serving_rate_limit(_: &mut impl PrecompileHandle, netuid: u16) -> EvmResult { - Ok(pallet_subtensor::ServingRateLimit::::get(NetUid::from( - netuid, - ))) + Ok(pallet_subtensor::Pallet::::get_serving_rate_limit( + NetUid::from(netuid), + )) } #[precompile::public("setServingRateLimit(uint16,uint64)")] @@ -153,9 +164,10 @@ where netuid: u16, serving_rate_limit: u64, ) -> EvmResult<()> { - let call = pallet_admin_utils::Call::::sudo_set_serving_rate_limit { - netuid: netuid.into(), - serving_rate_limit, + let call = pallet_rate_limiting::Call::::set_rate_limit { + target: RateLimitTarget::Group(subtensor_runtime_common::rate_limiting::GROUP_SERVE), + scope: Some(netuid.into()), + limit: RateLimitKind::Exact(serving_rate_limit.saturated_into()), }; handle.try_dispatch_runtime_call::( diff --git a/runtime/Cargo.toml b/runtime/Cargo.toml index b3aced2160..baec2c6f30 100644 --- a/runtime/Cargo.toml +++ b/runtime/Cargo.toml @@ -38,9 +38,12 @@ frame-system = { workspace = true } frame-try-runtime = { workspace = true, optional = true } pallet-timestamp.workspace = true pallet-transaction-payment.workspace = true +pallet-rate-limiting.workspace = true +pallet-rate-limiting-runtime-api.workspace = true pallet-subtensor-utility.workspace = true frame-executive.workspace = true frame-metadata-hash-extension.workspace = true +serde.workspace = true sp-api.workspace = true sp-block-builder.workspace = true sp-consensus-aura.workspace = true @@ -52,6 +55,7 @@ sp-inherents.workspace = true sp-offchain.workspace = true sp-runtime.workspace = true sp-session.workspace = true +sp-io.workspace = true sp-std.workspace = true sp-transaction-pool.workspace = true sp-version.workspace = true @@ -156,7 +160,6 @@ ethereum.workspace = true [dev-dependencies] frame-metadata.workspace = true -sp-io.workspace = true sp-tracing.workspace = true [build-dependencies] @@ -190,6 +193,8 @@ std = [ "pallet-timestamp/std", "pallet-transaction-payment-rpc-runtime-api/std", "pallet-transaction-payment/std", + "pallet-rate-limiting/std", + "pallet-rate-limiting-runtime-api/std", "pallet-subtensor-utility/std", "pallet-sudo/std", "pallet-multisig/std", @@ -198,6 +203,7 @@ std = [ "pallet-preimage/std", "pallet-commitments/std", "precompile-utils/std", + "serde/std", "sp-api/std", "sp-block-builder/std", "sp-core/std", @@ -333,6 +339,7 @@ try-runtime = [ "pallet-insecure-randomness-collective-flip/try-runtime", "pallet-timestamp/try-runtime", "pallet-transaction-payment/try-runtime", + "pallet-rate-limiting/try-runtime", "pallet-subtensor-utility/try-runtime", "pallet-safe-mode/try-runtime", "pallet-subtensor/try-runtime", diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index 52746675f9..65a61ebfd8 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -8,10 +8,12 @@ #[cfg(feature = "std")] include!(concat!(env!("OUT_DIR"), "/wasm_binary.rs")); +use core::marker::PhantomData; use core::num::NonZeroU64; pub mod check_nonce; mod migrations; +pub mod rate_limiting; pub mod sudo_wrapper; pub mod transaction_payment_wrapper; @@ -45,6 +47,7 @@ use pallet_subtensor_proxy as pallet_proxy; use pallet_subtensor_swap_runtime_api::SimSwapResult; use pallet_subtensor_utility as pallet_utility; use runtime_common::prod_or_fast; +use scale_info::TypeInfo; use sp_api::impl_runtime_apis; use sp_consensus_aura::sr25519::AuthorityId as AuraId; use sp_consensus_babe::BabeConfiguration; @@ -73,6 +76,13 @@ use sp_version::RuntimeVersion; use subtensor_precompiles::Precompiles; use subtensor_runtime_common::{AlphaCurrency, TaoCurrency, time::*, *}; use subtensor_swap_interface::{Order, SwapHandler}; +use subtensor_transaction_fee::{SubtensorTxFeeHandler, TransactionFeeHandler}; +// Frontier +use fp_rpc::TransactionStatus; +use pallet_ethereum::{Call::transact, PostLogContent, Transaction as EthereumTransaction}; +use pallet_evm::{ + Account as EVMAccount, BalanceConverter, EvmBalance, FeeCalculator, Runner, SubstrateBalance, +}; // A few exports that help ease life for downstream crates. pub use frame_support::{ @@ -94,21 +104,13 @@ pub use pallet_balances::Call as BalancesCall; use pallet_commitments::GetCommitments; pub use pallet_timestamp::Call as TimestampCall; use pallet_transaction_payment::{ConstFeeMultiplier, Multiplier}; +pub use rate_limiting::{ + ScopeResolver as RuntimeScopeResolver, UsageResolver as RuntimeUsageResolver, +}; #[cfg(any(feature = "std", test))] pub use sp_runtime::BuildStorage; pub use sp_runtime::{Perbill, Permill}; -use subtensor_transaction_fee::{SubtensorTxFeeHandler, TransactionFeeHandler}; - -use core::marker::PhantomData; - -use scale_info::TypeInfo; - -// Frontier -use fp_rpc::TransactionStatus; -use pallet_ethereum::{Call::transact, PostLogContent, Transaction as EthereumTransaction}; -use pallet_evm::{ - Account as EVMAccount, BalanceConverter, EvmBalance, FeeCalculator, Runner, SubstrateBalance, -}; +pub use subtensor_runtime_common::rate_limiting::RateLimitUsageKey; // Drand impl pallet_drand::Config for Runtime { @@ -179,8 +181,11 @@ impl frame_system::offchain::CreateSignedTransaction SudoTransactionExtension::::new(), pallet_subtensor::transaction_extension::SubtensorTransactionExtension::::new( ), - pallet_drand::drand_priority::DrandPriority::::new(), - frame_metadata_hash_extension::CheckMetadataHash::::new(true), + ( + pallet_drand::drand_priority::DrandPriority::::new(), + frame_metadata_hash_extension::CheckMetadataHash::::new(true), + ), + pallet_rate_limiting::RateLimitTransactionExtension::::new(), ); let raw_payload = SignedPayload::new(call.clone(), extra.clone()).ok()?; @@ -244,7 +249,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion { spec_version: 365, impl_version: 1, apis: RUNTIME_API_VERSIONS, - transaction_version: 1, + transaction_version: 2, system_version: 1, }; @@ -1028,7 +1033,6 @@ parameter_types! { pub const SubtensorInitialWeightsVersionKey: u64 = 0; pub const SubtensorInitialMinDifficulty: u64 = 10_000_000; pub const SubtensorInitialMaxDifficulty: u64 = u64::MAX / 4; - pub const SubtensorInitialServingRateLimit: u64 = 50; pub const SubtensorInitialBurn: u64 = 100_000_000; // 0.1 tao pub const SubtensorInitialMinBurn: u64 = 500_000; // 500k RAO pub const SubtensorInitialMaxBurn: u64 = 100_000_000_000; // 100 tao @@ -1101,7 +1105,6 @@ impl pallet_subtensor::Config for Runtime { type InitialWeightsVersionKey = SubtensorInitialWeightsVersionKey; type InitialMaxDifficulty = SubtensorInitialMaxDifficulty; type InitialMinDifficulty = SubtensorInitialMinDifficulty; - type InitialServingRateLimit = SubtensorInitialServingRateLimit; type InitialBurn = SubtensorInitialBurn; type InitialMaxBurn = SubtensorInitialMaxBurn; type InitialMinBurn = SubtensorInitialMinBurn; @@ -1137,9 +1140,32 @@ impl pallet_subtensor::Config for Runtime { type GetCommitments = GetCommitmentsStruct; type MaxImmuneUidsPercentage = MaxImmuneUidsPercentage; type CommitmentsInterface = CommitmentsI; + type RateLimiting = RateLimiting; type EvmKeyAssociateRateLimit = EvmKeyAssociateRateLimit; } +parameter_types! { + pub const RateLimitingMaxGroupMembers: u32 = 64; + pub const RateLimitingMaxGroupNameLength: u32 = 64; +} + +impl pallet_rate_limiting::Config for Runtime { + type RuntimeCall = RuntimeCall; + type AdminOrigin = EnsureRoot; + type LimitSettingRule = rate_limiting::LimitSettingRule; + type DefaultLimitSettingRule = rate_limiting::DefaultLimitSettingRule; + type LimitSettingOrigin = rate_limiting::LimitSettingOrigin; + type LimitScope = NetUid; + type LimitScopeResolver = RuntimeScopeResolver; + type UsageKey = RateLimitUsageKey; + type UsageResolver = RuntimeUsageResolver; + type GroupId = subtensor_runtime_common::rate_limiting::GroupId; + type MaxGroupMembers = RateLimitingMaxGroupMembers; + type MaxGroupNameLength = RateLimitingMaxGroupNameLength; + #[cfg(feature = "runtime-benchmarks")] + type BenchmarkHelper = (); +} + parameter_types! { pub const SwapProtocolId: PalletId = PalletId(*b"ten/swap"); pub const SwapMaxFeeRate: u16 = 10000; // 15.26% @@ -1642,6 +1668,7 @@ construct_runtime!( Swap: pallet_subtensor_swap = 28, Contracts: pallet_contracts = 29, MevShield: pallet_shield = 30, + RateLimiting: pallet_rate_limiting = 31, } ); @@ -1652,6 +1679,8 @@ pub type Header = generic::Header; // Block type as expected by this runtime. pub type Block = generic::Block; // The extensions to the basic transaction logic. +// Note: The SDK only implements TransactionExtension for tuples up to 12 items, so we nest the last +// two extensions to keep order/encoding while staying under the limit. pub type TransactionExtensions = ( frame_system::CheckNonZeroSender, frame_system::CheckSpecVersion, @@ -1663,8 +1692,11 @@ pub type TransactionExtensions = ( ChargeTransactionPaymentWrapper, SudoTransactionExtension, pallet_subtensor::transaction_extension::SubtensorTransactionExtension, - pallet_drand::drand_priority::DrandPriority, - frame_metadata_hash_extension::CheckMetadataHash, + ( + pallet_drand::drand_priority::DrandPriority, + frame_metadata_hash_extension::CheckMetadataHash, + ), + pallet_rate_limiting::RateLimitTransactionExtension, ); type Migrations = ( @@ -1673,6 +1705,7 @@ type Migrations = ( pallet_subtensor::migrations::migrate_init_total_issuance::initialise_total_issuance::Migration< Runtime, >, + rate_limiting::migration::Migration, ); // Unchecked extrinsic type as expected by this runtime. @@ -2222,6 +2255,50 @@ impl_runtime_apis! { } } + impl pallet_rate_limiting_runtime_api::RateLimitingRuntimeApi for Runtime { + fn get_rate_limit( + pallet: Vec, + extrinsic: Vec, + ) -> Option { + use pallet_rate_limiting::{Pallet as RateLimiting, RateLimit}; + use pallet_rate_limiting_runtime_api::RateLimitRpcResponse; + + let pallet_name = sp_std::str::from_utf8(&pallet).ok()?; + let extrinsic_name = sp_std::str::from_utf8(&extrinsic).ok()?; + + let identifier = RateLimiting::::identifier_for_call_names( + pallet_name, + extrinsic_name, + )?; + let target = + RateLimiting::::config_target(&identifier).ok()?; + let limits = + pallet_rate_limiting::Limits::::get(target)?; + let default_limit = + pallet_rate_limiting::DefaultLimit::::get(); + let resolved = + RateLimiting::::resolved_limit(&target, &None); + + let (global, contextual) = match limits { + RateLimit::Global(kind) => (Some(kind), sp_std::vec::Vec::new()), + RateLimit::Scoped(entries) => ( + None, + entries + .into_iter() + .map(|(scope, kind)| (scope.encode(), kind)) + .collect(), + ), + }; + + Some(RateLimitRpcResponse { + global, + contextual, + default_limit, + resolved, + }) + } + } + impl pallet_contracts::ContractsApi for Runtime { diff --git a/runtime/src/rate_limiting/legacy.rs b/runtime/src/rate_limiting/legacy.rs new file mode 100644 index 0000000000..cab3d37719 --- /dev/null +++ b/runtime/src/rate_limiting/legacy.rs @@ -0,0 +1,305 @@ +use codec::{Decode, Encode}; +use frame_support::{Identity, migration::storage_key_iter}; +use runtime_common::prod_or_fast; +use scale_info::TypeInfo; +use sp_io::{ + hashing::twox_128, + storage::{self as io_storage, next_key}, +}; +use sp_std::{collections::btree_map::BTreeMap, vec::Vec}; +use subtensor_runtime_common::NetUid; + +use super::AccountId; +use crate::{ + SubtensorInitialNetworkRateLimit, SubtensorInitialTxChildKeyTakeRateLimit, + SubtensorInitialTxDelegateTakeRateLimit, SubtensorInitialTxRateLimit, +}; + +pub use types::{Hyperparameter, RateLimitKey, TransactionType}; + +const PALLET_PREFIX: &[u8] = b"SubtensorModule"; +const BLAKE2_128_PREFIX_LEN: usize = 16; + +pub mod storage { + use super::*; + + pub fn serving_rate_limits() -> (BTreeMap, u64) { + let items: Vec<_> = + storage_key_iter::(PALLET_PREFIX, b"ServingRateLimit").collect(); + let reads = items.len() as u64; + (items.into_iter().collect(), reads) + } + + pub fn weights_set_rate_limits() -> (BTreeMap, u64) { + let items: Vec<_> = + storage_key_iter::(PALLET_PREFIX, b"WeightsSetRateLimit") + .collect(); + let reads = items.len() as u64; + (items.into_iter().collect(), reads) + } + + pub fn tx_rate_limit() -> (u64, u64) { + value_with_default(b"TxRateLimit", defaults::tx_rate_limit()) + } + + pub fn tx_delegate_take_rate_limit() -> (u64, u64) { + value_with_default( + b"TxDelegateTakeRateLimit", + defaults::tx_delegate_take_rate_limit(), + ) + } + + pub fn tx_childkey_take_rate_limit() -> (u64, u64) { + value_with_default( + b"TxChildkeyTakeRateLimit", + defaults::tx_childkey_take_rate_limit(), + ) + } + + pub fn network_rate_limit() -> (u64, u64) { + value_with_default(b"NetworkRateLimit", defaults::network_rate_limit()) + } + + pub fn owner_hyperparam_rate_limit() -> (u64, u64) { + let (value, reads) = value_with_default::( + b"OwnerHyperparamRateLimit", + defaults::owner_hyperparam_rate_limit(), + ); + (u64::from(value), reads) + } + + pub fn weights_version_key_rate_limit() -> (u64, u64) { + value_with_default( + b"WeightsVersionKeyRateLimit", + defaults::weights_version_key_rate_limit(), + ) + } + + pub fn last_rate_limited_blocks() -> (Vec<(RateLimitKey, u64)>, u64) { + let entries: Vec<_> = storage_key_iter::, u64, Identity>( + PALLET_PREFIX, + b"LastRateLimitedBlock", + ) + .collect(); + let reads = entries.len() as u64; + (entries, reads) + } + + pub fn transaction_key_last_block() -> (Vec<((AccountId, NetUid, u16), u64)>, u64) { + let prefix = storage_prefix(PALLET_PREFIX, b"TransactionKeyLastBlock"); + let mut cursor = prefix.clone(); + let mut entries = Vec::new(); + + while let Some(next) = next_key(&cursor) { + if !next.starts_with(&prefix) { + break; + } + if let Some(value) = io_storage::get(&next) { + let key_bytes = &next[prefix.len()..]; + if let (Some(key), Some(decoded_value)) = ( + decode_transaction_key(key_bytes), + decode_value::(&value), + ) { + entries.push((key, decoded_value)); + } + } + cursor = next; + } + + let reads = entries.len() as u64; + (entries, reads) + } + + fn storage_prefix(pallet: &[u8], storage: &[u8]) -> Vec { + [twox_128(pallet), twox_128(storage)].concat() + } + + fn value_with_default(storage_name: &[u8], default: V) -> (V, u64) { + let key = storage_prefix(PALLET_PREFIX, storage_name); + let value = io_storage::get(&key) + .and_then(|bytes| Decode::decode(&mut &bytes[..]).ok()) + .unwrap_or(default); + (value, 1) + } + + fn decode_value(bytes: &[u8]) -> Option { + Decode::decode(&mut &bytes[..]).ok() + } + + fn decode_transaction_key( + encoded: &[u8], + ) -> Option<(AccountId, NetUid, u16)> { + if encoded.len() < BLAKE2_128_PREFIX_LEN { + return None; + } + let mut slice = &encoded[BLAKE2_128_PREFIX_LEN..]; + let account = AccountId::decode(&mut slice).ok()?; + let netuid = NetUid::decode(&mut slice).ok()?; + let tx_kind = u16::decode(&mut slice).ok()?; + + Some((account, netuid, tx_kind)) + } +} + +pub mod defaults { + use super::*; + + pub fn serving_rate_limit() -> u64 { + // SubtensorInitialServingRateLimit::get() + 50 + } + + pub fn weights_set_rate_limit() -> u64 { + 100 + } + + pub fn tx_rate_limit() -> u64 { + SubtensorInitialTxRateLimit::get() + } + + pub fn tx_delegate_take_rate_limit() -> u64 { + SubtensorInitialTxDelegateTakeRateLimit::get() + } + + pub fn tx_childkey_take_rate_limit() -> u64 { + SubtensorInitialTxChildKeyTakeRateLimit::get() + } + + pub fn network_rate_limit() -> u64 { + if cfg!(feature = "pow-faucet") { + 0 + } else { + SubtensorInitialNetworkRateLimit::get() + } + } + + pub fn owner_hyperparam_rate_limit() -> u16 { + 2 + } + + pub fn weights_version_key_rate_limit() -> u64 { + 5 + } + + pub fn sn_owner_hotkey_rate_limit() -> u64 { + 50_400 + } + + pub fn mechanism_count_rate_limit() -> u64 { + prod_or_fast!(7_200, 1) + } + + pub fn mechanism_emission_rate_limit() -> u64 { + prod_or_fast!(7_200, 1) + } + + pub fn max_uids_trimming_rate_limit() -> u64 { + prod_or_fast!(30 * 7_200, 1) + } +} + +pub mod types { + use super::*; + + #[derive(Encode, Decode, Clone, PartialEq, Eq, Debug, TypeInfo)] + pub enum RateLimitKey { + #[codec(index = 0)] + SetSNOwnerHotkey(NetUid), + #[codec(index = 1)] + OwnerHyperparamUpdate(NetUid, Hyperparameter), + #[codec(index = 2)] + NetworkLastRegistered, + #[codec(index = 3)] + LastTxBlock(AccountId), + #[codec(index = 4)] + LastTxBlockChildKeyTake(AccountId), + #[codec(index = 5)] + LastTxBlockDelegateTake(AccountId), + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + #[non_exhaustive] + pub enum TransactionType { + SetChildren, + SetChildkeyTake, + Unknown, + RegisterNetwork, + SetWeightsVersionKey, + SetSNOwnerHotkey, + OwnerHyperparamUpdate(Hyperparameter), + MechanismCountUpdate, + MechanismEmission, + MaxUidsTrimming, + } + + impl From for TransactionType { + fn from(value: u16) -> Self { + match value { + 0 => TransactionType::SetChildren, + 1 => TransactionType::SetChildkeyTake, + 3 => TransactionType::RegisterNetwork, + 4 => TransactionType::SetWeightsVersionKey, + 5 => TransactionType::SetSNOwnerHotkey, + 6 => TransactionType::OwnerHyperparamUpdate(Hyperparameter::Unknown), + 7 => TransactionType::MechanismCountUpdate, + 8 => TransactionType::MechanismEmission, + 9 => TransactionType::MaxUidsTrimming, + _ => TransactionType::Unknown, + } + } + } + + impl From for u16 { + fn from(tx_type: TransactionType) -> Self { + match tx_type { + TransactionType::SetChildren => 0, + TransactionType::SetChildkeyTake => 1, + TransactionType::Unknown => 2, + TransactionType::RegisterNetwork => 3, + TransactionType::SetWeightsVersionKey => 4, + TransactionType::SetSNOwnerHotkey => 5, + TransactionType::OwnerHyperparamUpdate(_) => 6, + TransactionType::MechanismCountUpdate => 7, + TransactionType::MechanismEmission => 8, + TransactionType::MaxUidsTrimming => 9, + } + } + } + + #[derive(Encode, Decode, Clone, Copy, PartialEq, Eq, Debug, TypeInfo)] + #[non_exhaustive] + pub enum Hyperparameter { + Unknown = 0, + ServingRateLimit = 1, + MaxDifficulty = 2, + AdjustmentAlpha = 3, + MaxWeightLimit = 4, + ImmunityPeriod = 5, + MinAllowedWeights = 6, + Kappa = 7, + Rho = 8, + ActivityCutoff = 9, + PowRegistrationAllowed = 10, + MinBurn = 11, + MaxBurn = 12, + BondsMovingAverage = 13, + BondsPenalty = 14, + CommitRevealEnabled = 15, + LiquidAlphaEnabled = 16, + AlphaValues = 17, + WeightCommitInterval = 18, + TransferEnabled = 19, + AlphaSigmoidSteepness = 20, + Yuma3Enabled = 21, + BondsResetEnabled = 22, + ImmuneNeuronLimit = 23, + RecycleOrBurn = 24, + MaxAllowedUids = 25, + } + + impl From for TransactionType { + fn from(param: Hyperparameter) -> Self { + Self::OwnerHyperparamUpdate(param) + } + } +} diff --git a/runtime/src/rate_limiting/migration.rs b/runtime/src/rate_limiting/migration.rs new file mode 100644 index 0000000000..360933bb48 --- /dev/null +++ b/runtime/src/rate_limiting/migration.rs @@ -0,0 +1,1145 @@ +use core::{convert::TryFrom, marker::PhantomData}; + +use frame_support::{BoundedBTreeSet, BoundedVec, weights::Weight}; +use frame_system::pallet_prelude::BlockNumberFor; +use log::{info, warn}; +use pallet_rate_limiting::{ + GroupSharing, RateLimit, RateLimitGroup, RateLimitKind, RateLimitTarget, TransactionIdentifier, +}; +use pallet_subtensor::{ + self, AssociatedEvmAddress, Axons, Config as SubtensorConfig, HasMigrationRun, LastUpdate, + Pallet, Prometheus, +}; +use sp_runtime::traits::SaturatedConversion; +use sp_std::{ + collections::{btree_map::BTreeMap, btree_set::BTreeSet}, + vec, + vec::Vec, +}; +use subtensor_runtime_common::{ + NetUid, + rate_limiting::{ + GROUP_DELEGATE_TAKE, GROUP_OWNER_HPARAMS, GROUP_REGISTER_NETWORK, GROUP_SERVE, + GROUP_STAKING_OPS, GROUP_SWAP_KEYS, GROUP_WEIGHTS_SUBNET, GroupId, + }, +}; + +use super::{ + AccountId, LimitSettingRule, RateLimitUsageKey, Runtime, + legacy::{ + Hyperparameter, RateLimitKey, TransactionType, defaults as legacy_defaults, + storage as legacy_storage, + }, +}; + +type GroupNameOf = BoundedVec::MaxGroupNameLength>; +type GroupMembersOf = + BoundedBTreeSet::MaxGroupMembers>; + +// Pallet index assigned to `pallet_subtensor` in `construct_runtime!`. +const SUBTENSOR_PALLET_INDEX: u8 = 7; +// Pallet index assigned to `pallet_admin_utils` in `construct_runtime!`. +const ADMIN_UTILS_PALLET_INDEX: u8 = 19; + +/// Marker stored in `pallet_subtensor::HasMigrationRun` once the migration finishes. +pub const MIGRATION_NAME: &[u8] = b"migrate_rate_limiting"; + +// `set_children` is rate-limited to once every 150 blocks, it's hard-coded in the legacy code. +const SET_CHILDREN_RATE_LIMIT: u64 = 150; + +// Hyperparameter extrinsics routed through owner-or-root rate limiting. +const HYPERPARAMETERS: &[Hyperparameter] = &[ + Hyperparameter::ServingRateLimit, + Hyperparameter::MaxDifficulty, + Hyperparameter::AdjustmentAlpha, + Hyperparameter::ImmunityPeriod, + Hyperparameter::MinAllowedWeights, + Hyperparameter::MaxAllowedUids, + Hyperparameter::Kappa, + Hyperparameter::Rho, + Hyperparameter::ActivityCutoff, + Hyperparameter::PowRegistrationAllowed, + Hyperparameter::MinBurn, + Hyperparameter::MaxBurn, + Hyperparameter::BondsMovingAverage, + Hyperparameter::BondsPenalty, + Hyperparameter::CommitRevealEnabled, + Hyperparameter::LiquidAlphaEnabled, + Hyperparameter::AlphaValues, + Hyperparameter::WeightCommitInterval, + Hyperparameter::TransferEnabled, + Hyperparameter::AlphaSigmoidSteepness, + Hyperparameter::Yuma3Enabled, + Hyperparameter::BondsResetEnabled, + Hyperparameter::ImmuneNeuronLimit, + Hyperparameter::RecycleOrBurn, +]; + +/// Runtime hook that executes the rate-limiting migration. +pub struct Migration(PhantomData); + +impl frame_support::traits::OnRuntimeUpgrade for Migration +where + T: SubtensorConfig + pallet_rate_limiting::Config, + RateLimitUsageKey: Into<::UsageKey>, +{ + fn on_runtime_upgrade() -> Weight { + migrate_rate_limiting() + } +} + +pub fn migrate_rate_limiting() -> Weight { + let mut weight = ::DbWeight::get().reads(1); + if HasMigrationRun::::get(MIGRATION_NAME) { + info!("Rate-limiting migration already executed. Skipping."); + return weight; + } + + let (groups, commits, reads) = commits(); + weight = weight.saturating_add(::DbWeight::get().reads(reads)); + + let (limit_commits, last_seen_commits) = commits.into_iter().fold( + (Vec::new(), Vec::new()), + |(mut limits, mut seen), commit| { + match commit.kind { + CommitKind::Limit(limit) => limits.push((commit.target, limit)), + CommitKind::LastSeen(ls) => seen.push((commit.target, ls)), + } + (limits, seen) + }, + ); + + let (group_writes, group_count) = migrate_grouping(&groups); + let (limit_writes, limits_len) = migrate_limits(limit_commits); + let (last_seen_writes, last_seen_len) = migrate_last_seen(last_seen_commits); + + let mut writes = group_writes + .saturating_add(limit_writes) + .saturating_add(last_seen_writes); + + // Legacy parity: serving-rate-limit configuration is allowed for root OR subnet owner. + // Everything else remains default (`AdminOrigin` / root in this runtime). + pallet_rate_limiting::LimitSettingRules::::insert( + RateLimitTarget::Group(GROUP_SERVE), + LimitSettingRule::RootOrSubnetOwnerAdminWindow, + ); + writes += 1; + + HasMigrationRun::::insert(MIGRATION_NAME, true); + writes += 1; + + weight = + weight.saturating_add(::DbWeight::get().writes(writes)); + + info!( + "New migration wrote {} limits, {} last-seen entries, and {} groups into pallet-rate-limiting", + limits_len, last_seen_len, group_count + ); + + weight +} + +// Main entrypoint: build all groups and commits, along with storage reads. +fn commits() -> (Vec, Vec, u64) { + let mut groups = Vec::new(); + let mut commits = Vec::new(); + + // grouped + let mut reads = build_serving(&mut groups, &mut commits); + reads = reads.saturating_add(build_delegate_take(&mut groups, &mut commits)); + reads = reads.saturating_add(build_weights(&mut groups, &mut commits)); + reads = reads.saturating_add(build_register_network(&mut groups, &mut commits)); + reads = reads.saturating_add(build_owner_hparams(&mut groups, &mut commits)); + reads = reads.saturating_add(build_staking_ops(&mut groups, &mut commits)); + reads = reads.saturating_add(build_swap_keys(&mut groups, &mut commits)); + + // standalone + reads = reads.saturating_add(build_childkey_take(&mut commits)); + reads = reads.saturating_add(build_set_children(&mut commits)); + reads = reads.saturating_add(build_weights_version_key(&mut commits)); + reads = reads.saturating_add(build_sn_owner_hotkey(&mut commits)); + reads = reads.saturating_add(build_associate_evm(&mut commits)); + reads = reads.saturating_add(build_mechanism_count(&mut commits)); + reads = reads.saturating_add(build_mechanism_emission(&mut commits)); + reads = reads.saturating_add(build_trim_max_uids(&mut commits)); + + (groups, commits, reads) +} + +fn migrate_grouping(groups: &[GroupConfig]) -> (u64, usize) { + let mut writes: u64 = 0; + let mut max_group_id: Option = None; + + for group in groups { + let Ok(name) = GroupNameOf::::try_from(group.name.clone()) else { + warn!( + "rate-limiting migration: group name exceeds bounds, skipping id {}", + group.id + ); + continue; + }; + + pallet_rate_limiting::Groups::::insert( + group.id, + RateLimitGroup { + id: group.id, + name: name.clone(), + sharing: group.sharing, + }, + ); + pallet_rate_limiting::GroupNameIndex::::insert(name, group.id); + writes += 2; + + let mut member_set = BTreeSet::new(); + for call in &group.members { + member_set.insert(call.identifier()); + pallet_rate_limiting::CallGroups::::insert(call.identifier(), group.id); + writes += 1; + if call.read_only { + pallet_rate_limiting::CallReadOnly::::insert(call.identifier(), true); + writes += 1; + } + } + let Ok(bounded) = GroupMembersOf::::try_from(member_set) else { + warn!( + "rate-limiting migration: group {} has too many members, skipping assignment", + group.id + ); + continue; + }; + pallet_rate_limiting::GroupMembers::::insert(group.id, bounded); + writes += 1; + + max_group_id = Some(max_group_id.map_or(group.id, |current| current.max(group.id))); + } + + let next_group_id = max_group_id.map_or(0, |id| id.saturating_add(1)); + pallet_rate_limiting::NextGroupId::::put(next_group_id); + writes += 1; + + (writes, groups.len()) +} + +fn migrate_limits(limit_commits: Vec<(RateLimitTarget, MigratedLimit)>) -> (u64, usize) { + let mut writes: u64 = 0; + let mut limits: BTreeMap, RateLimit>> = + BTreeMap::new(); + + for (target, MigratedLimit { span, scope }) in limit_commits { + let entry = limits.entry(target).or_insert_with(|| match scope { + Some(s) => RateLimit::scoped_single(s, RateLimitKind::Exact(span)), + None => RateLimit::global(RateLimitKind::Exact(span)), + }); + + if let Some(netuid) = scope { + match entry { + RateLimit::Global(_) => { + *entry = RateLimit::scoped_single(netuid, RateLimitKind::Exact(span)); + } + RateLimit::Scoped(map) => { + map.insert(netuid, RateLimitKind::Exact(span)); + } + } + } else { + *entry = RateLimit::global(RateLimitKind::Exact(span)); + } + } + + let len = limits.len(); + for (target, limit) in limits { + pallet_rate_limiting::Limits::::insert(target, limit); + writes += 1; + } + + (writes, len) +} + +fn migrate_last_seen( + last_seen_commits: Vec<(RateLimitTarget, MigratedLastSeen)>, +) -> (u64, usize) { + let mut writes: u64 = 0; + let mut last_seen: BTreeMap< + ( + RateLimitTarget, + Option>, + ), + BlockNumberFor, + > = BTreeMap::new(); + + for (target, MigratedLastSeen { block, usage }) in last_seen_commits { + let key = (target, usage); + last_seen + .entry(key) + .and_modify(|existing| { + if block > *existing { + *existing = block; + } + }) + .or_insert(block); + } + + let len = last_seen.len(); + for ((target, usage), block) in last_seen { + pallet_rate_limiting::LastSeen::::insert(target, usage, block); + writes += 1; + } + + (writes, len) +} + +// Serving group (config+usage shared). +// scope: netuid +// usage: account+netuid, but different keys (endpoint value) for axon/prometheus +// legacy sources: ServingRateLimit (per netuid), Axons/Prometheus +fn build_serving(groups: &mut Vec, commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + // Create the group with all its members. + groups.push(GroupConfig { + id: GROUP_SERVE, + name: b"serving".to_vec(), + sharing: GroupSharing::ConfigAndUsage, + members: vec![ + MigratedCall::subtensor(4, false), // serve_axon + MigratedCall::subtensor(40, false), // serve_axon_tls + MigratedCall::subtensor(5, false), // serve_prometheus + ], + }); + + let (serving_limits, serving_reads) = legacy_storage::serving_rate_limits(); + reads = reads.saturating_add(serving_reads); + // Limits per netuid (written to the group target). + // Merge live subnets (which may rely on default rate-limit values) with any legacy entries that + // exist only in storage, so we migrate both current and previously stored netuids without + // duplicates. + let mut netuids = Pallet::::get_all_subnet_netuids(); + for (&netuid, _) in &serving_limits { + if !netuids.contains(&netuid) { + netuids.push(netuid); + } + } + let default_limit = legacy_defaults::serving_rate_limit(); + for netuid in netuids { + reads = reads.saturating_add(1); + push_limit_commit_if_non_zero( + commits, + RateLimitTarget::Group(GROUP_SERVE), + serving_limits + .get(&netuid) + .copied() + .unwrap_or(default_limit), + Some(netuid), + ); + } + + // Axon last-seen (group-shared usage). + for (netuid, hotkey, axon) in Axons::::iter() { + reads = reads.saturating_add(1); + if let Some(block) = block_number::(axon.block) { + commits.push(Commit { + target: RateLimitTarget::Group(GROUP_SERVE), + kind: CommitKind::LastSeen(MigratedLastSeen { + block, + usage: Some(RateLimitUsageKey::AccountSubnetServing { + account: hotkey.clone(), + netuid, + endpoint: crate::rate_limiting::ServingEndpoint::Axon, + }), + }), + }); + } + } + + // Prometheus last-seen (group-shared usage). + for (netuid, hotkey, prom) in Prometheus::::iter() { + reads = reads.saturating_add(1); + if let Some(block) = block_number::(prom.block) { + commits.push(Commit { + target: RateLimitTarget::Group(GROUP_SERVE), + kind: CommitKind::LastSeen(MigratedLastSeen { + block, + usage: Some(RateLimitUsageKey::AccountSubnetServing { + account: hotkey, + netuid, + endpoint: crate::rate_limiting::ServingEndpoint::Prometheus, + }), + }), + }); + } + } + + reads +} + +// Delegate take group (config + usage shared). +// usage: account +// legacy sources: TxDelegateTakeRateLimit, LastTxBlockDelegateTake +fn build_delegate_take(groups: &mut Vec, commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + groups.push(GroupConfig { + id: GROUP_DELEGATE_TAKE, + name: b"delegate-take".to_vec(), + sharing: GroupSharing::ConfigAndUsage, + members: vec![ + MigratedCall::subtensor(66, false), // increase_take + MigratedCall::subtensor(65, false), // decrease_take + ], + }); + + let target = RateLimitTarget::Group(GROUP_DELEGATE_TAKE); + let (delegate_take_limit, delegate_reads) = legacy_storage::tx_delegate_take_rate_limit(); + reads = reads.saturating_add(delegate_reads); + push_limit_commit_if_non_zero(commits, target, delegate_take_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_last_rate_limited_block( + commits, + |key| match key { + RateLimitKey::LastTxBlockDelegateTake(account) => { + Some((target, Some(RateLimitUsageKey::Account(account)))) + } + _ => None, + }, + ), + ); + + reads +} + +// Weights group (config + usage shared). +// scope: netuid +// usage: netuid+neuron/netuid+mechanism+neuron +// legacy source: SubnetWeightsSetRateLimit, LastUpdate (subnet/mechanism) +fn build_weights(groups: &mut Vec, commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + groups.push(GroupConfig { + id: GROUP_WEIGHTS_SUBNET, + name: b"weights".to_vec(), + sharing: GroupSharing::ConfigAndUsage, + members: vec![ + MigratedCall::subtensor(0, false), // set_weights + MigratedCall::subtensor(96, false), // commit_weights + MigratedCall::subtensor(100, false), // batch_commit_weights + MigratedCall::subtensor(113, false), // commit_timelocked_weights + MigratedCall::subtensor(97, false), // reveal_weights + MigratedCall::subtensor(98, false), // batch_reveal_weights + MigratedCall::subtensor(119, false), // set_mechanism_weights + MigratedCall::subtensor(115, false), // commit_mechanism_weights + MigratedCall::subtensor(117, false), // commit_crv3_mechanism_weights + MigratedCall::subtensor(118, false), // commit_timelocked_mechanism_weights + MigratedCall::subtensor(116, false), // reveal_mechanism_weights + ], + }); + + let (weights_limits, weights_reads) = legacy_storage::weights_set_rate_limits(); + reads = reads.saturating_add(weights_reads); + let default_limit = legacy_defaults::weights_set_rate_limit(); + for netuid in Pallet::::get_all_subnet_netuids() { + reads = reads.saturating_add(1); + push_limit_commit_if_non_zero( + commits, + RateLimitTarget::Group(GROUP_WEIGHTS_SUBNET), + weights_limits + .get(&netuid) + .copied() + .unwrap_or(default_limit), + Some(netuid), + ); + } + + for (index, blocks) in LastUpdate::::iter() { + reads = reads.saturating_add(1); + let (netuid, mecid) = + Pallet::::get_netuid_and_subid(index).unwrap_or((NetUid::ROOT, 0.into())); + for (uid, last_block) in blocks.into_iter().enumerate() { + let Some(block) = block_number::(last_block) else { + continue; + }; + let Ok(uid_u16) = u16::try_from(uid) else { + continue; + }; + let usage = if mecid == 0.into() { + RateLimitUsageKey::SubnetNeuron { + netuid, + uid: uid_u16, + } + } else { + RateLimitUsageKey::SubnetMechanismNeuron { + netuid, + mecid, + uid: uid_u16, + } + }; + commits.push(Commit { + target: RateLimitTarget::Group(GROUP_WEIGHTS_SUBNET), + kind: CommitKind::LastSeen(MigratedLastSeen { + block, + usage: Some(usage), + }), + }); + } + } + + reads +} + +// Register network group (config + usage shared). +// legacy sources: NetworkRateLimit, NetworkLastRegistered +fn build_register_network(groups: &mut Vec, commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + groups.push(GroupConfig { + id: GROUP_REGISTER_NETWORK, + name: b"register-network".to_vec(), + sharing: GroupSharing::ConfigAndUsage, + members: vec![ + MigratedCall::subtensor(59, false), // register_network + MigratedCall::subtensor(79, false), // register_network_with_identity + ], + }); + + let target = RateLimitTarget::Group(GROUP_REGISTER_NETWORK); + let (network_rate_limit, network_reads) = legacy_storage::network_rate_limit(); + reads = reads.saturating_add(network_reads); + push_limit_commit_if_non_zero(commits, target, network_rate_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_last_rate_limited_block( + commits, + |key| match key { + RateLimitKey::NetworkLastRegistered => Some((target, None)), + _ => None, + }, + ), + ); + + reads +} + +// Owner hyperparameter group (config shared, usage per call). +// usage: netuid +// legacy sources: OwnerHyperparamRateLimit * tempo, LastRateLimitedBlock per OwnerHyperparamUpdate +fn build_owner_hparams(groups: &mut Vec, commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + groups.push(GroupConfig { + id: GROUP_OWNER_HPARAMS, + name: b"owner-hparams".to_vec(), + sharing: GroupSharing::ConfigOnly, + members: HYPERPARAMETERS + .iter() + .filter_map(|h| identifier_for_hyperparameter(*h)) + .collect(), + }); + + let group_target = RateLimitTarget::Group(GROUP_OWNER_HPARAMS); + let (owner_limit, owner_reads) = legacy_storage::owner_hyperparam_rate_limit(); + reads = reads.saturating_add(owner_reads); + push_limit_commit_if_non_zero(commits, group_target, owner_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_last_rate_limited_block( + commits, + |key| match key { + RateLimitKey::OwnerHyperparamUpdate(netuid, hyper) => { + let Some(identifier) = identifier_for_hyperparameter(hyper) else { + return None; + }; + Some(( + RateLimitTarget::Transaction(identifier.identifier()), + Some(RateLimitUsageKey::Subnet(netuid)), + )) + } + _ => None, + }, + ), + ); + + reads +} + +// Staking ops group (config + usage shared, all ops 1 block). +// usage: coldkey+hotkey+netuid +// legacy sources: TxRateLimit (reset every block for staking ops), StakingOperationRateLimiter +fn build_staking_ops(groups: &mut Vec, commits: &mut Vec) -> u64 { + groups.push(GroupConfig { + id: GROUP_STAKING_OPS, + name: b"staking-ops".to_vec(), + sharing: GroupSharing::ConfigAndUsage, + members: vec![ + MigratedCall::subtensor(2, false), // add_stake + MigratedCall::subtensor(88, false), // add_stake_limit + MigratedCall::subtensor(3, true), // remove_stake + MigratedCall::subtensor(89, true), // remove_stake_limit + MigratedCall::subtensor(103, true), // remove_stake_full_limit + MigratedCall::subtensor(85, false), // move_stake + MigratedCall::subtensor(86, true), // transfer_stake + MigratedCall::subtensor(87, false), // swap_stake + MigratedCall::subtensor(90, false), // swap_stake_limit + ], + }); + + push_limit_commit_if_non_zero(commits, RateLimitTarget::Group(GROUP_STAKING_OPS), 1, None); + + // we don't need to migrate last-seen since the limiter is reset every block. + + 0 +} + +// Swap hotkey/coldkey share the lock and usage; swap_coldkey bypasses enforcement but records +// usage. +// usage: account (coldkey) +// legacy sources: TxRateLimit, LastRateLimitedBlock per LastTxBlock +fn build_swap_keys(groups: &mut Vec, commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + groups.push(GroupConfig { + id: GROUP_SWAP_KEYS, + name: b"swap-keys".to_vec(), + sharing: GroupSharing::ConfigAndUsage, + members: vec![ + MigratedCall::subtensor(70, false), // swap_hotkey + MigratedCall::subtensor(71, false), // swap_coldkey + ], + }); + + let target = RateLimitTarget::Group(GROUP_SWAP_KEYS); + let (tx_rate_limit, tx_reads) = legacy_storage::tx_rate_limit(); + reads = reads.saturating_add(tx_reads); + push_limit_commit_if_non_zero(commits, target, tx_rate_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_last_rate_limited_block( + commits, + |key| match key { + RateLimitKey::LastTxBlock(account) => { + Some((target, Some(RateLimitUsageKey::Account(account)))) + } + _ => None, + }, + ), + ); + + reads +} + +// Standalone set_childkey_take. +// usage: account+netuid +// legacy sources: TxChildkeyTakeRateLimit, TransactionKeyLastBlock per SetChildkeyTake +fn build_childkey_take(commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + let target = + RateLimitTarget::Transaction(TransactionIdentifier::new(SUBTENSOR_PALLET_INDEX, 75)); + let (childkey_limit, childkey_reads) = legacy_storage::tx_childkey_take_rate_limit(); + reads = reads.saturating_add(childkey_reads); + push_limit_commit_if_non_zero(commits, target, childkey_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_transaction_key_last_block( + commits, + target, + TransactionType::SetChildkeyTake, + ), + ); + + reads +} + +// Standalone set_children. +// usage: account+netuid +// legacy sources: SET_CHILDREN_RATE_LIMIT (constant 150), TransactionKeyLastBlock per SetChildren +fn build_set_children(commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + let target = + RateLimitTarget::Transaction(TransactionIdentifier::new(SUBTENSOR_PALLET_INDEX, 67)); + push_limit_commit_if_non_zero(commits, target, SET_CHILDREN_RATE_LIMIT, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_transaction_key_last_block( + commits, + target, + TransactionType::SetChildren, + ), + ); + + reads +} + +// Standalone set_weights_version_key. +// scope: netuid +// usage: account+netuid +// legacy sources: WeightsVersionKeyRateLimit * tempo, +// TransactionKeyLastBlock per SetWeightsVersionKey +fn build_weights_version_key(commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + let target = + RateLimitTarget::Transaction(TransactionIdentifier::new(ADMIN_UTILS_PALLET_INDEX, 6)); + let (weights_version_limit, weights_version_reads) = + legacy_storage::weights_version_key_rate_limit(); + reads = reads.saturating_add(weights_version_reads); + push_limit_commit_if_non_zero(commits, target, weights_version_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_transaction_key_last_block( + commits, + target, + TransactionType::SetWeightsVersionKey, + ), + ); + + reads +} + +// Standalone set_sn_owner_hotkey. +// usage: netuid +// legacy sources: DefaultSetSNOwnerHotkeyRateLimit, LastRateLimitedBlock per SetSNOwnerHotkey +fn build_sn_owner_hotkey(commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + let target = + RateLimitTarget::Transaction(TransactionIdentifier::new(ADMIN_UTILS_PALLET_INDEX, 67)); + let sn_owner_limit = legacy_defaults::sn_owner_hotkey_rate_limit(); + reads += 1; + push_limit_commit_if_non_zero(commits, target, sn_owner_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_last_rate_limited_block( + commits, + |key| match key { + RateLimitKey::SetSNOwnerHotkey(netuid) => { + Some((target, Some(RateLimitUsageKey::Subnet(netuid)))) + } + _ => None, + }, + ), + ); + + reads +} + +// Standalone associate_evm_key. +// usage: netuid+neuron +// legacy sources: EvmKeyAssociateRateLimit, AssociatedEvmAddress +fn build_associate_evm(commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + let target = + RateLimitTarget::Transaction(TransactionIdentifier::new(SUBTENSOR_PALLET_INDEX, 93)); + reads += 1; + push_limit_commit_if_non_zero( + commits, + target, + ::EvmKeyAssociateRateLimit::get(), + None, + ); + + for (netuid, uid, (_, block)) in AssociatedEvmAddress::::iter() { + reads = reads.saturating_add(1); + let Some(block) = block_number::(block) else { + continue; + }; + commits.push(Commit { + target, + kind: CommitKind::LastSeen(MigratedLastSeen { + block, + usage: Some(RateLimitUsageKey::SubnetNeuron { netuid, uid }), + }), + }); + } + + reads +} + +// Standalone mechanism count. +// usage: account+netuid +// legacy sources: MechanismCountSetRateLimit, TransactionKeyLastBlock per MechanismCountUpdate +// sudo_set_mechanism_count +fn build_mechanism_count(commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + let target = + RateLimitTarget::Transaction(TransactionIdentifier::new(ADMIN_UTILS_PALLET_INDEX, 76)); + let mechanism_limit = legacy_defaults::mechanism_count_rate_limit(); + push_limit_commit_if_non_zero(commits, target, mechanism_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_transaction_key_last_block( + commits, + target, + TransactionType::MechanismCountUpdate, + ), + ); + + reads +} + +// Standalone mechanism emission. +// usage: account+netuid +// legacy sources: MechanismEmissionRateLimit, TransactionKeyLastBlock per MechanismEmission +// sudo_set_mechanism_emission_split +fn build_mechanism_emission(commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + let target = + RateLimitTarget::Transaction(TransactionIdentifier::new(ADMIN_UTILS_PALLET_INDEX, 77)); + let emission_limit = legacy_defaults::mechanism_emission_rate_limit(); + push_limit_commit_if_non_zero(commits, target, emission_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_transaction_key_last_block( + commits, + target, + TransactionType::MechanismEmission, + ), + ); + + reads +} + +// Standalone trim_to_max_allowed_uids. +// usage: account+netuid +// legacy sources: MaxUidsTrimmingRateLimit, TransactionKeyLastBlock per MaxUidsTrimming +// sudo_trim_to_max_allowed_uids +fn build_trim_max_uids(commits: &mut Vec) -> u64 { + let mut reads: u64 = 0; + let target = + RateLimitTarget::Transaction(TransactionIdentifier::new(ADMIN_UTILS_PALLET_INDEX, 78)); + let trim_limit = legacy_defaults::max_uids_trimming_rate_limit(); + push_limit_commit_if_non_zero(commits, target, trim_limit, None); + + reads = reads.saturating_add( + last_seen_helpers::collect_last_seen_from_transaction_key_last_block( + commits, + target, + TransactionType::MaxUidsTrimming, + ), + ); + + reads +} + +struct Commit { + target: RateLimitTarget, + kind: CommitKind, +} + +enum CommitKind { + Limit(MigratedLimit), + LastSeen(MigratedLastSeen), +} + +struct MigratedLimit { + span: BlockNumberFor, + scope: Option, +} + +struct MigratedLastSeen { + block: BlockNumberFor, + usage: Option>, +} + +struct GroupConfig { + id: GroupId, + name: Vec, + sharing: GroupSharing, + members: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct MigratedCall { + identifier: TransactionIdentifier, + read_only: bool, +} + +impl MigratedCall { + const fn new(pallet_index: u8, call_index: u8, read_only: bool) -> Self { + Self { + identifier: TransactionIdentifier::new(pallet_index, call_index), + read_only, + } + } + + const fn subtensor(call_index: u8, read_only: bool) -> Self { + Self::new(SUBTENSOR_PALLET_INDEX, call_index, read_only) + } + + const fn admin(call_index: u8, read_only: bool) -> Self { + Self::new(ADMIN_UTILS_PALLET_INDEX, call_index, read_only) + } + + pub fn identifier(&self) -> TransactionIdentifier { + self.identifier + } +} + +fn push_limit_commit_if_non_zero( + commits: &mut Vec, + target: RateLimitTarget, + span: u64, + scope: Option, +) { + if let Some(span) = block_number::(span) { + commits.push(Commit { + target, + kind: CommitKind::Limit(MigratedLimit { span, scope }), + }); + } +} + +mod last_seen_helpers { + use core::mem::discriminant; + + use super::*; + + pub(super) fn collect_last_seen_from_last_rate_limited_block( + commits: &mut Vec, + map: impl Fn( + RateLimitKey, + ) -> Option<( + RateLimitTarget, + Option>, + )>, + ) -> u64 { + let mut reads: u64 = 0; + + let (entries, iter_reads) = legacy_storage::last_rate_limited_blocks(); + reads = reads.saturating_add(iter_reads); + for (key, block) in entries { + let Some((target, usage)) = map(key) else { + continue; + }; + let Some(block) = block_number::(block) else { + continue; + }; + commits.push(Commit { + target, + kind: CommitKind::LastSeen(MigratedLastSeen { block, usage }), + }); + } + + reads + } + + pub(super) fn collect_last_seen_from_transaction_key_last_block( + commits: &mut Vec, + target: RateLimitTarget, + tx_filter: TransactionType, + ) -> u64 { + let mut reads: u64 = 0; + + let (entries, iter_reads) = legacy_storage::transaction_key_last_block(); + reads = reads.saturating_add(iter_reads); + for ((account, netuid, tx_kind), block) in entries { + let tx = TransactionType::from(tx_kind); + if discriminant(&tx) != discriminant(&tx_filter) { + continue; + } + let Some(usage) = usage_key_from_transaction_type(tx, &account, netuid) else { + continue; + }; + let Some(block) = block_number::(block) else { + continue; + }; + commits.push(Commit { + target, + kind: CommitKind::LastSeen(MigratedLastSeen { + block, + usage: Some(usage), + }), + }); + } + + reads + } +} + +// Produces the usage key for a `TransactionType` that was stored in `TransactionKeyLastBlock`. +fn usage_key_from_transaction_type( + tx: TransactionType, + account: &AccountId, + netuid: NetUid, +) -> Option> { + match tx { + TransactionType::MechanismCountUpdate + | TransactionType::MaxUidsTrimming + | TransactionType::MechanismEmission + | TransactionType::SetChildkeyTake + | TransactionType::SetChildren + | TransactionType::SetWeightsVersionKey => Some(RateLimitUsageKey::AccountSubnet { + account: account.clone(), + netuid, + }), + TransactionType::SetSNOwnerHotkey | TransactionType::OwnerHyperparamUpdate(_) => { + Some(RateLimitUsageKey::Subnet(netuid)) + } + _ => None, + } +} + +// Returns the migrated call wrapper for the admin-utils extrinsic that controls `hparam`. +// +// Only hyperparameters that are currently rate-limited (i.e. routed through +// `ensure_sn_owner_or_root_with_limits`) are mapped; others return `None`. +fn identifier_for_hyperparameter(hparam: Hyperparameter) -> Option { + use Hyperparameter::*; + + let identifier = match hparam { + ServingRateLimit => MigratedCall::admin(3, false), + MaxDifficulty => MigratedCall::admin(5, false), + AdjustmentAlpha => MigratedCall::admin(9, false), + ImmunityPeriod => MigratedCall::admin(13, false), + MinAllowedWeights => MigratedCall::admin(14, false), + MaxAllowedUids => MigratedCall::admin(15, false), + Kappa => MigratedCall::admin(16, false), + Rho => MigratedCall::admin(17, false), + ActivityCutoff => MigratedCall::admin(18, false), + PowRegistrationAllowed => MigratedCall::admin(20, false), + MinBurn => MigratedCall::admin(22, false), + MaxBurn => MigratedCall::admin(23, false), + BondsMovingAverage => MigratedCall::admin(26, false), + BondsPenalty => MigratedCall::admin(60, false), + CommitRevealEnabled => MigratedCall::admin(49, false), + LiquidAlphaEnabled => MigratedCall::admin(50, false), + AlphaValues => MigratedCall::admin(51, false), + WeightCommitInterval => MigratedCall::admin(57, false), + TransferEnabled => MigratedCall::admin(61, false), + AlphaSigmoidSteepness => MigratedCall::admin(68, false), + Yuma3Enabled => MigratedCall::admin(69, false), + BondsResetEnabled => MigratedCall::admin(70, false), + ImmuneNeuronLimit => MigratedCall::admin(72, false), + RecycleOrBurn => MigratedCall::admin(80, false), + _ => return None, + }; + + Some(identifier) +} + +fn block_number(value: u64) -> Option> { + if value == 0 { + return None; + } + Some(value.saturated_into::>()) +} + +#[cfg(test)] +mod tests { + use codec::Encode; + use sp_io::TestExternalities; + use sp_io::{hashing::twox_128, storage}; + use sp_runtime::traits::{SaturatedConversion, Zero}; + + use super::*; + use crate::BuildStorage; + + const ACCOUNT: [u8; 32] = [7u8; 32]; + const DELEGATE_TAKE_GROUP_ID: GroupId = GROUP_DELEGATE_TAKE; + const PALLET_PREFIX: &[u8] = b"SubtensorModule"; + + fn new_test_ext() -> TestExternalities { + sp_tracing::try_init_simple(); + let mut ext: TestExternalities = crate::RuntimeGenesisConfig::default() + .build_storage() + .expect("runtime storage") + .into(); + ext.execute_with(|| crate::System::set_block_number(1)); + ext + } + + #[test] + fn maps_hyperparameters() { + assert_eq!( + identifier_for_hyperparameter(Hyperparameter::ServingRateLimit), + Some(MigratedCall::admin(3, false)) + ); + assert!(identifier_for_hyperparameter(Hyperparameter::MaxWeightLimit).is_none()); + } + + #[test] + fn migration_populates_limits_last_seen_and_groups() { + new_test_ext().execute_with(|| { + let account: AccountId = ACCOUNT.into(); + pallet_subtensor::HasMigrationRun::::remove(MIGRATION_NAME); + + put_legacy_value(b"TxRateLimit", 10u64); + put_legacy_value(b"TxDelegateTakeRateLimit", 3u64); + put_last_rate_limited_block(RateLimitKey::LastTxBlock(account.clone()), 5); + + let weight = migrate_rate_limiting(); + assert!(!weight.is_zero()); + assert!(pallet_subtensor::HasMigrationRun::::get( + MIGRATION_NAME + )); + + let tx_target = RateLimitTarget::Group(GROUP_SWAP_KEYS); + let delegate_group = RateLimitTarget::Group(DELEGATE_TAKE_GROUP_ID); + + assert_eq!( + pallet_rate_limiting::Limits::::get(tx_target), + Some(RateLimit::Global(RateLimitKind::Exact( + 10u64.saturated_into() + ))) + ); + assert_eq!( + pallet_rate_limiting::Limits::::get(delegate_group), + Some(RateLimit::Global(RateLimitKind::Exact( + 3u64.saturated_into() + ))) + ); + + let usage_key = RateLimitUsageKey::Account(account.clone()); + assert_eq!( + pallet_rate_limiting::LastSeen::::get(tx_target, Some(usage_key.clone())), + Some(5u64.saturated_into()) + ); + + let group = pallet_rate_limiting::Groups::::get(DELEGATE_TAKE_GROUP_ID) + .expect("group stored"); + assert_eq!(group.id, DELEGATE_TAKE_GROUP_ID); + assert_eq!(group.name.as_slice(), b"delegate-take"); + assert_eq!( + pallet_rate_limiting::CallGroups::::get( + MigratedCall::subtensor(66, false).identifier() + ), + Some(DELEGATE_TAKE_GROUP_ID) + ); + assert_eq!(pallet_rate_limiting::NextGroupId::::get(), 7); + + let serve_target = RateLimitTarget::Group(GROUP_SERVE); + assert!(pallet_rate_limiting::LimitSettingRules::::contains_key(serve_target)); + assert_eq!( + pallet_rate_limiting::LimitSettingRules::::get(serve_target), + crate::rate_limiting::LimitSettingRule::RootOrSubnetOwnerAdminWindow + ); + }); + } + + #[test] + fn migration_skips_when_already_run() { + new_test_ext().execute_with(|| { + pallet_subtensor::HasMigrationRun::::insert(MIGRATION_NAME, true); + put_legacy_value(b"TxRateLimit", 99u64); + + let base_weight = ::DbWeight::get().reads(1); + let weight = migrate_rate_limiting(); + + assert_eq!(weight, base_weight); + assert!( + pallet_rate_limiting::Limits::::iter() + .next() + .is_none() + ); + assert!( + pallet_rate_limiting::LastSeen::::iter() + .next() + .is_none() + ); + }); + } + + fn put_legacy_value(storage_name: &[u8], value: impl Encode) { + let key = storage_key(storage_name); + storage::set(&key, &value.encode()); + } + + fn put_last_rate_limited_block(key: RateLimitKey, block: u64) { + let mut storage_key = storage_key(b"LastRateLimitedBlock"); + storage_key.extend(key.encode()); + storage::set(&storage_key, &block.encode()); + } + + fn storage_key(storage_name: &[u8]) -> Vec { + [twox_128(PALLET_PREFIX), twox_128(storage_name)].concat() + } +} diff --git a/runtime/src/rate_limiting/mod.rs b/runtime/src/rate_limiting/mod.rs new file mode 100644 index 0000000000..9b370229fe --- /dev/null +++ b/runtime/src/rate_limiting/mod.rs @@ -0,0 +1,393 @@ +//! Runtime-level rate limiting wiring and resolvers. +//! +//! `pallet-rate-limiting` supports multiple independent instances, and is intended to be deployed +//! as “one instance per pallet” with pallet-specific scope/usage-key types and resolvers. +//! +//! This runtime module is centralized today because `pallet-subtensor` is currently centralized and +//! coupled with `pallet-admin-utils`; both share a single `pallet-rate-limiting` instance and a +//! single resolver implementation. +//! +//! For new pallets, do not reuse or extend the centralized scope/usage-key types or resolvers. +//! Prefer defining pallet-local types/resolvers and using a dedicated `pallet-rate-limiting` +//! instance. +//! +//! Long-term, we should refactor `pallet-subtensor` into smaller pallets and move to dedicated +//! `pallet-rate-limiting` instances per pallet. + +use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen}; +use frame_support::traits::Get; +use frame_system::RawOrigin; +use pallet_admin_utils::Call as AdminUtilsCall; +use pallet_rate_limiting::{ + BypassDecision, EnsureLimitSettingRule, RateLimitScopeResolver, RateLimitUsageResolver, +}; +use pallet_subtensor::{Call as SubtensorCall, Tempo}; +use scale_info::TypeInfo; +use serde::{Deserialize, Serialize}; +use sp_runtime::DispatchError; +use sp_std::{vec, vec::Vec}; +use subtensor_runtime_common::{ + BlockNumber, NetUid, + rate_limiting::{RateLimitUsageKey, ServingEndpoint}, +}; + +use crate::{AccountId, Runtime, RuntimeCall, RuntimeOrigin}; + +mod legacy; +pub mod migration; + +/// Authorization rules for configuring rate limits via `pallet-rate-limiting::set_rate_limit`. +/// +/// Legacy note: historically, all rate-limit setters were `Root`-only except +/// `admin-utils::sudo_set_serving_rate_limit` (subnet-owner-or-root). We preserve that behavior by +/// requiring a `scope` value when using the [`LimitSettingRule::RootOrSubnetOwnerAdminWindow`] rule and +/// validating subnet ownership against that `scope` (`netuid`). +#[derive( + Encode, + Decode, + DecodeWithMemTracking, + Serialize, + Deserialize, + Clone, + PartialEq, + Eq, + TypeInfo, + MaxEncodedLen, + Debug, +)] +pub enum LimitSettingRule { + /// Require `Root`. + Root, + /// Allow `Root` or the subnet owner for the provided `netuid` scope. + /// + /// This rule requires `scope == Some(netuid)`. + RootOrSubnetOwnerAdminWindow, +} + +pub struct DefaultLimitSettingRule; + +impl Get for DefaultLimitSettingRule { + fn get() -> LimitSettingRule { + LimitSettingRule::Root + } +} + +pub struct LimitSettingOrigin; + +impl EnsureLimitSettingRule for LimitSettingOrigin { + fn ensure_origin( + origin: RuntimeOrigin, + rule: &LimitSettingRule, + scope: &Option, + ) -> frame_support::dispatch::DispatchResult { + match rule { + LimitSettingRule::Root => frame_system::ensure_root(origin).map_err(Into::into), + LimitSettingRule::RootOrSubnetOwnerAdminWindow => { + let netuid = scope.ok_or(DispatchError::BadOrigin)?; + pallet_subtensor::Pallet::::ensure_admin_window_open(netuid)?; + pallet_subtensor::Pallet::::ensure_subnet_owner_or_root(origin, netuid) + .map(|_| ()) + .map_err(Into::into) + } + } + } +} + +#[derive(Default)] +pub struct ScopeResolver; + +impl RateLimitScopeResolver for ScopeResolver { + fn context(_origin: &RuntimeOrigin, call: &RuntimeCall) -> Option { + match call { + RuntimeCall::SubtensorModule(inner) => match inner { + SubtensorCall::serve_axon { netuid, .. } + | SubtensorCall::serve_axon_tls { netuid, .. } + | SubtensorCall::serve_prometheus { netuid, .. } + | SubtensorCall::set_weights { netuid, .. } + | SubtensorCall::commit_weights { netuid, .. } + | SubtensorCall::reveal_weights { netuid, .. } + | SubtensorCall::batch_reveal_weights { netuid, .. } + | SubtensorCall::commit_timelocked_weights { netuid, .. } + | SubtensorCall::set_mechanism_weights { netuid, .. } + | SubtensorCall::commit_mechanism_weights { netuid, .. } + | SubtensorCall::reveal_mechanism_weights { netuid, .. } + | SubtensorCall::commit_crv3_mechanism_weights { netuid, .. } + | SubtensorCall::commit_timelocked_mechanism_weights { netuid, .. } => { + Some(*netuid) + } + _ => None, + }, + _ => None, + } + } + + fn should_bypass(origin: &RuntimeOrigin, call: &RuntimeCall) -> BypassDecision { + if let RuntimeCall::SubtensorModule(inner) = call { + if matches!(origin.clone().into(), Ok(RawOrigin::Root)) { + // swap_coldkey should record last-seen but never fail; other root calls skip. + if matches!(inner, SubtensorCall::swap_coldkey { .. }) { + return BypassDecision::bypass_and_record(); + } + return BypassDecision::bypass_and_skip(); + } + + match inner { + SubtensorCall::set_childkey_take { + hotkey, + netuid, + take, + .. + } => { + let current = + pallet_subtensor::Pallet::::get_childkey_take(hotkey, *netuid); + return if *take <= current { + BypassDecision::bypass_and_record() + } else { + BypassDecision::enforce_and_record() + }; + } + SubtensorCall::add_stake { .. } + | SubtensorCall::add_stake_limit { .. } + | SubtensorCall::decrease_take { .. } + | SubtensorCall::swap_coldkey { .. } => { + return BypassDecision::bypass_and_record(); + } + SubtensorCall::reveal_weights { netuid, .. } + | SubtensorCall::batch_reveal_weights { netuid, .. } + | SubtensorCall::reveal_mechanism_weights { netuid, .. } => { + if pallet_subtensor::Pallet::::get_commit_reveal_weights_enabled( + *netuid, + ) { + // Legacy: reveals are not rate-limited while commit-reveal is enabled. + return BypassDecision::bypass_and_skip(); + } + } + _ => {} + } + } + + BypassDecision::enforce_and_record() + } + + fn adjust_span(_origin: &RuntimeOrigin, call: &RuntimeCall, span: BlockNumber) -> BlockNumber { + match call { + RuntimeCall::AdminUtils(inner) => { + if let Some(netuid) = owner_hparam_netuid(inner) { + if span == 0 { + return span; + } + let tempo = BlockNumber::from(Tempo::::get(netuid) as u32); + span.saturating_mul(tempo) + } else if let AdminUtilsCall::sudo_set_weights_version_key { netuid, .. } = inner { + if span == 0 { + return span; + } + let tempo = BlockNumber::from(Tempo::::get(netuid) as u32); + span.saturating_mul(tempo) + } else { + span + } + } + _ => span, + } + } +} + +#[derive(Default)] +pub struct UsageResolver; + +impl RateLimitUsageResolver> + for UsageResolver +{ + fn context( + origin: &RuntimeOrigin, + call: &RuntimeCall, + ) -> Option>> { + match call { + RuntimeCall::SubtensorModule(inner) => match inner { + SubtensorCall::swap_coldkey { new_coldkey, .. } => { + Some(vec![RateLimitUsageKey::::Account( + new_coldkey.clone(), + )]) + } + SubtensorCall::swap_hotkey { new_hotkey, .. } => { + // Record against the coldkey (enforcement) and the new hotkey to mirror legacy + // writes. + let coldkey = signed_origin(origin)?; + Some(vec![ + RateLimitUsageKey::::Account(coldkey), + RateLimitUsageKey::::Account(new_hotkey.clone()), + ]) + } + SubtensorCall::increase_take { hotkey, .. } => { + Some(vec![RateLimitUsageKey::::Account( + hotkey.clone(), + )]) + } + SubtensorCall::set_childkey_take { hotkey, netuid, .. } + | SubtensorCall::set_children { hotkey, netuid, .. } => { + Some(vec![RateLimitUsageKey::::AccountSubnet { + account: hotkey.clone(), + netuid: *netuid, + }]) + } + SubtensorCall::set_weights { netuid, .. } + | SubtensorCall::commit_weights { netuid, .. } + | SubtensorCall::reveal_weights { netuid, .. } + | SubtensorCall::batch_reveal_weights { netuid, .. } + | SubtensorCall::commit_timelocked_weights { netuid, .. } => { + let (_, uid) = neuron_identity(origin, *netuid)?; + Some(vec![RateLimitUsageKey::::SubnetNeuron { + netuid: *netuid, + uid, + }]) + } + // legacy implementation still used netuid only, but it was recalculating it using + // mecid, so switching to netuid AND mecid is logical here + SubtensorCall::set_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::reveal_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_crv3_mechanism_weights { netuid, mecid, .. } + | SubtensorCall::commit_timelocked_mechanism_weights { netuid, mecid, .. } => { + let (_, uid) = neuron_identity(origin, *netuid)?; + Some(vec![ + RateLimitUsageKey::::SubnetMechanismNeuron { + netuid: *netuid, + mecid: *mecid, + uid, + }, + ]) + } + SubtensorCall::serve_axon { netuid, .. } + | SubtensorCall::serve_axon_tls { netuid, .. } => { + let hotkey = signed_origin(origin)?; + Some(vec![RateLimitUsageKey::::AccountSubnetServing { + account: hotkey, + netuid: *netuid, + endpoint: ServingEndpoint::Axon, + }]) + } + SubtensorCall::serve_prometheus { netuid, .. } => { + let hotkey = signed_origin(origin)?; + Some(vec![RateLimitUsageKey::::AccountSubnetServing { + account: hotkey, + netuid: *netuid, + endpoint: ServingEndpoint::Prometheus, + }]) + } + SubtensorCall::associate_evm_key { netuid, .. } => { + let hotkey = signed_origin(origin)?; + let uid = pallet_subtensor::Pallet::::get_uid_for_net_and_hotkey( + *netuid, &hotkey, + ) + .ok()?; + Some(vec![RateLimitUsageKey::::SubnetNeuron { + netuid: *netuid, + uid, + }]) + } + // Staking calls share a group lock; only add_* write usage, the rest are read-only. + // Keep the usage key granular so the lock applies per (coldkey, hotkey, netuid). + SubtensorCall::add_stake { hotkey, netuid, .. } + | SubtensorCall::add_stake_limit { hotkey, netuid, .. } + | SubtensorCall::remove_stake { hotkey, netuid, .. } + | SubtensorCall::remove_stake_limit { hotkey, netuid, .. } + | SubtensorCall::remove_stake_full_limit { hotkey, netuid, .. } + | SubtensorCall::transfer_stake { + hotkey, + origin_netuid: netuid, + .. + } + | SubtensorCall::swap_stake { + hotkey, + origin_netuid: netuid, + .. + } + | SubtensorCall::swap_stake_limit { + hotkey, + origin_netuid: netuid, + .. + } + | SubtensorCall::move_stake { + origin_hotkey: hotkey, + origin_netuid: netuid, + .. + } => { + let coldkey = signed_origin(origin)?; + Some(vec![RateLimitUsageKey::::ColdkeyHotkeySubnet { + coldkey, + hotkey: hotkey.clone(), + netuid: *netuid, + }]) + } + _ => None, + }, + RuntimeCall::AdminUtils(inner) => { + if let Some(netuid) = owner_hparam_netuid(inner) { + Some(vec![RateLimitUsageKey::::Subnet(netuid)]) + } else { + match inner { + AdminUtilsCall::sudo_set_sn_owner_hotkey { netuid, .. } => { + Some(vec![RateLimitUsageKey::::Subnet(*netuid)]) + } + AdminUtilsCall::sudo_set_weights_version_key { netuid, .. } + | AdminUtilsCall::sudo_set_mechanism_count { netuid, .. } + | AdminUtilsCall::sudo_set_mechanism_emission_split { netuid, .. } + | AdminUtilsCall::sudo_trim_to_max_allowed_uids { netuid, .. } => { + let who = signed_origin(origin)?; + Some(vec![RateLimitUsageKey::::AccountSubnet { + account: who, + netuid: *netuid, + }]) + } + _ => None, + } + } + } + _ => None, + } + } +} + +fn neuron_identity(origin: &RuntimeOrigin, netuid: NetUid) -> Option<(AccountId, u16)> { + let hotkey = signed_origin(origin)?; + let uid = + pallet_subtensor::Pallet::::get_uid_for_net_and_hotkey(netuid, &hotkey).ok()?; + Some((hotkey, uid)) +} + +fn signed_origin(origin: &RuntimeOrigin) -> Option { + match origin.clone().into() { + Ok(RawOrigin::Signed(who)) => Some(who), + _ => None, + } +} + +fn owner_hparam_netuid(call: &AdminUtilsCall) -> Option { + match call { + AdminUtilsCall::sudo_set_activity_cutoff { netuid, .. } + | AdminUtilsCall::sudo_set_adjustment_alpha { netuid, .. } + | AdminUtilsCall::sudo_set_alpha_sigmoid_steepness { netuid, .. } + | AdminUtilsCall::sudo_set_alpha_values { netuid, .. } + | AdminUtilsCall::sudo_set_bonds_moving_average { netuid, .. } + | AdminUtilsCall::sudo_set_bonds_penalty { netuid, .. } + | AdminUtilsCall::sudo_set_bonds_reset_enabled { netuid, .. } + | AdminUtilsCall::sudo_set_commit_reveal_weights_enabled { netuid, .. } + | AdminUtilsCall::sudo_set_commit_reveal_weights_interval { netuid, .. } + | AdminUtilsCall::sudo_set_immunity_period { netuid, .. } + | AdminUtilsCall::sudo_set_liquid_alpha_enabled { netuid, .. } + | AdminUtilsCall::sudo_set_max_allowed_uids { netuid, .. } + | AdminUtilsCall::sudo_set_max_burn { netuid, .. } + | AdminUtilsCall::sudo_set_max_difficulty { netuid, .. } + | AdminUtilsCall::sudo_set_min_allowed_weights { netuid, .. } + | AdminUtilsCall::sudo_set_min_burn { netuid, .. } + | AdminUtilsCall::sudo_set_network_pow_registration_allowed { netuid, .. } + | AdminUtilsCall::sudo_set_owner_immune_neuron_limit { netuid, .. } + | AdminUtilsCall::sudo_set_recycle_or_burn { netuid, .. } + | AdminUtilsCall::sudo_set_rho { netuid, .. } + | AdminUtilsCall::sudo_set_serving_rate_limit { netuid, .. } + | AdminUtilsCall::sudo_set_toggle_transfer { netuid, .. } + | AdminUtilsCall::sudo_set_yuma3_enabled { netuid, .. } => Some(*netuid), + _ => None, + } +} diff --git a/runtime/tests/rate_limiting_behavior.rs b/runtime/tests/rate_limiting_behavior.rs new file mode 100644 index 0000000000..132fa1bbe3 --- /dev/null +++ b/runtime/tests/rate_limiting_behavior.rs @@ -0,0 +1,438 @@ +#![allow(clippy::unwrap_used)] + +use codec::Encode; +use frame_support::traits::OnRuntimeUpgrade; +use frame_system::pallet_prelude::BlockNumberFor; +use node_subtensor_runtime::{ + BuildStorage, Runtime, RuntimeCall, RuntimeGenesisConfig, RuntimeOrigin, RuntimeScopeResolver, + RuntimeUsageResolver, SubtensorModule, System, rate_limiting::migration::Migration, +}; +use pallet_rate_limiting::{RateLimitScopeResolver, RateLimitUsageResolver}; +use pallet_rate_limiting::{RateLimitTarget, TransactionIdentifier}; +use pallet_subtensor::Call as SubtensorCall; +use pallet_subtensor::{ + AxonInfo, HasMigrationRun, LastRateLimitedBlock, LastUpdate, NetworksAdded, PrometheusInfo, + RateLimitKey, TransactionKeyLastBlock, WeightsSetRateLimit, WeightsVersionKeyRateLimit, + utils::rate_limiting::TransactionType, +}; +use sp_core::{H160, ecdsa}; +use sp_io::{hashing::twox_128, storage}; +use sp_runtime::traits::SaturatedConversion; +use subtensor_runtime_common::{ + NetUid, NetUidStorageIndex, + rate_limiting::{GroupId, RateLimitUsageKey}, +}; + +type AccountId = ::AccountId; +type UsageKey = RateLimitUsageKey; + +const MIGRATION_NAME: &[u8] = b"migrate_rate_limiting"; + +fn new_ext() -> sp_io::TestExternalities { + sp_tracing::try_init_simple(); + let mut ext: sp_io::TestExternalities = RuntimeGenesisConfig::default() + .build_storage() + .unwrap() + .into(); + ext.execute_with(|| System::set_block_number(1)); + ext +} + +fn account(n: u8) -> AccountId { + AccountId::from([n; 32]) +} + +fn resolve_target(identifier: TransactionIdentifier) -> RateLimitTarget { + if let Some(group) = pallet_rate_limiting::CallGroups::::get(identifier) { + RateLimitTarget::Group(group) + } else { + RateLimitTarget::Transaction(identifier) + } +} + +fn exact_span(span: u64) -> BlockNumberFor { + span.saturated_into::>() +} + +fn clear_rate_limiting_storage() { + let limit = u32::MAX; + let _ = pallet_rate_limiting::Limits::::clear(limit, None); + let _ = pallet_rate_limiting::LastSeen::::clear(limit, None); + let _ = pallet_rate_limiting::Groups::::clear(limit, None); + let _ = pallet_rate_limiting::GroupMembers::::clear(limit, None); + let _ = pallet_rate_limiting::GroupNameIndex::::clear(limit, None); + let _ = pallet_rate_limiting::CallGroups::::clear(limit, None); + pallet_rate_limiting::NextGroupId::::kill(); +} + +fn set_legacy_serving_rate_limit(netuid: NetUid, span: u64) { + let mut key = twox_128(b"SubtensorModule").to_vec(); + key.extend(twox_128(b"ServingRateLimit")); + key.extend(netuid.encode()); + storage::set(&key, &span.encode()); +} + +fn parity_check( + now: u64, + call: RuntimeCall, + origin: RuntimeOrigin, + usage_override: Option>, + scope_override: Option, + legacy_check: F, +) where + F: Fn() -> bool, +{ + System::set_block_number(now.saturated_into()); + HasMigrationRun::::remove(MIGRATION_NAME); + clear_rate_limiting_storage(); + + // Run migration to hydrate pallet-rate-limiting state. + Migration::::on_runtime_upgrade(); + + let identifier = TransactionIdentifier::from_call(&call).expect("identifier for call"); + let scope = scope_override.or_else(|| RuntimeScopeResolver::context(&origin, &call)); + let usage: Option::UsageKey>> = + usage_override.or_else(|| RuntimeUsageResolver::context(&origin, &call)); + let target = resolve_target(identifier); + + // Use the runtime-adjusted span (handles tempo scaling for admin-utils). + let span = pallet_rate_limiting::Pallet::::effective_span( + &origin.clone().into(), + &call, + &target, + &scope, + ) + .unwrap_or_default(); + let span_u64: u64 = span.saturated_into(); + + let usage_keys: Vec::UsageKey>> = match usage { + None => vec![None], + Some(keys) => keys.into_iter().map(Some).collect(), + }; + + let within = usage_keys.iter().all(|key| { + pallet_rate_limiting::Pallet::::is_within_limit( + &origin.clone().into(), + &call, + &identifier, + &scope, + key, + ) + .expect("pallet rate limit result") + }); + assert_eq!(within, legacy_check(), "parity at now for {:?}", identifier); + + // Advance beyond the span and re-check (span==0 treated as allow). + let advance: BlockNumberFor = span.saturating_add(exact_span(1)); + System::set_block_number(System::block_number().saturating_add(advance)); + + let within_after = usage_keys.iter().all(|key| { + pallet_rate_limiting::Pallet::::is_within_limit( + &origin.clone().into(), + &call, + &identifier, + &scope, + key, + ) + .expect("pallet rate limit result (after)") + }); + assert!( + within_after || span_u64 == 0, + "parity after window for {:?}", + identifier + ); +} + +#[test] +fn register_network_parity() { + new_ext().execute_with(|| { + let now = 100u64; + let cold = account(1); + let hot = account(2); + let span = 5u64; + LastRateLimitedBlock::::insert(RateLimitKey::NetworkLastRegistered, now - 1); + pallet_subtensor::NetworkRateLimit::::put(span); + + let call = RuntimeCall::SubtensorModule(SubtensorCall::register_network { hotkey: hot }); + let origin = RuntimeOrigin::signed(cold.clone()); + let legacy = || TransactionType::RegisterNetwork.passes_rate_limit::(&cold); + parity_check(now, call, origin, None, None, legacy); + }); +} + +#[test] +fn swap_hotkey_parity() { + new_ext().execute_with(|| { + let now = 200u64; + let cold = account(10); + let old_hot = account(11); + let new_hot = account(12); + let span = 10u64; + LastRateLimitedBlock::::insert(RateLimitKey::LastTxBlock(cold.clone()), now - 1); + pallet_subtensor::TxRateLimit::::put(span); + + let call = RuntimeCall::SubtensorModule(SubtensorCall::swap_hotkey { + hotkey: old_hot, + new_hotkey: new_hot, + netuid: None, + }); + let origin = RuntimeOrigin::signed(cold.clone()); + let legacy = || !SubtensorModule::exceeds_tx_rate_limit(now - 1, now); + parity_check(now, call, origin, None, None, legacy); + }); +} + +#[test] +fn increase_take_parity() { + new_ext().execute_with(|| { + let now = 300u64; + let hot = account(20); + let span = 3u64; + LastRateLimitedBlock::::insert( + RateLimitKey::LastTxBlockDelegateTake(hot.clone()), + now - 1, + ); + pallet_subtensor::TxDelegateTakeRateLimit::::put(span); + + let call = RuntimeCall::SubtensorModule(SubtensorCall::increase_take { + hotkey: hot.clone(), + take: 5, + }); + let origin = RuntimeOrigin::signed(account(21)); + let legacy = || !SubtensorModule::exceeds_tx_delegate_take_rate_limit(now - 1, now); + parity_check(now, call, origin, None, None, legacy); + }); +} + +#[test] +fn set_childkey_take_parity() { + new_ext().execute_with(|| { + let now = 400u64; + let hot = account(30); + let netuid = NetUid::from(1u16); + let span = 7u64; + let tx_kind: u16 = TransactionType::SetChildkeyTake.into(); + TransactionKeyLastBlock::::insert((hot.clone(), netuid, tx_kind), now - 1); + pallet_subtensor::TxChildkeyTakeRateLimit::::put(span); + + let call = RuntimeCall::SubtensorModule(SubtensorCall::set_childkey_take { + hotkey: hot.clone(), + netuid, + take: 1, + }); + let origin = RuntimeOrigin::signed(account(31)); + let legacy = || { + TransactionType::SetChildkeyTake.passes_rate_limit_on_subnet::(&hot, netuid) + }; + parity_check(now, call, origin, None, None, legacy); + }); +} + +#[test] +fn set_children_parity() { + new_ext().execute_with(|| { + let now = 500u64; + let hot = account(40); + let netuid = NetUid::from(2u16); + let tx_kind: u16 = TransactionType::SetChildren.into(); + TransactionKeyLastBlock::::insert((hot.clone(), netuid, tx_kind), now - 1); + + let call = RuntimeCall::SubtensorModule(SubtensorCall::set_children { + hotkey: hot.clone(), + netuid, + children: Vec::new(), + }); + let origin = RuntimeOrigin::signed(account(41)); + let legacy = + || TransactionType::SetChildren.passes_rate_limit_on_subnet::(&hot, netuid); + parity_check(now, call, origin, None, None, legacy); + }); +} + +#[test] +fn serving_parity() { + new_ext().execute_with(|| { + let now = 600u64; + let hot = account(50); + let netuid = NetUid::from(3u16); + let span = 5u64; + set_legacy_serving_rate_limit(netuid, span); + pallet_subtensor::Axons::::insert( + netuid, + hot.clone(), + AxonInfo { + block: now - 1, + ..Default::default() + }, + ); + pallet_subtensor::Prometheus::::insert( + netuid, + hot.clone(), + PrometheusInfo { + block: now - 1, + ..Default::default() + }, + ); + + // Axon + let axon_call = RuntimeCall::SubtensorModule(SubtensorCall::serve_axon { + netuid, + version: 1, + ip: 0, + port: 0, + ip_type: 4, + protocol: 0, + placeholder1: 0, + placeholder2: 0, + }); + let origin = RuntimeOrigin::signed(hot.clone()); + let legacy_axon = || { + let info = AxonInfo { + block: now.saturating_sub(1), + ..Default::default() + }; + now.saturating_sub(info.block) >= span + }; + parity_check(now, axon_call, origin.clone(), None, None, legacy_axon); + + // Prometheus + let prom_call = RuntimeCall::SubtensorModule(SubtensorCall::serve_prometheus { + netuid, + version: 1, + ip: 0, + port: 0, + ip_type: 4, + }); + let legacy_prom = || { + let info = PrometheusInfo { + block: now.saturating_sub(1), + ..Default::default() + }; + now.saturating_sub(info.block) >= span + }; + parity_check(now, prom_call, origin, None, None, legacy_prom); + }); +} + +#[test] +fn weights_and_hparam_parity() { + new_ext().execute_with(|| { + let now = 700u64; + let hot = account(60); + let netuid = NetUid::from(4u16); + let uid: u16 = 0; + let weights_span = 4u64; + let tempo = 3u16; + // Ensure subnet exists so LastUpdate is imported. + NetworksAdded::::insert(netuid, true); + SubtensorModule::set_tempo(netuid, tempo); + WeightsSetRateLimit::::insert(netuid, weights_span); + LastUpdate::::insert(NetUidStorageIndex::from(netuid), vec![now - 1]); + + let weights_call = RuntimeCall::SubtensorModule(SubtensorCall::set_weights { + netuid, + dests: Vec::new(), + weights: Vec::new(), + version_key: 0, + }); + let origin = RuntimeOrigin::signed(hot.clone()); + let scope = Some(netuid); + let usage = Some(vec![UsageKey::SubnetNeuron { netuid, uid }]); + + let legacy_weights = || SubtensorModule::check_rate_limit(netuid.into(), uid, now); + parity_check( + now, + weights_call, + origin.clone(), + usage, + scope, + legacy_weights, + ); + + // Hyperparam (activity_cutoff) with tempo scaling. + let hparam_span_epochs = 2u16; + pallet_subtensor::OwnerHyperparamRateLimit::::put(hparam_span_epochs); + LastRateLimitedBlock::::insert( + RateLimitKey::OwnerHyperparamUpdate( + netuid, + pallet_subtensor::utils::rate_limiting::Hyperparameter::ActivityCutoff, + ), + now - 1, + ); + let hparam_call = + RuntimeCall::AdminUtils(pallet_admin_utils::Call::sudo_set_activity_cutoff { + netuid, + activity_cutoff: 1, + }); + let hparam_origin = RuntimeOrigin::signed(hot); + let legacy_hparam = || { + let span = (tempo as u64) * (hparam_span_epochs as u64); + let last = now - 1; + // same logic as TransactionType::OwnerHyperparamUpdate in legacy: passes if delta >= span. + let delta = now.saturating_sub(last); + delta >= span + }; + parity_check(now, hparam_call, hparam_origin, None, None, legacy_hparam); + }); +} + +#[test] +fn weights_version_parity() { + new_ext().execute_with(|| { + let now = 800u64; + let hot = account(70); + let netuid = NetUid::from(5u16); + NetworksAdded::::insert(netuid, true); + SubtensorModule::set_tempo(netuid, 4); + WeightsVersionKeyRateLimit::::put(2u64); + let tx_kind_wvk: u16 = TransactionType::SetWeightsVersionKey.into(); + TransactionKeyLastBlock::::insert((hot.clone(), netuid, tx_kind_wvk), now - 1); + + let wvk_call = + RuntimeCall::AdminUtils(pallet_admin_utils::Call::sudo_set_weights_version_key { + netuid, + weights_version_key: 0, + }); + let origin = RuntimeOrigin::signed(hot.clone()); + let legacy_wvk = || { + let limit = SubtensorModule::get_tempo(netuid) as u64 + * WeightsVersionKeyRateLimit::::get(); + let delta = now.saturating_sub(now - 1); + delta >= limit + }; + parity_check(now, wvk_call, origin, None, None, legacy_wvk); + }); +} + +#[test] +fn associate_evm_key_parity() { + new_ext().execute_with(|| { + let now = 900u64; + let hot = account(80); + let netuid = NetUid::from(6u16); + let uid: u16 = 0; + NetworksAdded::::insert(netuid, true); + pallet_subtensor::AssociatedEvmAddress::::insert( + netuid, + uid, + (H160::zero(), now - 1), + ); + + let call = RuntimeCall::SubtensorModule(SubtensorCall::associate_evm_key { + netuid, + evm_key: H160::zero(), + block_number: now, + signature: ecdsa::Signature::from_raw([0u8; 65]), + }); + let origin = RuntimeOrigin::signed(hot.clone()); + let usage = Some(vec![UsageKey::SubnetNeuron { netuid, uid }]); + let scope = Some(netuid); + let limit = ::EvmKeyAssociateRateLimit::get(); + let legacy = || { + let last = now - 1; + let delta = now.saturating_sub(last); + delta >= limit + }; + parity_check(now, call, origin, usage, scope, legacy); + }); +} diff --git a/runtime/tests/rate_limiting_migration.rs b/runtime/tests/rate_limiting_migration.rs new file mode 100644 index 0000000000..9e08f489b9 --- /dev/null +++ b/runtime/tests/rate_limiting_migration.rs @@ -0,0 +1,77 @@ +#![allow(clippy::unwrap_used)] + +use frame_support::traits::OnRuntimeUpgrade; +use frame_system::pallet_prelude::BlockNumberFor; +use pallet_rate_limiting::{RateLimit, RateLimitKind, RateLimitTarget, TransactionIdentifier}; +use pallet_subtensor::{HasMigrationRun, LastRateLimitedBlock, RateLimitKey}; +use sp_runtime::traits::SaturatedConversion; +use subtensor_runtime_common::{ + NetUid, + rate_limiting::{GROUP_REGISTER_NETWORK, RateLimitUsageKey}, +}; + +use node_subtensor_runtime::{ + BuildStorage, Runtime, RuntimeGenesisConfig, SubtensorModule, System, + rate_limiting::migration::{MIGRATION_NAME, Migration}, +}; + +type AccountId = ::AccountId; +type UsageKey = RateLimitUsageKey; + +fn new_test_ext() -> sp_io::TestExternalities { + sp_tracing::try_init_simple(); + let mut ext: sp_io::TestExternalities = RuntimeGenesisConfig::default() + .build_storage() + .unwrap() + .into(); + ext.execute_with(|| System::set_block_number(1)); + ext +} + +#[test] +fn migrates_global_register_network_last_seen() { + new_test_ext().execute_with(|| { + HasMigrationRun::::remove(MIGRATION_NAME); + + // Seed legacy global register rate-limit state. + LastRateLimitedBlock::::insert(RateLimitKey::NetworkLastRegistered, 10u64); + System::set_block_number(12); + + // Run migration. + Migration::::on_runtime_upgrade(); + + let target = RateLimitTarget::Group(GROUP_REGISTER_NETWORK); + + // LastSeen preserved globally (usage = None). + let stored = pallet_rate_limiting::LastSeen::::get(target, None::) + .expect("last seen entry"); + assert_eq!(stored, 10u64.saturated_into::>()); + }); +} + +#[test] +fn sn_owner_hotkey_limit_not_tempo_scaled_and_last_seen_preserved() { + new_test_ext().execute_with(|| { + HasMigrationRun::::remove(MIGRATION_NAME); + + let netuid = NetUid::from(1); + // Give the subnet a non-1 tempo to catch accidental scaling. + SubtensorModule::set_tempo(netuid, 5); + LastRateLimitedBlock::::insert(RateLimitKey::SetSNOwnerHotkey(netuid), 100u64); + + Migration::::on_runtime_upgrade(); + + let target = RateLimitTarget::Transaction(TransactionIdentifier::new(19, 67)); + + // Limit should remain the fixed default (50400 blocks), not tempo-scaled. + let limit = pallet_rate_limiting::Limits::::get(target).expect("limit stored"); + assert!(matches!(limit, RateLimit::Global(kind) if kind == RateLimitKind::Exact(50_400))); + + // LastSeen preserved per subnet. + let usage: Option<::UsageKey> = + Some(UsageKey::Subnet(netuid).into()); + let stored = + pallet_rate_limiting::LastSeen::::get(target, usage).expect("last seen entry"); + assert_eq!(stored, 100u64.saturated_into::>()); + }); +}