Merge pull request #59 from CosmWasm/duplicate-contract-code

Duplicate contract code

src/app.rs

@@ -728,6 +728,62 @@ where
         self.init_modules(|router, _, _| router.wasm.store_code(creator, code))
     }
 
+    /// Duplicates the contract code identified by `code_id` and returns
+    /// the identifier of the newly created copy of the contract code.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use cosmwasm_std::Addr;
+    /// use cw_multi_test::App;
+    ///
+    /// // contract implementation
+    /// mod echo {
+    ///   // contract entry points not shown here
+    /// #  use std::todo;
+    /// #  use cosmwasm_std::{Binary, Deps, DepsMut, Empty, Env, MessageInfo, Response, StdError, SubMsg, WasmMsg};
+    /// #  use serde::{Deserialize, Serialize};
+    /// #  use cw_multi_test::{Contract, ContractWrapper};
+    /// #
+    /// #  fn instantiate(_: DepsMut, _: Env, _: MessageInfo, _: Empty) -> Result<Response, StdError> {  
+    /// #    todo!()
+    /// #  }
+    /// #
+    /// #  fn execute(_: DepsMut, _: Env, _info: MessageInfo, msg: WasmMsg) -> Result<Response, StdError> {
+    /// #    todo!()
+    /// #  }
+    /// #
+    /// #  fn query(_deps: Deps, _env: Env, _msg: Empty) -> Result<Binary, StdError> {
+    /// #    todo!()
+    /// #  }
+    /// #  
+    ///   pub fn contract() -> Box<dyn Contract<Empty>> {
+    ///     // should return the contract
+    /// #   Box::new(ContractWrapper::new(execute, instantiate, query))
+    ///   }
+    /// }
+    ///
+    /// let mut app = App::default();
+    ///
+    /// // store a new contract, save the code id
+    /// # #[cfg(not(feature = "multitest_api_1_0"))]
+    /// let code_id = app.store_code_with_creator(Addr::unchecked("creator"), echo::contract());
+    /// # #[cfg(feature = "multitest_api_1_0")]
+    /// # let code_id = app.store_code(Addr::unchecked("creator"), echo::contract());
+    ///
+    /// // duplicate the existing contract, duplicated contract has different code id
+    /// assert_ne!(code_id, app.duplicate_code(code_id).unwrap());
+    ///
+    /// // zero is an invalid identifier for contract code, returns an error
+    /// assert_eq!("code id: invalid", app.duplicate_code(0).unwrap_err().to_string());
+    ///
+    /// // there is no contract code with identifier 100 stored yet, returns an error
+    /// assert_eq!("code id 100: no such code", app.duplicate_code(100).unwrap_err().to_string());
+    /// ```
+    pub fn duplicate_code(&mut self, code_id: u64) -> AnyResult<u64> {
+        self.init_modules(|router, _, _| router.wasm.duplicate_code(code_id))
+    }
+
     /// This allows to get `ContractData` for specific contract
     pub fn contract_data(&self, address: &Addr) -> AnyResult<ContractData> {
         self.read_module(|router, _, storage| router.wasm.load_contract(storage, address))
@@ -1136,6 +1192,30 @@ mod test {
     use crate::test_helpers::{CustomMsg, EmptyMsg};
     use crate::transactions::StorageTransaction;
 
+    #[test]
+    #[cfg(feature = "cosmwasm_1_2")]
+    fn duplicate_contract_code() {
+        // set up application
+        let mut app = App::default();
+
+        // store original contract code
+        #[cfg(not(feature = "multitest_api_1_0"))]
+        let original_code_id = app.store_code(payout::contract());
+        #[cfg(feature = "multitest_api_1_0")]
+        let original_code_id = app.store_code(Addr::unchecked("creator"), payout::contract());
+
+        // duplicate contract code
+        let duplicate_code_id = app.duplicate_code(original_code_id).unwrap();
+        assert_ne!(original_code_id, duplicate_code_id);
+
+        // query and compare code info of both contracts
+        let original_response = app.wrap().query_wasm_code_info(original_code_id).unwrap();
+        let duplicate_response = app.wrap().query_wasm_code_info(duplicate_code_id).unwrap();
+        assert_ne!(original_response.code_id, duplicate_response.code_id);
+        assert_eq!(original_response.creator, duplicate_response.creator);
+        assert_eq!(original_response.checksum, duplicate_response.checksum);
+    }
+
     fn get_balance<BankT, ApiT, StorageT, CustomT, WasmT>(
         app: &App<BankT, ApiT, StorageT, CustomT, WasmT>,
         addr: &Addr,
@@ -2744,7 +2824,7 @@ mod test {
 
         #[test]
         #[cfg(feature = "cosmwasm_1_2")]
-        fn query_contract_info() {
+        fn query_existing_code_info() {
             use super::*;
             let mut app = App::default();
             #[cfg(not(feature = "multitest_api_1_0"))]
@@ -2756,6 +2836,21 @@ mod test {
             assert_eq!("creator", code_info_response.creator);
             assert!(!code_info_response.checksum.is_empty());
         }
+
+        #[test]
+        #[cfg(feature = "cosmwasm_1_2")]
+        fn query_non_existing_code_info() {
+            use super::*;
+            let app = App::default();
+            assert_eq!(
+                "Generic error: Querier contract error: code id: invalid",
+                app.wrap().query_wasm_code_info(0).unwrap_err().to_string()
+            );
+            assert_eq!(
+                "Generic error: Querier contract error: code id 1: no such code",
+                app.wrap().query_wasm_code_info(1).unwrap_err().to_string()
+            );
+        }
     }
 
     mod custom_messages {

src/error.rs

@@ -21,7 +21,10 @@ pub enum Error {
     #[error("Unsupported wasm message: {0:?}")]
     UnsupportedWasmMsg(WasmMsg),
 
-    #[error("Unregistered code id")]
+    #[error("code id: invalid")]
+    InvalidCodeId,
+
+    #[error("code id {0}: no such code")]
     UnregisteredCodeId(u64),
 }
 

src/wasm.rs

@@ -66,13 +66,17 @@ pub struct ContractData {
 }
 
 /// Contract code base data.
-struct CodeData<ExecC, QueryC> {
-    /// Address of an account that initially created the contract.
+struct CodeData {
+    /// Address of an account that initially stored the contract code.
     //FIXME Remove this feature flag when the default flag is `cosmwasm_1_2` or higher.
     #[cfg_attr(feature = "cosmwasm_1_1", allow(dead_code))]
     creator: Addr,
-    /// The contract code base.
-    code: Box<dyn Contract<ExecC, QueryC>>,
+    /// Seed used to generate a checksum for underlying code base.
+    //FIXME Remove this feature flag when the default flag is `cosmwasm_1_2` or higher.
+    #[cfg_attr(feature = "cosmwasm_1_1", allow(dead_code))]
+    seed: usize,
+    /// Identifier of the code base where the contract code is stored in memory.
+    code_base_id: usize,
 }
 
 pub trait Wasm<ExecC, QueryC> {
@@ -110,9 +114,10 @@ pub trait Wasm<ExecC, QueryC> {
 }
 
 pub struct WasmKeeper<ExecC, QueryC> {
-    /// `codes` is in-memory lookup that stands in for wasm code,
-    /// this can only be edited on the WasmRouter, and just read in caches.
-    codes: Vec<CodeData<ExecC, QueryC>>,
+    /// Contract codes that stand for wasm code in real-life blockchain.
+    code_base: Vec<Box<dyn Contract<ExecC, QueryC>>>,
+    /// Code data with code base identifier and additional attributes.  
+    code_data: Vec<CodeData>,
     /// Just markers to make type elision fork when using it as `Wasm` trait
     _p: std::marker::PhantomData<QueryC>,
     generator: Box<dyn AddressGenerator>,
@@ -139,9 +144,11 @@ impl AddressGenerator for SimpleAddressGenerator {
 }
 
 impl<ExecC, QueryC> Default for WasmKeeper<ExecC, QueryC> {
+    /// Returns the default value for [WasmKeeper].
     fn default() -> Self {
         Self {
-            codes: Vec::default(),
+            code_base: Vec::default(),
+            code_data: Vec::default(),
             _p: std::marker::PhantomData,
             generator: Box::new(SimpleAddressGenerator()),
         }
@@ -181,15 +188,12 @@ where
             }
             #[cfg(feature = "cosmwasm_1_2")]
             WasmQuery::CodeInfo { code_id } => {
-                let code_data = self
-                    .codes
-                    .get((code_id - 1) as usize)
-                    .ok_or(Error::UnregisteredCodeId(code_id))?;
+                let code_data = self.code_data(code_id)?;
                 let mut res = cosmwasm_std::CodeInfoResponse::default();
                 res.code_id = code_id;
                 res.creator = code_data.creator.to_string();
                 res.checksum = cosmwasm_std::HexBinary::from(
-                    Sha256::digest(format!("contract code {}", res.code_id)).to_vec(),
+                    Sha256::digest(format!("contract code {}", code_data.seed)).to_vec(),
                 );
                 to_binary(&res).map_err(Into::into)
             }
@@ -234,11 +238,45 @@ impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
     /// Stores contract code in the in-memory lookup table.
     /// Returns an identifier of the stored contract code.
     pub fn store_code(&mut self, creator: Addr, code: Box<dyn Contract<ExecC, QueryC>>) -> u64 {
-        let code_id = self.codes.len() + 1;
-        self.codes.push(CodeData::<ExecC, QueryC> { creator, code });
+        let code_base_id = self.code_base.len();
+        self.code_base.push(code);
+        let code_id = self.code_data.len() + 1;
+        self.code_data.push(CodeData {
+            creator,
+            seed: code_id,
+            code_base_id,
+        });
         code_id as u64
     }
 
+    /// Duplicates contract code with specified identifier.
+    pub fn duplicate_code(&mut self, code_id: u64) -> AnyResult<u64> {
+        let code_data = self.code_data(code_id)?;
+        self.code_data.push(CodeData {
+            creator: code_data.creator.clone(),
+            seed: code_data.seed,
+            code_base_id: code_data.code_base_id,
+        });
+        Ok(code_id + 1)
+    }
+
+    /// Returns a handler to code of the contract with specified code id.
+    pub fn contract_code(&self, code_id: u64) -> AnyResult<&dyn Contract<ExecC, QueryC>> {
+        let code_data = self.code_data(code_id)?;
+        Ok(self.code_base[code_data.code_base_id].borrow())
+    }
+
+    /// Returns code data of the contract with specified code id.
+    fn code_data(&self, code_id: u64) -> AnyResult<&CodeData> {
+        if code_id < 1 {
+            bail!(Error::InvalidCodeId);
+        }
+        Ok(self
+            .code_data
+            .get((code_id - 1) as usize)
+            .ok_or(Error::UnregisteredCodeId(code_id))?)
+    }
+
     pub fn load_contract(&self, storage: &dyn Storage, address: &Addr) -> AnyResult<ContractData> {
         CONTRACTS
             .load(&prefixed_read(storage, NAMESPACE_WASM), address)
@@ -330,11 +368,9 @@ where
     }
 
     pub fn new_with_custom_address_generator(generator: impl AddressGenerator + 'static) -> Self {
-        let default = Self::default();
         Self {
-            codes: default.codes,
-            _p: default._p,
             generator: Box::new(generator),
+            ..Default::default()
         }
     }
 
@@ -534,7 +570,7 @@ where
                 let contract_addr = api.addr_validate(&contract_addr)?;
 
                 // check admin status and update the stored code_id
-                if new_code_id as usize > self.codes.len() {
+                if new_code_id as usize > self.code_data.len() {
                     bail!("Cannot migrate contract to unregistered code id");
                 }
                 let mut data = self.load_contract(storage, &contract_addr)?;
@@ -742,7 +778,7 @@ where
         label: String,
         created: u64,
     ) -> AnyResult<Addr> {
-        if code_id as usize > self.codes.len() {
+        if code_id as usize > self.code_data.len() {
             bail!("Cannot init contract with unregistered code id");
         }
 
@@ -876,15 +912,10 @@ where
         action: F,
     ) -> AnyResult<T>
     where
-        F: FnOnce(&Box<dyn Contract<ExecC, QueryC>>, Deps<QueryC>, Env) -> AnyResult<T>,
+        F: FnOnce(&dyn Contract<ExecC, QueryC>, Deps<QueryC>, Env) -> AnyResult<T>,
     {
         let contract = self.load_contract(storage, &address)?;
-        let handler = self
-            .codes
-            .get((contract.code_id - 1) as usize)
-            .ok_or(Error::UnregisteredCodeId(contract.code_id))?
-            .code
-            .borrow();
+        let handler = self.contract_code(contract.code_id)?;
         let storage = self.contract_storage_readonly(storage, &address);
         let env = self.get_env(address, block);
 
@@ -906,16 +937,11 @@ where
         action: F,
     ) -> AnyResult<T>
     where
-        F: FnOnce(&Box<dyn Contract<ExecC, QueryC>>, DepsMut<QueryC>, Env) -> AnyResult<T>,
+        F: FnOnce(&dyn Contract<ExecC, QueryC>, DepsMut<QueryC>, Env) -> AnyResult<T>,
         ExecC: DeserializeOwned,
     {
         let contract = self.load_contract(storage, &address)?;
-        let handler = self
-            .codes
-            .get((contract.code_id - 1) as usize)
-            .ok_or(Error::UnregisteredCodeId(contract.code_id))?
-            .code
-            .borrow();
+        let handler = self.contract_code(contract.code_id)?;
 
         // We don't actually need a transaction here, as it is already embedded in a transactional.
         // execute_submsg or App.execute_multi.