contracts: Allow ChainExtension::call() to access &mut self (#11874)

* Give chain extensions the ability to store some temporary values

* Update frame/contracts/src/wasm/runtime.rs

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

* Rename func_id -> id

* Replace `id` param by two functions on `env`

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

frame/contracts/fixtures/chain_extension.wat

@@ -31,14 +31,14 @@
 
 		;; the chain extension passes through the input and returns it as output
 		(call $seal_call_chain_extension
-			(i32.load (i32.const 4))	;; func_id
+			(i32.load (i32.const 4))	;; id
 			(i32.const 4)				;; input_ptr
 			(i32.load (i32.const 0))	;; input_len
 			(i32.const 16)				;; output_ptr
 			(i32.const 12)				;; output_len_ptr
 		)
 
-		;; the chain extension passes through the func_id
+		;; the chain extension passes through the id
 		(call $assert (i32.eq (i32.load (i32.const 4))))
 
 		(call $seal_return (i32.const 0) (i32.const 16) (i32.load (i32.const 12)))

frame/contracts/fixtures/chain_extension_temp_storage.wat

@@ -0,0 +1,85 @@
+;; Call chain extension two times with the specified func_ids
+;; It then calls itself once
+(module
+	(import "seal0" "seal_call_chain_extension"
+		(func $seal_call_chain_extension (param i32 i32 i32 i32 i32) (result i32))
+	)
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "seal0" "seal_address" (func $seal_address (param i32 i32)))
+	(import "seal1" "seal_call" (func $seal_call (param i32 i32 i64 i32 i32 i32 i32 i32) (result i32)))
+	(import "env" "memory" (memory 16 16))
+
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok (get_local 0))
+			(unreachable)
+		)
+	)
+
+	;; [0, 4) len of input buffer: 8 byte (func_ids) + 1byte (stop_recurse)
+	(data (i32.const 0) "\09")
+
+	;; [4, 16) buffer for input
+
+	;; [16, 48] buffer for self address
+
+	;; [48, 52] len of self address buffer
+	(data (i32.const 48) "\20")
+
+	(func (export "deploy"))
+
+	(func (export "call")
+		;; input: (func_id1: i32, func_id2: i32, stop_recurse: i8)
+		(call $seal_input (i32.const 4) (i32.const 0))
+
+		(call $seal_call_chain_extension
+			(i32.load (i32.const 4))	;; id
+			(i32.const 0)				;; input_ptr
+			(i32.const 0)				;; input_len
+			(i32.const 0xffffffff) 		;; u32 max sentinel value: do not copy output
+			(i32.const 0)				;; output_len_ptr
+		)
+		drop
+
+		(call $seal_call_chain_extension
+			(i32.load (i32.const 8))	;; _id
+			(i32.const 0)				;; input_ptr
+			(i32.const 0)				;; input_len
+			(i32.const 0xffffffff) 		;; u32 max sentinel value: do not copy output
+			(i32.const 0)				;; output_len_ptr
+		)
+		drop
+
+		(if (i32.eqz (i32.load8_u (i32.const 12)))
+			(then
+				;; stop recursion
+				(i32.store8 (i32.const 12) (i32.const 1))
+
+				;; load own address into buffer
+				(call $seal_address (i32.const 16) (i32.const 48))
+
+				;; call function 2 + 3 of chainext 3 next time
+				;; (3 << 16) | 2
+				;; (3 << 16) | 3
+				(i32.store (i32.const 4) (i32.const 196610))
+				(i32.store (i32.const 8) (i32.const 196611))
+
+				;; call self
+				(call $seal_call
+					(i32.const 8) ;; Set ALLOW_REENTRY
+					(i32.const 16)  ;; Pointer to "callee" address.
+					(i64.const 0)  ;; How much gas to devote for the execution. 0 = all.
+					(i32.const 512) ;; Pointer to the buffer with value to transfer
+					(i32.const 4) ;; Pointer to input data buffer address
+					(i32.load (i32.const 0))  ;; Length of input data buffer
+					(i32.const 4294967295) ;; u32 max value is the sentinel value: do not copy output
+					(i32.const 0) ;; Length is ignored in this case
+				)
+
+				;; check that call succeeded of call
+				(call $assert (i32.eqz))
+			)
+			(else)
+		)
+	)
+)

frame/contracts/src/chain_extension.rs

@@ -33,7 +33,7 @@
 //!
 //! Often there is a need for having multiple chain extensions. This is often the case when
 //! some generally useful off-the-shelf extensions should be included. To have multiple chain
-//! extensions they can be put into a tuple which is then passed to `[Config::ChainExtension]` like
+//! extensions they can be put into a tuple which is then passed to [`Config::ChainExtension`] like
 //! this `type Extensions = (ExtensionA, ExtensionB)`.
 //!
 //! However, only extensions implementing [`RegisteredChainExtension`] can be put into a tuple.
@@ -94,6 +94,12 @@ pub type Result<T> = sp_std::result::Result<T, DispatchError>;
 /// In order to create a custom chain extension this trait must be implemented and supplied
 /// to the pallet contracts configuration trait as the associated type of the same name.
 /// Consult the [module documentation](self) for a general explanation of chain extensions.
+///
+/// # Lifetime
+///
+/// The extension will be [`Default`] initialized at the beginning of each call
+/// (**not** per call stack) and dropped afterwards. Hence any value held inside the extension
+/// can be used as a per-call scratch buffer.
 pub trait ChainExtension<C: Config> {
 	/// Call the chain extension logic.
 	///
@@ -102,16 +108,14 @@ pub trait ChainExtension<C: Config> {
 	/// imported wasm function.
 	///
 	/// # Parameters
-	/// - `func_id`: The first argument to `seal_call_chain_extension`. Usually used to determine
-	///   which function to realize.
 	/// - `env`: Access to the remaining arguments and the execution environment.
 	///
 	/// # Return
 	///
 	/// In case of `Err` the contract execution is immediately suspended and the passed error
 	/// is returned to the caller. Otherwise the value of [`RetVal`] determines the exit
 	/// behaviour.
-	fn call<E>(func_id: u32, env: Environment<E, InitState>) -> Result<RetVal>
+	fn call<E>(&mut self, env: Environment<E, InitState>) -> Result<RetVal>
 	where
 		E: Ext<T = C>,
 		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>;
@@ -132,7 +136,7 @@ pub trait ChainExtension<C: Config> {
 ///
 /// An extension that implements this trait can be put in a tuple in order to have multiple
 /// extensions available. The tuple implementation routes requests based on the first two
-/// most significant bytes of the `func_id` passed to `call`.
+/// most significant bytes of the `id` passed to `call`.
 ///
 /// If this extensions is to be used by multiple runtimes consider
 /// [registering it](https://github.com/paritytech/chainextension-registry) to ensure that there
@@ -150,15 +154,15 @@ pub trait RegisteredChainExtension<C: Config>: ChainExtension<C> {
 #[impl_trait_for_tuples::impl_for_tuples(10)]
 #[tuple_types_custom_trait_bound(RegisteredChainExtension<C>)]
 impl<C: Config> ChainExtension<C> for Tuple {
-	fn call<E>(func_id: u32, mut env: Environment<E, InitState>) -> Result<RetVal>
+	fn call<E>(&mut self, mut env: Environment<E, InitState>) -> Result<RetVal>
 	where
 		E: Ext<T = C>,
 		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
 	{
 		for_tuples!(
 			#(
-				if (Tuple::ID == (func_id >> 16) as u16) && Tuple::enabled() {
-					return Tuple::call(func_id, env);
+				if (Tuple::ID == env.ext_id()) && Tuple::enabled() {
+					return Tuple.call(env);
 				}
 			)*
 		);
@@ -206,6 +210,22 @@ impl<'a, 'b, E: Ext, S: state::State> Environment<'a, 'b, E, S>
 where
 	<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
 {
+	/// The function id within the `id` passed by a contract.
+	///
+	/// It returns the two least significant bytes of the `id` passed by a contract as the other
+	/// two bytes represent the chain extension itself (the code which is calling this function).
+	pub fn func_id(&self) -> u16 {
+		(self.inner.id & 0x00FF) as u16
+	}
+
+	/// The chain extension id within the `id` passed by a contract.
+	///
+	/// It returns the two most significant bytes of the `id` passed by a contract which represent
+	/// the chain extension itself (the code which is calling this function).
+	pub fn ext_id(&self) -> u16 {
+		(self.inner.id >> 16) as u16
+	}
+
 	/// Charge the passed `amount` of weight from the overall limit.
 	///
 	/// It returns `Ok` when there the remaining weight budget is larger than the passed
@@ -251,13 +271,14 @@ impl<'a, 'b, E: Ext> Environment<'a, 'b, E, state::Init> {
 	/// ever create this type. Chain extensions merely consume it.
 	pub(crate) fn new(
 		runtime: &'a mut Runtime<'b, E>,
+		id: u32,
 		input_ptr: u32,
 		input_len: u32,
 		output_ptr: u32,
 		output_len_ptr: u32,
 	) -> Self {
 		Environment {
-			inner: Inner { runtime, input_ptr, input_len, output_ptr, output_len_ptr },
+			inner: Inner { runtime, id, input_ptr, input_len, output_ptr, output_len_ptr },
 			phantom: PhantomData,
 		}
 	}
@@ -406,6 +427,8 @@ struct Inner<'a, 'b, E: Ext> {
 	/// The runtime contains all necessary functions to interact with the running contract.
 	runtime: &'a mut Runtime<'b, E>,
 	/// Verbatim argument passed to `seal_call_chain_extension`.
+	id: u32,
+	/// Verbatim argument passed to `seal_call_chain_extension`.
 	input_ptr: u32,
 	/// Verbatim argument passed to `seal_call_chain_extension`.
 	input_len: u32,

frame/contracts/src/lib.rs

@@ -280,7 +280,7 @@ pub mod pallet {
 		type WeightInfo: WeightInfo;
 
 		/// Type that allows the runtime authors to add new host functions for a contract to call.
-		type ChainExtension: chain_extension::ChainExtension<Self>;
+		type ChainExtension: chain_extension::ChainExtension<Self> + Default;
 
 		/// Cost schedule and limits.
 		#[pallet::constant]