Advanced test environment

crates/engine/src/database.rs

@@ -39,7 +39,7 @@ pub fn storage_of_contract_key(who: &[u8], key: &[u8]) -> [u8; 32] {
 ///
 /// Everything is stored in here: accounts, balances, contract storage, etc..
 /// Just like in Substrate a prefix hash is computed for every contract.
-#[derive(Default)]
+#[derive(Default, Clone)]
 pub struct Database {
     hmap: HashMap<Vec<u8>, Vec<u8>>,
 }

crates/engine/src/exec_context.rs

@@ -21,7 +21,7 @@ use super::types::{
 
 /// The context of a contract execution.
 #[cfg_attr(test, derive(Debug, PartialEq, Eq))]
-#[derive(Default)]
+#[derive(Default, Clone)]
 pub struct ExecContext {
     /// The caller of the contract execution. Might be user or another contract.
     ///
@@ -43,6 +43,14 @@ pub struct ExecContext {
     pub block_number: BlockNumber,
     /// The current block timestamp.
     pub block_timestamp: BlockTimestamp,
+    /// The input of the call.
+    pub input: Vec<u8>,
+    /// The output buffer of the call.
+    pub output: Vec<u8>,
+    /// Is contract reverted.
+    pub reverted: bool,
+    /// Origin of the call.
+    pub origin: Option<Vec<u8>>,
 }
 
 impl ExecContext {

crates/engine/src/ext.rs

@@ -32,7 +32,12 @@ use crate::{
     },
 };
 use scale::Encode;
-use std::panic::panic_any;
+use std::{
+    cell::RefCell,
+    collections::HashMap,
+    panic::panic_any,
+    rc::Rc,
+};
 
 type Result = core::result::Result<(), Error>;
 
@@ -110,7 +115,66 @@ impl ReturnCode {
     }
 }
 
+#[derive(Default, Clone)]
+pub struct ContractStorage {
+    pub instantiated: HashMap<Vec<u8>, Vec<u8>>,
+    pub entrance_count: HashMap<Vec<u8>, u32>,
+    pub reentrancy_allowed: HashMap<Vec<u8>, bool>,
+    pub deployed: HashMap<Vec<u8>, Contract>,
+}
+
+impl ContractStorage {
+    pub fn get_entrance_count(&self, callee: Vec<u8>) -> u32 {
+        *self.entrance_count.get(&callee).unwrap_or(&0)
+    }
+
+    pub fn get_allow_reentry(&self, callee: Vec<u8>) -> bool {
+        *self.reentrancy_allowed.get(&callee).unwrap_or(&false)
+    }
+
+    pub fn set_allow_reentry(&mut self, callee: Vec<u8>, allow: bool) {
+        if allow {
+            self.reentrancy_allowed.insert(callee, allow);
+        } else {
+            self.reentrancy_allowed.remove(&callee);
+        }
+    }
+
+    pub fn increase_entrance_count(&mut self, callee: Vec<u8>) -> Result {
+        let entrance_count = self
+            .entrance_count
+            .get(&callee)
+            .map_or(1, |count| count + 1);
+        self.entrance_count.insert(callee, entrance_count);
+
+        Ok(())
+    }
+
+    pub fn decrease_entrance_count(&mut self, callee: Vec<u8>) -> Result {
+        let entrance_count = self.entrance_count.get(&callee).map_or_else(
+            || Err(Error::CalleeTrapped),
+            |count| {
+                if *count == 0 {
+                    Err(Error::CalleeTrapped)
+                } else {
+                    Ok(count - 1)
+                }
+            },
+        )?;
+
+        self.entrance_count.insert(callee, entrance_count);
+        Ok(())
+    }
+}
+
+#[derive(Clone)]
+pub struct Contract {
+    pub deploy: fn(),
+    pub call: fn(),
+}
+
 /// The off-chain engine.
+#[derive(Clone)]
 pub struct Engine {
     /// The environment database.
     pub database: Database,
@@ -123,10 +187,13 @@ pub struct Engine {
     /// The chain specification.
     pub chain_spec: ChainSpec,
     /// Handler for registered chain extensions.
-    pub chain_extension_handler: ChainExtensionHandler,
+    pub chain_extension_handler: Rc<RefCell<ChainExtensionHandler>>,
+    /// Contracts' store.
+    pub contracts: ContractStorage,
 }
 
 /// The chain specification.
+#[derive(Clone)]
 pub struct ChainSpec {
     /// The current gas price.
     pub gas_price: Balance,
@@ -162,7 +229,8 @@ impl Engine {
             exec_context: ExecContext::new(),
             debug_info: DebugInfo::new(),
             chain_spec: ChainSpec::default(),
-            chain_extension_handler: ChainExtensionHandler::new(),
+            chain_extension_handler: Rc::new(RefCell::new(ChainExtensionHandler::new())),
+            contracts: ContractStorage::default(),
         }
     }
 }
@@ -323,8 +391,8 @@ impl Engine {
             .caller
             .as_ref()
             .expect("no caller has been set")
-            .as_bytes();
-        set_output(output, caller);
+            .clone();
+        set_output(output, caller.as_bytes());
     }
 
     /// Returns the balance of the executed contract.
@@ -333,7 +401,8 @@ impl Engine {
             .exec_context
             .callee
             .as_ref()
-            .expect("no callee has been set");
+            .expect("no callee has been set")
+            .clone();
 
         let balance_in_storage = self
             .database
@@ -357,8 +426,8 @@ impl Engine {
             .callee
             .as_ref()
             .expect("no callee has been set")
-            .as_bytes();
-        set_output(output, callee)
+            .clone();
+        set_output(output, callee.as_bytes())
     }
 
     /// Records the given debug message and appends to stdout.
@@ -453,8 +522,8 @@ impl Engine {
         output: &mut &mut [u8],
     ) {
         let encoded_input = input.encode();
-        let (status_code, out) = self
-            .chain_extension_handler
+        let mut chain_extension_hanler = self.chain_extension_handler.borrow_mut();
+        let (status_code, out) = chain_extension_hanler
             .eval(func_id, &encoded_input)
             .unwrap_or_else(|error| {
                 panic!(
@@ -510,6 +579,79 @@ impl Engine {
             Err(_) => Err(Error::EcdsaRecoveryFailed),
         }
     }
+
+    /// Register the contract into the local storage.
+    pub fn register_contract(
+        &mut self,
+        hash: &[u8],
+        deploy: fn(),
+        call: fn(),
+    ) -> Option<Contract> {
+        self.contracts
+            .deployed
+            .insert(hash.to_vec(), Contract { deploy, call })
+    }
+
+    /// Apply call flags for the call and return the input that might be changed
+    pub fn apply_code_flags_before_call(
+        &mut self,
+        caller: Option<AccountId>,
+        callee: Vec<u8>,
+        call_flags: u32,
+        input: Vec<u8>,
+    ) -> core::result::Result<Vec<u8>, Error> {
+        let forward_input = (call_flags & 1) != 0;
+        let clone_input = ((call_flags & 2) >> 1) != 0;
+        let allow_reentry = ((call_flags & 8) >> 3) != 0;
+
+        // Allow/deny reentrancy to the caller
+        if let Some(caller) = caller {
+            self.contracts
+                .set_allow_reentry(caller.as_bytes().to_vec(), allow_reentry);
+        }
+
+        // Check if reentrancy that is not allowed is encountered
+        if !self.contracts.get_allow_reentry(callee.clone())
+            && self.contracts.get_entrance_count(callee.clone()) > 0
+        {
+            return Err(Error::CalleeTrapped)
+        }
+
+        self.contracts.increase_entrance_count(callee)?;
+
+        let new_input = if forward_input {
+            let previous_input = self.exec_context.input.clone();
+
+            // delete the input because we will forward it
+            self.exec_context.input.clear();
+
+            previous_input
+        } else if clone_input {
+            self.exec_context.input.clone()
+        } else {
+            input
+        };
+
+        Ok(new_input)
+    }
+
+    /// Apply call flags after the call
+    pub fn apply_code_flags_after_call(
+        &mut self,
+        caller: Option<AccountId>,
+        callee: Vec<u8>,
+        _call_flags: u32,
+        _output: Vec<u8>,
+    ) -> core::result::Result<(), Error> {
+        self.contracts.decrease_entrance_count(callee)?;
+
+        if let Some(caller) = caller {
+            self.contracts
+                .reentrancy_allowed
+                .remove(&caller.as_bytes().to_vec());
+        }
+        Ok(())
+    }
 }
 
 /// Copies the `slice` into `output`.
@@ -525,3 +667,36 @@ fn set_output(output: &mut &mut [u8], slice: &[u8]) {
     );
     output[..slice.len()].copy_from_slice(slice);
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    pub fn contract_storage_works() {
+        let mut storage = ContractStorage::default();
+
+        let account = [0u8; 32].to_vec();
+
+        assert!(!storage.get_allow_reentry(account.clone()));
+        storage.set_allow_reentry(account.clone(), true);
+        assert!(storage.get_allow_reentry(account.clone()));
+
+        assert_eq!(storage.increase_entrance_count(account.clone()), Ok(()));
+        assert_eq!(storage.get_entrance_count(account.clone()), 1);
+        assert_eq!(storage.decrease_entrance_count(account.clone()), Ok(()));
+        assert_eq!(storage.get_entrance_count(account), 0);
+    }
+
+    #[test]
+    pub fn decrease_entrance_count_fails() {
+        let mut storage = ContractStorage::default();
+
+        let account = [0u8; 32].to_vec();
+
+        assert_eq!(
+            storage.decrease_entrance_count(account),
+            Err(Error::CalleeTrapped)
+        );
+    }
+}

crates/engine/src/lib.rs

@@ -22,7 +22,7 @@ pub mod test_api;
 
 mod chain_extension;
 mod database;
-mod exec_context;
+pub mod exec_context;
 mod hashing;
 mod types;
 

crates/engine/src/test_api.rs

@@ -72,6 +72,7 @@ impl IntoIterator for RecordedDebugMessages {
 }
 
 /// Recorder for relevant interactions with this crate.
+#[derive(Clone)]
 pub struct DebugInfo {
     /// Emitted events recorder.
     emitted_events: Vec<EmittedEvent>,
@@ -181,9 +182,9 @@ impl Engine {
     /// Returns the total number of reads and writes of the contract's storage.
     pub fn get_contract_storage_rw(&self, account_id: Vec<u8>) -> (usize, usize) {
         let account_id = AccountId::from(account_id);
-        let reads = self.debug_info.count_reads.get(&account_id).unwrap_or(&0);
-        let writes = self.debug_info.count_writes.get(&account_id).unwrap_or(&0);
-        (*reads, *writes)
+        let reads = *self.debug_info.count_reads.get(&account_id).unwrap_or(&0);
+        let writes = *self.debug_info.count_writes.get(&account_id).unwrap_or(&0);
+        (reads, writes)
     }
 
     /// Returns the total number of reads executed.
@@ -210,14 +211,15 @@ impl Engine {
     ///
     /// Returns `None` if the `account_id` is non-existent.
     pub fn count_used_storage_cells(&self, account_id: &[u8]) -> Result<usize, Error> {
-        let cells = self
+        let cells_len = self
             .debug_info
             .cells_per_account
             .get(&account_id.to_owned().into())
             .ok_or_else(|| {
                 Error::Account(AccountError::NoAccountForId(account_id.to_vec()))
-            })?;
-        Ok(cells.len())
+            })?
+            .len();
+        Ok(cells_len)
     }
 
     /// Advances the chain by a single block.

crates/engine/src/types.rs

@@ -27,7 +27,7 @@ pub type BlockTimestamp = u64;
 pub type Balance = u128;
 
 /// The Account Id type used by this crate.
-#[derive(Debug, From, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(Default, Debug, From, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[repr(transparent)]
 pub struct AccountId(Vec<u8>);
 

crates/env/src/api.rs

@@ -423,7 +423,7 @@ where
 /// # Note
 ///
 /// This function  stops the execution of the contract immediately.
-pub fn return_value<R>(return_flags: ReturnFlags, return_value: &R) -> !
+pub fn return_value<R>(return_flags: ReturnFlags, return_value: &R)
 where
     R: scale::Encode,
 {