diff --git a/Cargo.lock b/Cargo.lock index aee74c67..19a1b6c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6209,6 +6209,7 @@ dependencies = [ "pallet-timestamp", "parity-scale-codec", "qp-high-security", + "qp-wormhole", "scale-info", "sp-core", "sp-io", diff --git a/frame/metadata/src/lib.rs b/frame/metadata/src/lib.rs index eb626b68..5791968c 100644 --- a/frame/metadata/src/lib.rs +++ b/frame/metadata/src/lib.rs @@ -83,11 +83,31 @@ pub mod v16; pub const META_RESERVED: u32 = 0x6174656d; // 'meta' warning for endianness. /// Metadata prefixed by a u32 for reserved usage +/// +/// # Decoding untrusted input +/// +/// [`Decode`] rejects blobs whose prefix is not [`META_RESERVED`] after reading only the first +/// four bytes. Beyond that, decoding materializes unbounded owned containers (`Vec`s, maps, +/// strings) whose size is proportional to the input, with no schema-level caps: callers decoding +/// metadata from an untrusted source must bound the input size *before* decoding. #[derive(Eq, Encode, PartialEq, Debug)] -#[cfg_attr(feature = "decode", derive(Decode))] #[cfg_attr(feature = "serde_full", derive(Serialize))] pub struct RuntimeMetadataPrefixed(pub u32, pub RuntimeMetadata); +#[cfg(feature = "decode")] +impl Decode for RuntimeMetadataPrefixed { + fn decode(input: &mut I) -> Result { + // Validate the reserved prefix *before* decoding the payload, so a blob that is not + // runtime metadata is rejected after reading four bytes instead of after fully + // materializing an attacker-sized metadata structure. + let prefix = u32::decode(input)?; + if prefix != META_RESERVED { + return Err("Invalid metadata prefix: expected `META_RESERVED` ('meta')".into()) + } + Ok(Self(prefix, RuntimeMetadata::decode(input)?)) + } +} + impl From for Vec { fn from(value: RuntimeMetadataPrefixed) -> Self { value.encode() @@ -272,4 +292,19 @@ mod test { Decode::decode(&mut load_metadata(14).as_slice()).unwrap(); assert!(matches!(meta.1, RuntimeMetadata::V14(_))); } + + #[test] + fn should_reject_invalid_prefix_before_decoding_payload() { + // A wrong magic followed by a huge claimed payload: the prefix check must fail after + // the first four bytes, without attempting to materialize the payload. + let mut blob = 0xdeadbeef_u32.encode(); + blob.extend_from_slice(&load_metadata(14)[4..]); + let res = RuntimeMetadataPrefixed::decode(&mut blob.as_slice()); + assert!(res.is_err()); + + // A valid prefix still decodes. + let meta: RuntimeMetadataPrefixed = + Decode::decode(&mut load_metadata(14).as_slice()).unwrap(); + assert_eq!(meta.0, META_RESERVED); + } } diff --git a/frame/support-procedural/src/construct_runtime/expand/call.rs b/frame/support-procedural/src/construct_runtime/expand/call.rs index f2ec5594..44daf4cb 100644 --- a/frame/support-procedural/src/construct_runtime/expand/call.rs +++ b/frame/support-procedural/src/construct_runtime/expand/call.rs @@ -162,7 +162,10 @@ pub fn expand_outer_dispatch( <<#pallet_names as Callable<#runtime>>::RuntimeCall as GetCallName>::get_call_names(), )* - _ => unreachable!(), + // An unknown module name is a recoverable "not found" rather than an + // impossible state: this is public metadata API reachable with + // caller-supplied strings, so return an empty slice instead of panicking. + _ => &[], } } } diff --git a/frame/support-procedural/src/construct_runtime/expand/task.rs b/frame/support-procedural/src/construct_runtime/expand/task.rs index b1cf4f85..6dc16b2d 100644 --- a/frame/support-procedural/src/construct_runtime/expand/task.rs +++ b/frame/support-procedural/src/construct_runtime/expand/task.rs @@ -98,7 +98,9 @@ pub fn expand_outer_task( #[automatically_derived] impl #scrate::traits::Task for RuntimeTask { - type Enumeration = #prelude::IntoIter; + type Enumeration = #scrate::__private::Box< + dyn #prelude::Iterator + >; fn is_valid(&self) -> bool { match self { @@ -141,12 +143,18 @@ pub fn expand_outer_task( } fn iter() -> Self::Enumeration { - let mut all_tasks = Vec::new(); + // Stream the pallet task iterators lazily instead of collecting them into + // heap vectors: task sets can be derived from attacker-growable state, so + // eager materialization allows unbounded allocation before the first yield. + let mut all_tasks = + #scrate::__private::Vec::::new(); #( #cfg_attrs - all_tasks.extend(<#task_types>::iter().map(RuntimeTask::from).collect::>()); + all_tasks.push(#scrate::__private::Box::new( + <#task_types>::iter().map(RuntimeTask::from), + )); )* - all_tasks.into_iter() + #scrate::__private::Box::new(all_tasks.into_iter().flatten()) } } diff --git a/frame/support-procedural/src/pallet/expand/hooks.rs b/frame/support-procedural/src/pallet/expand/hooks.rs index c31ddd8a..e1a2310a 100644 --- a/frame/support-procedural/src/pallet/expand/hooks.rs +++ b/frame/support-procedural/src/pallet/expand/hooks.rs @@ -214,24 +214,16 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream { for #pallet_ident<#type_use_gen> #where_clause { fn before_all_runtime_migrations() -> #frame_support::weights::Weight { - use #frame_support::traits::{Get, PalletInfoAccess}; - use #frame_support::__private::hashing::twox_128; - use #frame_support::storage::unhashed::contains_prefixed_key; #frame_support::__private::sp_tracing::enter_span!( #frame_support::__private::sp_tracing::trace_span!("before_all") ); - // Check if the pallet has any keys set, including the storage version. If there are - // no keys set, the pallet was just added to the runtime and needs to have its - // version initialized. - let pallet_hashed_prefix = ::name_hash(); - let exists = contains_prefixed_key(&pallet_hashed_prefix); - if !exists { - #initialize_on_chain_storage_version - ::DbWeight::get().reads_writes(1, 1) - } else { - ::DbWeight::get().reads(1) - } + // The on-chain storage version of a newly added pallet is initialized + // *after* all migrations ran (see `on_runtime_upgrade` below). Seeding it + // here would run before the runtime's custom migrations and cause + // version-gated migrations targeting a new (still empty) pallet prefix, + // e.g. renames or imports from another namespace, to be silently skipped. + #frame_support::weights::Weight::zero() } } @@ -240,6 +232,7 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream { for #pallet_ident<#type_use_gen> #where_clause { fn on_runtime_upgrade() -> #frame_support::weights::Weight { + use #frame_support::traits::Get; #frame_support::__private::sp_tracing::enter_span!( #frame_support::__private::sp_tracing::trace_span!("on_runtime_update") ); @@ -247,11 +240,28 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream { // log info about the upgrade. #log_runtime_upgrade - < + let mut weight = < Self as #frame_support::traits::Hooks< #frame_system::pallet_prelude::BlockNumberFor:: > - >::on_runtime_upgrade() + >::on_runtime_upgrade(); + + // If the storage-version key is still absent after every migration had a + // chance to run, the pallet is genuinely new: initialize its on-chain + // storage version. Doing this after the migrations (instead of in + // `before_all_runtime_migrations`) keeps version-gated migrations that + // target a new pallet prefix from being skipped. + if !#frame_support::traits::StorageVersion::exists::() { + #initialize_on_chain_storage_version + weight = weight.saturating_add( + ::DbWeight::get().reads_writes(1, 1) + ); + } else { + weight = weight + .saturating_add(::DbWeight::get().reads(1)); + } + + weight } #frame_support::try_runtime_enabled! { diff --git a/frame/support-procedural/src/pallet/expand/storage.rs b/frame/support-procedural/src/pallet/expand/storage.rs index a0951a25..96799731 100644 --- a/frame/support-procedural/src/pallet/expand/storage.rs +++ b/frame/support-procedural/src/pallet/expand/storage.rs @@ -66,7 +66,10 @@ fn check_prefix_duplicates( return Err(err); } - if let Metadata::CountedMap { .. } = storage_def.metadata { + // Counted maps and counted N-maps both generate an auxiliary counter storage instance under + // the derived `CounterFor...` prefix, so that namespace must be reserved for both kinds; + // otherwise another storage item could silently alias the counter's physical key. + if let Metadata::CountedMap { .. } | Metadata::CountedNMap { .. } = storage_def.metadata { let counter_prefix = counter_prefix(&prefix); let counter_dup_err = syn::Error::new( storage_def.prefix_span(), diff --git a/frame/support-procedural/src/pallet/expand/tasks.rs b/frame/support-procedural/src/pallet/expand/tasks.rs index 6cd2436c..f2f0f777 100644 --- a/frame/support-procedural/src/pallet/expand/tasks.rs +++ b/frame/support-procedural/src/pallet/expand/tasks.rs @@ -166,15 +166,22 @@ impl TasksDef { impl #impl_generics #frame_support::traits::Task for #enum_use { - type Enumeration = #frame_support::__private::IntoIter<#enum_use>; + type Enumeration = #frame_support::__private::Box< + dyn #frame_support::traits::tasks::__private::Iterator + >; fn iter() -> Self::Enumeration { - let mut all_tasks = #frame_support::__private::vec![]; - #(all_tasks - .extend(#task_iters.map(|(#(#task_arg_names),*)| #enum_ident::#task_fn_idents { #(#task_arg_names: #task_arg_names.clone()),* }) - .collect::<#frame_support::__private::Vec<_>>()); + // Stream the task lists lazily instead of collecting them into heap + // vectors: task lists can be derived from attacker-growable state, so + // eager materialization allows unbounded allocation before the first + // yield. + let mut all_tasks = + #frame_support::__private::Vec::::new(); + #(all_tasks.push(#frame_support::__private::Box::new( + #task_iters.map(|(#(#task_arg_names),*)| #enum_ident::#task_fn_idents { #(#task_arg_names: #task_arg_names.clone()),* }) + )); )* - all_tasks.into_iter() + #frame_support::__private::Box::new(all_tasks.into_iter().flatten()) } fn task_index(&self) -> u32 { diff --git a/frame/support-procedural/src/pallet/parse/call.rs b/frame/support-procedural/src/pallet/parse/call.rs index 851b75e3..02642ef2 100644 --- a/frame/support-procedural/src/pallet/parse/call.rs +++ b/frame/support-procedural/src/pallet/parse/call.rs @@ -144,12 +144,26 @@ impl syn::parse::Parse for FunctionAttr { content.parse::()?; content.parse::()?; + // Reject attributes with unparsed trailing tokens instead of silently ignoring + // them: attributes like `#[pallet::weight(...)]` define security-critical dispatch + // metadata, so accepting only a leading prefix of the attribute could silently + // underprice a call. + fn ensure_empty(stream: syn::parse::ParseStream) -> syn::Result<()> { + if stream.is_empty() { + Ok(()) + } else { + Err(stream.error("unexpected trailing tokens in pallet attribute")) + } + } + let lookahead = content.lookahead1(); - if lookahead.peek(keyword::weight) { + let attr = if lookahead.peek(keyword::weight) { content.parse::()?; let weight_content; syn::parenthesized!(weight_content in content); - Ok(FunctionAttr::Weight(weight_content.parse::()?)) + let weight = weight_content.parse::()?; + ensure_empty(&weight_content)?; + FunctionAttr::Weight(weight) } else if lookahead.peek(keyword::call_index) { content.parse::()?; let call_index_content; @@ -159,33 +173,40 @@ impl syn::parse::Parse for FunctionAttr { let msg = "Number literal must not have a suffix"; return Err(syn::Error::new(index.span(), msg)); } - Ok(FunctionAttr::CallIndex(index.base10_parse()?)) + ensure_empty(&call_index_content)?; + FunctionAttr::CallIndex(index.base10_parse()?) } else if lookahead.peek(keyword::feeless_if) { content.parse::()?; let closure_content; syn::parenthesized!(closure_content in content); - Ok(FunctionAttr::FeelessIf( - closure_content.span(), - closure_content.parse::().map_err(|e| { - let msg = "Invalid feeless_if attribute: expected a closure"; - let mut err = syn::Error::new(closure_content.span(), msg); - err.combine(e); - err - })?, - )) + let closure = closure_content.parse::().map_err(|e| { + let msg = "Invalid feeless_if attribute: expected a closure"; + let mut err = syn::Error::new(closure_content.span(), msg); + err.combine(e); + err + })?; + ensure_empty(&closure_content)?; + FunctionAttr::FeelessIf(closure_content.span(), closure) } else if lookahead.peek(keyword::authorize) { content.parse::()?; let closure_content; syn::parenthesized!(closure_content in content); - Ok(FunctionAttr::Authorize(closure_content.parse::()?)) + let expr = closure_content.parse::()?; + ensure_empty(&closure_content)?; + FunctionAttr::Authorize(expr) } else if lookahead.peek(keyword::weight_of_authorize) { content.parse::()?; let closure_content; syn::parenthesized!(closure_content in content); - Ok(FunctionAttr::WeightOfAuthorize(closure_content.parse::()?)) + let expr = closure_content.parse::()?; + ensure_empty(&closure_content)?; + FunctionAttr::WeightOfAuthorize(expr) } else { - Err(lookahead.error()) - } + return Err(lookahead.error()); + }; + + ensure_empty(&content)?; + Ok(attr) } } diff --git a/frame/support-procedural/src/pallet/parse/tests/calls.rs b/frame/support-procedural/src/pallet/parse/tests/calls.rs new file mode 100644 index 00000000..320d2543 --- /dev/null +++ b/frame/support-procedural/src/pallet/parse/tests/calls.rs @@ -0,0 +1,92 @@ +// This file is part of Substrate. + +// Copyright (C) Parity Technologies (UK) Ltd. +// SPDX-License-Identifier: Apache-2.0 + +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use syn::parse_quote; + +#[test] +fn test_weight_with_trailing_tokens_is_rejected() { + assert_pallet_parse_error! { + #[manifest_dir("../examples/basic")] + #[error_regex("unexpected trailing tokens in pallet attribute")] + #[frame_support::pallet] + pub mod pallet { + #[pallet::config] + pub trait Config: frame_system::Config {} + + #[pallet::pallet] + pub struct Pallet(_); + + #[pallet::call] + impl Pallet { + #[pallet::call_index(0)] + #[pallet::weight(Weight::zero(), DispatchClass::Operational)] + pub fn noop(origin: OriginFor) -> DispatchResult { + Ok(()) + } + } + } + } +} + +#[test] +fn test_call_index_with_trailing_tokens_is_rejected() { + assert_pallet_parse_error! { + #[manifest_dir("../examples/basic")] + #[error_regex("unexpected trailing tokens in pallet attribute")] + #[frame_support::pallet] + pub mod pallet { + #[pallet::config] + pub trait Config: frame_system::Config {} + + #[pallet::pallet] + pub struct Pallet(_); + + #[pallet::call] + impl Pallet { + #[pallet::call_index(0, 1)] + #[pallet::weight(Weight::zero())] + pub fn noop(origin: OriginFor) -> DispatchResult { + Ok(()) + } + } + } + } +} + +#[test] +fn test_well_formed_call_attributes_parse() { + assert_pallet_parses! { + #[manifest_dir("../examples/basic")] + #[frame_support::pallet] + pub mod pallet { + #[pallet::config] + pub trait Config: frame_system::Config {} + + #[pallet::pallet] + pub struct Pallet(_); + + #[pallet::call] + impl Pallet { + #[pallet::call_index(0)] + #[pallet::weight(Weight::zero())] + pub fn noop(origin: OriginFor) -> DispatchResult { + Ok(()) + } + } + } + }; +} diff --git a/frame/support-procedural/src/pallet/parse/tests/mod.rs b/frame/support-procedural/src/pallet/parse/tests/mod.rs index 86caf1e9..ad0f3720 100644 --- a/frame/support-procedural/src/pallet/parse/tests/mod.rs +++ b/frame/support-procedural/src/pallet/parse/tests/mod.rs @@ -220,6 +220,7 @@ pub fn simulate_manifest_dir, F: FnOnce() + std::panic result.unwrap(); } +mod calls; mod tasks; #[test] diff --git a/frame/support/src/generate_genesis_config.rs b/frame/support/src/generate_genesis_config.rs index 283840d7..c54281e1 100644 --- a/frame/support/src/generate_genesis_config.rs +++ b/frame/support/src/generate_genesis_config.rs @@ -64,19 +64,52 @@ impl<'a> InitializedField<'a> { } } +/// Converts a single `snake_case` identifier segment to `camelCase`, exactly matching +/// serde's `#[serde(rename_all = "camelCase")]` field-renaming rule. +/// +/// Serde first produces `PascalCase` (uppercase the first character of every +/// underscore-separated word, dropping the underscores) and then lowercases the first +/// character of the result. +fn to_camel_case_segment(segment: &str) -> String { + let mut pascal = String::new(); + let mut capitalize = true; + for ch in segment.chars() { + if ch == '_' { + capitalize = true; + } else if capitalize { + pascal.push(ch.to_ascii_uppercase()); + capitalize = false; + } else { + pascal.push(ch); + } + } + match pascal.char_indices().nth(1) { + Some((split, _)) => pascal[..split].to_ascii_lowercase() + &pascal[split..], + None => pascal.to_ascii_lowercase(), + } +} + +/// Converts a dotted field path (e.g. `outer.inner_field`) to its `camelCase` form, +/// converting each `.`-separated segment individually. +fn to_camel_case_path(path: &str) -> String { + let mut out = String::new(); + for (i, segment) in path.split('.').enumerate() { + if i > 0 { + out.push('.'); + } + out.push_str(&to_camel_case_segment(segment)); + } + out +} + impl PartialEq for InitializedField<'_> { fn eq(&self, other: &String) -> bool { - #[inline] - /// We need to respect the `camelCase` naming for field names. This means that - /// `"camelCaseKey"` should be considered equal to `"camel_case_key"`. This - /// function implements this comparison. - fn compare_keys(ident_chars: core::str::Chars, camel_chars: core::str::Chars) -> bool { - ident_chars - .filter(|c| *c != '_') - .map(|c| c.to_ascii_uppercase()) - .eq(camel_chars.map(|c| c.to_ascii_uppercase())) - } - *self.1 == *other || compare_keys(self.1.chars(), other.chars()) + // A JSON key matches the recorded identifier either exactly (structs without a + // `rename_all` attribute serialize field names verbatim) or as its `camelCase` + // rename. The camelCase comparison uses serde's exact rule rather than an + // underscore-insensitive match, so distinct identifiers such as `a_b` (=> `aB`) + // and `ab` (=> `ab`) never compare equal and retention stays injective. + *self.1 == *other || to_camel_case_path(&self.1) == *other } } @@ -1183,53 +1216,67 @@ mod retain_keys_test { use super::*; use serde_json::json; - macro_rules! check_initialized_field_eq_cc( - ( $s:literal ) => { - let field = InitializedField::full($s); - let cc = inflector::cases::camelcase::to_camel_case($s); - assert_eq!(field,cc); - } ; - ( &[ $f:literal $(, $r:literal)* ]) => { - let field = InitializedField::full( - concat!( $f $(,".",$r)+ ) - ); - let cc = [ $f $(,$r)+ ].into_iter() - .map(|s| inflector::cases::camelcase::to_camel_case(s)) - .collect::>() - .join("."); - assert_eq!(field,cc); - } ; - ); + // Ground truth is serde's own `rename_all = "camelCase"` rule (see + // `serde_derive`'s `RenameRule::CamelCase`): PascalCase, then lowercase the first char. + #[track_caller] + fn assert_field_eq_cc(ident: &str, expected_camel: &str) { + let field = InitializedField::full(ident); + // A field is equal to its camelCase rename... + assert_eq!(field, expected_camel.to_string(), "{ident} should equal {expected_camel}"); + // ...and to its verbatim (non-renamed) form. + assert_eq!(field, ident.to_string(), "{ident} should equal itself"); + } #[test] fn test_initialized_field_eq_cc_string() { - check_initialized_field_eq_cc!("a_"); - check_initialized_field_eq_cc!("abc"); - check_initialized_field_eq_cc!("aBc"); - check_initialized_field_eq_cc!("aBC"); - check_initialized_field_eq_cc!("ABC"); - check_initialized_field_eq_cc!("2abs"); - check_initialized_field_eq_cc!("2Abs"); - check_initialized_field_eq_cc!("2ABs"); - check_initialized_field_eq_cc!("2aBs"); - check_initialized_field_eq_cc!("AlreadyCamelCase"); - check_initialized_field_eq_cc!("alreadyCamelCase"); - check_initialized_field_eq_cc!("C"); - check_initialized_field_eq_cc!("1a"); - check_initialized_field_eq_cc!("_1a"); - check_initialized_field_eq_cc!("a_b"); - check_initialized_field_eq_cc!("_a_b"); - check_initialized_field_eq_cc!("a___b"); - check_initialized_field_eq_cc!("__a_b"); - check_initialized_field_eq_cc!("_a___b_C"); - check_initialized_field_eq_cc!("__A___B_C"); - check_initialized_field_eq_cc!(&["a_b", "b_c"]); - check_initialized_field_eq_cc!(&["al_pha", "_a___b_C"]); - check_initialized_field_eq_cc!(&["al_pha_", "_a___b_C"]); - check_initialized_field_eq_cc!(&["first_field", "al_pha_", "_a___b_C"]); - check_initialized_field_eq_cc!(&["al_pha_", "__2nd_field", "_a___b_C"]); - check_initialized_field_eq_cc!(&["al_pha_", "__2nd3and_field", "_a___b_C"]); - check_initialized_field_eq_cc!(&["_a1", "_a2", "_a3_"]); + assert_field_eq_cc("a_", "a"); + assert_field_eq_cc("abc", "abc"); + assert_field_eq_cc("aBc", "aBc"); + assert_field_eq_cc("aBC", "aBC"); + assert_field_eq_cc("ABC", "aBC"); + assert_field_eq_cc("2abs", "2abs"); + assert_field_eq_cc("2Abs", "2Abs"); + assert_field_eq_cc("2ABs", "2ABs"); + assert_field_eq_cc("2aBs", "2aBs"); + assert_field_eq_cc("AlreadyCamelCase", "alreadyCamelCase"); + assert_field_eq_cc("alreadyCamelCase", "alreadyCamelCase"); + assert_field_eq_cc("C", "c"); + assert_field_eq_cc("1a", "1a"); + assert_field_eq_cc("_1a", "1a"); + assert_field_eq_cc("a_b", "aB"); + assert_field_eq_cc("_a_b", "aB"); + assert_field_eq_cc("a___b", "aB"); + assert_field_eq_cc("__a_b", "aB"); + assert_field_eq_cc("_a___b_C", "aBC"); + assert_field_eq_cc("__A___B_C", "aBC"); + assert_field_eq_cc("a_b.b_c", "aB.bC"); + assert_field_eq_cc("al_pha._a___b_C", "alPha.aBC"); + assert_field_eq_cc("al_pha_._a___b_C", "alPha.aBC"); + assert_field_eq_cc("first_field.al_pha_._a___b_C", "firstField.alPha.aBC"); + } + + #[test] + fn test_camel_case_comparison_is_injective() { + // The reported collision: `a_b` (=> `aB`) must not match the JSON key `ab`, and + // `ab` must not match the JSON key `aB`. + assert_ne!(InitializedField::full("a_b"), "ab".to_string()); + assert_ne!(InitializedField::full("ab"), "aB".to_string()); + // Sanity: each still matches its own canonical camelCase / verbatim form. + assert_eq!(InitializedField::full("a_b"), "aB".to_string()); + assert_eq!(InitializedField::full("ab"), "ab".to_string()); + } + + #[test] + fn retain_does_not_keep_colliding_unselected_field() { + // Serialized default object exposes both `aB` (from `a_b`) and `ab` (from `ab`). + let mut v = json!({ + "aB": 7, + "ab": true, + }); + // Only `a_b` was initialized. + retain_initialized_fields(&mut v, &[InitializedField::full("a_b")], String::default()); + // The unrelated `ab` field must be pruned, not retained by an ambiguous match. + assert_eq!(v, json!({ "aB": 7 })); } #[test] diff --git a/frame/support/src/storage/bounded_vec.rs b/frame/support/src/storage/bounded_vec.rs index 06c58915..dc13116e 100644 --- a/frame/support/src/storage/bounded_vec.rs +++ b/frame/support/src/storage/bounded_vec.rs @@ -30,6 +30,10 @@ impl> StorageTryAppend for BoundedVec { fn bound() -> usize { S::get() as usize } + + fn len(value: &Self) -> usize { + value.len() + } } #[cfg(test)] diff --git a/frame/support/src/storage/child.rs b/frame/support/src/storage/child.rs index 7109e921..4dd868cb 100644 --- a/frame/support/src/storage/child.rs +++ b/frame/support/src/storage/child.rs @@ -80,8 +80,13 @@ pub fn put(child_info: &ChildInfo, key: &[u8], value: &T) { /// Remove `key` from storage, returning its value if it had an explicit entry or `None` otherwise. pub fn take(child_info: &ChildInfo, key: &[u8]) -> Option { + // Remove any explicit entry, even one whose bytes fail to decode (in which case `get` + // returns `None`). Conditioning the removal on a successful decode would leave a + // corrupted entry in place and break the `take_or*` contract that no explicit entry + // remains on return. + let had_explicit_entry = exists(child_info, key); let r = get(child_info, key); - if r.is_some() { + if had_explicit_entry { kill(child_info, key); } r @@ -176,11 +181,53 @@ pub fn kill_storage(child_info: &ChildInfo, limit: Option) -> KillStorageRe /// /// Please note that keys which are residing in the overlay for the child are deleted without /// counting towards the `limit`. +/// +/// When a `limit` is given, the pre/post key walks used to attribute overlay-only removals are +/// themselves bounded (to `limit` plus a fixed allowance), so a child trie holding many more +/// keys than the limit cannot force unbounded counting work. If the trie holds more keys than +/// the bound, overlay-only removals may be under-reported (never over-reported) in +/// `unique`/`loops`; the `backend` count is always exact. pub fn clear_storage( child_info: &ChildInfo, maybe_limit: Option, _maybe_cursor: Option<&[u8]>, ) -> MultiRemovalResults { + // The legacy `storage_kill` host function only reports keys removed from the backend; it + // also deletes overlay-resident keys (those written earlier in the current block) but does + // not count them. Count the visible keys before and after so those overlay-only removals + // are still reflected in `unique`/`loops`, which callers may use for weight or cleanup + // accounting. + // + // Both walks stop after `up_to` keys: without a bound, a limited clear of a large + // (potentially user-growable) child trie would perform O(total keys) `next_key` host calls + // just for accounting — the very unbounded-work pattern the `limit` exists to prevent. + fn key_count(child_info: &ChildInfo, up_to: u32) -> u32 { + let mut count: u32 = if exists(child_info, &[]) { 1 } else { 0 }; + let mut previous_key = Vec::new(); + while count < up_to { + match sp_io::default_child_storage::next_key(child_info.storage_key(), &previous_key) { + Some(next_key) => { + count = count.saturating_add(1); + previous_key = next_key; + }, + None => break, + } + } + count + } + + // Extra headroom on top of `limit` for the counting walks. Overlay-resident keys (written + // earlier in the current block) do not count towards `limit`, so allow enumerating this many + // keys beyond it before giving up on exact overlay attribution. When both walks saturate at + // the cap the computed overlay removals collapse to zero rather than going negative. + const OVERLAY_ATTRIBUTION_CAP: u32 = 1024; + // With no limit the host call itself is already O(total keys), so exact counting does not + // change the asymptotic cost. + let count_cap = + maybe_limit.map_or(u32::MAX, |limit| limit.saturating_add(OVERLAY_ATTRIBUTION_CAP)); + + let keys_before = key_count(child_info, count_cap); + // TODO: Once the network has upgraded to include the new host functions, this code can be // enabled. // sp_io::default_child_storage::storage_kill(prefix, maybe_limit, maybe_cursor) @@ -193,7 +240,14 @@ pub fn clear_storage( AllRemoved(db) => (None, db), SomeRemaining(db) => (Some(child_info.storage_key().to_vec()), db), }; - MultiRemovalResults { maybe_cursor, backend, unique: backend, loops: backend } + + // Overlay-only removals = total keys deleted (before - after) minus those the host call + // already accounted for in `backend`. + let keys_after = key_count(child_info, count_cap); + let overlay = keys_before.saturating_sub(keys_after).saturating_sub(backend); + let total = backend.saturating_add(overlay); + + MultiRemovalResults { maybe_cursor, backend, unique: total, loops: total } } /// Ensure `key` has no explicit entry in storage. @@ -237,3 +291,103 @@ pub fn len(child_info: &ChildInfo, key: &[u8]) -> Option { }, } } + +#[cfg(test)] +mod tests { + use super::*; + use sp_io::TestExternalities; + + #[test] + fn take_removes_undecodable_explicit_entry() { + TestExternalities::new_empty().execute_with(|| { + let child_info = ChildInfo::new_default(b"test-child"); + let key = b"slot"; + // `[0x02]` is not a valid SCALE `bool` (only `0x00`/`0x01` decode). + put_raw(&child_info, key, &[0x02]); + assert!(exists(&child_info, key)); + + // `take` on a corrupted entry returns `None` (decode fails) but must still + // clear the explicit entry. + let taken = take::(&child_info, key); + assert_eq!(taken, None); + assert!(!exists(&child_info, key), "corrupt explicit entry must be removed by take"); + }); + } + + #[test] + fn take_or_clears_corrupt_entry_and_returns_default() { + TestExternalities::new_empty().execute_with(|| { + let child_info = ChildInfo::new_default(b"test-child"); + let key = b"one-shot"; + put_raw(&child_info, key, &[0x02]); + + // `take_or` must honor its contract: no explicit entry remains afterwards. + assert!(!take_or::(&child_info, key, false)); + assert!(!exists(&child_info, key)); + // A subsequent emptiness check now sees the slot as free. + assert!(!exists(&child_info, key)); + }); + } + + #[test] + fn clear_storage_counts_overlay_only_removals() { + TestExternalities::new_empty().execute_with(|| { + let child_info = ChildInfo::new_default(b"overlay-child"); + // Stage several keys in the overlay within this block. + for i in 0u8..5 { + put_raw(&child_info, &[i], &[i]); + } + + // Even with a backend limit of zero, the overlay keys are deleted and must be + // reported in `unique`/`loops` instead of being counted as zero work. + let res = clear_storage(&child_info, Some(0), None); + assert_eq!(res.unique, 5); + assert_eq!(res.loops, 5); + for i in 0u8..5 { + assert!(!exists(&child_info, &[i])); + } + }); + } + + #[test] + fn clear_storage_counting_is_bounded_for_large_tries() { + let child_info = ChildInfo::new_default(b"big-child"); + let mut ext = TestExternalities::new_empty(); + // Commit more backend keys than the counting cap (limit + 1024) can enumerate. + ext.execute_with(|| { + for i in 0u32..1_500 { + put_raw(&child_info, &i.to_le_bytes(), &[1]); + } + }); + ext.commit_all().unwrap(); + + ext.execute_with(|| { + // Stage a few overlay-only keys on top. + for i in 0u8..5 { + put_raw(&child_info, &[0xff, i], &[i]); + } + + // The counting walks saturate at the cap, so overlay attribution degrades to zero + // (documented under-reporting) instead of going negative or walking the whole trie. + let res = clear_storage(&child_info, Some(10), None); + assert_eq!(res.backend, 10, "backend removals are reported exactly"); + assert_eq!(res.unique, 10, "saturated counting must not inflate `unique`"); + for i in 0u8..5 { + assert!(!exists(&child_info, &[0xff, i]), "overlay keys are still deleted"); + } + }); + } + + #[test] + fn take_still_returns_and_clears_valid_entry() { + TestExternalities::new_empty().execute_with(|| { + let child_info = ChildInfo::new_default(b"test-child"); + let key = b"slot"; + put(&child_info, key, &123u32); + assert_eq!(take::(&child_info, key), Some(123)); + assert!(!exists(&child_info, key)); + // Taking an absent key returns `None` and remains absent. + assert_eq!(take::(&child_info, key), None); + }); + } +} diff --git a/frame/support/src/storage/generator/double_map.rs b/frame/support/src/storage/generator/double_map.rs index a9116f1f..fbbf1f07 100644 --- a/frame/support/src/storage/generator/double_map.rs +++ b/frame/support/src/storage/generator/double_map.rs @@ -318,7 +318,7 @@ where V: StorageAppend, { let final_key = Self::storage_double_map_final_key(k1, k2); - sp_io::storage::append(&final_key, item.encode()); + V::append(&final_key, item); } fn migrate_keys< @@ -396,6 +396,7 @@ where let mut key_material = G::Hasher2::reverse(raw_key_without_prefix); K2::decode(&mut key_material) }, + phantom: Default::default(), } } @@ -450,6 +451,7 @@ where let k2 = K2::decode(&mut k2_material)?; Ok((k1, k2)) }, + phantom: Default::default(), } } diff --git a/frame/support/src/storage/generator/map.rs b/frame/support/src/storage/generator/map.rs index 2d1f6c9f..b18b502a 100644 --- a/frame/support/src/storage/generator/map.rs +++ b/frame/support/src/storage/generator/map.rs @@ -114,6 +114,7 @@ where let mut key_material = G::Hasher::reverse(raw_key_without_prefix); K::decode(&mut key_material) }, + phantom: Default::default(), } } @@ -295,7 +296,7 @@ impl> storage::StorageMap V: StorageAppend, { let key = Self::storage_map_final_key(key); - sp_io::storage::append(&key, item.encode()); + V::append(&key, item); } fn migrate_key>(key: KeyArg) -> Option { diff --git a/frame/support/src/storage/generator/nmap.rs b/frame/support/src/storage/generator/nmap.rs index 9083aba9..32814247 100755 --- a/frame/support/src/storage/generator/nmap.rs +++ b/frame/support/src/storage/generator/nmap.rs @@ -287,7 +287,7 @@ where V: StorageAppend, { let final_key = Self::storage_n_map_final_key::(key); - sp_io::storage::append(&final_key, item.encode()); + V::append(&final_key, item); } fn migrate_keys(key: KArg, hash_fns: K::HArg) -> Option @@ -356,6 +356,7 @@ impl> previous_key: prefix, drain: false, closure: K::decode_partial_key, + phantom: Default::default(), } } @@ -412,6 +413,7 @@ impl> let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?; Ok(final_key) }, + phantom: Default::default(), } } diff --git a/frame/support/src/storage/generator/value.rs b/frame/support/src/storage/generator/value.rs index 21166b39..aa85d881 100644 --- a/frame/support/src/storage/generator/value.rs +++ b/frame/support/src/storage/generator/value.rs @@ -156,6 +156,6 @@ impl> storage::StorageValue for G { T: StorageAppend, { let key = Self::storage_value_final_key(); - sp_io::storage::append(&key, item.encode()); + T::append(&key, item); } } diff --git a/frame/support/src/storage/migration.rs b/frame/support/src/storage/migration.rs index 67bccfb2..3af7baa9 100644 --- a/frame/support/src/storage/migration.rs +++ b/frame/support/src/storage/migration.rs @@ -416,6 +416,17 @@ pub fn move_prefix_bounded( return MovePrefixResult { moved: 0, maybe_cursor: None } } + // The source and destination prefixes must be disjoint. If either is a prefix of the + // other, the destination keys written during the drain can fall back inside the source + // traversal domain, causing the move to reprocess its own output (unbounded work) or to + // terminate leaving values under the source prefix (silent partial migration). A bad + // prefix pair is a migration-author error, so fail loudly rather than corrupt storage. + assert!( + !to_prefix.starts_with(from_prefix) && !from_prefix.starts_with(to_prefix), + "move_prefix: `from_prefix` and `to_prefix` must be disjoint (neither may be a prefix \ + of the other), otherwise the move would overlap its own traversal domain", + ); + let previous_key = maybe_cursor.map(|c| c.to_vec()).unwrap_or_else(|| from_prefix.to_vec()); let mut iter = PrefixIterator::<(Vec, Vec)>::new( from_prefix.to_vec(), @@ -511,6 +522,26 @@ mod tests { }) } + #[test] + #[should_panic(expected = "must be disjoint")] + fn move_prefix_rejects_destination_nested_in_source() { + TestExternalities::new_empty().execute_with(|| { + // `to_prefix` starts with `from_prefix`: moved keys would re-enter the source + // traversal domain and get reprocessed. This must be rejected. + crate::storage::unhashed::put_raw(b"aab", b"v"); + move_prefix(b"a", b"aa"); + }) + } + + #[test] + #[should_panic(expected = "must be disjoint")] + fn move_prefix_rejects_source_nested_in_destination() { + TestExternalities::new_empty().execute_with(|| { + crate::storage::unhashed::put_raw(b"aab", b"v"); + move_prefix(b"aa", b"a"); + }) + } + #[test] fn test_move_prefix_bounded_moves_in_chunks() { TestExternalities::new_empty().execute_with(|| { diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index 61939256..85e77c55 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -1030,12 +1030,17 @@ impl PrefixIterator { /// Trait for specialising on removal logic of [`PrefixIterator`]. pub trait PrefixIteratorOnRemoval { + /// Whether [`Self::on_removal`] needs the removed value. When `false`, key-only drain + /// iterators skip the per-key value read and pass an empty slice instead, avoiding a host + /// call that would exist only to feed the hook. + const NEEDS_VALUE: bool = true; /// This function is called whenever a key/value is removed. fn on_removal(key: &[u8], value: &[u8]); } /// No-op implementation. impl PrefixIteratorOnRemoval for () { + const NEEDS_VALUE: bool = false; fn on_removal(_key: &[u8], _value: &[u8]) {} } @@ -1132,7 +1137,9 @@ impl Iterator for PrefixIterator { +/// +/// If draining, then the hook `OnRemoval::on_removal` is called after each removal. +pub struct KeyPrefixIterator { prefix: Vec, previous_key: Vec, /// If true then value are removed while iterating @@ -1140,6 +1147,20 @@ pub struct KeyPrefixIterator { /// Function that take `raw_key_without_prefix` and decode `T`. /// `raw_key_without_prefix` is the raw storage key without the prefix iterated on. closure: fn(&[u8]) -> Result, + phantom: core::marker::PhantomData, +} + +impl KeyPrefixIterator { + /// Converts to the same iterator but with the different 'OnRemoval' type + pub fn convert_on_removal(self) -> KeyPrefixIterator { + KeyPrefixIterator:: { + prefix: self.prefix, + previous_key: self.previous_key, + drain: self.drain, + closure: self.closure, + phantom: Default::default(), + } + } } impl KeyPrefixIterator { @@ -1150,14 +1171,25 @@ impl KeyPrefixIterator { /// a `Result` containing the decoded key type `T` if successful, and a `codec::Error` on /// failure. The `&[u8]` argument represents the raw, undecoded key without the prefix of the /// current item. + /// + /// The returned iterator uses the no-op `OnRemoval` hook; use [`Self::convert_on_removal`] to + /// attach a different one. pub fn new( prefix: Vec, previous_key: Vec, decode_fn: fn(&[u8]) -> Result, ) -> Self { - KeyPrefixIterator { prefix, previous_key, drain: false, closure: decode_fn } + KeyPrefixIterator { + prefix, + previous_key, + drain: false, + closure: decode_fn, + phantom: Default::default(), + } } +} +impl KeyPrefixIterator { /// Get the last key that has been iterated upon and return it. pub fn last_raw_key(&self) -> &[u8] { &self.previous_key @@ -1180,7 +1212,7 @@ impl KeyPrefixIterator { } } -impl Iterator for KeyPrefixIterator { +impl Iterator for KeyPrefixIterator { type Item = T; fn next(&mut self) -> Option { @@ -1191,7 +1223,27 @@ impl Iterator for KeyPrefixIterator { if let Some(next) = maybe_next { self.previous_key = next; if self.drain { - unhashed::kill(&self.previous_key); + if OnRemoval::NEEDS_VALUE { + // The raw value is read before removal so the `OnRemoval` hook can + // observe what was deleted. + let raw_value = match unhashed::get_raw(&self.previous_key) { + Some(raw_value) => raw_value, + None => { + log::error!( + "next_key returned a key with no value at {:?}", + self.previous_key, + ); + continue + }, + }; + unhashed::kill(&self.previous_key); + OnRemoval::on_removal(&self.previous_key, &raw_value); + } else { + // The hook does not inspect the value (e.g. the counted map counter + // update), so skip the read that would exist only to feed it. + unhashed::kill(&self.previous_key); + OnRemoval::on_removal(&self.previous_key, &[]); + } } let raw_key_without_prefix = &self.previous_key[self.prefix.len()..]; @@ -1465,10 +1517,19 @@ pub trait StoragePrefixedMap { } } -/// Marker trait that will be implemented for types that support the `storage::append` api. +/// Trait that will be implemented for types that support the `storage::append` api. /// /// This trait is sealed. -pub trait StorageAppend: private::Sealed {} +pub trait StorageAppend: private::Sealed { + /// Append `item` to the storage entry at `key`. + /// + /// The default implementation performs a raw SCALE append, which is only sound for + /// sequence-like encodings whose raw representation cannot violate the container's + /// invariants (e.g. `Vec`). Containers with uniqueness invariants must override this. + fn append>(key: &[u8], item: EncodeLikeItem) { + sp_io::storage::append(key, item.encode()); + } +} /// It is expected that the length is at the beginning of the encoded object /// and that the length is a `Compact`. @@ -1562,7 +1623,20 @@ mod private { impl StorageAppend for Vec {} impl StorageDecodeLength for Vec {} -impl StorageAppend for BTreeSet {} +impl StorageAppend for BTreeSet { + fn append>(key: &[u8], item: EncodeLikeItem) { + // A raw SCALE append can store duplicate raw entries that a decoded `BTreeSet` + // silently deduplicates, letting repeated appends of the same element grow state and + // decode cost without changing the logical set. Instead, decode the stored set, + // insert the item, and write the canonical deduplicated representation back. + let mut value: Self = unhashed::get(key).unwrap_or_default(); + let item = item + .using_encoded(|mut encoded| T::decode(&mut encoded)) + .expect("an `EncodeLike` value must encode to a decodable `T`; qed"); + value.insert(item); + unhashed::put(key, &value); + } +} impl StorageDecodeNonDedupLength for BTreeSet {} // Blanket implementation StorageDecodeNonDedupLength for all types that are StorageDecodeLength. @@ -1583,6 +1657,33 @@ impl StorageAppend for Digest {} /// This trait is sealed. pub trait StorageTryAppend: StorageDecodeLength + private::Sealed { fn bound() -> usize; + + /// The number of items in `value`. + /// + /// Used by the bounded `try_append` implementations to enforce the bound against the + /// effective query value (e.g. a non-empty `OnEmpty` default) when no explicit entry is + /// present in storage and `decode_len` therefore returns `None`. + fn len(value: &Self) -> usize; +} + +/// Compute the length of the *effective* value of a bounded storage entry that has no explicit +/// (decodable) trie entry, given the query-kind conversions of its storage type. +/// +/// A missing explicit entry is not necessarily an empty collection: `ValueQuery` storage falls +/// back to its `OnEmpty` default, which may be non-empty or even at capacity. Bounded appends +/// must enforce their bound against that effective default instead of assuming zero. +fn effective_absent_len( + from_optional_value_to_query: FromOptional, + from_query_to_optional_value: ToOptional, +) -> usize +where + T: StorageTryAppend, + FromOptional: FnOnce(Option) -> Query, + ToOptional: FnOnce(Query) -> Option, +{ + from_query_to_optional_value(from_optional_value_to_query(None)) + .map(|value| >::len(&value)) + .unwrap_or(0) } /// Storage value that is capable of [`StorageTryAppend`]. @@ -1601,7 +1702,15 @@ where { fn try_append>(item: LikeI) -> Result<(), ()> { let bound = T::bound(); - let current = Self::decode_len().unwrap_or_default(); + // `decode_len` only sees explicitly stored entries; with no (decodable) entry the + // effective value is the query-kind default, which may be non-empty. Enforce the + // bound against that effective value instead of assuming it is empty. + let current = Self::decode_len().unwrap_or_else(|| { + effective_absent_len::( + Self::from_optional_value_to_query, + Self::from_query_to_optional_value, + ) + }); if current < bound { // NOTE: we cannot reuse the implementation for `Vec` here because we never want to // mark `BoundedVec` as `StorageAppend`. @@ -1637,7 +1746,14 @@ where item: LikeI, ) -> Result<(), ()> { let bound = T::bound(); - let current = Self::decode_len(key.clone()).unwrap_or_default(); + // See `TryAppendValue::try_append`: a missing entry may still have a non-empty + // query-kind default, so the bound is enforced against the effective value. + let current = Self::decode_len(key.clone()).unwrap_or_else(|| { + effective_absent_len::( + Self::from_optional_value_to_query, + Self::from_query_to_optional_value, + ) + }); if current < bound { let key = Self::storage_map_final_key(key); sp_io::storage::append(&key, item.encode()); @@ -1682,7 +1798,14 @@ where item: LikeI, ) -> Result<(), ()> { let bound = T::bound(); - let current = Self::decode_len(key1.clone(), key2.clone()).unwrap_or_default(); + // See `TryAppendValue::try_append`: a missing entry may still have a non-empty + // query-kind default, so the bound is enforced against the effective value. + let current = Self::decode_len(key1.clone(), key2.clone()).unwrap_or_else(|| { + effective_absent_len::( + Self::from_optional_value_to_query, + Self::from_query_to_optional_value, + ) + }); if current < bound { let double_map_key = Self::storage_double_map_final_key(key1, key2); sp_io::storage::append(&double_map_key, item.encode()); @@ -1722,7 +1845,14 @@ where item: LikeI, ) -> Result<(), ()> { let bound = T::bound(); - let current = Self::decode_len(key.clone()).unwrap_or_default(); + // See `TryAppendValue::try_append`: a missing entry may still have a non-empty + // query-kind default, so the bound is enforced against the effective value. + let current = Self::decode_len(key.clone()).unwrap_or_else(|| { + effective_absent_len::( + Self::from_optional_value_to_query, + Self::from_query_to_optional_value, + ) + }); if current < bound { let key = Self::storage_n_map_final_key::(key); sp_io::storage::append(&key, item.encode()); @@ -2098,6 +2228,61 @@ mod test { }); } + pub struct FullDefault; + impl crate::traits::Get>> for FullDefault { + fn get() -> BoundedVec> { + vec![1, 2, 3].try_into().expect("fits the bound") + } + } + + struct WithDefaultPrefix; + impl crate::traits::StorageInstance for WithDefaultPrefix { + const STORAGE_PREFIX: &'static str = "with_default"; + fn pallet_prefix() -> &'static str { + "test" + } + } + + // `#[storage_alias]` does not support an `OnEmpty` parameter, so define these directly. + type FooWithDefault = types::StorageValue< + WithDefaultPrefix, + BoundedVec>, + types::ValueQuery, + FullDefault, + >; + type FooMapWithDefault = types::StorageMap< + WithDefaultPrefix, + Twox128, + u32, + BoundedVec>, + types::ValueQuery, + FullDefault, + >; + + #[test] + fn try_append_respects_non_empty_default_of_missing_entry() { + TestExternalities::default().execute_with(|| { + // No explicit entry exists, but the effective `ValueQuery` value is the full + // `OnEmpty` default: `try_append` must respect the bound of that effective value + // instead of treating the missing entry as an empty collection. + assert_eq!(FooWithDefault::get().len(), 3); + assert!(FooWithDefault::try_append(4).is_err()); + + assert_eq!(FooMapWithDefault::get(1).len(), 3); + assert!(FooMapWithDefault::try_append(1, 4).is_err()); + + // Explicitly stored values below the bound can still be appended to. + let bounded: BoundedVec> = vec![1].try_into().unwrap(); + FooWithDefault::put(bounded.clone()); + assert_ok!(FooWithDefault::try_append(2)); + assert_eq!(FooWithDefault::get().into_inner(), vec![1, 2]); + + FooMapWithDefault::insert(1, bounded); + assert_ok!(FooMapWithDefault::try_append(1, 2)); + assert_eq!(FooMapWithDefault::get(1).into_inner(), vec![1, 2]); + }); + } + #[test] fn try_append_works() { TestExternalities::default().execute_with(|| { @@ -2205,6 +2390,20 @@ mod test { }); } + #[test] + fn btree_set_append_deduplicates() { + TestExternalities::default().execute_with(|| { + FooSet::append(4); + FooSet::append(4); // duplicate value + FooSet::append(5); + + // The append path canonicalizes the set: the duplicate raw entry is not stored, so + // repeated appends of the same element cannot grow state. + assert_eq!(FooSet::decode_non_dedup_len().unwrap(), 2); + assert_eq!(FooSet::get().unwrap(), BTreeSet::from([4, 5])); + }); + } + #[docify::export] #[test] fn btree_set_decode_non_dedup_len() { @@ -2212,13 +2411,16 @@ mod test { type Store = StorageValue>; TestExternalities::default().execute_with(|| { - Store::append(4); - Store::append(4); // duplicate value - Store::append(5); + // A raw (non-canonical) encoding, e.g. written by a legacy runtime, can contain + // duplicate entries that a decoded `BTreeSet` deduplicates. + sp_io::storage::append(&Store::hashed_key(), 4u32.encode()); + sp_io::storage::append(&Store::hashed_key(), 4u32.encode()); // duplicate value + sp_io::storage::append(&Store::hashed_key(), 5u32.encode()); let length_with_dup_items = 3; assert_eq!(Store::decode_non_dedup_len().unwrap(), length_with_dup_items); + assert_eq!(Store::get().unwrap().len(), 2); }); } } diff --git a/frame/support/src/storage/types/counted_map.rs b/frame/support/src/storage/types/counted_map.rs index 872c6679..deac72e8 100644 --- a/frame/support/src/storage/types/counted_map.rs +++ b/frame/support/src/storage/types/counted_map.rs @@ -132,6 +132,8 @@ pub struct OnRemovalCounterUpdate(core::marker::PhantomData); impl crate::storage::PrefixIteratorOnRemoval for OnRemovalCounterUpdate { + // Only the counter is touched; key-only drains need not read the removed value. + const NEEDS_VALUE: bool = false; fn on_removal(_key: &[u8], _value: &[u8]) { CounterFor::::mutate(|value| value.saturating_dec()); } @@ -325,7 +327,18 @@ where pub fn migrate_key>( key: KeyArg, ) -> Option { - ::Map::migrate_key::(key) + // If the same logical key already has an entry under the current hasher, migrating the + // old-hasher entry overwrites it and collapses two stored keys into one, so the counter + // must be reconciled against the actual pre-migration existence of both entries. + let old_hashed = key.using_encoded(OldHasher::hash); + let new_hashed = key.using_encoded(Hasher::hash); + let collapses = old_hashed.as_ref() != new_hashed.as_ref() && + ::Map::contains_key(Ref::from(&key)); + let value = ::Map::migrate_key::(key); + if value.is_some() && collapses { + CounterFor::::mutate(|value| value.saturating_dec()); + } + value } /// Remove all values in the map. @@ -410,9 +423,20 @@ where Value: StorageTryAppend, { let bound = Value::bound(); - let current = ::Map::decode_len(Ref::from(&key)).unwrap_or_default(); + // See `TryAppendValue::try_append`: a missing entry may still have a non-empty + // query-kind default, so the bound is enforced against the effective value. + let current = ::Map::decode_len(Ref::from(&key)).unwrap_or_else(|| { + crate::storage::effective_absent_len::( + QueryKind::from_optional_value_to_query, + QueryKind::from_query_to_optional_value, + ) + }); if current < bound { - CounterFor::::mutate(|value| value.saturating_inc()); + // The counter tracks the number of keys in the map, so it must only grow when the + // append creates a previously absent key, not on every successful append. + if !::Map::contains_key(Ref::from(&key)) { + CounterFor::::mutate(|value| value.saturating_inc()); + } let key = ::Map::hashed_key_for(key); sp_io::storage::append(&key, item.encode()); Ok(()) @@ -492,8 +516,8 @@ where /// Enumerate all keys in the counted map. /// /// If you alter the map while doing this, you'll get undefined results. - pub fn iter_keys() -> crate::storage::KeyPrefixIterator { - ::Map::iter_keys() + pub fn iter_keys() -> crate::storage::KeyPrefixIterator> { + ::Map::iter_keys().convert_on_removal() } } @@ -1129,6 +1153,32 @@ mod test { }) } + #[test] + fn try_append_respects_non_empty_default_of_missing_entry() { + struct FullBoundedDefault; + impl crate::traits::Get>> for FullBoundedDefault { + fn get() -> BoundedVec> { + vec![1, 2, 3].try_into().expect("fits the bound") + } + } + type B = CountedStorageMap< + Prefix, + Twox64Concat, + u16, + BoundedVec>, + ValueQuery, + FullBoundedDefault, + >; + + TestExternalities::default().execute_with(|| { + // No explicit entry exists, but the effective `ValueQuery` value is the full + // `OnEmpty` default: appending must be rejected and the counter untouched. + assert_eq!(B::get(0).len(), 3); + assert!(B::try_append(0, 4).is_err()); + assert_eq!(B::count(), 0); + }) + } + #[test] fn migrate_keys_works() { type A = CountedStorageMap; @@ -1140,6 +1190,67 @@ mod test { }) } + #[test] + fn iter_keys_drain_updates_count() { + type A = CountedStorageMap; + TestExternalities::default().execute_with(|| { + A::insert(1, 10); + A::insert(2, 20); + assert_eq!(A::count(), 2); + + assert_eq!(A::iter_keys().drain().count(), 2); + assert_eq!(A::count(), 0); + assert_eq!(A::iter().count(), 0); + }) + } + + #[test] + fn try_append_increments_counter_only_for_new_keys() { + type B = CountedStorageMap>>; + TestExternalities::default().execute_with(|| { + B::try_append(0, 1).unwrap(); + assert_eq!(B::count(), 1); + // Repeated appends to an existing key must not inflate the key count. + B::try_append(0, 2).unwrap(); + B::try_append(0, 3).unwrap(); + assert_eq!(B::count(), 1); + B::try_append(1, 1).unwrap(); + assert_eq!(B::count(), 2); + }) + } + + #[test] + fn migrate_key_reconciles_counter_when_entries_collapse() { + type A = CountedStorageMap; + type B = CountedStorageMap; + TestExternalities::default().execute_with(|| { + // The same logical key exists under both the old and the current hasher. + A::insert(1, 1); + B::insert(1, 2); + assert_eq!(B::count(), 2); + + // Migration overwrites the current-hasher entry: two stored keys collapse into one + // and the counter must follow. + assert_eq!(B::migrate_key::(1), Some(1)); + assert_eq!(B::get(1), Some(1)); + assert_eq!(B::count(), 1); + + // A degenerate migration with the current hasher rewrites the same physical key and + // must not change the count. + assert_eq!(B::migrate_key::(1), Some(1)); + assert_eq!(B::count(), 1); + + // Migrating a key that only exists under the old hasher moves it without changing + // the count, and migrating a missing key is a no-op. + A::insert(2, 4); + assert_eq!(B::count(), 2); + assert_eq!(B::migrate_key::(2), Some(4)); + assert_eq!(B::count(), 2); + assert_eq!(B::migrate_key::(3), None); + assert_eq!(B::count(), 2); + }) + } + #[test] fn translate_values() { type A = CountedStorageMap; diff --git a/frame/support/src/storage/types/counted_nmap.rs b/frame/support/src/storage/types/counted_nmap.rs index 267ffd92..2ff67e9c 100644 --- a/frame/support/src/storage/types/counted_nmap.rs +++ b/frame/support/src/storage/types/counted_nmap.rs @@ -126,6 +126,8 @@ pub struct OnRemovalCounterUpdate(core::marker::PhantomData); impl crate::storage::PrefixIteratorOnRemoval for OnRemovalCounterUpdate { + // Only the counter is touched; key-only drains need not read the removed value. + const NEEDS_VALUE: bool = false; fn on_removal(_key: &[u8], _value: &[u8]) { CounterFor::::mutate(|value| value.saturating_dec()); } @@ -181,16 +183,16 @@ where } /// Store or remove the value to be associated with `key` so that `get` returns the `query`. - /// It decrements the counter when the value is removed. + /// It updates the counter when a key is actually added to or removed from the map. pub fn set + TupleToEncodedIter>( key: KArg, query: QueryKind::Query, ) { let option = QueryKind::from_query_to_optional_value(query); - if option.is_none() { - CounterFor::::mutate(|value| value.saturating_dec()); - } - ::Map::set(key, QueryKind::from_optional_value_to_query(option)) + // The counter must reflect actual key existence transitions: writing a value can create + // a previously absent key and removing one only shrinks the map if the key existed. + // `try_mutate_exists` reconciles the counter against both states. + Self::mutate_exists(key, |value| *value = option); } /// Take a value from storage, removing it afterwards. @@ -269,19 +271,21 @@ where Key: HasKeyPrefix, { let result = ::Map::clear_prefix(partial_key, limit, maybe_cursor); - match result.maybe_cursor { - None => CounterFor::::kill(), - Some(_) => CounterFor::::mutate(|x| x.saturating_reduce(result.unique)), - } + // A `None` cursor only means that no items remain under *this* partial prefix; other + // prefixes may still hold entries, so the counter must be decremented by the number of + // removed items rather than killed (which would report a non-empty map as empty). + CounterFor::::mutate(|x| x.saturating_reduce(result.unique)); result } /// Iterate over values that share the first key. - pub fn iter_prefix_values(partial_key: KP) -> PrefixIterator + pub fn iter_prefix_values( + partial_key: KP, + ) -> PrefixIterator> where Key: HasKeyPrefix, { - ::Map::iter_prefix_values(partial_key) + ::Map::iter_prefix_values(partial_key).convert_on_removal() } /// Mutate the value under the given keys. @@ -396,7 +400,24 @@ where where KArg: EncodeLikeTuple + TupleToEncodedIter, { - ::Map::migrate_keys::<_>(key, hash_fns) + // This reimplements the raw map migration (rather than delegating) because the counter + // must be reconciled when the same logical key already has an entry under the current + // hashers: migrating then collapses two stored keys into one. The single-use `hash_fns` + // closures cannot be reused for both a collision check and a delegated call. + let old_key = { + let mut old_key = Self::map_storage_final_prefix(); + old_key.extend_from_slice(&Key::migrate_key(&key, hash_fns)); + old_key + }; + let new_key = ::Map::hashed_key_for(key); + let value = crate::storage::unhashed::take::(&old_key); + if let Some(value) = &value { + if old_key != new_key && crate::storage::unhashed::exists(&new_key) { + CounterFor::::mutate(|value| value.saturating_dec()); + } + crate::storage::unhashed::put(&new_key, value); + } + value } /// Attempt to remove all items from the map. @@ -434,8 +455,8 @@ where /// Iter over all value of the storage. /// /// NOTE: If a value failed to decode because storage is corrupted then it is skipped. - pub fn iter_values() -> crate::storage::PrefixIterator { - ::Map::iter_values() + pub fn iter_values() -> crate::storage::PrefixIterator> { + ::Map::iter_values().convert_on_removal() } /// Translate the values of all elements by a function `f`, in the map in no particular order. @@ -495,11 +516,14 @@ where /// undefined results. pub fn iter_prefix( kp: KP, - ) -> crate::storage::PrefixIterator<(>::Suffix, Value)> + ) -> crate::storage::PrefixIterator< + (>::Suffix, Value), + OnRemovalCounterUpdate, + > where Key: HasReversibleKeyPrefix, { - ::Map::iter_prefix(kp) + ::Map::iter_prefix(kp).convert_on_removal() } /// Enumerate all elements in the map with prefix key `kp` after a specified `starting_raw_key` @@ -526,11 +550,14 @@ where /// undefined results. pub fn iter_key_prefix( kp: KP, - ) -> crate::storage::KeyPrefixIterator<>::Suffix> + ) -> crate::storage::KeyPrefixIterator< + >::Suffix, + OnRemovalCounterUpdate, + > where Key: HasReversibleKeyPrefix, { - ::Map::iter_key_prefix(kp) + ::Map::iter_key_prefix(kp).convert_on_removal() } /// Enumerate all suffix keys in the map with prefix key `kp` after a specified @@ -541,11 +568,14 @@ where pub fn iter_key_prefix_from( kp: KP, starting_raw_key: Vec, - ) -> crate::storage::KeyPrefixIterator<>::Suffix> + ) -> crate::storage::KeyPrefixIterator< + >::Suffix, + OnRemovalCounterUpdate, + > where Key: HasReversibleKeyPrefix, { - ::Map::iter_key_prefix_from(kp, starting_raw_key) + ::Map::iter_key_prefix_from(kp, starting_raw_key).convert_on_removal() } /// Remove all elements from the map with prefix key `kp` and iterate through them in no @@ -585,8 +615,9 @@ where /// Enumerate all keys in the map in no particular order. /// /// If you add or remove values to the map while doing this, you'll get undefined results. - pub fn iter_keys() -> crate::storage::KeyPrefixIterator { - ::Map::iter_keys() + pub fn iter_keys() -> crate::storage::KeyPrefixIterator> + { + ::Map::iter_keys().convert_on_removal() } /// Enumerate all keys in the map after a specified `starting_raw_key` in no particular order. @@ -594,8 +625,8 @@ where /// If you add or remove values to the map while doing this, you'll get undefined results. pub fn iter_keys_from( starting_raw_key: Vec, - ) -> crate::storage::KeyPrefixIterator { - ::Map::iter_keys_from(starting_raw_key) + ) -> crate::storage::KeyPrefixIterator> { + ::Map::iter_keys_from(starting_raw_key).convert_on_removal() } /// Remove all elements from the map and iterate through them in no particular order. @@ -710,6 +741,123 @@ mod test { } } + #[test] + fn clear_prefix_keeps_count_of_other_prefixes() { + type A = CountedStorageNMap< + Prefix, + (NMapKey, NMapKey), + u32, + OptionQuery, + >; + + let mut ext = TestExternalities::default(); + ext.execute_with(|| { + A::insert((1, 10), 100); + A::insert((1, 11), 101); + A::insert((2, 20), 200); + assert_eq!(A::count(), 3); + }); + // Commit to the backend so the legacy `kill_prefix` host function reports the + // removals in `unique`. + ext.commit_all().unwrap(); + ext.execute_with(|| { + // Clearing one partial prefix to completion (cursor `None`) must only subtract the + // removed entries, not reset the global counter: entries under other prefixes + // remain in the map. + let res = A::clear_prefix((1,), u32::MAX, None); + assert!(res.maybe_cursor.is_none()); + assert_eq!(res.unique, 2); + assert_eq!(A::count(), 1); + assert_eq!(A::get((2, 20)), Some(200)); + }); + } + + #[test] + fn drainable_iterator_variants_update_count() { + type A = CountedStorageNMap< + Prefix, + (NMapKey, NMapKey), + u32, + OptionQuery, + >; + + TestExternalities::default().execute_with(|| { + A::insert((3, 30), 11); + A::insert((3, 31), 12); + A::insert((4, 40), 13); + assert_eq!(A::count(), 3); + + // Draining through a value-only prefix iterator must decrement the counter. + assert_eq!(A::iter_prefix_values((3,)).drain().count(), 2); + assert_eq!(A::count(), 1); + assert_eq!(A::iter().collect::>(), vec![((4, 40), 13)]); + + // Draining through the key iterator must decrement the counter as well. + assert_eq!(A::iter_keys().drain().collect::>(), vec![(4, 40)]); + assert_eq!(A::count(), 0); + assert_eq!(A::iter().count(), 0); + }); + } + + #[test] + fn set_reconciles_counter_with_actual_key_existence() { + type A = CountedStorageNMap, u32, OptionQuery>; + + TestExternalities::default().execute_with(|| { + // Setting a value on an absent key creates it and must increment the count. + A::set((1,), Some(10)); + assert_eq!(A::count(), 1); + // Overwriting an existing key must not change the count. + A::set((1,), Some(11)); + assert_eq!(A::get((1,)), Some(11)); + assert_eq!(A::count(), 1); + // Removing an absent key must not decrement the count. + A::set((2,), None); + assert_eq!(A::count(), 1); + // Removing an existing key decrements. + A::set((1,), None); + assert!(!A::contains_key((1,))); + assert_eq!(A::count(), 0); + }) + } + + #[test] + fn migrate_keys_reconciles_counter_when_entries_collapse() { + type A = CountedStorageNMap, u32, OptionQuery>; + type B = CountedStorageNMap, u32, ValueQuery>; + + TestExternalities::default().execute_with(|| { + // The same logical key exists under both the old and the current hashers. + B::insert((2,), 10); + A::insert((2,), 20); + assert_eq!(A::count(), 2); + + // Migration overwrites the current-hasher entry: two stored keys collapse into one + // and the counter must follow. + assert_eq!( + A::migrate_keys((2,), (Box::new(|key| Blake2_256::hash(key).to_vec()),)), + Some(10) + ); + assert_eq!(A::get((2,)), Some(10)); + assert_eq!(A::count(), 1); + + // A degenerate migration with the current hasher rewrites the same physical key and + // must not change the count. + assert_eq!( + A::migrate_keys((2,), (Box::new(|key| Blake2_128Concat::hash(key)),)), + Some(10) + ); + assert_eq!(A::count(), 1); + + // Migrating a missing key is a no-op. + assert_eq!( + A::migrate_keys((3,), (Box::new(|key| Blake2_256::hash(key).to_vec()),)), + None + ); + assert_eq!(A::count(), 1); + }) + } + #[test] fn test_1_key() { type A = CountedStorageNMap, u32, OptionQuery>; diff --git a/frame/support/src/storage/unhashed.rs b/frame/support/src/storage/unhashed.rs index 495c50ca..6a9541e6 100644 --- a/frame/support/src/storage/unhashed.rs +++ b/frame/support/src/storage/unhashed.rs @@ -140,12 +140,41 @@ pub fn kill_prefix(prefix: &[u8], limit: Option) -> sp_io::KillStorageResul pub fn clear_prefix( prefix: &[u8], maybe_limit: Option, - _maybe_cursor: Option<&[u8]>, + maybe_cursor: Option<&[u8]>, ) -> sp_io::MultiRemovalResults { + use sp_io::{KillStorageResult::*, MultiRemovalResults}; + + // The legacy `kill_prefix` host function has no cursor support and, when called repeatedly + // with a limit in the same block, keeps re-reporting the same backend keys. To honour the + // documented cursor contract, resumed calls walk `next_key` from the supplied cursor + // (which observes the overlay, so keys already deleted this block are skipped) and delete + // keys one by one, returning a cursor that reflects actual progress. + if let Some(cursor) = maybe_cursor { + let limit = maybe_limit.unwrap_or(u32::MAX); + let mut previous_key = cursor.to_vec(); + let mut unique = 0u32; + + while unique < limit { + match sp_io::storage::next_key(&previous_key) { + Some(next_key) if next_key.starts_with(prefix) => { + sp_io::storage::clear(&next_key); + previous_key = next_key; + unique = unique.saturating_add(1); + }, + _ => break, + } + } + + let maybe_cursor = match sp_io::storage::next_key(&previous_key) { + Some(next_key) if next_key.starts_with(prefix) => Some(previous_key), + _ => None, + }; + return MultiRemovalResults { maybe_cursor, backend: unique, unique, loops: unique } + } + // TODO: Once the network has upgraded to include the new host functions, this code can be // enabled. // sp_io::storage::clear_prefix(prefix, maybe_limit, maybe_cursor) - use sp_io::{KillStorageResult::*, MultiRemovalResults}; #[allow(deprecated)] let (maybe_cursor, i) = match kill_prefix(prefix, maybe_limit) { AllRemoved(i) => (None, i), @@ -177,3 +206,42 @@ pub fn get_raw(key: &[u8]) -> Option> { pub fn put_raw(key: &[u8], value: &[u8]) { sp_io::storage::set(key, value) } + +#[cfg(test)] +mod tests { + use super::*; + use sp_io::TestExternalities; + + #[test] + fn clear_prefix_resumes_from_cursor_instead_of_repeating_the_same_batch() { + let mut ext = TestExternalities::default(); + let prefix = b"test_prefix:"; + let key = |i: u8| [&prefix[..], &[i]].concat(); + + ext.execute_with(|| { + for i in 0u8..4 { + put_raw(&key(i), &[i]); + } + }); + // Move the keys into the backend so the limited first call leaves a remainder. + ext.commit_all().unwrap(); + + ext.execute_with(|| { + let first = clear_prefix(prefix, Some(2), None); + assert_eq!(first.unique, 2); + let first_cursor = + first.maybe_cursor.expect("entries remain after first partial clear"); + + // The resumed call must make real progress past the keys already deleted in + // this block rather than re-reporting the same backend batch. + let second = clear_prefix(prefix, Some(2), Some(&first_cursor)); + assert_eq!(second.unique, 2); + assert!(second.maybe_cursor.is_none(), "all entries are gone after the second call"); + + assert!(!contains_prefixed_key(prefix)); + for i in 0u8..4 { + assert!(!exists(&key(i)), "key {} must be deleted", i); + } + }); + } +} diff --git a/frame/support/src/storage/weak_bounded_vec.rs b/frame/support/src/storage/weak_bounded_vec.rs index 41d27b51..cf9d6cd8 100644 --- a/frame/support/src/storage/weak_bounded_vec.rs +++ b/frame/support/src/storage/weak_bounded_vec.rs @@ -30,6 +30,10 @@ impl> StorageTryAppend for WeakBoundedVec { fn bound() -> usize { S::get() as usize } + + fn len(value: &Self) -> usize { + value.len() + } } #[cfg(test)] diff --git a/frame/support/src/tests/mod.rs b/frame/support/src/tests/mod.rs index 43811ee4..151f71f0 100644 --- a/frame/support/src/tests/mod.rs +++ b/frame/support/src/tests/mod.rs @@ -215,6 +215,23 @@ pub mod frame_system { #[pallet::storage] pub type Numbers = StorageMap<_, Twox64Concat, u32, u32, OptionQuery>; + #[pallet::tasks_experimental] + impl Pallet { + /// Add a number into the total and remove it. + #[pallet::task_list(Numbers::::iter_keys())] + #[pallet::task_condition(|i| Numbers::::contains_key(i))] + #[pallet::task_weight(crate::weights::Weight::zero())] + #[pallet::task_index(0)] + pub fn add_number_into_total(i: u32) -> crate::dispatch::DispatchResult { + let v = Numbers::::take(i).ok_or(Error::::InvalidTask)?; + Total::::mutate(|(total_keys, total_values)| { + *total_keys += i; + *total_values += v; + }); + Ok(()) + } + } + pub mod pallet_prelude { pub type OriginFor = ::RuntimeOrigin; @@ -225,6 +242,27 @@ pub mod frame_system { } } +/// A minimal second pallet with a non-zero in-code storage version, used to test the +/// interaction between new-pallet storage-version initialization and version-gated +/// migrations. +#[pallet] +pub mod pallet2 { + use super::frame_system; + use crate::pallet_prelude::*; + + pub const STORAGE_VERSION: StorageVersion = StorageVersion::new(3); + + #[pallet::pallet] + #[pallet::storage_version(STORAGE_VERSION)] + pub struct Pallet(_); + + #[pallet::config] + pub trait Config: frame_system::Config {} + + #[pallet::storage] + pub type Values = StorageMap<_, Twox64Concat, u32, u32, OptionQuery>; +} + type BlockNumber = u32; type AccountId = u32; type Header = generic::Header; @@ -250,6 +288,9 @@ mod runtime { #[runtime::pallet_index(0)] pub type System = self::frame_system; + + #[runtime::pallet_index(1)] + pub type Example2 = self::pallet2; } #[crate::derive_impl(self::frame_system::config_preludes::TestDefaultConfig as self::frame_system::DefaultConfig)] @@ -259,10 +300,113 @@ impl Config for Runtime { type ExampleConstant = (); } +impl pallet2::Config for Runtime {} + fn new_test_ext() -> TestExternalities { RuntimeGenesisConfig::default().build_storage().unwrap().into() } +#[test] +fn new_pallet_storage_version_initialized_after_migrations_not_before() { + use crate::{ + migrations::VersionedMigration, + traits::{ + BeforeAllRuntimeMigrations, GetStorageVersion, OnRuntimeUpgrade, StorageVersion, + UncheckedOnRuntimeUpgrade, + }, + weights::Weight, + }; + + /// A migration importing state into the (new) `pallet2` prefix, gated on version 0. + pub struct ImportIntoPallet2; + impl UncheckedOnRuntimeUpgrade for ImportIntoPallet2 { + fn on_runtime_upgrade() -> Weight { + pallet2::Values::::insert(1, 100); + Weight::zero() + } + } + type Migration = VersionedMigration<0, 3, ImportIntoPallet2, pallet2::Pallet, ()>; + + // Empty state: the pallet was just added to the runtime, no genesis ran for it. + TestExternalities::default().execute_with(|| { + assert!(!StorageVersion::exists::>()); + + // Mirrors `Executive::execute_on_runtime_upgrade` ordering: + // 1. `before_all_runtime_migrations` must not pre-seed the storage version, otherwise the + // version-gated migration below is silently skipped. + as BeforeAllRuntimeMigrations>::before_all_runtime_migrations(); + assert!(!StorageVersion::exists::>()); + + // 2. The custom (version-gated) migration runs and observes version 0. + Migration::on_runtime_upgrade(); + assert_eq!(pallet2::Values::::get(1), Some(100), "migration must not be skipped"); + assert_eq!(pallet2::Pallet::::on_chain_storage_version(), 3); + + // 3. The pallet's own hook must not clobber the version set by the migration. + as OnRuntimeUpgrade>::on_runtime_upgrade(); + assert_eq!(pallet2::Pallet::::on_chain_storage_version(), 3); + }); + + // A genuinely new pallet with no migration still gets its version initialized once + // all migrations had the chance to run. + TestExternalities::default().execute_with(|| { + as BeforeAllRuntimeMigrations>::before_all_runtime_migrations(); + as OnRuntimeUpgrade>::on_runtime_upgrade(); + assert!(StorageVersion::exists::>()); + assert_eq!(pallet2::Pallet::::on_chain_storage_version(), 3); + }); +} + +#[test] +fn runtime_task_enumeration_is_lazy() { + use crate::traits::Task as TaskTrait; + + new_test_ext().execute_with(|| { + for i in 0..5u32 { + frame_system::Numbers::::insert(i, i); + } + + // Create the enumeration first, then clear the backing storage. A lazy + // enumeration observes the cleared state; an eager implementation would have + // collected all five tasks into heap vectors before yielding any of them. + let mut tasks = ::iter(); + let _ = frame_system::Numbers::::clear(u32::MAX, None); + assert_eq!(tasks.next(), None); + }); +} + +#[test] +fn runtime_task_enumeration_yields_pallet_tasks() { + use crate::traits::Task as TaskTrait; + + new_test_ext().execute_with(|| { + frame_system::Numbers::::insert(7, 3); + + let tasks = ::iter().collect::>(); + assert_eq!( + tasks, + vec![RuntimeTask::System(frame_system::Task::::AddNumberIntoTotal { i: 7 })] + ); + }); +} + +#[test] +fn get_call_names_unknown_module_returns_empty_not_panic() { + use crate::traits::GetCallMetadata; + + // A known module resolves to its call names. + let modules = ::get_module_names(); + assert!(modules.contains(&"System")); + assert!(!::get_call_names("System").is_empty()); + + // `get_call_names` is public metadata API reachable with caller-supplied strings; an + // unknown module must be a recoverable empty result, not a panic. + assert_eq!( + ::get_call_names("DefinitelyMissingPallet"), + &[] as &[&str], + ); +} + trait Sorted { fn sorted(self) -> Self; } diff --git a/frame/support/src/traits/misc.rs b/frame/support/src/traits/misc.rs index 16da4558..fc87956c 100644 --- a/frame/support/src/traits/misc.rs +++ b/frame/support/src/traits/misc.rs @@ -1119,11 +1119,29 @@ impl WrapperKeepOpaque { } /// Create from the given encoded `data`. + /// + /// # Warning + /// + /// This does not enforce the [`MaxEncodedLen`] bound advertised by this type: `data` longer + /// than [`Self::max_encoded_len`] produces a value that violates the bound and will fail to + /// decode again. Prefer [`Self::try_from_encoded`], which rejects oversized input. + #[deprecated(note = "use `try_from_encoded`, which enforces the `MaxEncodedLen` bound")] pub fn from_encoded(data: Vec) -> Self { Self { data, _phantom: core::marker::PhantomData } } } +impl WrapperKeepOpaque { + /// Create from the given encoded `data`, enforcing the advertised [`MaxEncodedLen`] bound. + /// + /// Returns `None` if `data` is longer than the maximum encoded length of `T`, since such a + /// value would exceed the `max_size` reported for bounded storage and PoV accounting. + pub fn try_from_encoded(data: Vec) -> Option { + (data.len() <= T::max_encoded_len()) + .then(|| Self { data, _phantom: core::marker::PhantomData }) + } +} + impl EncodeLike for WrapperKeepOpaque {} impl EncodeLike> for WrapperKeepOpaque {} @@ -1145,9 +1163,17 @@ impl Encode for WrapperKeepOpaque { } } -impl Decode for WrapperKeepOpaque { +impl Decode for WrapperKeepOpaque { fn decode(input: &mut I) -> Result { - Ok(Self { data: Vec::::decode(input)?, _phantom: core::marker::PhantomData }) + let data = Vec::::decode(input)?; + // Enforce the bound advertised via `MaxEncodedLen`: the opaque bytes are the encoding of + // `T`, so they can never legitimately exceed `T::max_encoded_len()`. Rejecting oversized + // input keeps decoded values (e.g. read back from storage or an extrinsic) within the + // `max_size` used for bounded-storage and PoV accounting. + if data.len() > T::max_encoded_len() { + return Err("WrapperKeepOpaque: encoded data exceeds `T::max_encoded_len()`".into()) + } + Ok(Self { data, _phantom: core::marker::PhantomData }) } fn skip(input: &mut I) -> Result<(), codec::Error> { @@ -1495,6 +1521,23 @@ mod test { WrapperOpaque::::decode(&mut &data[..]).unwrap(); } + #[test] + fn keep_opaque_wrapper_enforces_max_encoded_len() { + // The advertised bound is derived from `T`; opaque bytes longer than + // `T::max_encoded_len()` must be rejected on decode so decoded values never exceed the + // `max_size` used for bounded storage / PoV accounting. + let bound = >::max_encoded_len(); + + let oversized: Vec = vec![0x41; bound + 32]; + let encoded = oversized.encode(); + assert!(WrapperKeepOpaque::::decode(&mut &encoded[..]).is_err()); + + // The checked constructor rejects oversized bytes and accepts in-bound bytes. + assert!(WrapperKeepOpaque::::try_from_encoded(oversized).is_none()); + let within_bound = 3u32.encode(); + assert!(WrapperKeepOpaque::::try_from_encoded(within_bound).is_some()); + } + #[test] fn defensive_min_works() { assert_eq!(10, 10_u32.defensive_min(11_u32)); diff --git a/frame/support/src/traits/reality.rs b/frame/support/src/traits/reality.rs index c73a6d69..f8ef6087 100644 --- a/frame/support/src/traits/reality.rs +++ b/frame/support/src/traits/reality.rs @@ -290,6 +290,18 @@ pub enum Statement { /// Describes the location within the runtime of a callback, along with other type information such /// as parameters passed into the callback. +/// +/// # Warning: not stable across runtime upgrades +/// +/// A `Callback` records only the raw `(pallet_index, call_index)` pair, not the pallet/call name, +/// runtime version, or metadata hash. [`Self::curry`] resolves those indexes against the call +/// layout that exists *when it runs*, not when the callback was created. If a runtime upgrade +/// reorders pallets or call variants, a stored callback can decode into a different call than +/// intended. Implementations of asynchronous consumers such as [`StatementOracle`] must therefore +/// not persist a caller-provided `Callback` across a runtime upgrade (e.g. resolve pending +/// judgements before applying an upgrade, or bind the callback to a version/identity that is +/// revalidated on resolution). Callbacks selected by trusted code within a single runtime version +/// are unaffected. #[derive( CloneNoBound, PartialEqNoBound, EqNoBound, RuntimeDebug, Encode, Decode, MaxEncodedLen, TypeInfo, )] diff --git a/frame/support/src/traits/tokens/fungible/hold.rs b/frame/support/src/traits/tokens/fungible/hold.rs index 356a2625..e692806b 100644 --- a/frame/support/src/traits/tokens/fungible/hold.rs +++ b/frame/support/src/traits/tokens/fungible/hold.rs @@ -362,6 +362,14 @@ pub trait Mutate: ensure!(amount <= have, TokenError::FundsUnavailable); } + // In `OnHold` mode the destination keeps the funds on hold and must already exist: + // `can_deposit(.., Extant)` only concerns issuance accounting, not account existence, so + // without this check a held balance could be credited to a non-existent account, leaving + // hold state with no underlying account. + if mode == OnHold { + ensure!(!Self::total_balance(dest).is_zero(), TokenError::CannotCreate); + } + // We want to make sure we can deposit the amount in advance. If we can't then something is // very wrong. ensure!(Self::can_deposit(dest, amount, Extant) == Success, TokenError::CannotCreate); diff --git a/frame/support/src/traits/tokens/fungible/mod.rs b/frame/support/src/traits/tokens/fungible/mod.rs index 31f8e015..dd994d0e 100644 --- a/frame/support/src/traits/tokens/fungible/mod.rs +++ b/frame/support/src/traits/tokens/fungible/mod.rs @@ -359,10 +359,17 @@ impl< { fn new(who: &A, footprint: Fp) -> Result { ensure!(Fx::balance_frozen(&Rx::get(), who).is_zero(), DispatchError::Unavailable); - Fx::set_frozen(&Rx::get(), who, D::convert(footprint), Polite).map(|_| Self(PhantomData)) + // A zero-cost ticket leaves no freeze marker, so it is indistinguishable from a later + // paid ticket for the same account/reason. Since `drop` thaws the whole reason bucket, + // such a stale zero ticket could later clear a newer ticket's freeze. Reject it. + let amount = D::convert(footprint); + ensure!(!amount.is_zero(), DispatchError::Unavailable); + Fx::set_frozen(&Rx::get(), who, amount, Polite).map(|_| Self(PhantomData)) } fn update(self, who: &A, footprint: Fp) -> Result { - Fx::set_frozen(&Rx::get(), who, D::convert(footprint), Polite).map(|_| Self(PhantomData)) + let amount = D::convert(footprint); + ensure!(!amount.is_zero(), DispatchError::Unavailable); + Fx::set_frozen(&Rx::get(), who, amount, Polite).map(|_| Self(PhantomData)) } fn drop(self, who: &A) -> Result<(), DispatchError> { Fx::thaw(&Rx::get(), who).map(|_| ()) @@ -404,10 +411,18 @@ impl< { fn new(who: &A, footprint: Fp) -> Result { ensure!(F::balance_on_hold(&R::get(), who).is_zero(), DispatchError::Unavailable); - F::set_on_hold(&R::get(), who, D::convert(footprint)).map(|_| Self(PhantomData)) + // A zero-cost ticket leaves no hold marker, so it is indistinguishable from a later paid + // ticket for the same account/reason. Since `drop`/`burn` operate on the whole reason + // bucket, such a stale zero ticket could later release or burn a newer ticket's hold. + // Reject it. + let amount = D::convert(footprint); + ensure!(!amount.is_zero(), DispatchError::Unavailable); + F::set_on_hold(&R::get(), who, amount).map(|_| Self(PhantomData)) } fn update(self, who: &A, footprint: Fp) -> Result { - F::set_on_hold(&R::get(), who, D::convert(footprint)).map(|_| Self(PhantomData)) + let amount = D::convert(footprint); + ensure!(!amount.is_zero(), DispatchError::Unavailable); + F::set_on_hold(&R::get(), who, amount).map(|_| Self(PhantomData)) } fn drop(self, who: &A) -> Result<(), DispatchError> { F::release_all(&R::get(), who, BestEffort).map(|_| ()) diff --git a/frame/support/src/traits/tokens/fungibles/hold.rs b/frame/support/src/traits/tokens/fungibles/hold.rs index 026bfc87..18ef5567 100644 --- a/frame/support/src/traits/tokens/fungibles/hold.rs +++ b/frame/support/src/traits/tokens/fungibles/hold.rs @@ -403,6 +403,14 @@ pub trait Mutate: ensure!(amount <= have, TokenError::FundsUnavailable); } + // In `OnHold` mode the destination keeps the funds on hold and must already exist: + // `can_deposit(.., Extant)` only concerns issuance accounting, not account existence, so + // without this check a held balance could be credited to a non-existent account, leaving + // hold state with no underlying asset account. + if mode == OnHold { + ensure!(!Self::total_balance(asset.clone(), dest).is_zero(), TokenError::CannotCreate); + } + // We want to make sure we can deposit the amount in advance. If we can't then something is // very wrong. ensure!( diff --git a/frame/support/src/view_functions.rs b/frame/support/src/view_functions.rs index cba451b5..c08583b7 100644 --- a/frame/support/src/view_functions.rs +++ b/frame/support/src/view_functions.rs @@ -18,7 +18,7 @@ //! Traits for querying pallet view functions. use alloc::vec::Vec; -use codec::{Decode, DecodeAll, Encode, Output}; +use codec::{Decode, DecodeWithMemLimit, Encode, Output}; use scale_info::TypeInfo; use sp_runtime::RuntimeDebug; @@ -98,9 +98,16 @@ pub trait ViewFunctionIdSuffix { const SUFFIX: [u8; 16]; } +/// Maximum heap memory a single view-function input is allowed to allocate while decoding. +/// +/// View functions are reachable through the `RuntimeViewFunction` runtime API with +/// attacker-controlled input, so decoding is bounded to prevent a small crafted request from +/// forcing a disproportionately large allocation (a decode "bomb"). +pub const MAX_VIEW_FUNCTION_DECODE_MEM: usize = 16 * 1024 * 1024; + /// Automatically implemented for each pallet view function method by the macro /// [`pallet`](crate::pallet). -pub trait ViewFunction: DecodeAll { +pub trait ViewFunction: DecodeWithMemLimit { fn id() -> ViewFunctionId; type ReturnType: Encode; @@ -110,7 +117,12 @@ pub trait ViewFunction: DecodeAll { input: &mut &[u8], output: &mut O, ) -> Result<(), ViewFunctionDispatchError> { - let view_function = Self::decode_all(input)?; + // Use the mem-tracked decode (bounded by `MAX_VIEW_FUNCTION_DECODE_MEM`) instead of a + // plain `DecodeAll`, and still require the whole input to be consumed. + let view_function = Self::decode_with_mem_limit(input, MAX_VIEW_FUNCTION_DECODE_MEM)?; + if !input.is_empty() { + return Err(ViewFunctionDispatchError::Codec) + } let result = view_function.invoke(); Encode::encode_to(&result, output); Ok(()) @@ -132,3 +144,75 @@ pub mod runtime_api { } } } + +#[cfg(test)] +mod tests { + use super::*; + use codec::{Decode, DecodeWithMemTracking}; + + #[derive(Encode, Decode, DecodeWithMemTracking)] + struct DoubleIt { + value: u32, + } + + impl ViewFunction for DoubleIt { + fn id() -> ViewFunctionId { + ViewFunctionId { prefix: [0u8; 16], suffix: [0u8; 16] } + } + type ReturnType = u32; + + fn invoke(self) -> Self::ReturnType { + self.value.wrapping_mul(2) + } + } + + #[test] + fn execute_decodes_invokes_and_encodes() { + let input = 21u32.encode(); + let mut output = Vec::new(); + DoubleIt::execute(&mut &input[..], &mut output).unwrap(); + assert_eq!(u32::decode(&mut &output[..]).unwrap(), 42); + } + + #[test] + fn execute_rejects_trailing_bytes() { + // A valid `u32` followed by extra trailing bytes must be rejected rather than + // silently decoded from a prefix of the input. + let mut input = 21u32.encode(); + input.extend_from_slice(&[0xff, 0xff]); + let mut output = Vec::new(); + assert!(matches!( + DoubleIt::execute(&mut &input[..], &mut output), + Err(ViewFunctionDispatchError::Codec) + )); + } + + #[test] + fn execute_bounds_decode_allocation() { + // A `Vec` view function argument whose declared length would allocate far more + // than the mem limit is rejected instead of being decoded. + #[derive(Encode, Decode, DecodeWithMemTracking)] + struct TakesVec { + data: Vec, + } + impl ViewFunction for TakesVec { + fn id() -> ViewFunctionId { + ViewFunctionId { prefix: [1u8; 16], suffix: [1u8; 16] } + } + type ReturnType = u32; + fn invoke(self) -> Self::ReturnType { + self.data.len() as u32 + } + } + + // Encode a compact length larger than the mem limit, then provide that many bytes. + let len = (MAX_VIEW_FUNCTION_DECODE_MEM + 1) as u32; + let mut input = codec::Compact(len).encode(); + input.extend(core::iter::repeat(0u8).take(len as usize)); + let mut output = Vec::new(); + assert!(matches!( + TakesVec::execute(&mut &input[..], &mut output), + Err(ViewFunctionDispatchError::Codec) + )); + } +} diff --git a/pallets/balances/src/tests/fungible_tests.rs b/pallets/balances/src/tests/fungible_tests.rs index 2c27f81e..b7aa3018 100644 --- a/pallets/balances/src/tests/fungible_tests.rs +++ b/pallets/balances/src/tests/fungible_tests.rs @@ -23,7 +23,7 @@ use frame_support::traits::{ Fortitude::{Force, Polite}, Precision::{BestEffort, Exact}, Preservation::{Expendable, Preserve, Protect}, - Restriction::Free, + Restriction::{Free, OnHold}, }, Consideration, Footprint, LinearStoragePrice, MaybeConsideration, }; @@ -280,6 +280,32 @@ fn frozen_hold_balance_cannot_be_moved_without_force() { }); } +#[test] +fn transfer_on_hold_onhold_requires_existing_dest() { + ExtBuilder::default() + .existential_deposit(1) + .monied(true) + .build_and_execute_with(|| { + assert_ok!(Balances::hold(&TestId::Foo, &1, 7)); + let dest = 99u64; // does not exist + assert_eq!(Balances::total_balance(&dest), 0); + + // `OnHold` mode keeps the funds on hold in `dest`, which must already exist. A + // non-existent destination must be rejected rather than gaining orphan hold state. + assert_noop!( + Balances::transfer_on_hold(&TestId::Foo, &1, &dest, 5, Exact, OnHold, Force), + TokenError::CannotCreate, + ); + assert_eq!(Balances::balance_on_hold(&TestId::Foo, &1), 7); + assert_eq!(Balances::total_balance(&dest), 0); + + // `Free` mode may create the destination, so it still succeeds. + assert_ok!(Balances::transfer_on_hold(&TestId::Foo, &1, &dest, 5, Exact, Free, Force)); + assert_eq!(Balances::balance_on_hold(&TestId::Foo, &1), 2); + assert_eq!(Balances::free_balance(dest), 5); + }); +} + #[test] fn transfer_and_hold() { ExtBuilder::default() @@ -806,7 +832,12 @@ fn lone_freeze_consideration_works() { >; let who = 4; - let zero_ticket = Consideration::new(&who, Footprint::from_parts(0, 0)).unwrap(); + // A zero-footprint ticket is rejected: it would leave no freeze marker and could + // later thaw a newer ticket's freeze under the same reason. + assert_noop!( + Consideration::new(&who, Footprint::from_parts(0, 0)), + sp_runtime::DispatchError::Unavailable, + ); assert_eq!(Balances::balance_frozen(&TestId::Foo, &who), 0); let ticket = Consideration::new(&who, Footprint::from_parts(10, 1)).unwrap(); @@ -818,11 +849,12 @@ fn lone_freeze_consideration_works() { let ticket = ticket.update(&who, Footprint::from_parts(4, 1)).unwrap(); assert_eq!(Balances::balance_frozen(&TestId::Foo, &who), 4); - assert_eq!(ticket.update(&who, Footprint::from_parts(0, 0)).unwrap(), zero_ticket); - assert_eq!(Balances::balance_frozen(&TestId::Foo, &who), 0); - - let ticket = Consideration::new(&who, Footprint::from_parts(10, 1)).unwrap(); - assert_eq!(Balances::balance_frozen(&TestId::Foo, &who), 10); + // Updating to a zero footprint is likewise rejected and leaves the freeze intact. + assert_noop!( + ticket.clone().update(&who, Footprint::from_parts(0, 0)), + sp_runtime::DispatchError::Unavailable, + ); + assert_eq!(Balances::balance_frozen(&TestId::Foo, &who), 4); let _ = ticket.drop(&who).unwrap(); assert_eq!(Balances::balance_frozen(&TestId::Foo, &who), 0); @@ -844,7 +876,12 @@ fn lone_hold_consideration_works() { >; let who = 4; - let zero_ticket = Consideration::new(&who, Footprint::from_parts(0, 0)).unwrap(); + // A zero-footprint ticket is rejected: it would leave no hold marker and could + // later release/burn a newer ticket's hold under the same reason. + assert_noop!( + Consideration::new(&who, Footprint::from_parts(0, 0)), + sp_runtime::DispatchError::Unavailable, + ); assert_eq!(Balances::balance_on_hold(&TestId::Foo, &who), 0); let ticket = Consideration::new(&who, Footprint::from_parts(10, 1)).unwrap(); @@ -871,11 +908,12 @@ fn lone_hold_consideration_works() { let ticket = ticket.update(&who, Footprint::from_parts(4, 1)).unwrap(); assert_eq!(Balances::balance_on_hold(&TestId::Foo, &who), 4); - assert_eq!(ticket.update(&who, Footprint::from_parts(0, 0)).unwrap(), zero_ticket); - assert_eq!(Balances::balance_on_hold(&TestId::Foo, &who), 0); - - let ticket = Consideration::new(&who, Footprint::from_parts(10, 1)).unwrap(); - assert_eq!(Balances::balance_on_hold(&TestId::Foo, &who), 10); + // Updating to a zero footprint is likewise rejected and leaves the hold intact. + assert_noop!( + ticket.clone().update(&who, Footprint::from_parts(0, 0)), + sp_runtime::DispatchError::Unavailable, + ); + assert_eq!(Balances::balance_on_hold(&TestId::Foo, &who), 4); let _ = ticket.drop(&who).unwrap(); assert_eq!(Balances::balance_on_hold(&TestId::Foo, &who), 0); diff --git a/pallets/frame-system/src/extensions/check_weight.rs b/pallets/frame-system/src/extensions/check_weight.rs index 16522611..98675e6e 100644 --- a/pallets/frame-system/src/extensions/check_weight.rs +++ b/pallets/frame-system/src/extensions/check_weight.rs @@ -383,13 +383,13 @@ mod tests { fn operational_extrinsic_limited_by_operational_space_limit() { new_test_ext().execute_with(|| { let weights = block_weights(); - let operational_limit = weights + // The derived cap already accounts for `base_block`, the class base extrinsic weight + // and the average initialization cost, so it is the largest single-extrinsic weight + // admitted by validation. + let call_weight = weights .get(DispatchClass::Operational) - .max_total - .unwrap_or_else(|| weights.max_block); - let base_weight = weights.get(DispatchClass::Operational).base_extrinsic; - - let call_weight = operational_limit - base_weight; + .max_extrinsic + .expect("mock sets `avg_block_initialization`, so the cap is derived; qed"); let okay = DispatchInfo { call_weight, class: DispatchClass::Operational, @@ -609,14 +609,17 @@ mod tests { #[test] fn signed_ext_check_weight_works_normal_tx() { new_test_ext().execute_with(|| { - let normal_limit = normal_weight_limit(); + // The derived per-extrinsic cap accounts for `base_block`, the class base extrinsic + // weight and the average initialization cost. + let max_extrinsic = block_weights() + .get(DispatchClass::Normal) + .max_extrinsic + .expect("mock sets `avg_block_initialization`, so the cap is derived; qed"); let small = DispatchInfo { call_weight: Weight::from_parts(100, 0), ..Default::default() }; - let base_extrinsic = block_weights().get(DispatchClass::Normal).base_extrinsic; - let medium = - DispatchInfo { call_weight: normal_limit - base_extrinsic, ..Default::default() }; + let medium = DispatchInfo { call_weight: max_extrinsic, ..Default::default() }; let big = DispatchInfo { - call_weight: normal_limit - base_extrinsic + Weight::from_parts(1, 0), + call_weight: max_extrinsic + Weight::from_parts(1, 0), ..Default::default() }; let len = 0_usize; diff --git a/pallets/frame-system/src/extensions/weights.rs b/pallets/frame-system/src/extensions/weights.rs index a9c21e7b..5e872c41 100644 --- a/pallets/frame-system/src/extensions/weights.rs +++ b/pallets/frame-system/src/extensions/weights.rs @@ -50,6 +50,16 @@ use frame_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}}; use core::marker::PhantomData; /// Weight functions needed for `frame_system_extensions`. +/// +/// # Note on signer-dependent extensions +/// +/// [`Self::check_non_zero_sender`] and [`Self::check_nonce`] model a *constant* cost, yet their +/// implementations do work proportional to the SCALE-encoded length of the signer (scanning the +/// encoded `AccountId`, encoding it into transaction-pool dependency tags). That work is bounded +/// by `AccountId::max_encoded_len()`, so the constant model is sound only if the benchmarks that +/// generate these values are run with the runtime's actual `AccountId` type using a +/// maximal-length (worst-case) signer. Runtimes with variable-length account identifiers must +/// regenerate these weights accordingly rather than reuse values benchmarked with short signers. pub trait WeightInfo { fn check_genesis() -> Weight; fn check_mortality_mortal_transaction() -> Weight; diff --git a/pallets/frame-system/src/lib.rs b/pallets/frame-system/src/lib.rs index 7de51adb..898c653e 100644 --- a/pallets/frame-system/src/lib.rs +++ b/pallets/frame-system/src/lib.rs @@ -743,6 +743,7 @@ pub mod pallet { Self::can_set_code(&code, true).into_result()?; T::OnSetCode::set_code(code)?; // consume the rest of the block to prevent further transactions + Self::drain_remaining_block_weight(); Ok(Some(T::BlockWeights::get().max_block).into()) } @@ -759,6 +760,8 @@ pub mod pallet { ensure_root(origin)?; Self::can_set_code(&code, false).into_result()?; T::OnSetCode::set_code(code)?; + // consume the rest of the block to prevent further transactions + Self::drain_remaining_block_weight(); Ok(Some(T::BlockWeights::get().max_block).into()) } @@ -912,6 +915,9 @@ pub mod pallet { }; T::OnSetCode::set_code(code)?; + // consume the rest of the block to prevent further transactions + Self::drain_remaining_block_weight(); + Ok(PostDispatchInfo { // consume the rest of the block to prevent further transactions actual_weight: Some(T::BlockWeights::get().max_block), @@ -1932,6 +1938,20 @@ impl Pallet { }); } + /// Consume all remaining weight of the current block, preventing any further transaction + /// from being included after the current one. + /// + /// Used by the runtime-upgrade dispatchables: returning `max_block` as post-dispatch + /// `actual_weight` is not enough, since post-dispatch weight is capped at the (much smaller) + /// static pre-dispatch weight by `calc_actual_weight`. Registering the missing weight + /// directly ensures the block really is drained. + fn drain_remaining_block_weight() { + let max_block = T::BlockWeights::get().max_block; + let remaining = max_block.saturating_sub(Self::block_weight().total()); + // `Operational` so the consumption also exceeds any operational `reserved` allowance. + Self::register_extra_weight_unchecked(remaining, DispatchClass::Operational); + } + /// Start the execution of a particular block. /// /// # Panics diff --git a/pallets/frame-system/src/limits.rs b/pallets/frame-system/src/limits.rs index ab5a98a6..dfa77709 100644 --- a/pallets/frame-system/src/limits.rs +++ b/pallets/frame-system/src/limits.rs @@ -414,12 +414,18 @@ impl BlockWeightsBuilder { } // compute max size of single extrinsic if let Some(init_weight) = init_cost.map(|rate| rate * weights.max_block) { + let base_block = weights.base_block; for class in DispatchClass::all() { let per_class = weights.per_class.get_mut(*class); if per_class.max_extrinsic.is_none() && init_cost.is_some() { + // Every block starts with `base_block` already consumed (in addition to the + // average initialization cost), so it must be subtracted as well; otherwise an + // extrinsic at the derived cap validates against `max_extrinsic` but can never + // be included, as `base_block + base_extrinsic + weight` exceeds `max_block`. per_class.max_extrinsic = per_class .max_total .map(|x| x.saturating_sub(init_weight)) + .map(|x| x.saturating_sub(base_block)) .map(|x| x.saturating_sub(per_class.base_extrinsic)); } } @@ -447,4 +453,39 @@ mod tests { fn default_weights_are_valid() { BlockWeights::default().validate().unwrap(); } + + #[test] + fn derived_max_extrinsic_accounts_for_base_block() { + let base_block = Weight::from_parts(100, 5); + let base_extrinsic = Weight::from_parts(10, 1); + let max_total = Weight::from_parts(1_000, 100); + let init_rate = Perbill::from_percent(1); + + let weights = BlockWeights::builder() + .base_block(base_block) + .for_class(DispatchClass::all(), |w| { + w.base_extrinsic = base_extrinsic; + }) + .for_class(DispatchClass::non_mandatory(), |w| { + w.max_total = Some(max_total); + }) + .avg_block_initialization(init_rate) + .build() + .unwrap(); + + // A block starts with `base_block` (and on average `init_rate * max_block` of + // initialization) already consumed, so an extrinsic at the derived cap must still fit + // within `max_block` on top of those. + let max_extrinsic = weights.get(DispatchClass::Normal).max_extrinsic.unwrap(); + let consumed = weights + .base_block + .saturating_add(init_rate * weights.max_block) + .saturating_add(base_extrinsic) + .saturating_add(max_extrinsic); + assert!( + consumed.all_lte(weights.max_block), + "an extrinsic at the derived cap must fit in the block: {consumed:?} > {:?}", + weights.max_block, + ); + } } diff --git a/pallets/frame-system/src/tests.rs b/pallets/frame-system/src/tests.rs index 710879f8..c05a9c42 100644 --- a/pallets/frame-system/src/tests.rs +++ b/pallets/frame-system/src/tests.rs @@ -659,6 +659,38 @@ fn set_code_checks_works() { } } +#[test] +fn set_code_drains_remaining_block_weight() { + struct ReadRuntimeVersion(Vec); + + impl sp_core::traits::ReadRuntimeVersion for ReadRuntimeVersion { + fn read_runtime_version( + &self, + _wasm_code: &[u8], + _ext: &mut dyn sp_externalities::Externalities, + ) -> Result, String> { + Ok(self.0.clone()) + } + } + + let version = + RuntimeVersion { spec_name: "test".into(), spec_version: 2, ..Default::default() }; + let read_runtime_version = ReadRuntimeVersion(version.encode()); + + let mut ext = new_test_ext(); + ext.register_extension(sp_core::traits::ReadRuntimeVersionExt::new(read_runtime_version)); + ext.execute_with(|| { + let max_block = ::BlockWeights::get().max_block; + assert!(System::block_weight().total().all_lt(max_block)); + + // The post-dispatch `actual_weight` returned by `set_code` is capped at the static + // pre-dispatch weight, so the block must be drained by direct weight registration. + assert_ok!(System::set_code(RawOrigin::Root.into(), vec![1, 2, 3, 4])); + + assert_eq!(System::block_weight().total(), max_block); + }); +} + fn assert_runtime_updated_digest(num: usize) { assert_eq!( System::digest() diff --git a/pallets/multisig/src/mock.rs b/pallets/multisig/src/mock.rs index fe2066ee..63bcb79b 100644 --- a/pallets/multisig/src/mock.rs +++ b/pallets/multisig/src/mock.rs @@ -374,6 +374,7 @@ impl pallet_utility::Config for Test { type PalletsOrigin = OriginCaller; type WeightInfo = (); type HighSecurity = (); + type AddressRevealer = (); } pub fn new_test_ext() -> sp_io::TestExternalities { diff --git a/pallets/reversible-transfers/src/tests/mock.rs b/pallets/reversible-transfers/src/tests/mock.rs index bd2f4e5b..1daab458 100644 --- a/pallets/reversible-transfers/src/tests/mock.rs +++ b/pallets/reversible-transfers/src/tests/mock.rs @@ -366,6 +366,7 @@ impl pallet_utility::Config for Test { type PalletsOrigin = OriginCaller; type WeightInfo = (); type HighSecurity = (); + type AddressRevealer = (); } // Build genesis storage according to the mock runtime. diff --git a/pallets/scheduler/src/lib.rs b/pallets/scheduler/src/lib.rs index 55dcef8e..e0971094 100644 --- a/pallets/scheduler/src/lib.rs +++ b/pallets/scheduler/src/lib.rs @@ -754,10 +754,26 @@ impl Pallet { origin: T::PalletsOrigin, call: BoundedCallOf, ) -> Result, DispatchError> { - let when = Self::resolve_time(when)?; + // `call` may already carry a preimage noted by `StorePreimage::bound` (called by the + // scheduler entrypoints before this). If any later fallible step fails, drop that + // preimage so a failed schedule cannot leave an unowned, system-requested preimage + // behind as state bloat. This mirrors the cleanup performed on cancellation. + let when = match Self::resolve_time(when) { + Ok(when) => when, + Err(e) => { + T::Preimages::drop(&call); + return Err(e) + }, + }; let lookup_hash = call.lookup_hash(); let task = Scheduled { maybe_id: None, priority, call, origin, _phantom: PhantomData }; - let res = Self::place_task(when, task).map_err(|x| x.0)?; + let res = match Self::place_task(when, task) { + Ok(res) => res, + Err((e, task)) => { + T::Preimages::drop(&task.call); + return Err(e) + }, + }; if let Some(hash) = lookup_hash { T::Preimages::request(&hash); } @@ -826,14 +842,29 @@ impl Pallet { origin: T::PalletsOrigin, call: BoundedCallOf, ) -> Result, DispatchError> { + // See `do_schedule`: drop any preimage noted by `StorePreimage::bound` on every fallible + // path so a failed schedule cannot leave an unowned, system-requested preimage behind. if Lookup::::contains_key(id) { + T::Preimages::drop(&call); return Err(Error::::FailedToSchedule.into()); } - let when = Self::resolve_time(when)?; + let when = match Self::resolve_time(when) { + Ok(when) => when, + Err(e) => { + T::Preimages::drop(&call); + return Err(e) + }, + }; let lookup_hash = call.lookup_hash(); let task = Scheduled { maybe_id: Some(id), priority, call, origin, _phantom: Default::default() }; - let res = Self::place_task(when, task).map_err(|x| x.0)?; + let res = match Self::place_task(when, task) { + Ok(res) => res, + Err((e, task)) => { + T::Preimages::drop(&task.call); + return Err(e) + }, + }; if let Some(hash) = lookup_hash { T::Preimages::request(&hash); } diff --git a/pallets/scheduler/src/tests.rs b/pallets/scheduler/src/tests.rs index 9768af7c..4ccc10fb 100644 --- a/pallets/scheduler/src/tests.rs +++ b/pallets/scheduler/src/tests.rs @@ -752,6 +752,34 @@ fn reschedule_named_works() { }); } +#[test] +fn failed_schedule_drops_noted_preimage() { + new_test_ext().execute_with(|| { + run_to_block(1); + + // A call large enough to be stored as a `Lookup` preimage rather than inline. + let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![0u8; 1024] }); + let bounded = Preimage::bound(call).unwrap(); + assert!(bounded.lookup_needed(), "the call must be noted as a preimage"); + // `bound` has already noted the preimage. + assert!(Preimage::have(&bounded)); + + // Scheduling at a past target fails in `resolve_time`, after the preimage was noted. + // Note: this deliberately mutates storage (it drops the preimage), so `assert_noop!` + // is not applicable here. + assert_err!( + Scheduler::do_schedule(DispatchTime::At(1), 127, root(), bounded.clone()), + Error::::TargetBlockNumberInPast, + ); + + // The failed schedule must not leave the noted preimage behind. + assert!( + !Preimage::have(&bounded), + "a failed schedule must drop the preimage noted by `bound`", + ); + }); +} + #[test] fn cancel_named_scheduling_works_with_normal_cancel() { new_test_ext().execute_with(|| { diff --git a/pallets/utility/Cargo.toml b/pallets/utility/Cargo.toml index 7be28a14..916aaac0 100644 --- a/pallets/utility/Cargo.toml +++ b/pallets/utility/Cargo.toml @@ -21,6 +21,7 @@ frame-benchmarking = { optional = true, workspace = true } frame-support.workspace = true frame-system.workspace = true qp-high-security = { path = "../../primitives/high-security", default-features = false } +qp-wormhole = { workspace = true, default-features = false } scale-info = { features = ["derive"], workspace = true } sp-core.workspace = true sp-io.workspace = true @@ -46,6 +47,7 @@ std = [ "frame-support/std", "frame-system/std", "qp-high-security/std", + "qp-wormhole/std", "scale-info/std", "sp-core/std", "sp-io/std", diff --git a/pallets/utility/src/lib.rs b/pallets/utility/src/lib.rs index 6152de57..207d1706 100644 --- a/pallets/utility/src/lib.rs +++ b/pallets/utility/src/lib.rs @@ -114,8 +114,22 @@ pub mod pallet { Self::AccountId, ::RuntimeCall, >; + + /// Handle into the wormhole soundness counter. An `as_derivative` pseudonym spends + /// through its controller and never signs an extrinsic itself, so it would otherwise stay + /// "ambiguous" (nonce 0) forever while receiving tracked credits. On first use of a + /// pseudonym it is revealed here (deducting its balance from the potential wormhole + /// pool) and remembered in [`KnownDerivatives`], so it can be excluded from the ambiguous + /// set afterwards. + type AddressRevealer: qp_wormhole::AddressRevealer; } + /// Derivative pseudonym accounts that have been used with [`Pallet::as_derivative`] at least + /// once and have therefore been revealed to the wormhole soundness counter. + #[pallet::storage] + pub type KnownDerivatives = + StorageMap<_, Blake2_128Concat, T::AccountId, (), OptionQuery>; + #[pallet::event] #[pallet::generate_deposit(pub(super) fn deposit_event)] pub enum Event { @@ -267,6 +281,9 @@ pub mod pallet { T::WeightInfo::as_derivative() // AccountData for inner call origin accountdata. .saturating_add(T::DbWeight::get().reads_writes(1, 1)) + // First-use derivative reveal: `KnownDerivatives` read + insert and the + // wormhole pool write. + .saturating_add(T::DbWeight::get().reads_writes(1, 2)) .saturating_add(dispatch_info.call_weight), dispatch_info.class, ) @@ -283,12 +300,25 @@ pub mod pallet { T::HighSecurity::is_call_allowed(&pseudonym, &call), Error::::CallNotAllowedForHighSecurity ); + // A pseudonym spends via its controller and never signs itself, so it never reveals + // itself to the wormhole soundness counter the way regular accounts do on their first + // signature. Reveal it on first use here (deducting its balance from the potential + // wormhole pool) and remember it so the runtime can exclude it from the ambiguous set + // afterwards; otherwise funds bounced between derivative pseudonyms would inflate the + // pool without bound. + if !KnownDerivatives::::contains_key(&pseudonym) { + KnownDerivatives::::insert(&pseudonym, ()); + >::reveal_address( + pseudonym.clone(), + ); + } origin.set_caller_from(frame_system::RawOrigin::Signed(pseudonym)); let info = call.get_dispatch_info(); let result = call.dispatch(origin); // Always take into account the base weight of this call. let mut weight = T::WeightInfo::as_derivative() - .saturating_add(T::DbWeight::get().reads_writes(1, 1)); + .saturating_add(T::DbWeight::get().reads_writes(1, 1)) + .saturating_add(T::DbWeight::get().reads_writes(1, 2)); // Add the real weight of the dispatch. weight = weight.saturating_add(extract_actual_weight(&result, &info)); result @@ -599,6 +629,12 @@ pub mod pallet { } impl Pallet { + /// Whether `account` is a known [`Pallet::as_derivative`] pseudonym, i.e. it has been + /// used (and revealed to the wormhole soundness counter) at least once. + pub fn is_derivative(account: &T::AccountId) -> bool { + KnownDerivatives::::contains_key(account) + } + /// Get the accumulated `weight` and the dispatch class for the given `calls`. fn weight_and_dispatch_class( calls: &[::RuntimeCall], diff --git a/pallets/utility/src/tests.rs b/pallets/utility/src/tests.rs index a1337947..4ed4a403 100644 --- a/pallets/utility/src/tests.rs +++ b/pallets/utility/src/tests.rs @@ -22,6 +22,7 @@ use super::*; use crate as utility; +use core::cell::RefCell; use frame_support::{ assert_err_ignore_postinfo, assert_noop, assert_ok, derive_impl, dispatch::{DispatchErrorWithPostInfo, Pays}, @@ -225,6 +226,20 @@ impl Config for Test { type PalletsOrigin = OriginCaller; type WeightInfo = (); type HighSecurity = qp_high_security::testing::TestHighSecurity; + type AddressRevealer = MockAddressRevealer; +} + +thread_local! { + /// Records addresses passed to `AddressRevealer::reveal_address`, so tests can assert a + /// derivative pseudonym is revealed to the wormhole soundness counter on first use. + pub static REVEALED_ADDRESSES: RefCell> = const { RefCell::new(Vec::new()) }; +} + +pub struct MockAddressRevealer; +impl qp_wormhole::AddressRevealer for MockAddressRevealer { + fn reveal_address(account: u64) { + REVEALED_ADDRESSES.with(|a| a.borrow_mut().push(account)); + } } /// High-security accounts in tests may only dispatch `System::remark`. @@ -299,6 +314,41 @@ fn as_derivative_works() { }); } +#[test] +fn as_derivative_reveals_pseudonym_on_first_use_only() { + new_test_ext().execute_with(|| { + REVEALED_ADDRESSES.with(|a| a.borrow_mut().clear()); + let sub_1_0 = derivative_account_id(1u64, 0); + assert!(!Utility::is_derivative(&sub_1_0)); + + // First use reveals the pseudonym to the wormhole soundness counter and marks it known. + assert_ok!(Utility::as_derivative( + RuntimeOrigin::signed(1), + 0, + Box::new(RuntimeCall::System(frame_system::Call::remark { remark: vec![] })), + )); + assert!(Utility::is_derivative(&sub_1_0)); + assert_eq!(REVEALED_ADDRESSES.with(|a| a.borrow().clone()), vec![sub_1_0]); + + // A second use must not reveal again. + assert_ok!(Utility::as_derivative( + RuntimeOrigin::signed(1), + 0, + Box::new(RuntimeCall::System(frame_system::Call::remark { remark: vec![] })), + )); + assert_eq!(REVEALED_ADDRESSES.with(|a| a.borrow().clone()), vec![sub_1_0]); + + // A different index is a different pseudonym and reveals separately. + let sub_1_1 = derivative_account_id(1u64, 1); + assert_ok!(Utility::as_derivative( + RuntimeOrigin::signed(1), + 1, + Box::new(RuntimeCall::System(frame_system::Call::remark { remark: vec![] })), + )); + assert_eq!(REVEALED_ADDRESSES.with(|a| a.borrow().clone()), vec![sub_1_0, sub_1_1]); + }); +} + #[test] fn as_derivative_handles_weight_refund() { new_test_ext().execute_with(|| { diff --git a/pallets/wormhole/src/lib.rs b/pallets/wormhole/src/lib.rs index 2ce139ad..c8560cce 100644 --- a/pallets/wormhole/src/lib.rs +++ b/pallets/wormhole/src/lib.rs @@ -806,4 +806,12 @@ pub mod pallet { Self::reveal_account(&account.into()); } } + + // Narrow reveal-only handle for pallets without balance/asset types (e.g. `pallet-utility` + // revealing `as_derivative` pseudonyms). + impl qp_wormhole::AddressRevealer<::AccountId> for Pallet { + fn reveal_address(account: ::AccountId) { + Self::reveal_account(&account); + } + } } diff --git a/primitives/wormhole/src/lib.rs b/primitives/wormhole/src/lib.rs index 66f2644f..22483ef8 100644 --- a/primitives/wormhole/src/lib.rs +++ b/primitives/wormhole/src/lib.rs @@ -94,6 +94,20 @@ pub trait TransferProofRecorder { fn reveal_address(account: AccountId); } +/// Narrow handle into the wormhole soundness counter for pallets that only need to reveal +/// addresses (see [`TransferProofRecorder::reveal_address`]) and have no notion of balances or +/// assets, e.g. `pallet-utility` revealing `as_derivative` pseudonyms on first use. +pub trait AddressRevealer { + /// Reveal `account` to the wormhole soundness counter, removing its current balance from the + /// pool of value that could be exited via the wormhole. + fn reveal_address(account: AccountId); +} + +/// No-op revealer for tests and runtimes without a wormhole. +impl AddressRevealer for () { + fn reveal_address(_account: AccountId) {} +} + /// Derive a wormhole address from a 32-byte inner_digest (already hashed). /// /// This hashes the inner_digest using Poseidon to get the wormhole account address. diff --git a/runtime/src/configs/mod.rs b/runtime/src/configs/mod.rs index 1c921ea9..762b04bd 100644 --- a/runtime/src/configs/mod.rs +++ b/runtime/src/configs/mod.rs @@ -464,6 +464,7 @@ impl pallet_utility::Config for Runtime { type PalletsOrigin = OriginCaller; type WeightInfo = pallet_utility::weights::SubstrateWeight; type HighSecurity = HighSecurityConfig; + type AddressRevealer = Wormhole; } parameter_types! { @@ -706,6 +707,13 @@ impl Contains for NonWormholeAccounts { return true; } + // `as_derivative` pseudonyms likewise spend through their controller and never sign. + // The utility pallet reveals a pseudonym on its first use (deducting its balance from + // the pool); afterwards it is excluded here so later receipts don't inflate the pool. + if pallet_utility::Pallet::::is_derivative(account) { + return true; + } + // The configured treasury account receives a block reward every block but is a keyless // governance account (a multisig that need not be registered in the multisig pallet), not // a wormhole deposit. In this fork the treasury address is configurable, so we read the diff --git a/runtime/src/transaction_extensions.rs b/runtime/src/transaction_extensions.rs index 26b8f9c2..8cc4955d 100644 --- a/runtime/src/transaction_extensions.rs +++ b/runtime/src/transaction_extensions.rs @@ -106,12 +106,32 @@ impl WormholeProofRecorderExtension Self(PhantomData) } + /// Weight charged per recorded transfer proof. + /// + /// Per recorded transfer, `record_transfer` touches (worst case): + /// reads: TransferCount (1) + the `is_ambiguous_account` lookup on the recipient, + /// which reads the recipient nonce (1) and, when nonce == 0, the + /// `NonWormholeAccounts` membership checks `Multisigs` (1), the utility + /// pallet's `KnownDerivatives` (1) and the configured `TreasuryAccount` (1) + /// => 5 reads. + /// writes: TransferProof / ZK-tree leaf (1) + TransferCount (1) + the conditional + /// `PotentialWormholeBalance` deposit add (1) => 3 writes. + fn per_transfer_weight() -> Weight { + T::DbWeight::get().reads_writes(5, 3) + } + fn count_transfers(call: &RuntimeCall) -> u64 { // NOTE: this must stay in sync with the events matched by `record_proofs_from_events_since` // — we only weight calls whose emitted events we actually record. In particular // `Balances::force_set_balance` is deliberately NOT counted here: it emits `BalanceSet` // (an absolute set, not a transfer/mint), which we cannot turn into a transfer proof and // therefore never record. See the soundness-counter caveat on `reduce_potential_balance`. + // + // Wrappers whose inner call is stored on-chain rather than in the submitted call + // (`Multisig::execute`, `ReversibleTransfers::recover_funds`, ...) cannot be counted + // statically. Proof-recording work they trigger is reconciled in `post_dispatch`, which + // registers any weight shortfall against the block via + // `register_extra_weight_unchecked`. match call { RuntimeCall::Balances(pallet_balances::Call::transfer_keep_alive { .. }) | RuntimeCall::Balances(pallet_balances::Call::transfer_allow_death { .. }) | @@ -130,6 +150,7 @@ impl WormholeProofRecorderExtension RuntimeCall::Utility(pallet_utility::Call::dispatch_as { call, .. }) | RuntimeCall::Utility(pallet_utility::Call::with_weight { call, .. }) | + RuntimeCall::Utility(pallet_utility::Call::as_derivative { call, .. }) | RuntimeCall::Recovery(pallet_recovery::Call::as_recovered { call, .. }) => Self::count_transfers(call), @@ -143,7 +164,10 @@ impl WormholeProofRecorderExtension /// /// `event_count_before` is the value from `frame_system::Pallet::event_count()` /// captured in `prepare()`. - fn record_proofs_from_events_since(event_count_before: u32) { + /// + /// Returns the number of transfer proofs recorded, so callers can reconcile the actual + /// proof-recording work against the statically charged weight. + fn record_proofs_from_events_since(event_count_before: u32) -> u64 { // IMPORTANT: We must collect all transfers FIRST before calling record_transfer_proof, // because record_transfer_proof deposits new events which would invalidate the // stream_iter iterator (causing "Corrupted state" errors). @@ -191,18 +215,23 @@ impl WormholeProofRecorderExtension .collect(); // Now record the proofs - this is safe because we're no longer iterating over Events + let recorded = transfers_to_record.len() as u64; for (asset_id, from, to, amount) in transfers_to_record { >::record_transfer_proof( asset_id, from, to, amount, ); } + recorded } } impl TransactionExtension for WormholeProofRecorderExtension { - type Pre = u32; + /// `(event_count_snapshot, statically_charged_transfer_count)`. The snapshot bounds the + /// event scan in `post_dispatch`; the charged count lets `post_dispatch` reconcile actual + /// proof-recording work against the weight reserved by `weight()`. + type Pre = (u32, u64); /// Carries the pre-fee balance captured in `validate` for a first-time signer that must be /// revealed (subtracted from the potential wormhole pool). `None` when there is nothing to /// reveal. The actual subtraction is committed in `prepare`. @@ -213,24 +242,14 @@ impl TransactionEx fn weight(&self, call: &RuntimeCall) -> Weight { let n = Self::count_transfers(call); - let transfer_weight = if n > 0 { - // Per recorded transfer, `record_transfer` touches (worst case): - // reads: TransferCount (1) + the `is_ambiguous_account` lookup on the recipient, - // which reads the recipient nonce (1) and, when nonce == 0, the - // `NonWormholeAccounts` membership checks `Multisigs` (1) and the configured - // `TreasuryAccount` (1) => 4 reads. - // writes: TransferProof / ZkTree leaf (1) + TransferCount (1) + the conditional - // `PotentialWormholeBalance` deposit add (1) => 3 writes. - T::DbWeight::get().reads_writes(4 * n, 3 * n) - } else { - Weight::zero() - }; + let transfer_weight = + if n > 0 { Self::per_transfer_weight().saturating_mul(n) } else { Weight::zero() }; // Soundness reveal bookkeeping for the signer. `validate` runs `is_ambiguous_account` on - // the signer — nonce (1) + `Multisigs` (1) + `TreasuryAccount` (1) — and, when ambiguous, - // reads the signer's balance (1): 4 reads worst case. `prepare` then writes - // `PotentialWormholeBalance` once. - let reveal_weight = T::DbWeight::get().reads_writes(4, 1); + // the signer — nonce (1) + `Multisigs` (1) + `KnownDerivatives` (1) + `TreasuryAccount` + // (1) — and, when ambiguous, reads the signer's balance (1): 5 reads worst case. + // `prepare` then writes `PotentialWormholeBalance` once. + let reveal_weight = T::DbWeight::get().reads_writes(5, 1); transfer_weight.saturating_add(reveal_weight) } @@ -239,7 +258,7 @@ impl TransactionEx self, val: Self::Val, _origin: &sp_runtime::traits::DispatchOriginOf, - _call: &RuntimeCall, + call: &RuntimeCall, _info: &sp_runtime::traits::DispatchInfoOf, _len: usize, ) -> Result { @@ -251,8 +270,9 @@ impl TransactionEx } // Snapshot current event count so we only process events added by this tx - // (and any events from previous txs in the same block). - Ok(frame_system::Pallet::::event_count()) + // (and any events from previous txs in the same block), and remember how many transfer + // proofs were statically charged for so post_dispatch can reconcile. + Ok((frame_system::Pallet::::event_count(), Self::count_transfers(call))) } fn validate( @@ -292,7 +312,7 @@ impl TransactionEx fn post_dispatch( pre: Self::Pre, - _info: &DispatchInfoOf, + info: &DispatchInfoOf, _post_info: &mut PostDispatchInfoOf, _len: usize, result: &DispatchResult, @@ -300,7 +320,21 @@ impl TransactionEx // Only record proofs if the transaction succeeded. // Use the event count snapshot from prepare() to avoid duplicate recording. if result.is_ok() { - Self::record_proofs_from_events_since(pre); + let (event_count_before, charged_transfers) = pre; + let recorded = Self::record_proofs_from_events_since(event_count_before); + + // Wrappers that dispatch inner calls stored on-chain (`Multisig::execute`, + // `ReversibleTransfers::recover_funds`, ...) can emit transfer events the static + // `count_transfers` matcher cannot see, so the proof-recording work above may exceed + // the weight reserved by `weight()`. Register the shortfall against the block so + // block-weight based DoS protection stays sound even when the static count drifts. + if recorded > charged_transfers { + frame_system::Pallet::::register_extra_weight_unchecked( + Self::per_transfer_weight() + .saturating_mul(recorded.saturating_sub(charged_transfers)), + info.class, + ); + } } Ok(()) @@ -683,10 +717,10 @@ mod tests { let ext = WormholeProofRecorderExtension::::new(); // Even non-transfer calls carry the constant soundness reveal-bookkeeping overhead - // (read signer nonce + multisig/treasury membership + balance, possibly write the - // pool), so the base weight is non-zero. + // (read signer nonce + multisig/derivative/treasury membership + balance, possibly + // write the pool), so the base weight is non-zero. let reveal_weight = - ::DbWeight::get().reads_writes(4, 1); + ::DbWeight::get().reads_writes(5, 1); let non_transfer = RuntimeCall::System(frame_system::Call::remark { remark: vec![1, 2, 3] }); let base_weight = as TransactionExtension< @@ -723,6 +757,64 @@ mod tests { }); } + #[test] + fn wormhole_proof_recorder_counts_as_derivative_wrapped_transfers() { + new_test_ext().execute_with(|| { + let ext = WormholeProofRecorderExtension::::new(); + let reveal_weight = + ::DbWeight::get().reads_writes(5, 1); + + // A transfer hidden behind `as_derivative` (possibly batched) must be charged the + // same per-transfer weight as a direct transfer. + let wrapped = RuntimeCall::Utility(pallet_utility::Call::as_derivative { + index: 0, + call: boxed(RuntimeCall::Utility(pallet_utility::Call::batch { + calls: vec![non_whitelisted_transfer(), non_whitelisted_transfer()], + })), + }); + let weight = as TransactionExtension< + RuntimeCall, + >>::weight(&ext, &wrapped); + assert_eq!( + weight, + WormholeProofRecorderExtension::::per_transfer_weight() + .saturating_mul(2) + .saturating_add(reveal_weight), + "as_derivative-wrapped transfers must be statically counted" + ); + }); + } + + #[test] + fn wormhole_proof_recorder_registers_extra_weight_for_uncounted_transfers() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + + // The presented call is opaque to the static matcher (like `Multisig::execute` or + // `ReversibleTransfers::recover_funds`, whose inner call lives on-chain), but the + // dispatch emits a real transfer event that post_dispatch must record. + let opaque_call = RuntimeCall::System(frame_system::Call::remark { remark: vec![1] }); + assert_eq!(WormholeProofRecorderExtension::::count_transfers(&opaque_call), 0); + + let weight_before = frame_system::Pallet::::block_weight().total(); + + run_lifecycle(&alice(), opaque_call, || { + assert_ok!(Balances::transfer_keep_alive( + RuntimeOrigin::signed(alice()), + MultiAddress::Id(bob()), + EXISTENTIAL_DEPOSIT * 50, + )); + }); + + let weight_after = frame_system::Pallet::::block_weight().total(); + assert_eq!( + weight_after.saturating_sub(weight_before), + WormholeProofRecorderExtension::::per_transfer_weight(), + "the uncounted recorded transfer must be registered as extra block weight" + ); + }); + } + #[test] fn wormhole_proof_recorder_extension_prepare_succeeds() { new_test_ext().execute_with(|| { @@ -1578,4 +1670,75 @@ mod tests { ); }); } + + // --- derivative pseudonyms must not inflate the pool without bound --- + // An `as_derivative` pseudonym never signs (its controller does), so before the reveal-on-use + // fix an attacker could bounce fixed capital between derivative addresses, adding the amount + // to `PotentialWormholeBalance` on every hop while never subtracting anything. + #[test] + fn counter_derivative_hops_do_not_inflate_pool() { + new_test_ext().execute_with(|| { + System::set_block_number(1); + let controller = intg_account(70); + fund(&controller, SENDER_BAL); + mark_revealed(&controller); + pallet_wormhole::PotentialWormholeBalance::::put(POOL_BASE); + + let d0 = pallet_utility::derivative_account_id(controller.clone(), 0u16); + let d1 = pallet_utility::derivative_account_id(controller.clone(), 1u16); + + // Fund derivative 0: it has never been used, so it looks ambiguous and is counted. + let funding = 100 * UNIT; + run_transfer(&controller, &d0, funding); + assert_eq!(pool(), POOL_BASE + funding, "credit to an unused pseudonym is counted"); + + // Hop funds from derivative 0 to derivative 1. First use of derivative 0 reveals it + // (subtracting its whole balance), while the credit to (still unused) derivative 1 + // is counted — so the hop nets to the hopped amount, not a double count. + let hop = 40 * UNIT; + let inner = transfer_call(&d1, hop); + let call = RuntimeCall::Utility(pallet_utility::Call::as_derivative { + index: 0, + call: alloc::boxed::Box::new(inner.clone()), + }); + run_lifecycle(&controller, call, || { + assert_ok!(Utility::as_derivative( + RuntimeOrigin::signed(controller.clone()), + 0, + alloc::boxed::Box::new(inner), + )); + }); + assert_eq!( + pool(), + POOL_BASE + hop, + "first use reveals the pseudonym: its counted funding is subtracted" + ); + assert!( + !pallet_wormhole::Pallet::::is_ambiguous_account(&d0), + "a used pseudonym must no longer be treated as ambiguous" + ); + + // Hop back from derivative 1 to derivative 0. Derivative 1 reveals on its first use + // (-hop), and the credit to the already-revealed derivative 0 is NOT counted, so the + // pool returns to its baseline instead of growing on every bounce. + let hop_back = 20 * UNIT; + let inner = transfer_call(&d0, hop_back); + let call = RuntimeCall::Utility(pallet_utility::Call::as_derivative { + index: 1, + call: alloc::boxed::Box::new(inner.clone()), + }); + run_lifecycle(&controller, call, || { + assert_ok!(Utility::as_derivative( + RuntimeOrigin::signed(controller.clone()), + 1, + alloc::boxed::Box::new(inner), + )); + }); + assert_eq!( + pool(), + POOL_BASE, + "bouncing fixed capital between pseudonyms must not inflate the pool" + ); + }); + } }