[pallet_contracts] Add support for transient storage in contracts host functions

substrate/bin/node/runtime/src/lib.rs

@@ -1378,6 +1378,7 @@ impl pallet_contracts::Config for Runtime {
 	type UploadOrigin = EnsureSigned<Self::AccountId>;
 	type InstantiateOrigin = EnsureSigned<Self::AccountId>;
 	type MaxDebugBufferLen = ConstU32<{ 2 * 1024 * 1024 }>;
+	type MaxTransientStorageLen = ConstU32<{ 1 * 1024 * 1024 }>;
 	type RuntimeHoldReason = RuntimeHoldReason;
 	#[cfg(not(feature = "runtime-benchmarks"))]
 	type Migrations = ();

substrate/frame/contracts/src/exec.rs

@@ -20,6 +20,7 @@ use crate::{
 	gas::GasMeter,
 	primitives::{ExecReturnValue, StorageDeposit},
 	storage::{self, meter::Diff, WriteOutcome},
+	transient_storage::TransientStorage,
 	BalanceOf, CodeHash, CodeInfo, CodeInfoOf, Config, ContractInfo, ContractInfoOf,
 	DebugBufferVec, Determinism, Error, Event, Nonce, Origin, Pallet as Contracts, Schedule,
 	LOG_TARGET,
@@ -209,6 +210,27 @@ pub trait Ext: sealing::Sealed {
 		take_old: bool,
 	) -> Result<WriteOutcome, DispatchError>;
 
+	/// Returns the storage entry of the executing account by the given `key`.
+	///
+	/// Returns `None` if the `key` wasn't previously set by `set_storage` or
+	/// was deleted.
+	fn get_transient_storage(&self, key: &Key<Self::T>) -> Option<Vec<u8>>;
+
+	/// 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.
+	fn get_transient_storage_size(&self, key: &Key<Self::T>) -> Option<u32>;
+
+	/// Sets the storage entry by the given key to the specified value. If `value` is `None` then
+	/// the storage entry is deleted.
+	fn set_transient_storage(
+		&mut self,
+		key: &Key<Self::T>,
+		value: Option<Vec<u8>>,
+		take_old: bool,
+	) -> Result<WriteOutcome, DispatchError>;
+
 	/// Returns the caller.
 	fn caller(&self) -> Origin<Self::T>;
 
@@ -473,6 +495,8 @@ pub struct Stack<'a, T: Config, E> {
 	debug_message: Option<&'a mut DebugBufferVec<T>>,
 	/// The determinism requirement of this call stack.
 	determinism: Determinism,
+	/// Transient storage
+	transient_storage: TransientStorage<T>,
 	/// No executable is held by the struct but influences its behaviour.
 	_phantom: PhantomData<E>,
 }
@@ -784,6 +808,7 @@ where
 			false,
 		)?;
 
+		let transient_storage_limit = T::MaxTransientStorageLen::get();
 		let stack = Self {
 			origin,
 			schedule,
@@ -796,6 +821,11 @@ where
 			frames: Default::default(),
 			debug_message,
 			determinism,
+			transient_storage: TransientStorage::new(
+				transient_storage_limit,
+				transient_storage_limit
+					.saturating_div((T::CallStack::size() as u32).saturating_add(1)),
+			),
 			_phantom: Default::default(),
 		};
 
@@ -926,6 +956,9 @@ where
 		let entry_point = frame.entry_point;
 		let delegated_code_hash =
 			if frame.delegate_caller.is_some() { Some(*executable.code_hash()) } else { None };
+
+		self.transient_storage.start_transaction();
+
 		let do_transaction = || {
 			// We need to charge the storage deposit before the initial transfer so that
 			// it can create the account in case the initial transfer is < ed.
@@ -1046,6 +1079,12 @@ where
 			Err(error) => (false, Err(error.into())),
 		};
 
+		if success {
+			self.transient_storage.commit_transaction();
+		} else {
+			self.transient_storage.rollback_transaction();
+		}
+
 		self.pop_frame(success);
 		output
 	}
@@ -1326,22 +1365,24 @@ where
 			return Err(Error::<T>::TerminatedWhileReentrant.into())
 		}
 		let frame = self.top_frame_mut();
+		let account_id = frame.account_id.clone();
 		let info = frame.terminate();
 		frame.nested_storage.terminate(&info, beneficiary.clone());
 
 		info.queue_trie_for_deletion();
-		ContractInfoOf::<T>::remove(&frame.account_id);
+		ContractInfoOf::<T>::remove(&account_id);
 		Self::decrement_refcount(info.code_hash);
 
 		for (code_hash, deposit) in info.delegate_dependencies() {
 			Self::decrement_refcount(*code_hash);
 			frame
 				.nested_storage
-				.charge_deposit(frame.account_id.clone(), StorageDeposit::Refund(*deposit));
+				.charge_deposit(account_id.clone(), StorageDeposit::Refund(*deposit));
 		}
 
+		self.transient_storage.remove(&account_id);
 		Contracts::<T>::deposit_event(Event::Terminated {
-			contract: frame.account_id.clone(),
+			contract: account_id,
 			beneficiary: beneficiary.clone(),
 		});
 		Ok(())
@@ -1374,6 +1415,24 @@ where
 		)
 	}
 
+	fn get_transient_storage(&self, key: &Key<T>) -> Option<Vec<u8>> {
+		self.transient_storage.read(self.address(), key)
+	}
+
+	fn get_transient_storage_size(&self, key: &Key<T>) -> Option<u32> {
+		self.transient_storage.read(self.address(), key).map(|value| value.len() as _)
+	}
+
+	fn set_transient_storage(
+		&mut self,
+		key: &Key<T>,
+		value: Option<Vec<u8>>,
+		take_old: bool,
+	) -> Result<WriteOutcome, DispatchError> {
+		let account_id = self.address().clone();
+		self.transient_storage.write(&account_id, key, value, take_old)
+	}
+
 	fn address(&self) -> &T::AccountId {
 		&self.top_frame().account_id
 	}
@@ -3826,6 +3885,142 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn transient_storage_works() {
+		// Call stack: BOB -> CHARLIE(success) -> BOB' (success)
+		let storage_key_1 = &Key::Fix([1; 32]);
+		let storage_key_2 = &Key::Fix([2; 32]);
+		let code_bob = MockLoader::insert(Call, |ctx, _| {
+			if ctx.input_data[0] == 0 {
+				assert_eq!(
+					ctx.ext.set_transient_storage(storage_key_1, Some(vec![1, 2]), false),
+					Ok(WriteOutcome::New)
+				);
+				assert_eq!(
+					ctx.ext.call(
+						Weight::zero(),
+						BalanceOf::<Test>::zero(),
+						CHARLIE,
+						0,
+						vec![],
+						true,
+						false,
+					),
+					exec_success()
+				);
+				assert_eq!(ctx.ext.get_transient_storage(storage_key_1), Some(vec![3]));
+				assert_eq!(ctx.ext.get_transient_storage(storage_key_2), Some(vec![4]));
+			} else {
+				assert_eq!(
+					ctx.ext.set_transient_storage(storage_key_1, Some(vec![3]), true),
+					Ok(WriteOutcome::Taken(vec![1, 2]))
+				);
+				assert_eq!(
+					ctx.ext.set_transient_storage(storage_key_2, Some(vec![4]), false),
+					Ok(WriteOutcome::New)
+				);
+			}
+			exec_success()
+		});
+		let code_charlie = MockLoader::insert(Call, |ctx, _| {
+			assert!(ctx
+				.ext
+				.call(Weight::zero(), BalanceOf::<Test>::zero(), BOB, 0, vec![99], true, false)
+				.is_ok());
+			// CHARLIE can not read BOB`s storage.
+			assert_eq!(ctx.ext.get_transient_storage(storage_key_1), None);
+			exec_success()
+		});
+
+		// This one tests passing the input data into a contract via call.
+		ExtBuilder::default().build().execute_with(|| {
+			let schedule = <Test as Config>::Schedule::get();
+			place_contract(&BOB, code_bob);
+			place_contract(&CHARLIE, code_charlie);
+			let contract_origin = Origin::from_account_id(ALICE);
+			let mut storage_meter =
+				storage::meter::Meter::new(&contract_origin, Some(0), 0).unwrap();
+
+			let result = MockStack::run_call(
+				contract_origin,
+				BOB,
+				&mut GasMeter::<Test>::new(GAS_LIMIT),
+				&mut storage_meter,
+				&schedule,
+				0,
+				vec![0],
+				None,
+				Determinism::Enforced,
+			);
+			assert_matches!(result, Ok(_));
+		});
+	}
+
+	#[test]
+	fn transient_storage_rollback_works() {
+		// Call stack: BOB -> CHARLIE (trap) -> BOB' (success)
+		let storage_key = &Key::Fix([1; 32]);
+		let code_bob = MockLoader::insert(Call, |ctx, _| {
+			if ctx.input_data[0] == 0 {
+				assert_eq!(
+					ctx.ext.set_transient_storage(storage_key, Some(vec![1, 2]), false),
+					Ok(WriteOutcome::New)
+				);
+				assert_eq!(
+					ctx.ext.call(
+						Weight::zero(),
+						BalanceOf::<Test>::zero(),
+						CHARLIE,
+						0,
+						vec![],
+						true,
+						false
+					),
+					exec_trapped()
+				);
+				assert_eq!(ctx.ext.get_transient_storage(storage_key), Some(vec![1, 2]));
+			} else {
+				let overwritten_length = ctx.ext.get_transient_storage_size(storage_key).unwrap();
+				assert_eq!(
+					ctx.ext.set_transient_storage(storage_key, Some(vec![3]), false),
+					Ok(WriteOutcome::Overwritten(overwritten_length))
+				);
+				assert_eq!(ctx.ext.get_transient_storage(storage_key), Some(vec![3]));
+			}
+			exec_success()
+		});
+		let code_charlie = MockLoader::insert(Call, |ctx, _| {
+			assert!(ctx
+				.ext
+				.call(Weight::zero(), BalanceOf::<Test>::zero(), BOB, 0, vec![99], true, false)
+				.is_ok());
+			exec_trapped()
+		});
+
+		// This one tests passing the input data into a contract via call.
+		ExtBuilder::default().build().execute_with(|| {
+			let schedule = <Test as Config>::Schedule::get();
+			place_contract(&BOB, code_bob);
+			place_contract(&CHARLIE, code_charlie);
+			let contract_origin = Origin::from_account_id(ALICE);
+			let mut storage_meter =
+				storage::meter::Meter::new(&contract_origin, Some(0), 0).unwrap();
+
+			let result = MockStack::run_call(
+				contract_origin,
+				BOB,
+				&mut GasMeter::<Test>::new(GAS_LIMIT),
+				&mut storage_meter,
+				&schedule,
+				0,
+				vec![0],
+				None,
+				Determinism::Enforced,
+			);
+			assert_matches!(result, Ok(_));
+		});
+	}
+
 	#[test]
 	fn ecdsa_to_eth_address_returns_proper_value() {
 		let bob_ch = MockLoader::insert(Call, |ctx, _| {

substrate/frame/contracts/src/lib.rs

@@ -96,6 +96,7 @@ pub use primitives::*;
 
 mod schedule;
 mod storage;
+mod transient_storage;
 mod wasm;
 
 pub mod chain_extension;
@@ -409,6 +410,10 @@ pub mod pallet {
 		#[pallet::constant]
 		type MaxDebugBufferLen: Get<u32>;
 
+		/// The maximum length of the transient storage in bytes.
+		#[pallet::constant]
+		type MaxTransientStorageLen: Get<u32>;
+
 		/// Origin allowed to upload code.
 		///
 		/// By default, it is safe to set this to `EnsureSigned`, allowing anyone to upload contract
@@ -554,6 +559,7 @@ pub mod pallet {
 			type MaxDebugBufferLen = ConstU32<{ 2 * 1024 * 1024 }>;
 			type MaxDelegateDependencies = MaxDelegateDependencies;
 			type MaxStorageKeyLen = ConstU32<128>;
+			type MaxTransientStorageLen = ConstU32<{ 1 * 1024 * 1024 }>;
 			type Migrations = ();
 			type Time = Self;
 			type Randomness = Self;
@@ -605,7 +611,11 @@ pub mod pallet {
 			// Max call depth is CallStack::size() + 1
 			let max_call_depth = u32::try_from(T::CallStack::size().saturating_add(1))
 				.expect("CallStack size is too big");
-
+			// Transient storage uses a BTreeMap, which has overhead compared to the raw size of
+			// key-value data. To ensure safety, a margin of 2x the raw key-value size is used.
+			let max_transient_storage_len = T::MaxTransientStorageLen::get()
+				.checked_mul(2)
+				.expect("MaxTransientStorageLen to big");
 			// Check that given configured `MaxCodeLen`, runtime heap memory limit can't be broken.
 			//
 			// In worst case, the decoded Wasm contract code would be `x16` times larger than the
@@ -615,7 +625,7 @@ pub mod pallet {
 			// Next, the pallet keeps the Wasm blob for each
 			// contract, hence we add up `MaxCodeLen` to the safety margin.
 			//
-			// Finally, the inefficiencies of the freeing-bump allocator
+			// The inefficiencies of the freeing-bump allocator
 			// being used in the client for the runtime memory allocations, could lead to possible
 			// memory allocations for contract code grow up to `x4` times in some extreme cases,
 			// which gives us total multiplier of `17*4` for `MaxCodeLen`.
@@ -624,17 +634,20 @@ pub mod pallet {
 			// memory should be available. Note that maximum allowed heap memory and stack size per
 			// each contract (stack frame) should also be counted.
 			//
+			// The pallet holds transient storage with a size up to `max_transient_storage_size`.
+			//
 			// Finally, we allow 50% of the runtime memory to be utilized by the contracts call
 			// stack, keeping the rest for other facilities, such as PoV, etc.
 			//
 			// This gives us the following formula:
 			//
-			// `(MaxCodeLen * 17 * 4 + MAX_STACK_SIZE + max_heap_size) * max_call_depth <
-			// max_runtime_mem/2`
+			// `(MaxCodeLen * 17 * 4 + MAX_STACK_SIZE + max_heap_size) * max_call_depth +
+			// max_transient_storage_size < max_runtime_mem/2`
 			//
 			// Hence the upper limit for the `MaxCodeLen` can be defined as follows:
 			let code_len_limit = max_runtime_mem
 				.saturating_div(2)
+				.saturating_sub(max_transient_storage_len)
 				.saturating_div(max_call_depth)
 				.saturating_sub(max_heap_size)
 				.saturating_sub(MAX_STACK_SIZE)
@@ -1235,6 +1248,8 @@ pub mod pallet {
 		DelegateDependencyAlreadyExists,
 		/// Can not add a delegate dependency to the code hash of the contract itself.
 		CannotAddSelfAsDelegateDependency,
+		/// Can not add more data to transient storage.
+		OutOfStorage,
 	}
 
 	/// A reason for the pallet contracts placing a hold on funds.