Implement chain extension testing in experimental off-chain env (#1152)

* Implement chain extension testing in experimental off-chain env

* Make `clippy` happy

crates/engine/src/chain_extension.rs

@@ -0,0 +1,100 @@
+// Copyright 2018-2022 Parity Technologies (UK) Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::Error;
+use derive_more::From;
+use std::collections::{
+    hash_map::Entry,
+    HashMap,
+};
+
+/// Chain extension registry.
+///
+/// Allows to register chain extension methods and call them.
+pub struct ChainExtensionHandler {
+    /// The currently registered runtime call handler.
+    registered: HashMap<ExtensionId, Box<dyn ChainExtension>>,
+    /// The output buffer used and reused for chain extension method call results.
+    output: Vec<u8>,
+}
+
+/// The unique ID of the registered chain extension method.
+#[derive(
+    Debug, From, scale::Encode, scale::Decode, PartialEq, Eq, PartialOrd, Ord, Hash,
+)]
+pub struct ExtensionId(u32);
+
+/// Types implementing this trait can be used as chain extensions.
+///
+/// This trait is only useful for testing contract via the off-chain environment.
+pub trait ChainExtension {
+    /// The static function ID of the chain extension.
+    ///
+    /// # Note
+    ///
+    /// This is expected to return a constant value.
+    fn func_id(&self) -> u32;
+
+    /// Calls the chain extension with the given input.
+    ///
+    /// Returns an error code and may fill the `output` buffer with a SCALE encoded result.
+    #[allow(clippy::ptr_arg)]
+    fn call(&mut self, input: &[u8], output: &mut Vec<u8>) -> u32;
+}
+
+impl Default for ChainExtensionHandler {
+    fn default() -> Self {
+        ChainExtensionHandler::new()
+    }
+}
+
+impl ChainExtensionHandler {
+    /// Creates a new chain extension handler.
+    ///
+    /// Initialized with an empty set of chain extensions.
+    pub fn new() -> Self {
+        Self {
+            registered: HashMap::new(),
+            output: Vec::new(),
+        }
+    }
+
+    /// Resets the chain extension handler to uninitialized state.
+    pub fn reset(&mut self) {
+        self.registered.clear();
+        self.output.clear();
+    }
+
+    /// Register a new chain extension.
+    pub fn register(&mut self, extension: Box<dyn ChainExtension>) {
+        let func_id = extension.func_id();
+        self.registered
+            .insert(ExtensionId::from(func_id), extension);
+    }
+
+    /// Evaluates the chain extension with the given parameters.
+    ///
+    /// Upon success returns the values returned by the evaluated chain extension.
+    pub fn eval(&mut self, func_id: u32, input: &[u8]) -> Result<(u32, &[u8]), Error> {
+        self.output.clear();
+        let extension_id = ExtensionId::from(func_id);
+        match self.registered.entry(extension_id) {
+            Entry::Occupied(occupied) => {
+                let status_code = occupied.into_mut().call(input, &mut self.output);
+                Ok((status_code, &mut self.output))
+            }
+            Entry::Vacant(_vacant) => Err(Error::UnregisteredChainExtension),
+        }
+    }
+}

crates/engine/src/ext.rs

@@ -18,6 +18,7 @@
 //! for more information.
 
 use crate::{
+    chain_extension::ChainExtensionHandler,
     database::Database,
     exec_context::ExecContext,
     test_api::{
@@ -125,6 +126,8 @@ pub struct Engine {
     pub(crate) debug_info: DebugInfo,
     /// The chain specification.
     pub chain_spec: ChainSpec,
+    /// Handler for registered chain extensions.
+    pub chain_extension_handler: ChainExtensionHandler,
 }
 
 /// The chain specification.
@@ -165,6 +168,7 @@ impl Engine {
             exec_context: ExecContext::new(),
             debug_info: DebugInfo::new(),
             chain_spec: ChainSpec::default(),
+            chain_extension_handler: ChainExtensionHandler::new(),
         }
     }
 }
@@ -445,15 +449,26 @@ impl Engine {
         set_output(output, &rng_bytes[..])
     }
 
+    /// Calls the chain extension method registered at `func_id` with `input`.
     pub fn call_chain_extension(
         &mut self,
-        _func_id: u32,
-        _input: &[u8],
-        _output: &mut &mut [u8],
-    ) -> u32 {
-        unimplemented!(
-            "off-chain environment does not yet support `call_chain_extension`"
-        );
+        func_id: u32,
+        input: &[u8],
+        output: &mut &mut [u8],
+    ) {
+        let encoded_input = input.encode();
+        let (status_code, out) = self
+            .chain_extension_handler
+            .eval(func_id, &encoded_input)
+            .unwrap_or_else(|error| {
+                panic!(
+                    "Encountered unexpected missing chain extension method: {:?}",
+                    error
+                );
+            });
+        let res = (status_code, out);
+        let decoded: Vec<u8> = scale::Encode::encode(&res);
+        set_output(output, &decoded[..])
     }
 
     /// Recovers the compressed ECDSA public key for given `signature` and `message_hash`,

crates/engine/src/lib.rs

@@ -15,6 +15,7 @@
 pub mod ext;
 pub mod test_api;
 
+mod chain_extension;
 mod database;
 mod exec_context;
 mod hashing;
@@ -23,6 +24,7 @@ mod types;
 #[cfg(test)]
 mod tests;
 
+pub use chain_extension::ChainExtension;
 pub use types::AccountError;
 
 use derive_more::From;

crates/env/src/engine/experimental_off_chain/impls.rs

@@ -305,12 +305,18 @@ impl EnvBackend for EnvInstance {
         let enc_input = &scale::Encode::encode(input)[..];
         let mut output: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE];
 
-        status_to_result(self.engine.call_chain_extension(
-            func_id,
-            enc_input,
-            &mut &mut output[..],
-        ))?;
-        let decoded = decode_to_result(&output[..])?;
+        self.engine
+            .call_chain_extension(func_id, enc_input, &mut &mut output[..]);
+        let (status, out): (u32, Vec<u8>) = scale::Decode::decode(&mut &output[..])
+            .unwrap_or_else(|error| {
+                panic!(
+                    "could not decode `call_chain_extension` output: {:?}",
+                    error
+                )
+            });
+
+        status_to_result(status)?;
+        let decoded = decode_to_result(&out[..])?;
         Ok(decoded)
     }
 }

crates/env/src/engine/experimental_off_chain/test_api.rs

@@ -26,6 +26,8 @@ use core::fmt::Debug;
 use ink_engine::test_api::RecordedDebugMessages;
 use std::panic::UnwindSafe;
 
+pub use ink_engine::ChainExtension;
+
 /// Record for an emitted event.
 #[derive(Clone)]
 pub struct EmittedEvent {
@@ -82,6 +84,19 @@ where
     })
 }
 
+/// Registers a new chain extension.
+pub fn register_chain_extension<E>(extension: E)
+where
+    E: ink_engine::ChainExtension + 'static,
+{
+    <EnvInstance as OnInstance>::on_instance(|instance| {
+        instance
+            .engine
+            .chain_extension_handler
+            .register(Box::new(extension));
+    })
+}
+
 /// Set the entropy hash of the current block.
 ///
 /// # Note

examples/rand-extension/Cargo.toml

@@ -30,3 +30,4 @@ std = [
     "scale-info/std",
 ]
 ink-as-dependency = []
+ink-experimental-engine = ["ink_env/ink-experimental-engine"]

examples/rand-extension/lib.rs

@@ -153,4 +153,53 @@ mod rand_extension {
             assert_eq!(rand_extension.get(), [1; 32]);
         }
     }
+
+    /// Unit tests in Rust are normally defined within such a `#[cfg(test)]`
+    #[cfg(test)]
+    #[cfg(feature = "ink-experimental-engine")]
+    mod tests_experimental_engine {
+        /// Imports all the definitions from the outer scope so we can use them here.
+        use super::*;
+        use ink_lang as ink;
+
+        /// We test if the default constructor does its job.
+        #[ink::test]
+        fn default_works() {
+            let rand_extension = RandExtension::default();
+            assert_eq!(rand_extension.get(), [0; 32]);
+        }
+
+        #[ink::test]
+        fn chain_extension_works() {
+            // given
+            struct MockedExtension;
+            impl ink_env::test::ChainExtension for MockedExtension {
+                /// The static function id of the chain extension.
+                fn func_id(&self) -> u32 {
+                    1101
+                }
+
+                /// The chain extension is called with the given input.
+                ///
+                /// Returns an error code and may fill the `output` buffer with a
+                /// SCALE encoded result. The error code is taken from the
+                /// `ink_env::chain_extension::FromStatusCode` implementation for
+                /// `RandomReadErr`.
+                fn call(&mut self, _input: &[u8], output: &mut Vec<u8>) -> u32 {
+                    let ret: [u8; 32] = [1; 32];
+                    scale::Encode::encode_to(&ret, output);
+                    0
+                }
+            }
+            ink_env::test::register_chain_extension(MockedExtension);
+            let mut rand_extension = RandExtension::default();
+            assert_eq!(rand_extension.get(), [0; 32]);
+
+            // when
+            rand_extension.update([0_u8; 32]).expect("update must work");
+
+            // then
+            assert_eq!(rand_extension.get(), [1; 32]);
+        }
+    }
 }