[pallet_contracts] Modify the storage host function benchmarks to run on an unbalanced storage trie.

prdoc/pr_5036.prdoc

@@ -0,0 +1,16 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: "[pallet_contracts] Modify the storage host function benchmarks to be run on an unbalanced storage trie"
+
+doc:
+ - audience: Runtime User
+ description: |
+ This PR modifies the storage host function benchmarks. Previously, they were run
+ on an empty storage trie. Now, they are run on an unbalanced storage trie
+ to reflect the worst-case scenario. This approach increases the storage host
+ function weights and decreases the probability of DoS attacks.
+
+crates:
+ - name: pallet-contracts
+ bump: patch

substrate/frame/contracts/src/benchmarking/mod.rs

@@ -41,6 +41,7 @@ use frame_benchmarking::v2::*;
 use frame_support::{
  self, assert_ok,
  pallet_prelude::StorageVersion,
+ storage::child,
  traits::{fungible::InspectHold, Currency},
  weights::{Weight, WeightMeter},
 };
@@ -63,6 +64,9 @@ const API_BENCHMARK_RUNS: u32 = 1600;
 /// benchmarks are faster.
 const INSTR_BENCHMARK_RUNS: u32 = 5000;
 
+/// Number of layers in a Radix16 unbalanced trie.
+const UNBALANCED_TRIE_LAYERS: u32 = 20;
+
 /// An instantiated and deployed contract.
 #[derive(Clone)]
 struct Contract<T: Config> {
@@ -152,6 +156,36 @@ where
  Ok(())
  }
 
+ /// Create a new contract with the specified unbalanced storage trie.
+ fn with_unbalanced_storage_trie(code: WasmModule<T>, key: &[u8]) -> Result<Self, &'static str> {
+ if (key.len() as u32) < (UNBALANCED_TRIE_LAYERS + 1) / 2 {
+ return Err("Key size too small to create the specified trie");
+ }
+
+ let value = vec![16u8; T::Schedule::get().limits.payload_len as usize];
+ let contract = Contract::<T>::new(code, vec![])?;
+ let info = contract.info()?;
+ let child_trie_info = info.child_trie_info();
+ child::put_raw(&child_trie_info, &key, &value);
+ for l in 0..UNBALANCED_TRIE_LAYERS {
+ let pos = l as usize / 2;
+ let mut key_new = key.to_vec();
+ for i in 0u8..16 {
+ key_new[pos] = if l % 2 == 0 {
+ (key_new[pos] & 0xF0) | i
+ } else {
+ (key_new[pos] & 0x0F) | (i << 4)
+ };
+
+ if key == &key_new {
+ continue
+ }
+ child::put_raw(&child_trie_info, &key_new, &value);
+ }
+ }
+ Ok(contract)
+ }
+
  /// Get the `ContractInfo` of the `addr` or an error if it no longer exists.
  fn address_info(addr: &T::AccountId) -> Result<ContractInfo<T>, &'static str> {
  ContractInfoOf::<T>::get(addr).ok_or("Expected contract to exist at this point.")
@@ -1014,6 +1048,102 @@ mod benchmarks {
  assert_eq!(setup.debug_message().unwrap().len() as u32, i);
  }
 
+ #[benchmark(skip_meta, pov_mode = Measured)]
+ fn get_storage_empty() -> Result<(), BenchmarkError> {
+ let max_key_len = T::MaxStorageKeyLen::get();
+ let key = vec![0u8; max_key_len as usize];
+ let max_value_len = T::Schedule::get().limits.payload_len as usize;
+ let value = vec![1u8; max_value_len];
+
+ let instance = Contract::<T>::new(WasmModule::dummy(), vec![])?;
+ let info = instance.info()?;
+ let child_trie_info = info.child_trie_info();
+ info.bench_write_raw(&key, Some(value.clone()), false)
+ .map_err(|_| "Failed to write to storage during setup.")?;
+
+ let result;
+ #[block]
+ {
+ result = child::get_raw(&child_trie_info, &key);
+ }
+
+ assert_eq!(result, Some(value));
+ Ok(())
+ }
+
+ #[benchmark(skip_meta, pov_mode = Measured)]
+ fn get_storage_full() -> Result<(), BenchmarkError> {
+ let max_key_len = T::MaxStorageKeyLen::get();
+ let key = vec![0u8; max_key_len as usize];
+ let max_value_len = T::Schedule::get().limits.payload_len;
+ let value = vec![1u8; max_value_len as usize];
+
+ let instance = Contract::<T>::with_unbalanced_storage_trie(WasmModule::dummy(), &key)?;
+ let info = instance.info()?;
+ let child_trie_info = info.child_trie_info();
+ info.bench_write_raw(&key, Some(value.clone()), false)
+ .map_err(|_| "Failed to write to storage during setup.")?;
+
+ let result;
+ #[block]
+ {
+ result = child::get_raw(&child_trie_info, &key);
+ }
+
+ assert_eq!(result, Some(value));
+ Ok(())
+ }
+
+ #[benchmark(skip_meta, pov_mode = Measured)]
+ fn set_storage_empty() -> Result<(), BenchmarkError> {
+ let max_key_len = T::MaxStorageKeyLen::get();
+ let key = vec![0u8; max_key_len as usize];
+ let max_value_len = T::Schedule::get().limits.payload_len as usize;
+ let value = vec![1u8; max_value_len];
+
+ let instance = Contract::<T>::new(WasmModule::dummy(), vec![])?;
+ let info = instance.info()?;
+ let child_trie_info = info.child_trie_info();
+ info.bench_write_raw(&key, Some(vec![42u8; max_value_len]), false)
+ .map_err(|_| "Failed to write to storage during setup.")?;
+
+ let val = Some(value.clone());
+ let result;
+ #[block]
+ {
+ result = info.bench_write_raw(&key, val, true);
+ }
+
+ assert_ok!(result);
+ assert_eq!(child::get_raw(&child_trie_info, &key).unwrap(), value);
+ Ok(())
+ }
+
+ #[benchmark(skip_meta, pov_mode = Measured)]
+ fn set_storage_full() -> Result<(), BenchmarkError> {
+ let max_key_len = T::MaxStorageKeyLen::get();
+ let key = vec![0u8; max_key_len as usize];
+ let max_value_len = T::Schedule::get().limits.payload_len;
+ let value = vec![1u8; max_value_len as usize];
+
+ let instance = Contract::<T>::with_unbalanced_storage_trie(WasmModule::dummy(), &key)?;
+ let info = instance.info()?;
+ let child_trie_info = info.child_trie_info();
+ info.bench_write_raw(&key, Some(vec![42u8; max_value_len as usize]), false)
+ .map_err(|_| "Failed to write to storage during setup.")?;
+
+ let val = Some(value.clone());
+ let result;
+ #[block]
+ {
+ result = info.bench_write_raw(&key, val, true);
+ }
+
+ assert_ok!(result);
+ assert_eq!(child::get_raw(&child_trie_info, &key).unwrap(), value);
+ Ok(())
+ }
+
  // n: new byte size
  // o: old byte size
  #[benchmark(skip_meta, pov_mode = Measured)]

substrate/frame/contracts/src/storage.rs

@@ -164,13 +164,35 @@ impl<T: Config> ContractInfo<T> {
  storage_meter: Option<&mut meter::NestedMeter<T>>,
  take: bool,
  ) -> Result<WriteOutcome, DispatchError> {
- let child_trie_info = &self.child_trie_info();
  let hashed_key = key.hash();
+ self.write_raw(&hashed_key, new_value, storage_meter, take)
+ }
+
+ /// Update a storage entry into a contract's kv storage.
+ /// Function used in benchmarks, which can simulate prefix collision in keys.
+ #[cfg(feature = "runtime-benchmarks")]
+ pub fn bench_write_raw(
+ &self,
+ key: &[u8],
+ new_value: Option<Vec<u8>>,
+ take: bool,
+ ) -> Result<WriteOutcome, DispatchError> {
+ self.write_raw(key, new_value, None, take)
+ }
+
+ fn write_raw(
+ &self,
+ key: &[u8],
+ new_value: Option<Vec<u8>>,
+ storage_meter: Option<&mut meter::NestedMeter<T>>,
+ take: bool,
+ ) -> Result<WriteOutcome, DispatchError> {
+ let child_trie_info = &self.child_trie_info();
  let (old_len, old_value) = if take {
- let val = child::get_raw(child_trie_info, &hashed_key);
+ let val = child::get_raw(child_trie_info, key);
  (val.as_ref().map(|v| v.len() as u32), val)
  } else {
- (child::len(child_trie_info, &hashed_key), None)
+ (child::len(child_trie_info, key), None)
  };
 
  if let Some(storage_meter) = storage_meter {
@@ -196,8 +218,8 @@ impl<T: Config> ContractInfo<T> {
  }
 
  match &new_value {
- Some(new_value) => child::put_raw(child_trie_info, &hashed_key, new_value),
- None => child::kill(child_trie_info, &hashed_key),
+ Some(new_value) => child::put_raw(child_trie_info, key, new_value),
+ None => child::kill(child_trie_info, key),
  }
 
  Ok(match (old_len, old_value) {

substrate/frame/contracts/src/wasm/runtime.rs

@@ -255,31 +255,34 @@ pub enum RuntimeCosts {
  UnlockDelegateDependency,
 }
 
-// For the function that modifies the storage, the benchmarks are done with one item in the
-// transient_storage (BTreeMap). To consider the worst-case scenario, the weight of the overhead of
-// writing to a full BTreeMap should be included. On top of that, the rollback weight is added,
-// which is the worst scenario.
-macro_rules! cost_write {
- // cost_write!(name, a, b, c) -> T::WeightInfo::name(a, b, c).saturating_add(T::WeightInfo::rollback_transient_storage())
- // .saturating_add(T::WeightInfo::set_transient_storage_full().saturating_sub(T::WeightInfo::set_transient_storage_empty())
- ($name:ident $(, $arg:expr )*) => {
- (T::WeightInfo::$name($( $arg ),*).saturating_add(T::WeightInfo::rollback_transient_storage()).saturating_add(cost_write!(@cost_storage)))
- };
+/// For functions that modify storage, benchmarks are performed with one item in the
+/// storage. To account for the worst-case scenario, the weight of the overhead of
+/// writing to or reading from full storage is included. For transient storage writes,
+/// the rollback weight is added to reflect the worst-case scenario for this operation.
+macro_rules! cost_storage {
+ (write_transient, $name:ident $(, $arg:expr )*) => {
+ T::WeightInfo::$name($( $arg ),*)
+ .saturating_add(T::WeightInfo::rollback_transient_storage())
+ .saturating_add(T::WeightInfo::set_transient_storage_full()
+ .saturating_sub(T::WeightInfo::set_transient_storage_empty()))
+ };
 
- (@cost_storage) => {
- T::WeightInfo::set_transient_storage_full().saturating_sub(T::WeightInfo::set_transient_storage_empty())
+ (read_transient, $name:ident $(, $arg:expr )*) => {
+ T::WeightInfo::$name($( $arg ),*)
+ .saturating_add(T::WeightInfo::get_transient_storage_full()
+ .saturating_sub(T::WeightInfo::get_transient_storage_empty()))
  };
-}
 
-macro_rules! cost_read {
- // cost_read!(name, a, b, c) -> T::WeightInfo::name(a, b, c).saturating_add(T::WeightInfo::get_transient_storage_full()
- // .saturating_sub(T::WeightInfo::get_transient_storage_empty())
- ($name:ident $(, $arg:expr )*) => {
- (T::WeightInfo::$name($( $arg ),*).saturating_add(cost_read!(@cost_storage)))
- };
+ (write, $name:ident $(, $arg:expr )*) => {
+ T::WeightInfo::$name($( $arg ),*)
+ .saturating_add(T::WeightInfo::set_storage_full()
+ .saturating_sub(T::WeightInfo::set_storage_empty()))
+ };
 
- (@cost_storage) => {
- T::WeightInfo::get_transient_storage_full().saturating_sub(T::WeightInfo::get_transient_storage_empty())
+ (read, $name:ident $(, $arg:expr )*) => {
+ T::WeightInfo::$name($( $arg ),*)
+ .saturating_add(T::WeightInfo::get_storage_full()
+ .saturating_sub(T::WeightInfo::get_storage_empty()))
  };
 }
 
@@ -329,17 +332,21 @@ impl<T: Config> Token<T> for RuntimeCosts {
  DepositEvent { num_topic, len } => T::WeightInfo::seal_deposit_event(num_topic, len),
  DebugMessage(len) => T::WeightInfo::seal_debug_message(len),
  SetStorage { new_bytes, old_bytes } =>
- T::WeightInfo::seal_set_storage(new_bytes, old_bytes),
- ClearStorage(len) => T::WeightInfo::seal_clear_storage(len),
- ContainsStorage(len) => T::WeightInfo::seal_contains_storage(len),
- GetStorage(len) => T::WeightInfo::seal_get_storage(len),
- TakeStorage(len) => T::WeightInfo::seal_take_storage(len),
+ cost_storage!(write, seal_set_storage, new_bytes, old_bytes),
+ ClearStorage(len) => cost_storage!(write, seal_clear_storage, len),
+ ContainsStorage(len) => cost_storage!(read, seal_contains_storage, len),
+ GetStorage(len) => cost_storage!(read, seal_get_storage, len),
+ TakeStorage(len) => cost_storage!(write, seal_take_storage, len),
  SetTransientStorage { new_bytes, old_bytes } =>
- cost_write!(seal_set_transient_storage, new_bytes, old_bytes),
- ClearTransientStorage(len) => cost_write!(seal_clear_transient_storage, len),
- ContainsTransientStorage(len) => cost_read!(seal_contains_transient_storage, len),
- GetTransientStorage(len) => cost_read!(seal_get_transient_storage, len),
- TakeTransientStorage(len) => cost_write!(seal_take_transient_storage, len),
+ cost_storage!(write_transient, seal_set_transient_storage, new_bytes, old_bytes),
+ ClearTransientStorage(len) =>
+ cost_storage!(write_transient, seal_clear_transient_storage, len),
+ ContainsTransientStorage(len) =>
+ cost_storage!(read_transient, seal_contains_transient_storage, len),
+ GetTransientStorage(len) =>
+ cost_storage!(read_transient, seal_get_transient_storage, len),
+ TakeTransientStorage(len) =>
+ cost_storage!(write_transient, seal_take_transient_storage, len),
  Transfer => T::WeightInfo::seal_transfer(),
  CallBase => T::WeightInfo::seal_call(0, 0),
  DelegateCallBase => T::WeightInfo::seal_delegate_call(),