contracts: Composable ChainExtension (paritytech#11816)

* Add `RegisteredChainExtension`

* Add tests

* Update frame/contracts/src/chain_extension.rs

Co-authored-by: Michael Müller <michi@parity.io>

* Add more docs

* Remove debugging leftover

* Make ChainExtension-registry lowercase

* Apply suggestions from code review

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

* Improve clarity of test inputs

Co-authored-by: Michael Müller <michi@parity.io>
Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

frame/contracts/Cargo.toml

@@ -26,6 +26,7 @@ smallvec = { version = "1", default-features = false, features = [
 	"const_generics",
 ] }
 wasmi-validation = { version = "0.4", default-features = false }
+impl-trait-for-tuples = "0.2"
 
 # Only used in benchmarking to generate random contract code
 rand = { version = "0.8", optional = true, default-features = false }

frame/contracts/fixtures/chain_extension.wat

@@ -15,12 +15,12 @@
 	)
 
 	;; [0, 4) len of input output
-	(data (i32.const 0) "\02")
+	(data (i32.const 0) "\08")
 
 	;; [4, 12) buffer for input
 
-	;; [12, 16) len of output buffer
-	(data (i32.const 12) "\02")
+	;; [12, 48) len of output buffer
+	(data (i32.const 12) "\20")
 
 	;; [16, inf) buffer for output
 
@@ -31,15 +31,15 @@
 
 		;; the chain extension passes through the input and returns it as output
 		(call $seal_call_chain_extension
-			(i32.load8_u (i32.const 4))	;; func_id
+			(i32.load (i32.const 4))	;; func_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
-		(call $assert (i32.eq (i32.load8_u (i32.const 4))))
+		(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/src/chain_extension.rs

@@ -29,6 +29,22 @@
 //! required for this endeavour are defined or re-exported in this module. There is an
 //! implementation on `()` which can be used to signal that no chain extension is available.
 //!
+//! # Using multiple chain extensions
+//!
+//! 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
+//! this `type Extensions = (ExtensionA, ExtensionB)`.
+//!
+//! However, only extensions implementing [`RegisteredChainExtension`] can be put into a tuple.
+//! This is because the [`RegisteredChainExtension::ID`] is used to decide which of those extensions
+//! should should be used when the contract calls a chain extensions. Extensions which are generally
+//! useful should claim their `ID` with [the registry](https://github.com/paritytech/chainextension-registry)
+//! so that no collisions with other vendors will occur.
+//!
+//! **Chain specific extensions must use the reserved `ID = 0` so that they can't be registered with
+//! the registry.**
+//!
 //! # Security
 //!
 //! The chain author alone is responsible for the security of the chain extension.
@@ -112,20 +128,51 @@ pub trait ChainExtension<C: Config> {
 	}
 }
 
-/// Implementation that indicates that no chain extension is available.
-impl<C: Config> ChainExtension<C> for () {
-	fn call<E>(_func_id: u32, mut _env: Environment<E, InitState>) -> Result<RetVal>
+/// A [`ChainExtension`] that can be composed with other extensions using a tuple.
+///
+/// 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`.
+///
+/// If this extensions is to be used by multiple runtimes consider
+/// [registering it](https://github.com/paritytech/chainextension-registry) to ensure that there
+/// are no collisions with other vendors.
+///
+/// # Note
+///
+/// Currently, we support tuples of up to ten registred chain extensions. If more chain extensions
+/// are needed consider opening an issue.
+pub trait RegisteredChainExtension<C: Config>: ChainExtension<C> {
+	/// The extensions globally unique identifier.
+	const ID: u16;
+}
+
+#[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>
 	where
 		E: Ext<T = C>,
 		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
 	{
-		// Never called since [`Self::enabled()`] is set to `false`. Because we want to
-		// avoid panics at all costs we supply a sensible error value here instead
-		// of an `unimplemented!`.
+		for_tuples!(
+			#(
+				if (Tuple::ID == (func_id >> 16) as u16) && Tuple::enabled() {
+					return Tuple::call(func_id, env);
+				}
+			)*
+		);
 		Err(Error::<E::T>::NoChainExtension.into())
 	}
 
 	fn enabled() -> bool {
+		for_tuples!(
+			#(
+				if Tuple::enabled() {
+					return true;
+				}
+			)*
+		);
 		false
 	}
 }

frame/contracts/src/tests.rs

@@ -17,8 +17,8 @@
 
 use crate::{
 	chain_extension::{
-		ChainExtension, Environment, Ext, InitState, Result as ExtensionResult, RetVal,
-		ReturnFlags, SysConfig, UncheckedFrom,
+		ChainExtension, Environment, Ext, InitState, RegisteredChainExtension,
+		Result as ExtensionResult, RetVal, ReturnFlags, SysConfig, UncheckedFrom,
 	},
 	exec::{FixSizedKey, Frame},
 	storage::Storage,
@@ -118,6 +118,10 @@ pub struct TestExtension {
 	last_seen_inputs: (u32, u32, u32, u32),
 }
 
+pub struct RevertingExtension;
+
+pub struct DisabledExtension;
+
 impl TestExtension {
 	fn disable() {
 		TEST_EXTENSION.with(|e| e.borrow_mut().enabled = false)
@@ -147,7 +151,7 @@ impl ChainExtension<Test> for TestExtension {
 		match func_id {
 			0 => {
 				let mut env = env.buf_in_buf_out();
-				let input = env.read(2)?;
+				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))
@@ -162,7 +166,7 @@ impl ChainExtension<Test> for TestExtension {
 			},
 			2 => {
 				let mut env = env.buf_in_buf_out();
-				let weight = env.read(2)?[1].into();
+				let weight = env.read(5)?[4].into();
 				env.charge_weight(weight)?;
 				Ok(RetVal::Converging(func_id))
 			},
@@ -178,6 +182,46 @@ impl ChainExtension<Test> for TestExtension {
 	}
 }
 
+impl RegisteredChainExtension<Test> for TestExtension {
+	const ID: u16 = 0;
+}
+
+impl ChainExtension<Test> for RevertingExtension {
+	fn call<E>(_func_id: u32, _env: Environment<E, InitState>) -> ExtensionResult<RetVal>
+	where
+		E: Ext<T = Test>,
+		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
+	{
+		Ok(RetVal::Diverging { flags: ReturnFlags::REVERT, data: vec![0x4B, 0x1D] })
+	}
+
+	fn enabled() -> bool {
+		TEST_EXTENSION.with(|e| e.borrow().enabled)
+	}
+}
+
+impl RegisteredChainExtension<Test> for RevertingExtension {
+	const ID: u16 = 1;
+}
+
+impl ChainExtension<Test> for DisabledExtension {
+	fn call<E>(_func_id: u32, _env: Environment<E, InitState>) -> ExtensionResult<RetVal>
+	where
+		E: Ext<T = Test>,
+		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
+	{
+		panic!("Disabled chain extensions are never called")
+	}
+
+	fn enabled() -> bool {
+		false
+	}
+}
+
+impl RegisteredChainExtension<Test> for DisabledExtension {
+	const ID: u16 = 2;
+}
+
 parameter_types! {
 	pub BlockWeights: frame_system::limits::BlockWeights =
 		frame_system::limits::BlockWeights::simple_max(2 * WEIGHT_PER_SECOND);
@@ -281,7 +325,7 @@ impl Config for Test {
 	type CallStack = [Frame<Self>; 31];
 	type WeightPrice = Self;
 	type WeightInfo = ();
-	type ChainExtension = TestExtension;
+	type ChainExtension = (TestExtension, DisabledExtension, RevertingExtension);
 	type DeletionQueueDepth = ConstU32<1024>;
 	type DeletionWeightLimit = ConstU64<500_000_000_000>;
 	type Schedule = MySchedule;
@@ -1523,6 +1567,23 @@ 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();
@@ -1543,31 +1604,107 @@ fn chain_extension_works() {
 		// 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 result =
-			Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, None, vec![0, 99], false);
-		let gas_consumed = result.gas_consumed;
-		assert_eq!(TestExtension::last_seen_buffer(), vec![0, 99]);
-		assert_eq!(result.result.unwrap().data, Bytes(vec![0, 99]));
+			Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, None, input.clone(), false);
+		assert_eq!(TestExtension::last_seen_buffer(), input);
+		assert_eq!(result.result.unwrap().data, Bytes(input));
 
 		// 1 = treat inputs as integer primitives and store the supplied integers
-		Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, None, vec![1], false)
-			.result
-			.unwrap();
+		Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			Input { extension_id: 0, func_id: 1, extra: &[] }.into(),
+			false,
+		)
+		.result
+		.unwrap();
 		// those values passed in the fixture
-		assert_eq!(TestExtension::last_seen_inputs(), (4, 1, 16, 12));
+		assert_eq!(TestExtension::last_seen_inputs(), (4, 4, 16, 12));
 
-		// 2 = charge some extra weight (amount supplied in second byte)
-		let result =
-			Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, None, vec![2, 42], false);
+		// 2 = charge some extra weight (amount supplied in the fifth byte)
+		let result = Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			Input { extension_id: 0, func_id: 2, extra: &[0] }.into(),
+			false,
+		);
+		assert_ok!(result.result);
+		let gas_consumed = result.gas_consumed;
+		let result = Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			Input { extension_id: 0, func_id: 2, extra: &[42] }.into(),
+			false,
+		);
 		assert_ok!(result.result);
 		assert_eq!(result.gas_consumed, gas_consumed + 42);
+		let result = Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			Input { extension_id: 0, func_id: 2, extra: &[95] }.into(),
+			false,
+		);
+		assert_ok!(result.result);
+		assert_eq!(result.gas_consumed, gas_consumed + 95);
 
 		// 3 = diverging chain extension call that sets flags to 0x1 and returns a fixed buffer
-		let result = Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, None, vec![3], false)
-			.result
-			.unwrap();
+		let result = Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			Input { extension_id: 0, func_id: 3, extra: &[] }.into(),
+			false,
+		)
+		.result
+		.unwrap();
 		assert_eq!(result.flags, ReturnFlags::REVERT);
 		assert_eq!(result.data, Bytes(vec![42, 99]));
+
+		// diverging to second chain extension that sets flags to 0x1 and returns a fixed buffer
+		// We set the MSB part to 1 (instead of 0) which routes the request into the second
+		// extension
+		let result = Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			Input { extension_id: 1, func_id: 0, extra: &[] }.into(),
+			false,
+		)
+		.result
+		.unwrap();
+		assert_eq!(result.flags, ReturnFlags::REVERT);
+		assert_eq!(result.data, Bytes(vec![0x4B, 0x1D]));
+
+		// Diverging to third chain extension that is disabled
+		// We set the MSB part to 2 (instead of 0) which routes the request into the third extension
+		assert_err_ignore_postinfo!(
+			Contracts::call(
+				Origin::signed(ALICE),
+				addr.clone(),
+				0,
+				GAS_LIMIT,
+				None,
+				Input { extension_id: 2, func_id: 0, extra: &[] }.into(),
+			),
+			Error::<Test>::NoChainExtension,
+		);
 	});
 }