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>

substrate/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)))

substrate/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)
+		)
+	)
+)

substrate/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,8 +108,6 @@ 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
@@ -111,7 +115,7 @@ pub trait ChainExtension<C: Config> {
 	/// 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,

substrate/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]

substrate/frame/contracts/src/tests.rs

@@ -118,10 +118,17 @@ pub struct TestExtension {
 	last_seen_inputs: (u32, u32, u32, u32),
 }
 
+#[derive(Default)]
 pub struct RevertingExtension;
 
+#[derive(Default)]
 pub struct DisabledExtension;
 
+#[derive(Default)]
+pub struct TempStorageExtension {
+	storage: u32,
+}
+
 impl TestExtension {
 	fn disable() {
 		TEST_EXTENSION.with(|e| e.borrow_mut().enabled = false)
@@ -143,18 +150,20 @@ impl Default for TestExtension {
 }
 
 impl ChainExtension<Test> for TestExtension {
-	fn call<E>(func_id: u32, env: Environment<E, InitState>) -> ExtensionResult<RetVal>
+	fn call<E>(&mut self, env: Environment<E, InitState>) -> ExtensionResult<RetVal>
 	where
 		E: Ext<T = Test>,
 		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
 	{
+		let func_id = env.func_id();
+		let id = env.ext_id() as u32 | func_id as u32;
 		match func_id {
 			0 => {
 				let mut env = env.buf_in_buf_out();
 				let input = env.read(8)?;
 				env.write(&input, false, None)?;
 				TEST_EXTENSION.with(|e| e.borrow_mut().last_seen_buffer = input);
-				Ok(RetVal::Converging(func_id))
+				Ok(RetVal::Converging(id))
 			},
 			1 => {
 				let env = env.only_in();
@@ -162,17 +171,17 @@ impl ChainExtension<Test> for TestExtension {
 					e.borrow_mut().last_seen_inputs =
 						(env.val0(), env.val1(), env.val2(), env.val3())
 				});
-				Ok(RetVal::Converging(func_id))
+				Ok(RetVal::Converging(id))
 			},
 			2 => {
 				let mut env = env.buf_in_buf_out();
 				let weight = env.read(5)?[4].into();
 				env.charge_weight(weight)?;
-				Ok(RetVal::Converging(func_id))
+				Ok(RetVal::Converging(id))
 			},
 			3 => Ok(RetVal::Diverging { flags: ReturnFlags::REVERT, data: vec![42, 99] }),
 			_ => {
-				panic!("Passed unknown func_id to test chain extension: {}", func_id);
+				panic!("Passed unknown id to test chain extension: {}", func_id);
 			},
 		}
 	}
@@ -187,7 +196,7 @@ impl RegisteredChainExtension<Test> for TestExtension {
 }
 
 impl ChainExtension<Test> for RevertingExtension {
-	fn call<E>(_func_id: u32, _env: Environment<E, InitState>) -> ExtensionResult<RetVal>
+	fn call<E>(&mut self, _env: Environment<E, InitState>) -> ExtensionResult<RetVal>
 	where
 		E: Ext<T = Test>,
 		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
@@ -205,7 +214,7 @@ impl RegisteredChainExtension<Test> for RevertingExtension {
 }
 
 impl ChainExtension<Test> for DisabledExtension {
-	fn call<E>(_func_id: u32, _env: Environment<E, InitState>) -> ExtensionResult<RetVal>
+	fn call<E>(&mut self, _env: Environment<E, InitState>) -> ExtensionResult<RetVal>
 	where
 		E: Ext<T = Test>,
 		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
@@ -222,6 +231,37 @@ impl RegisteredChainExtension<Test> for DisabledExtension {
 	const ID: u16 = 2;
 }
 
+impl ChainExtension<Test> for TempStorageExtension {
+	fn call<E>(&mut self, env: Environment<E, InitState>) -> ExtensionResult<RetVal>
+	where
+		E: Ext<T = Test>,
+		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
+	{
+		let func_id = env.func_id();
+		match func_id {
+			0 => self.storage = 42,
+			1 => assert_eq!(self.storage, 42, "Storage is preserved inside the same call."),
+			2 => {
+				assert_eq!(self.storage, 0, "Storage is different for different calls.");
+				self.storage = 99;
+			},
+			3 => assert_eq!(self.storage, 99, "Storage is preserved inside the same call."),
+			_ => {
+				panic!("Passed unknown id to test chain extension: {}", func_id);
+			},
+		}
+		Ok(RetVal::Converging(0))
+	}
+
+	fn enabled() -> bool {
+		TEST_EXTENSION.with(|e| e.borrow().enabled)
+	}
+}
+
+impl RegisteredChainExtension<Test> for TempStorageExtension {
+	const ID: u16 = 3;
+}
+
 parameter_types! {
 	pub BlockWeights: frame_system::limits::BlockWeights =
 		frame_system::limits::BlockWeights::simple_max(2 * WEIGHT_PER_SECOND);
@@ -325,7 +365,8 @@ impl Config for Test {
 	type CallStack = [Frame<Self>; 31];
 	type WeightPrice = Self;
 	type WeightInfo = ();
-	type ChainExtension = (TestExtension, DisabledExtension, RevertingExtension);
+	type ChainExtension =
+		(TestExtension, DisabledExtension, RevertingExtension, TempStorageExtension);
 	type DeletionQueueDepth = ConstU32<1024>;
 	type DeletionWeightLimit = ConstU64<500_000_000_000>;
 	type Schedule = MySchedule;
@@ -396,6 +437,29 @@ fn initialize_block(number: u64) {
 	System::initialize(&number, &[0u8; 32].into(), &Default::default());
 }
 
+struct ExtensionInput<'a> {
+	extension_id: u16,
+	func_id: u16,
+	extra: &'a [u8],
+}
+
+impl<'a> ExtensionInput<'a> {
+	fn to_vec(&self) -> Vec<u8> {
+		((self.extension_id as u32) << 16 | (self.func_id as u32))
+			.to_le_bytes()
+			.iter()
+			.chain(self.extra)
+			.cloned()
+			.collect()
+	}
+}
+
+impl<'a> From<ExtensionInput<'a>> for Vec<u8> {
+	fn from(input: ExtensionInput) -> Vec<u8> {
+		input.to_vec()
+	}
+}
+
 // Perform a call to a plain account.
 // The actual transfer fails because we can only call contracts.
 // Then we check that at least the base costs where charged (no runtime gas costs.)
@@ -1567,23 +1631,6 @@ fn disabled_chain_extension_errors_on_call() {
 
 #[test]
 fn chain_extension_works() {
-	struct Input<'a> {
-		extension_id: u16,
-		func_id: u16,
-		extra: &'a [u8],
-	}
-
-	impl<'a> From<Input<'a>> for Vec<u8> {
-		fn from(input: Input) -> Vec<u8> {
-			((input.extension_id as u32) << 16 | (input.func_id as u32))
-				.to_le_bytes()
-				.iter()
-				.chain(input.extra)
-				.cloned()
-				.collect()
-		}
-	}
-
 	let (code, hash) = compile_module::<Test>("chain_extension").unwrap();
 	ExtBuilder::default().existential_deposit(50).build().execute_with(|| {
 		let min_balance = <Test as Config>::Currency::minimum_balance();
@@ -1599,12 +1646,8 @@ fn chain_extension_works() {
 		),);
 		let addr = Contracts::contract_address(&ALICE, &hash, &[]);
 
-		// The contract takes a up to 2 byte buffer where the first byte passed is used as
-		// as func_id to the chain extension which behaves differently based on the
-		// func_id.
-
 		// 0 = read input buffer and pass it through as output
-		let input: Vec<u8> = Input { extension_id: 0, func_id: 0, extra: &[99] }.into();
+		let input: Vec<u8> = ExtensionInput { extension_id: 0, func_id: 0, extra: &[99] }.into();
 		let result =
 			Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, None, input.clone(), false);
 		assert_eq!(TestExtension::last_seen_buffer(), input);
@@ -1617,7 +1660,7 @@ fn chain_extension_works() {
 			0,
 			GAS_LIMIT,
 			None,
-			Input { extension_id: 0, func_id: 1, extra: &[] }.into(),
+			ExtensionInput { extension_id: 0, func_id: 1, extra: &[] }.into(),
 			false,
 		)
 		.result
@@ -1632,7 +1675,7 @@ fn chain_extension_works() {
 			0,
 			GAS_LIMIT,
 			None,
-			Input { extension_id: 0, func_id: 2, extra: &[0] }.into(),
+			ExtensionInput { extension_id: 0, func_id: 2, extra: &[0] }.into(),
 			false,
 		);
 		assert_ok!(result.result);
@@ -1643,7 +1686,7 @@ fn chain_extension_works() {
 			0,
 			GAS_LIMIT,
 			None,
-			Input { extension_id: 0, func_id: 2, extra: &[42] }.into(),
+			ExtensionInput { extension_id: 0, func_id: 2, extra: &[42] }.into(),
 			false,
 		);
 		assert_ok!(result.result);
@@ -1654,7 +1697,7 @@ fn chain_extension_works() {
 			0,
 			GAS_LIMIT,
 			None,
-			Input { extension_id: 0, func_id: 2, extra: &[95] }.into(),
+			ExtensionInput { extension_id: 0, func_id: 2, extra: &[95] }.into(),
 			false,
 		);
 		assert_ok!(result.result);
@@ -1667,7 +1710,7 @@ fn chain_extension_works() {
 			0,
 			GAS_LIMIT,
 			None,
-			Input { extension_id: 0, func_id: 3, extra: &[] }.into(),
+			ExtensionInput { extension_id: 0, func_id: 3, extra: &[] }.into(),
 			false,
 		)
 		.result
@@ -1684,7 +1727,7 @@ fn chain_extension_works() {
 			0,
 			GAS_LIMIT,
 			None,
-			Input { extension_id: 1, func_id: 0, extra: &[] }.into(),
+			ExtensionInput { extension_id: 1, func_id: 0, extra: &[] }.into(),
 			false,
 		)
 		.result
@@ -1701,13 +1744,46 @@ fn chain_extension_works() {
 				0,
 				GAS_LIMIT,
 				None,
-				Input { extension_id: 2, func_id: 0, extra: &[] }.into(),
+				ExtensionInput { extension_id: 2, func_id: 0, extra: &[] }.into(),
 			),
 			Error::<Test>::NoChainExtension,
 		);
 	});
 }
 
+#[test]
+fn chain_extension_temp_storage_works() {
+	let (code, hash) = compile_module::<Test>("chain_extension_temp_storage").unwrap();
+	ExtBuilder::default().existential_deposit(50).build().execute_with(|| {
+		let min_balance = <Test as Config>::Currency::minimum_balance();
+		let _ = Balances::deposit_creating(&ALICE, 1000 * min_balance);
+		assert_ok!(Contracts::instantiate_with_code(
+			Origin::signed(ALICE),
+			min_balance * 100,
+			GAS_LIMIT,
+			None,
+			code,
+			vec![],
+			vec![],
+		),);
+		let addr = Contracts::contract_address(&ALICE, &hash, &[]);
+
+		// Call func 0 and func 1 back to back.
+		let stop_recursion = 0u8;
+		let mut input: Vec<u8> = ExtensionInput { extension_id: 3, func_id: 0, extra: &[] }.into();
+		input.extend_from_slice(
+			ExtensionInput { extension_id: 3, func_id: 1, extra: &[stop_recursion] }
+				.to_vec()
+				.as_ref(),
+		);
+
+		assert_ok!(
+			Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, None, input.clone(), false)
+				.result
+		);
+	})
+}
+
 #[test]
 fn lazy_removal_works() {
 	let (code, hash) = compile_module::<Test>("self_destruct").unwrap();

substrate/frame/contracts/src/wasm/runtime.rs

@@ -65,7 +65,7 @@ impl KeyType {
 /// This enum can be extended in the future: New codes can be added but existing codes
 /// will not be changed or removed. This means that any contract **must not** exhaustively
 /// match return codes. Instead, contracts should prepare for unknown variants and deal with
-/// those errors gracefuly in order to be forward compatible.
+/// those errors gracefully in order to be forward compatible.
 #[repr(u32)]
 pub enum ReturnCode {
 	/// API call successful.
@@ -101,8 +101,9 @@ pub enum ReturnCode {
 }
 
 impl ConvertibleToWasm for ReturnCode {
-	type NativeType = Self;
 	const VALUE_TYPE: ValueType = ValueType::I32;
+	type NativeType = Self;
+
 	fn to_typed_value(self) -> sp_sandbox::Value {
 		sp_sandbox::Value::I32(self as i32)
 	}
@@ -439,6 +440,7 @@ pub struct Runtime<'a, E: Ext + 'a> {
 	input_data: Option<Vec<u8>>,
 	memory: sp_sandbox::default_executor::Memory,
 	trap_reason: Option<TrapReason>,
+	chain_extension: Option<Box<<E::T as Config>::ChainExtension>>,
 }
 
 impl<'a, E> Runtime<'a, E>
@@ -452,7 +454,13 @@ where
 		input_data: Vec<u8>,
 		memory: sp_sandbox::default_executor::Memory,
 	) -> Self {
-		Runtime { ext, input_data: Some(input_data), memory, trap_reason: None }
+		Runtime {
+			ext,
+			input_data: Some(input_data),
+			memory,
+			trap_reason: None,
+			chain_extension: Some(Box::new(Default::default())),
+		}
 	}
 
 	/// Converts the sandbox result and the runtime state into the execution outcome.
@@ -2006,7 +2014,7 @@ define_env!(Env, <E: Ext>,
 	// module error.
 	[seal0] seal_call_chain_extension(
 		ctx,
-		func_id: u32,
+		id: u32,
 		input_ptr: u32,
 		input_len: u32,
 		output_ptr: u32,
@@ -2016,14 +2024,20 @@ define_env!(Env, <E: Ext>,
 		if !<E::T as Config>::ChainExtension::enabled() {
 			return Err(Error::<E::T>::NoChainExtension.into());
 		}
-		let env = Environment::new(ctx, input_ptr, input_len, output_ptr, output_len_ptr);
-		match <E::T as Config>::ChainExtension::call(func_id, env)? {
+		let mut chain_extension = ctx.chain_extension.take().expect(
+			"Constructor initializes with `Some`. This is the only place where it is set to `None`.\
+			It is always reset to `Some` afterwards. qed"
+		);
+		let env = Environment::new(ctx, id, input_ptr, input_len, output_ptr, output_len_ptr);
+		let ret = match chain_extension.call(env)? {
 			RetVal::Converging(val) => Ok(val),
 			RetVal::Diverging{flags, data} => Err(TrapReason::Return(ReturnData {
 				flags: flags.bits(),
 				data,
 			})),
-		}
+		};
+		ctx.chain_extension = Some(chain_extension);
+		ret
 	},
 
 	// Emit a custom debug message.