[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(),