Pre-Charge max size when contracts access storage (paritytech#10691)

* Fix seal_get_storage

* Fix seal_take_storage

* Add more benchmarks

* cargo run --quiet --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_contracts --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/contracts/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* Fix seal_set_storage

* Fix seal_contains_storage and seal_clear_storage

* Fix benchmarks

* cargo run --quiet --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_contracts --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/contracts/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* Get rid of mem::size_of in benchmarks

* Fix up code loading

* Apply suggestions from code review

Co-authored-by: Hernando Castano <[email protected]>

* Fix test to call same function twice

* Replaced u32::MAX by SENTINEL const

* Fix seal_contains_storage benchmark

* cargo run --quiet --profile=production  --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_contracts --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/contracts/src/weights.rs --template=./.maintain/frame-weight-template.hbs

Co-authored-by: Parity Bot <[email protected]>
Co-authored-by: Hernando Castano <[email protected]>

frame/contracts/src/chain_extension.rs

@@ -329,7 +329,7 @@ where
 	///
 	/// If the contract supplied buffer is smaller than the passed `buffer` an `Err` is returned.
 	/// If `allow_skip` is set to true the contract is allowed to skip the copying of the buffer
-	/// by supplying the guard value of `u32::MAX` as `out_ptr`. The
+	/// by supplying the guard value of `pallet-contracts::SENTINEL` as `out_ptr`. The
 	/// `weight_per_byte` is only charged when the write actually happens and is not skipped or
 	/// failed due to a too small output buffer.
 	pub fn write(

frame/contracts/src/exec.rs

@@ -24,7 +24,7 @@ use crate::{
 use frame_support::{
 	dispatch::{DispatchError, DispatchResult, DispatchResultWithPostInfo, Dispatchable},
 	storage::{with_transaction, TransactionOutcome},
-	traits::{Contains, Currency, ExistenceRequirement, Get, OriginTrait, Randomness, Time},
+	traits::{Contains, Currency, ExistenceRequirement, OriginTrait, Randomness, Time},
 	weights::Weight,
 };
 use frame_system::RawOrigin;
@@ -140,11 +140,11 @@ pub trait Ext: sealing::Sealed {
 	/// was deleted.
 	fn get_storage(&mut self, key: &StorageKey) -> Option<Vec<u8>>;
 
-	/// Returns true iff some storage entry exists under the supplied `key`
+	/// Returns `Some(len)` (in bytes) if a storage item exists at `key`.
 	///
-	/// Returns `false` if the `key` wasn't previously set by `set_storage` or
+	/// Returns `None` if the `key` wasn't previously set by `set_storage` or
 	/// was deleted.
-	fn contains_storage(&mut self, key: &StorageKey) -> bool;
+	fn get_storage_size(&mut self, key: &StorageKey) -> Option<u32>;
 
 	/// Sets the storage entry by the given key to the specified value. If `value` is `None` then
 	/// the storage entry is deleted.
@@ -996,8 +996,8 @@ where
 		Storage::<T>::read(&self.top_frame_mut().contract_info().trie_id, key)
 	}
 
-	fn contains_storage(&mut self, key: &StorageKey) -> bool {
-		Storage::<T>::contains(&self.top_frame_mut().contract_info().trie_id, key)
+	fn get_storage_size(&mut self, key: &StorageKey) -> Option<u32> {
+		Storage::<T>::size(&self.top_frame_mut().contract_info().trie_id, key)
 	}
 
 	fn set_storage(
@@ -1056,7 +1056,7 @@ where
 	}
 
 	fn max_value_size(&self) -> u32 {
-		T::Schedule::get().limits.payload_len
+		self.schedule.limits.payload_len
 	}
 
 	fn get_weight_price(&self, weight: Weight) -> BalanceOf<Self::T> {
@@ -2432,16 +2432,16 @@ mod tests {
 	}
 
 	#[test]
-	fn contains_storage_works() {
+	fn get_storage_size_works() {
 		let code_hash = MockLoader::insert(Call, |ctx, _| {
 			assert_eq!(
 				ctx.ext.set_storage([1; 32], Some(vec![1, 2, 3]), false),
 				Ok(WriteOutcome::New)
 			);
 			assert_eq!(ctx.ext.set_storage([2; 32], Some(vec![]), false), Ok(WriteOutcome::New));
-			assert_eq!(ctx.ext.contains_storage(&[1; 32]), true);
-			assert_eq!(ctx.ext.contains_storage(&[1; 32]), true);
-			assert_eq!(ctx.ext.contains_storage(&[3; 32]), false);
+			assert_eq!(ctx.ext.get_storage_size(&[1; 32]), Some(3));
+			assert_eq!(ctx.ext.get_storage_size(&[2; 32]), Some(0));
+			assert_eq!(ctx.ext.get_storage_size(&[3; 32]), None);
 
 			exec_success()
 		});

frame/contracts/src/lib.rs

@@ -136,6 +136,14 @@ type BalanceOf<T> =
 /// The current storage version.
 const STORAGE_VERSION: StorageVersion = StorageVersion::new(6);
 
+/// Used as a sentinel value when reading and writing contract memory.
+///
+/// It is usually used to signal `None` to a contract when only a primitive is allowed
+/// and we don't want to go through encoding a full Rust type. Using `u32::Max` is a safe
+/// sentinel because contracts are never allowed to use such a large amount of resources
+/// that this value makes sense for a memory location or length.
+const SENTINEL: u32 = u32::MAX;
+
 /// Provides the contract address generation method.
 ///
 /// See [`DefaultAddressGenerator`] for the default implementation.
@@ -831,7 +839,7 @@ where
 		module: &mut PrefabWasmModule<T>,
 		schedule: &Schedule<T>,
 	) -> frame_support::dispatch::DispatchResult {
-		self::wasm::reinstrument(module, schedule)
+		self::wasm::reinstrument(module, schedule).map(|_| ())
 	}
 
 	/// Internal function that does the actual call.

frame/contracts/src/schedule.rs

@@ -316,15 +316,24 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight of calling `seal_set_storage`.
 	pub set_storage: Weight,
 
-	/// Weight per byte of an item stored with `seal_set_storage`.
-	pub set_storage_per_byte: Weight,
+	/// Weight per written byten of an item stored with `seal_set_storage`.
+	pub set_storage_per_new_byte: Weight,
+
+	/// Weight per overwritten byte of an item stored with `seal_set_storage`.
+	pub set_storage_per_old_byte: Weight,
 
 	/// Weight of calling `seal_clear_storage`.
 	pub clear_storage: Weight,
 
+	/// Weight of calling `seal_clear_storage` per byte of the stored item.
+	pub clear_storage_per_byte: Weight,
+
 	/// Weight of calling `seal_contains_storage`.
 	pub contains_storage: Weight,
 
+	/// Weight of calling `seal_contains_storage` per byte of the stored item.
+	pub contains_storage_per_byte: Weight,
+
 	/// Weight of calling `seal_get_storage`.
 	pub get_storage: Weight,
 
@@ -586,9 +595,12 @@ impl<T: Config> Default for HostFnWeights<T> {
 			),
 			debug_message: cost_batched!(seal_debug_message),
 			set_storage: cost_batched!(seal_set_storage),
-			set_storage_per_byte: cost_byte_batched!(seal_set_storage_per_kb),
+			set_storage_per_new_byte: cost_byte_batched!(seal_set_storage_per_new_kb),
+			set_storage_per_old_byte: cost_byte_batched!(seal_set_storage_per_old_kb),
 			clear_storage: cost_batched!(seal_clear_storage),
+			clear_storage_per_byte: cost_byte_batched!(seal_clear_storage_per_kb),
 			contains_storage: cost_batched!(seal_contains_storage),
+			contains_storage_per_byte: cost_byte_batched!(seal_contains_storage_per_kb),
 			get_storage: cost_batched!(seal_get_storage),
 			get_storage_per_byte: cost_byte_batched!(seal_get_storage_per_kb),
 			take_storage: cost_batched!(seal_take_storage),

frame/contracts/src/storage.rs

@@ -22,7 +22,7 @@ pub mod meter;
 use crate::{
 	exec::{AccountIdOf, StorageKey},
 	weights::WeightInfo,
-	BalanceOf, CodeHash, Config, ContractInfoOf, DeletionQueue, Error, TrieId,
+	BalanceOf, CodeHash, Config, ContractInfoOf, DeletionQueue, Error, TrieId, SENTINEL,
 };
 use codec::{Decode, Encode};
 use frame_support::{
@@ -87,6 +87,33 @@ pub enum WriteOutcome {
 	Taken(Vec<u8>),
 }
 
+impl WriteOutcome {
+	/// Extracts the size of the overwritten value or `0` if there
+	/// was no value in storage.
+	pub fn old_len(&self) -> u32 {
+		match self {
+			Self::New => 0,
+			Self::Overwritten(len) => *len,
+			Self::Taken(value) => value.len() as u32,
+		}
+	}
+
+	/// Extracts the size of the overwritten value or `SENTINEL` if there
+	/// was no value in storage.
+	///
+	/// # Note
+	///
+	/// We cannot use `0` as sentinel value because there could be a zero sized
+	/// storage entry which is different from a non existing one.
+	pub fn old_len_with_sentinel(&self) -> u32 {
+		match self {
+			Self::New => SENTINEL,
+			Self::Overwritten(len) => *len,
+			Self::Taken(value) => value.len() as u32,
+		}
+	}
+}
+
 pub struct Storage<T>(PhantomData<T>);
 
 impl<T> Storage<T>
@@ -102,9 +129,12 @@ where
 		child::get_raw(&child_trie_info(trie_id), &blake2_256(key))
 	}
 
-	/// Returns `true` iff the `key` exists in storage.
-	pub fn contains(trie_id: &TrieId, key: &StorageKey) -> bool {
-		child::exists(&child_trie_info(trie_id), &blake2_256(key))
+	/// Returns `Some(len)` (in bytes) if a storage item exists at `key`.
+	///
+	/// Returns `None` if the `key` wasn't previously set by `set_storage` or
+	/// was deleted.
+	pub fn size(trie_id: &TrieId, key: &StorageKey) -> Option<u32> {
+		child::len(&child_trie_info(trie_id), &blake2_256(key))
 	}
 
 	/// Update a storage entry into a contract's kv storage.

frame/contracts/src/tests.rs

@@ -294,7 +294,7 @@ pub const BOB: AccountId32 = AccountId32::new([2u8; 32]);
 pub const CHARLIE: AccountId32 = AccountId32::new([3u8; 32]);
 pub const DJANGO: AccountId32 = AccountId32::new([4u8; 32]);
 
-pub const GAS_LIMIT: Weight = 10_000_000_000;
+pub const GAS_LIMIT: Weight = 100_000_000_000;
 
 pub struct ExtBuilder {
 	existential_deposit: u64,

frame/contracts/src/wasm/code_cache.rs

@@ -38,7 +38,6 @@ use crate::{
 use frame_support::{
 	dispatch::{DispatchError, DispatchResult},
 	ensure,
-	storage::StorageMap,
 	traits::ReservableCurrency,
 };
 use sp_core::crypto::UncheckedFrom;
@@ -149,52 +148,38 @@ pub fn load<T: Config>(
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
 {
-	gas_meter.charge(CodeToken::Load(estimate_code_size::<T, CodeStorage<T>, _>(&code_hash)?))?;
+	let charged = gas_meter.charge(CodeToken::Load(schedule.limits.code_len))?;
 
 	let mut prefab_module =
 		<CodeStorage<T>>::get(code_hash).ok_or_else(|| Error::<T>::CodeNotFound)?;
+	gas_meter.adjust_gas(charged, CodeToken::Load(prefab_module.code.len() as u32));
 	prefab_module.code_hash = code_hash;
 
 	if prefab_module.instruction_weights_version < schedule.instruction_weights.version {
 		// The instruction weights have changed.
 		// We need to re-instrument the code with the new instruction weights.
-		gas_meter.charge(CodeToken::Reinstrument(estimate_code_size::<T, PristineCode<T>, _>(
-			&code_hash,
-		)?))?;
-		reinstrument(&mut prefab_module, schedule)?;
+		let charged = gas_meter.charge(CodeToken::Reinstrument(schedule.limits.code_len))?;
+		let code_size = reinstrument(&mut prefab_module, schedule)?;
+		gas_meter.adjust_gas(charged, CodeToken::Reinstrument(code_size));
 	}
 
 	Ok(prefab_module)
 }
 
 /// Instruments the passed prefab wasm module with the supplied schedule.
+///
+/// Returns the size in bytes of the uninstrumented code.
 pub fn reinstrument<T: Config>(
 	prefab_module: &mut PrefabWasmModule<T>,
 	schedule: &Schedule<T>,
-) -> Result<(), DispatchError> {
+) -> Result<u32, DispatchError> {
 	let original_code =
 		<PristineCode<T>>::get(&prefab_module.code_hash).ok_or_else(|| Error::<T>::CodeNotFound)?;
+	let original_code_len = original_code.len();
 	prefab_module.code = prepare::reinstrument_contract::<T>(original_code, schedule)?;
 	prefab_module.instruction_weights_version = schedule.instruction_weights.version;
 	<CodeStorage<T>>::insert(&prefab_module.code_hash, &*prefab_module);
-	Ok(())
-}
-
-/// Get the size of the code stored at `code_hash` without loading it.
-///
-/// The returned value is slightly too large when using it for the [`PrefabWasmModule`]
-/// because it has other fields in addition to the code itself. However, those are negligible
-/// when compared to the code size. Additionally, charging too much weight is completely safe.
-fn estimate_code_size<T, M, V>(code_hash: &CodeHash<T>) -> Result<u32, DispatchError>
-where
-	T: Config,
-	M: StorageMap<CodeHash<T>, V>,
-	V: codec::FullCodec,
-{
-	let key = M::hashed_key_for(code_hash);
-	let mut data = [0u8; 0];
-	let len = sp_io::storage::read(&key, &mut data, 0).ok_or_else(|| Error::<T>::CodeNotFound)?;
-	Ok(len)
+	Ok(original_code_len as u32)
 }
 
 /// Costs for operations that are related to code handling.

frame/contracts/src/wasm/mod.rs

@@ -385,8 +385,8 @@ mod tests {
 		fn get_storage(&mut self, key: &StorageKey) -> Option<Vec<u8>> {
 			self.storage.get(key).cloned()
 		}
-		fn contains_storage(&mut self, key: &StorageKey) -> bool {
-			self.storage.contains_key(key)
+		fn get_storage_size(&mut self, key: &StorageKey) -> Option<u32> {
+			self.storage.get(key).map(|val| val.len() as u32)
 		}
 		fn set_storage(
 			&mut self,
@@ -2023,7 +2023,7 @@ mod tests {
 		// value did not exist before -> sentinel returned
 		let input = ([1u8; 32], [42u8, 48]).encode();
 		let result = execute(CODE, input, &mut ext).unwrap();
-		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), u32::MAX);
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), crate::SENTINEL);
 		assert_eq!(ext.storage.get(&[1u8; 32]).unwrap(), &[42u8, 48]);
 
 		// value do exist -> length of old value returned
@@ -2083,7 +2083,7 @@ mod tests {
 
 		// value does not exist -> sentinel returned
 		let result = execute(CODE, [3u8; 32].encode(), &mut ext).unwrap();
-		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), u32::MAX);
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), crate::SENTINEL);
 		assert_eq!(ext.storage.get(&[3u8; 32]), None);
 
 		// value did exist -> length returned
@@ -2228,25 +2228,16 @@ mod tests {
 		ext.storage.insert([1u8; 32], vec![42u8]);
 		ext.storage.insert([2u8; 32], vec![]);
 
-		// value does not exist -> error returned
+		// value does not exist -> sentinel value returned
 		let result = execute(CODE, [3u8; 32].encode(), &mut ext).unwrap();
-		assert_eq!(
-			u32::from_le_bytes(result.data.0.try_into().unwrap()),
-			ReturnCode::KeyNotFound as u32
-		);
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), crate::SENTINEL);
 
 		// value did exist -> success
 		let result = execute(CODE, [1u8; 32].encode(), &mut ext).unwrap();
-		assert_eq!(
-			u32::from_le_bytes(result.data.0.try_into().unwrap()),
-			ReturnCode::Success as u32
-		);
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), 1,);
 
 		// value did exist -> success (zero sized type)
 		let result = execute(CODE, [2u8; 32].encode(), &mut ext).unwrap();
-		assert_eq!(
-			u32::from_le_bytes(result.data.0.try_into().unwrap()),
-			ReturnCode::Success as u32
-		);
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), 0,);
 	}
 }