Cap storage items count and adapt claim_surcharge to that

bin/node/runtime/src/lib.rs

@@ -679,6 +679,18 @@ parameter_types! {
 	pub const MaxDepth: u32 = 32;
 	pub const StorageSizeOffset: u32 = 8;
 	pub const MaxValueSize: u32 = 16 * 1024;
+	pub MaxItemCount: u32 = {
+		// We allow to accumulate 75% of the storage that is possible to remove in a
+		// single extrinsic.
+		use pallet_contracts::WeightInfo;
+		type SubWeight = pallet_contracts::weights::SubstrateWeight<Runtime>;
+		let weight_per_item =
+			(SubWeight::claim_surcharge(1) - SubWeight::claim_surcharge(0))
+				.max(SubWeight::seal_terminate_per_item(1) - SubWeight::seal_terminate_per_item(0))
+				.max(1);
+		let upper_limit: u32 = (MaximumExtrinsicWeight::get() / weight_per_item).saturated_into();
+		Perbill::from_percent(75) * upper_limit
+	};
 }
 
 impl pallet_contracts::Config for Runtime {
@@ -695,6 +707,7 @@ impl pallet_contracts::Config for Runtime {
 	type SurchargeReward = SurchargeReward;
 	type MaxDepth = MaxDepth;
 	type MaxValueSize = MaxValueSize;
+	type MaxItemCount = MaxItemCount;
 	type WeightPrice = pallet_transaction_payment::Module<Self>;
 	type WeightInfo = pallet_contracts::weights::SubstrateWeight<Self>;
 }

frame/contracts/src/exec.rs

@@ -64,7 +64,7 @@ pub trait Ext {
 
 	/// Sets the storage entry by the given key to the specified value. If `value` is `None` then
 	/// the storage entry is deleted.
-	fn set_storage(&mut self, key: StorageKey, value: Option<Vec<u8>>);
+	fn set_storage(&mut self, key: StorageKey, value: Option<Vec<u8>>) -> Result<(), DispatchError>;
 
 	/// Instantiate a contract from the given code.
 	///
@@ -510,23 +510,18 @@ where
 		Storage::<T>::read(trie_id, key)
 	}
 
-	fn set_storage(&mut self, key: StorageKey, value: Option<Vec<u8>>) {
+	fn set_storage(
+		&mut self,
+		key: StorageKey,
+		value: Option<Vec<u8>>
+	) -> Result<(), DispatchError> {
 		let trie_id = self.ctx.self_trie_id.as_ref().expect(
 			"`ctx.self_trie_id` points to an alive contract within the `CallContext`;\
 				it cannot be `None`;\
 				expect can't fail;\
 				qed",
 		);
-		if let Err(storage::ContractAbsentError) =
-			Storage::<T>::write(&self.ctx.self_account, trie_id, &key, value)
-		{
-			panic!(
-				"the contract must be in the alive state within the `CallContext`;\
-				the contract cannot be absent in storage;
-				write cannot return `None`;
-				qed"
-			);
-		}
+		Storage::<T>::write(&self.ctx.self_account, trie_id, &key, value)
 	}
 
 	fn instantiate(

frame/contracts/src/lib.rs

@@ -124,7 +124,7 @@ use frame_system::{ensure_signed, ensure_root};
 use pallet_contracts_primitives::{
 	RentProjectionResult, GetStorageResult, ContractAccessError, ContractExecResult, ExecResult,
 };
-use frame_support::weights::Weight;
+use frame_support::weights::{Weight, WithPostDispatchInfo};
 
 pub type CodeHash<T> = <T as frame_system::Config>::Hash;
 pub type TrieId = Vec<u8>;
@@ -312,6 +312,14 @@ pub trait Config: frame_system::Config {
 	/// The maximum size of a storage value and event payload in bytes.
 	type MaxValueSize: Get<u32>;
 
+	/// The maximum number of storage items a contract is allowed to have.
+	///
+	/// The weight of evicting or removing a contract is linearly dependend on the number
+	/// of storage items. Therefore it is necessary to prevent the creation of more items
+	/// than it is possibel to remove in a single block. Otherwise it becomes impossible
+	/// to remove contracts that are too large.
+	type MaxItemCount: Get<u32>;
+
 	/// Used to answer contracts's queries regarding the current weight price. This is **not**
 	/// used to calculate the actual fee and is only for informational purposes.
 	type WeightPrice: Convert<Weight, BalanceOf<Self>>;
@@ -378,6 +386,8 @@ decl_error! {
 		/// on the call stack. Those actions are contract self destruction and restoration
 		/// of a tombstone.
 		ReentranceDenied,
+		/// The count in [`T::MaxItemCount`] was exceeded.
+		TooManyStorageItems,
 	}
 }
 
@@ -431,6 +441,14 @@ decl_module! {
 		/// The maximum size of a storage value in bytes. A reasonable default is 16 KiB.
 		const MaxValueSize: u32 = T::MaxValueSize::get();
 
+		/// The maximum number of storage items a contract is allowed to have.
+		///
+		/// The weight of evicting or removing a contract is linearly dependend on the number
+		/// of storage items. Therefore it is necessary to prevent the creation of more items
+		/// than it is possibel to remove in a single block. Otherwise it becomes impossible
+		/// to remove contracts that are too large.
+		const MaxItemCount: u32 = T::MaxItemCount::get();
+
 		fn deposit_event() = default;
 
 		/// Updates the schedule for metering contracts.
@@ -533,17 +551,28 @@ decl_module! {
 		///
 		/// If contract is not evicted as a result of this call, no actions are taken and
 		/// the sender is not eligible for the reward.
-		#[weight = T::WeightInfo::claim_surcharge()]
-		fn claim_surcharge(origin, dest: T::AccountId, aux_sender: Option<T::AccountId>) {
-			let origin = origin.into();
-			let (signed, rewarded) = match (origin, aux_sender) {
+		///
+		/// A priori the weight for removing a maximum sized contract is assumed. Post dispatch
+		/// weight correction is used to account for the following:
+		///  1. In case of error only the basic overhead of calling this extrinsic is charged.
+		///	 2. The weight is corrected to the actual amount of storage a contract has.
+		///	 3. In case of succesful eviction no weight fees are payed for the removal.
+		#[weight = T::WeightInfo::claim_surcharge(T::MaxItemCount::get())]
+		fn claim_surcharge(
+			origin,
+			dest: T::AccountId,
+			aux_sender: Option<T::AccountId>
+		) -> DispatchResultWithPostInfo {
+			let (signed, rewarded) = match (origin.into(), aux_sender) {
 				(Ok(frame_system::RawOrigin::Signed(account)), None) => {
 					(true, account)
 				},
 				(Ok(frame_system::RawOrigin::None), Some(aux_sender)) => {
 					(false, aux_sender)
 				},
-				_ => Err(Error::<T>::InvalidSurchargeClaim)?,
+				_ => Err(
+					Error::<T>::InvalidSurchargeClaim.with_weight(T::WeightInfo::claim_surcharge(0))
+				)?,
 			};
 
 			// Add some advantage for block producers (who send unsigned extrinsics) by
@@ -556,8 +585,18 @@ decl_module! {
 			};
 
 			// If poking the contract has lead to eviction of the contract, give out the rewards.
-			if Rent::<T>::snitch_contract_should_be_evicted(&dest, handicap) {
-				T::Currency::deposit_into_existing(&rewarded, T::SurchargeReward::get())?;
+			if let Some(count) = Rent::<T>::snitch_contract_should_be_evicted(&dest, handicap) {
+				use frame_support::weights::Pays;
+				let actual_weight = T::WeightInfo::claim_surcharge(count);
+				T::Currency::deposit_into_existing(&rewarded, T::SurchargeReward::get())
+				.map_err(|e| {
+					let mut err = e.with_weight(actual_weight);
+					err.post_info.pays_fee = Pays::No;
+					err
+				})?;
+				Ok((Some(actual_weight.into()), Pays::No).into())
+			} else {
+				Ok(Some(T::WeightInfo::claim_surcharge(0)).into())
 			}
 		}
 	}

frame/contracts/src/rent.rs

@@ -305,8 +305,9 @@ where
 
 	/// Process a report that a contract under the given address should be evicted.
 	///
-	/// Enact the eviction right away if the contract should be evicted and return true.
-	/// Otherwise, **do nothing** and return false.
+	/// Enact the eviction right away if the contract should be evicted and returns the number
+	/// of storage items that where removed.
+	/// Otherwise, **do nothing** and return None.
 	///
 	/// The `handicap` parameter gives a way to check the rent to a moment in the past instead
 	/// of current block. E.g. if the contract is going to be evicted at the current block,
@@ -318,10 +319,10 @@ where
 	pub fn snitch_contract_should_be_evicted(
 		account: &T::AccountId,
 		handicap: T::BlockNumber,
-	) -> bool {
+	) -> Option<u32> {
 		let contract_info = <ContractInfoOf<T>>::get(account);
 		let alive_contract_info = match contract_info {
-			None | Some(ContractInfo::Tombstone(_)) => return false,
+			None | Some(ContractInfo::Tombstone(_)) => return None,
 			Some(ContractInfo::Alive(contract)) => contract,
 		};
 		let current_block_number = <frame_system::Module<T>>::block_number();
@@ -335,10 +336,11 @@ where
 		// Enact the verdict only if the contract gets removed.
 		match verdict {
 			Verdict::Kill | Verdict::Evict { .. } => {
+				let result = alive_contract_info.total_pair_count;
 				Self::enact_verdict(account, alive_contract_info, current_block_number, verdict);
-				true
+				Some(result)
 			}
-			_ => false,
+			_ => None,
 		}
 	}
 

frame/contracts/src/storage.rs

@@ -19,20 +19,26 @@
 use crate::{
 	exec::{AccountIdOf, StorageKey},
 	AliveContractInfo, BalanceOf, CodeHash, ContractInfo, ContractInfoOf, Config, TrieId,
-	AccountCounter,
+	AccountCounter, Error,
 };
 use sp_std::prelude::*;
 use sp_std::marker::PhantomData;
 use sp_io::hashing::blake2_256;
 use sp_runtime::traits::Bounded;
 use sp_core::crypto::UncheckedFrom;
-use frame_support::{storage::child, StorageMap};
+use frame_support::{storage::child, StorageMap, dispatch::DispatchError, traits::Get};
 
 /// An error that means that the account requested either doesn't exist or represents a tombstone
 /// account.
 #[cfg_attr(test, derive(PartialEq, Eq, Debug))]
 pub struct ContractAbsentError;
 
+impl From<ContractAbsentError> for DispatchError {
+	fn from(_: ContractAbsentError) -> Self {
+		DispatchError::Other("Called storage function on non-alive contract.")
+	}
+}
+
 pub struct Storage<T>(PhantomData<T>);
 
 impl<T> Storage<T>
@@ -62,10 +68,10 @@ where
 		trie_id: &TrieId,
 		key: &StorageKey,
 		opt_new_value: Option<Vec<u8>>,
-	) -> Result<(), ContractAbsentError> {
+	) -> Result<(), DispatchError> {
 		let mut new_info = match <ContractInfoOf<T>>::get(account) {
 			Some(ContractInfo::Alive(alive)) => alive,
-			None | Some(ContractInfo::Tombstone(_)) => return Err(ContractAbsentError),
+			None | Some(ContractInfo::Tombstone(_)) => Err(ContractAbsentError)?,
 		};
 
 		let hashed_key = blake2_256(key);
@@ -91,6 +97,9 @@ where
 				}
 			},
 			(None, Some(new_value)) => {
+				if new_info.total_pair_count >= T::MaxItemCount::get() {
+					return Err(Error::<T>::TooManyStorageItems.into());
+				}
 				new_info.total_pair_count += 1;
 				if new_value.is_empty() {
 					new_info.empty_pair_count += 1;

frame/contracts/src/tests.rs

@@ -164,6 +164,7 @@ parameter_types! {
 	pub const SurchargeReward: u64 = 150;
 	pub const MaxDepth: u32 = 100;
 	pub const MaxValueSize: u32 = 16_384;
+	pub const MaxItemCount: u32 = 4096;
 }
 
 parameter_types! {
@@ -190,6 +191,7 @@ impl Config for Test {
 	type SurchargeReward = SurchargeReward;
 	type MaxDepth = MaxDepth;
 	type MaxValueSize = MaxValueSize;
+	type MaxItemCount = MaxItemCount;
 	type WeightPrice = Self;
 	type WeightInfo = ();
 }

frame/contracts/src/wasm/mod.rs

@@ -228,8 +228,13 @@ mod tests {
 		fn get_storage(&self, key: &StorageKey) -> Option<Vec<u8>> {
 			self.storage.get(key).cloned()
 		}
-		fn set_storage(&mut self, key: StorageKey, value: Option<Vec<u8>>) {
+		fn set_storage(
+			&mut self,
+			key: StorageKey,
+			value: Option<Vec<u8>>
+		) -> Result<(), DispatchError> {
 			*self.storage.entry(key).or_insert(Vec::new()) = value.unwrap_or(Vec::new());
+			Ok(())
 		}
 		fn instantiate(
 			&mut self,
@@ -362,7 +367,11 @@ mod tests {
 		fn get_storage(&self, key: &[u8; 32]) -> Option<Vec<u8>> {
 			(**self).get_storage(key)
 		}
-		fn set_storage(&mut self, key: [u8; 32], value: Option<Vec<u8>>) {
+		fn set_storage(
+			&mut self,
+			key: [u8; 32],
+			value: Option<Vec<u8>>
+		) -> Result<(), DispatchError> {
 			(**self).set_storage(key, value)
 		}
 		fn instantiate(

frame/contracts/src/wasm/runtime.rs

@@ -643,8 +643,7 @@ define_env!(Env, <E: Ext>,
 		let mut key: StorageKey = [0; 32];
 		ctx.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
 		let value = Some(ctx.read_sandbox_memory(value_ptr, value_len)?);
-		ctx.ext.set_storage(key, value);
-		Ok(())
+		ctx.ext.set_storage(key, value).map_err(|e| ctx.store_err(e))
 	},
 
 	// Clear the value at the given key in the contract storage.
@@ -656,8 +655,7 @@ define_env!(Env, <E: Ext>,
 		ctx.charge_gas(RuntimeToken::ClearStorage)?;
 		let mut key: StorageKey = [0; 32];
 		ctx.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
-		ctx.ext.set_storage(key, None);
-		Ok(())
+		ctx.ext.set_storage(key, None).map_err(|e| ctx.store_err(e))
 	},
 
 	// Retrieve the value under the given key from storage.