Duplicate contract code

src/app.rs

@@ -728,6 +728,28 @@ where
         self.init_modules(|router, _, _| router.wasm.store_code(creator, code))
     }
 
+    /// Duplicates the contract code identified by `code_id` and returns
+    /// the identifier of newly created copy of the contract.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use cosmwasm_std::Addr;
+    /// use cw_multi_test::App;
+    ///
+    /// let mut app = App::default();
+    ///
+    /// // there is no contract code with identifier 100 stored yet, returns an error
+    /// assert_eq!("Unregistered code id: 100", app.duplicate_code(100).unwrap_err().to_string());
+    ///
+    /// // zero is an invalid identifier for contract code, returns an error
+    /// assert_eq!("Unregistered code id: 0", app.duplicate_code(0).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 +1158,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,

src/error.rs

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

src/wasm.rs

@@ -66,13 +66,17 @@
 }
 
 /// 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 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<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()),
         }
@@ -182,14 +189,14 @@
             #[cfg(feature = "cosmwasm_1_2")]
             WasmQuery::CodeInfo { code_id } => {
                 let code_data = self
-                    .codes
+                    .code_data
                     .get((code_id - 1) as usize)
                     .ok_or(Error::UnregisteredCodeId(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 +241,46 @@
     /// 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> {
+        if code_id < 1 {
+            bail!(Error::UnregisteredCodeId(code_id));
+        }
+        let code_data = self
+            .code_data
+            .get((code_id - 1) as usize)
+            .ok_or(Error::UnregisteredCodeId(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 get_contract_code(&self, code_id: u64) -> AnyResult<&dyn Contract<ExecC, QueryC>> {
+        if code_id < 1 {
+            bail!(Error::UnregisteredCodeId(code_id));
+        }
+        let code_data = self
+            .code_data
+            .get((code_id - 1) as usize)
+            .ok_or(Error::UnregisteredCodeId(code_id))?;
+        Ok(self.code_base[code_data.code_base_id].borrow())
+    }
+
     pub fn load_contract(&self, storage: &dyn Storage, address: &Addr) -> AnyResult<ContractData> {
         CONTRACTS
             .load(&prefixed_read(storage, NAMESPACE_WASM), address)
@@ -330,11 +372,9 @@
     }
 
     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 +574,7 @@
                 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 +782,7 @@
         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 +916,10 @@
         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.get_contract_code(contract.code_id)?;
         let storage = self.contract_storage_readonly(storage, &address);
         let env = self.get_env(address, block);
 
@@ -906,16 +941,11 @@
         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.get_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.