cw_multi_test/
wasm.rs

1use crate::addresses::{AddressGenerator, SimpleAddressGenerator};
2use crate::app::{CosmosRouter, RouterQuerier};
3use crate::checksums::{ChecksumGenerator, SimpleChecksumGenerator};
4use crate::contracts::Contract;
5use crate::error::{bail, AnyContext, AnyError, AnyResult, Error};
6use crate::executor::AppResponse;
7use crate::prefixed_storage::{prefixed, prefixed_read, PrefixedStorage, ReadonlyPrefixedStorage};
8use crate::transactions::transactional;
9use cosmwasm_std::testing::mock_wasmd_attr;
10#[cfg(feature = "stargate")]
11use cosmwasm_std::GovMsg;
12use cosmwasm_std::{
13    to_json_binary, Addr, Api, Attribute, BankMsg, Binary, BlockInfo, Checksum, Coin, ContractInfo,
14    ContractInfoResponse, CosmosMsg, CustomMsg, CustomQuery, Deps, DepsMut, Env, Event,
15    MessageInfo, MsgResponse, Order, Querier, QuerierWrapper, Record, Reply, ReplyOn, Response,
16    StdResult, Storage, SubMsg, SubMsgResponse, SubMsgResult, TransactionInfo, WasmMsg, WasmQuery,
17};
18#[cfg(feature = "staking")]
19use cosmwasm_std::{DistributionMsg, StakingMsg};
20use cw_storage_plus::Map;
21use prost::Message;
22use schemars::JsonSchema;
23use serde::de::DeserializeOwned;
24use serde::{Deserialize, Serialize};
25use std::borrow::Borrow;
26use std::collections::BTreeMap;
27use std::fmt::Debug;
28
29/// Contract state kept in storage, separate from the contracts themselves (contract code).
30const CONTRACTS: Map<&Addr, ContractData> = Map::new("contracts");
31
32/// Wasm module namespace.
33const NAMESPACE_WASM: &[u8] = b"wasm";
34
35/// Contract [address namespace].
36///
37/// [address namespace]: https://github.com/CosmWasm/wasmd/blob/96e2b91144c9a371683555f3c696f882583cc6a2/x/wasm/types/events.go#L59
38const CONTRACT_ATTR: &str = "_contract_address";
39
40/// A structure representing a privileged message.
41#[derive(Clone, Debug, PartialEq, Eq, JsonSchema)]
42pub struct WasmSudo {
43    /// Address of a contract the privileged action will be sent to.
44    pub contract_addr: Addr,
45    /// Message representing privileged action to be executed by contract `sudo` entry-point.
46    pub message: Binary,
47}
48
49impl WasmSudo {
50    /// Creates a new privileged message for specified contract address and action to be executed.
51    pub fn new<T: Serialize>(contract_addr: &Addr, msg: &T) -> StdResult<WasmSudo> {
52        Ok(WasmSudo {
53            contract_addr: contract_addr.clone(),
54            message: to_json_binary(msg)?,
55        })
56    }
57}
58
59/// Contract data includes information about contract,
60/// equivalent of `ContractInfo` in `wasmd` interface.
61#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
62pub struct ContractData {
63    /// Identifier of stored contract code
64    pub code_id: u64,
65    /// Address of account who initially instantiated the contract
66    pub creator: Addr,
67    /// Optional address of account who can execute migrations
68    pub admin: Option<Addr>,
69    /// Metadata passed while contract instantiation
70    pub label: String,
71    /// Blockchain height in the moment of instantiating the contract
72    pub created: u64,
73}
74
75/// Contract code base data.
76struct CodeData {
77    /// Address of an account that initially stored the contract code.
78    creator: Addr,
79    /// Checksum of the contract's code base.
80    checksum: Checksum,
81    /// Identifier of the _source_ code of the contract stored in wasm keeper.
82    source_id: usize,
83}
84
85/// This trait implements the interface of the Wasm module.
86pub trait Wasm<ExecC, QueryC> {
87    /// Handles all `WasmMsg` messages.
88    fn execute(
89        &self,
90        api: &dyn Api,
91        storage: &mut dyn Storage,
92        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
93        block: &BlockInfo,
94        sender: Addr,
95        msg: WasmMsg,
96    ) -> AnyResult<AppResponse>;
97
98    /// Handles all `WasmQuery` requests.
99    fn query(
100        &self,
101        api: &dyn Api,
102        storage: &dyn Storage,
103        querier: &dyn Querier,
104        block: &BlockInfo,
105        request: WasmQuery,
106    ) -> AnyResult<Binary>;
107
108    /// Handles all sudo messages, this is an admin interface and can not be called via `CosmosMsg`.
109    fn sudo(
110        &self,
111        api: &dyn Api,
112        storage: &mut dyn Storage,
113        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
114        block: &BlockInfo,
115        msg: WasmSudo,
116    ) -> AnyResult<AppResponse>;
117
118    /// Stores the contract's code and returns an identifier of the stored contract's code.
119    fn store_code(&mut self, creator: Addr, code: Box<dyn Contract<ExecC, QueryC>>) -> u64;
120
121    /// Stores the contract's code under specified identifier,
122    /// returns the same code identifier when successful.
123    fn store_code_with_id(
124        &mut self,
125        creator: Addr,
126        code_id: u64,
127        code: Box<dyn Contract<ExecC, QueryC>>,
128    ) -> AnyResult<u64>;
129
130    /// Duplicates the contract's code with specified identifier
131    /// and returns an identifier of the copy of the contract's code.
132    fn duplicate_code(&mut self, code_id: u64) -> AnyResult<u64>;
133
134    /// Returns `ContractData` for the contract with specified address.
135    fn contract_data(&self, storage: &dyn Storage, address: &Addr) -> AnyResult<ContractData>;
136
137    /// Returns a raw state dump of all key-values held by a contract with specified address.
138    fn dump_wasm_raw(&self, storage: &dyn Storage, address: &Addr) -> Vec<Record>;
139
140    /// Returns the namespace of the contract storage.
141    fn contract_namespace(&self, contract: &Addr) -> Vec<u8> {
142        let mut name = b"contract_data/".to_vec();
143        name.extend_from_slice(contract.as_bytes());
144        name
145    }
146
147    /// Returns **read-only** (not mutable) contract storage.
148    fn contract_storage<'a>(
149        &self,
150        storage: &'a dyn Storage,
151        address: &Addr,
152    ) -> Box<dyn Storage + 'a> {
153        // We double-namespace this, once from global storage -> wasm_storage
154        // then from wasm_storage -> the contracts subspace
155        let namespace = self.contract_namespace(address);
156        let storage = ReadonlyPrefixedStorage::multilevel(storage, &[NAMESPACE_WASM, &namespace]);
157        Box::new(storage)
158    }
159
160    /// Returns **read-write** (mutable) contract storage.
161    fn contract_storage_mut<'a>(
162        &self,
163        storage: &'a mut dyn Storage,
164        address: &Addr,
165    ) -> Box<dyn Storage + 'a> {
166        // We double-namespace this, once from global storage -> wasm_storage
167        // then from wasm_storage -> the contracts subspace
168        let namespace = self.contract_namespace(address);
169        let storage = PrefixedStorage::multilevel(storage, &[NAMESPACE_WASM, &namespace]);
170        Box::new(storage)
171    }
172}
173
174/// A structure representing a default wasm keeper.
175pub struct WasmKeeper<ExecC, QueryC> {
176    /// Contract codes that stand for wasm code in real-life blockchain.
177    code_base: Vec<Box<dyn Contract<ExecC, QueryC>>>,
178    /// Code data with code base identifier and additional attributes.
179    code_data: BTreeMap<u64, CodeData>,
180    /// Contract's address generator.
181    address_generator: Box<dyn AddressGenerator>,
182    /// Contract's code checksum generator.
183    checksum_generator: Box<dyn ChecksumGenerator>,
184    /// Just markers to make type elision fork when using it as `Wasm` trait
185    _p: std::marker::PhantomData<QueryC>,
186}
187
188impl<ExecC, QueryC> Default for WasmKeeper<ExecC, QueryC> {
189    /// Returns the default value for [WasmKeeper].
190    fn default() -> Self {
191        Self {
192            code_base: Vec::default(),
193            code_data: BTreeMap::default(),
194            address_generator: Box::new(SimpleAddressGenerator),
195            checksum_generator: Box::new(SimpleChecksumGenerator),
196            _p: std::marker::PhantomData,
197        }
198    }
199}
200
201impl<ExecC, QueryC> Wasm<ExecC, QueryC> for WasmKeeper<ExecC, QueryC>
202where
203    ExecC: CustomMsg + DeserializeOwned + 'static,
204    QueryC: CustomQuery + DeserializeOwned + 'static,
205{
206    fn execute(
207        &self,
208        api: &dyn Api,
209        storage: &mut dyn Storage,
210        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
211        block: &BlockInfo,
212        sender: Addr,
213        msg: WasmMsg,
214    ) -> AnyResult<AppResponse> {
215        self.execute_wasm(api, storage, router, block, sender.clone(), msg.clone())
216            .context(format!(
217                "Error executing WasmMsg:\n  sender: {}\n  {:?}",
218                sender, msg
219            ))
220    }
221
222    fn query(
223        &self,
224        api: &dyn Api,
225        storage: &dyn Storage,
226        querier: &dyn Querier,
227        block: &BlockInfo,
228        request: WasmQuery,
229    ) -> AnyResult<Binary> {
230        match request {
231            WasmQuery::Smart { contract_addr, msg } => {
232                let addr = api.addr_validate(&contract_addr)?;
233                self.query_smart(addr, api, storage, querier, block, msg.into())
234            }
235            WasmQuery::Raw { contract_addr, key } => {
236                let addr = api.addr_validate(&contract_addr)?;
237                Ok(self.query_raw(addr, storage, &key))
238            }
239            WasmQuery::ContractInfo { contract_addr } => {
240                let addr = api.addr_validate(&contract_addr)?;
241                let contract = self.contract_data(storage, &addr)?;
242                let res = ContractInfoResponse::new(
243                    contract.code_id,
244                    contract.creator,
245                    contract.admin,
246                    false,
247                    None,
248                );
249                to_json_binary(&res).map_err(Into::into)
250            }
251            #[cfg(feature = "cosmwasm_1_2")]
252            WasmQuery::CodeInfo { code_id } => {
253                let code_data = self.code_data(code_id)?;
254                let res = cosmwasm_std::CodeInfoResponse::new(
255                    code_id,
256                    code_data.creator.clone(),
257                    code_data.checksum,
258                );
259                to_json_binary(&res).map_err(Into::into)
260            }
261            _ => unimplemented!("{}", Error::unsupported_wasm_query(request)),
262        }
263    }
264
265    fn sudo(
266        &self,
267        api: &dyn Api,
268        storage: &mut dyn Storage,
269        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
270        block: &BlockInfo,
271        msg: WasmSudo,
272    ) -> AnyResult<AppResponse> {
273        let custom_event = Event::new("sudo").add_attribute(CONTRACT_ATTR, &msg.contract_addr);
274        let res = self.call_sudo(
275            msg.contract_addr.clone(),
276            api,
277            storage,
278            router,
279            block,
280            msg.message.to_vec(),
281        )?;
282        let (res, msgs) = self.build_app_response(&msg.contract_addr, custom_event, res);
283        self.process_response(api, router, storage, block, msg.contract_addr, res, msgs)
284    }
285
286    /// Stores the contract's code in the in-memory lookup table.
287    /// Returns an identifier of the stored contract code.
288    fn store_code(&mut self, creator: Addr, code: Box<dyn Contract<ExecC, QueryC>>) -> u64 {
289        let code_id = self
290            .next_code_id()
291            .unwrap_or_else(|| panic!("{}", Error::NoMoreCodeIdAvailable));
292        self.save_code(code_id, creator, code)
293    }
294
295    /// Stores the contract's code in the in-memory lookup table.
296    /// Returns an identifier of the stored contract code.
297    fn store_code_with_id(
298        &mut self,
299        creator: Addr,
300        code_id: u64,
301        code: Box<dyn Contract<ExecC, QueryC>>,
302    ) -> AnyResult<u64> {
303        // validate provided contract code identifier
304        if self.code_data.contains_key(&code_id) {
305            bail!(Error::duplicated_code_id(code_id));
306        } else if code_id == 0 {
307            bail!(Error::invalid_code_id());
308        }
309        Ok(self.save_code(code_id, creator, code))
310    }
311
312    /// Duplicates the contract's code with specified identifier.
313    /// Returns an identifier of the copy of the contract's code.
314    fn duplicate_code(&mut self, code_id: u64) -> AnyResult<u64> {
315        let code_data = self.code_data(code_id)?;
316        let new_code_id = self
317            .next_code_id()
318            .ok_or_else(Error::no_more_code_id_available)?;
319        self.code_data.insert(
320            new_code_id,
321            CodeData {
322                creator: code_data.creator.clone(),
323                checksum: code_data.checksum,
324                source_id: code_data.source_id,
325            },
326        );
327        Ok(new_code_id)
328    }
329
330    /// Returns `ContractData` for the contract with specified address.
331    fn contract_data(&self, storage: &dyn Storage, address: &Addr) -> AnyResult<ContractData> {
332        CONTRACTS
333            .load(&prefixed_read(storage, NAMESPACE_WASM), address)
334            .map_err(Into::into)
335    }
336
337    /// Returns a raw state dump of all key-values held by a contract with specified address.
338    fn dump_wasm_raw(&self, storage: &dyn Storage, address: &Addr) -> Vec<Record> {
339        let storage = self.contract_storage(storage, address);
340        storage.range(None, None, Order::Ascending).collect()
341    }
342}
343
344impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC>
345where
346    ExecC: CustomMsg + DeserializeOwned + 'static,
347    QueryC: CustomQuery + DeserializeOwned + 'static,
348{
349    /// Creates a wasm keeper with default settings.
350    ///
351    /// # Example
352    ///
353    /// ```
354    /// use cw_multi_test::{no_init, AppBuilder, WasmKeeper};
355    ///
356    /// // create wasm keeper
357    /// let wasm_keeper = WasmKeeper::new();
358    ///
359    /// // create and use the application with newly created wasm keeper
360    /// let mut app = AppBuilder::default().with_wasm(wasm_keeper).build(no_init);
361    /// ```
362    pub fn new() -> Self {
363        Self::default()
364    }
365
366    /// Populates an existing [WasmKeeper] with custom contract address generator.
367    ///
368    /// # Example
369    ///
370    /// ```
371    /// use cosmwasm_std::{Addr, Api, Storage};
372    /// use cw_multi_test::{no_init, AddressGenerator, AppBuilder, WasmKeeper};
373    /// use cw_multi_test::error::AnyResult;
374    /// # use cosmwasm_std::testing::MockApi;
375    ///
376    /// struct CustomAddressGenerator;
377    ///
378    /// impl AddressGenerator for CustomAddressGenerator {
379    ///     fn contract_address(
380    ///         &self,
381    ///         api: &dyn Api,
382    ///         storage: &mut dyn Storage,
383    ///         code_id: u64,
384    ///         instance_id: u64,
385    ///     ) -> AnyResult<Addr> {
386    ///         // here implement your address generation logic
387    /// #       Ok(MockApi::default().addr_make("test_address"))
388    ///     }
389    /// }
390    ///
391    /// // populate wasm with your custom address generator
392    /// let wasm_keeper = WasmKeeper::new().with_address_generator(CustomAddressGenerator);
393    ///
394    /// // create and use the application with customized wasm keeper
395    /// let mut app = AppBuilder::default().with_wasm(wasm_keeper).build(no_init);
396    /// ```
397    pub fn with_address_generator(
398        mut self,
399        address_generator: impl AddressGenerator + 'static,
400    ) -> Self {
401        self.address_generator = Box::new(address_generator);
402        self
403    }
404
405    /// Populates an existing [WasmKeeper] with custom checksum generator for the contract code.
406    ///
407    /// # Example
408    ///
409    /// ```
410    /// use cosmwasm_std::{Addr, Checksum};
411    /// use cw_multi_test::{no_init, AppBuilder, ChecksumGenerator, WasmKeeper};
412    ///
413    /// struct MyChecksumGenerator;
414    ///
415    /// impl ChecksumGenerator for MyChecksumGenerator {
416    ///     fn checksum(&self, creator: &Addr, code_id: u64) -> Checksum {
417    ///         // here implement your custom checksum generator
418    /// #       Checksum::from_hex("custom_checksum").unwrap()
419    ///     }
420    /// }
421    ///
422    /// // populate wasm keeper with your custom checksum generator
423    /// let wasm_keeper = WasmKeeper::new().with_checksum_generator(MyChecksumGenerator);
424    ///
425    /// // create and use the application with customized wasm keeper
426    /// let mut app = AppBuilder::default().with_wasm(wasm_keeper).build(no_init);
427    /// ```
428    pub fn with_checksum_generator(
429        mut self,
430        checksum_generator: impl ChecksumGenerator + 'static,
431    ) -> Self {
432        self.checksum_generator = Box::new(checksum_generator);
433        self
434    }
435
436    /// Returns a handler to code of the contract with specified code id.
437    pub fn contract_code(&self, code_id: u64) -> AnyResult<&dyn Contract<ExecC, QueryC>> {
438        let code_data = self.code_data(code_id)?;
439        Ok(self.code_base[code_data.source_id].borrow())
440    }
441
442    /// Returns code data of the contract with specified code id.
443    fn code_data(&self, code_id: u64) -> AnyResult<&CodeData> {
444        if code_id < 1 {
445            bail!(Error::invalid_code_id());
446        }
447        Ok(self
448            .code_data
449            .get(&code_id)
450            .ok_or_else(|| Error::unregistered_code_id(code_id))?)
451    }
452
453    /// Validates all attributes.
454    ///
455    /// In `wasmd`, before version v0.45.0 empty attribute values were not allowed.
456    /// Since `wasmd` v0.45.0 empty attribute values are allowed,
457    /// so the value is not validated anymore.
458    fn verify_attributes(attributes: &[Attribute]) -> AnyResult<()> {
459        for attr in attributes {
460            let key = attr.key.trim();
461            let val = attr.value.trim();
462            if key.is_empty() {
463                bail!(Error::empty_attribute_key(val));
464            }
465            if key.starts_with('_') {
466                bail!(Error::reserved_attribute_key(key));
467            }
468        }
469        Ok(())
470    }
471
472    fn verify_response<T>(response: Response<T>) -> AnyResult<Response<T>>
473    where
474        T: CustomMsg,
475    {
476        Self::verify_attributes(&response.attributes)?;
477
478        for event in &response.events {
479            Self::verify_attributes(&event.attributes)?;
480            let ty = event.ty.trim();
481            if ty.len() < 2 {
482                bail!(Error::event_type_too_short(ty));
483            }
484        }
485
486        Ok(response)
487    }
488
489    fn save_code(
490        &mut self,
491        code_id: u64,
492        creator: Addr,
493        code: Box<dyn Contract<ExecC, QueryC>>,
494    ) -> u64 {
495        // prepare the next identifier for the contract's code
496        let source_id = self.code_base.len();
497        // prepare the contract's Wasm blob checksum
498        let checksum = code
499            .checksum()
500            .unwrap_or(self.checksum_generator.checksum(&creator, code_id));
501        // store the 'source' code of the contract
502        self.code_base.push(code);
503        // store the additional code attributes like creator address and checksum
504        self.code_data.insert(
505            code_id,
506            CodeData {
507                creator,
508                checksum,
509                source_id,
510            },
511        );
512        code_id
513    }
514
515    /// Returns the next contract's code identifier.
516    fn next_code_id(&self) -> Option<u64> {
517        self.code_data.keys().last().unwrap_or(&0u64).checked_add(1)
518    }
519
520    /// Executes the contract's `query` entry-point.
521    pub fn query_smart(
522        &self,
523        address: Addr,
524        api: &dyn Api,
525        storage: &dyn Storage,
526        querier: &dyn Querier,
527        block: &BlockInfo,
528        msg: Vec<u8>,
529    ) -> AnyResult<Binary> {
530        self.with_storage_readonly(
531            api,
532            storage,
533            querier,
534            block,
535            address,
536            |handler, deps, env| handler.query(deps, env, msg),
537        )
538    }
539
540    /// Returns the value stored under specified key in contracts storage.
541    pub fn query_raw(&self, address: Addr, storage: &dyn Storage, key: &[u8]) -> Binary {
542        let storage = self.contract_storage(storage, &address);
543        let data = storage.get(key).unwrap_or_default();
544        data.into()
545    }
546
547    fn send<T>(
548        &self,
549        api: &dyn Api,
550        storage: &mut dyn Storage,
551        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
552        block: &BlockInfo,
553        sender: T,
554        recipient: String,
555        amount: &[Coin],
556    ) -> AnyResult<AppResponse>
557    where
558        T: Into<Addr>,
559    {
560        if !amount.is_empty() {
561            let msg: CosmosMsg<ExecC> = BankMsg::Send {
562                to_address: recipient,
563                amount: amount.to_vec(),
564            }
565            .into();
566            let res = router.execute(api, storage, block, sender.into(), msg)?;
567            Ok(res)
568        } else {
569            Ok(AppResponse::default())
570        }
571    }
572
573    /// unified logic for UpdateAdmin and ClearAdmin messages
574    fn update_admin(
575        &self,
576        api: &dyn Api,
577        storage: &mut dyn Storage,
578        sender: Addr,
579        contract_addr: &str,
580        new_admin: Option<String>,
581    ) -> AnyResult<AppResponse> {
582        let contract_addr = api.addr_validate(contract_addr)?;
583        let admin = new_admin.map(|a| api.addr_validate(&a)).transpose()?;
584
585        // check admin status
586        let mut contract_data = self.contract_data(storage, &contract_addr)?;
587        if contract_data.admin != Some(sender) {
588            bail!(
589                "Only admin can update the contract admin: {:?}",
590                contract_data.admin
591            );
592        }
593        // update admin field
594        contract_data.admin = admin;
595        self.save_contract(storage, &contract_addr, &contract_data)?;
596
597        // No custom events or data here.
598        Ok(AppResponse::default())
599    }
600
601    // this returns the contract address as well, so we can properly resend the data
602    fn execute_wasm(
603        &self,
604        api: &dyn Api,
605        storage: &mut dyn Storage,
606        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
607        block: &BlockInfo,
608        sender: Addr,
609        msg: WasmMsg,
610    ) -> AnyResult<AppResponse> {
611        match msg {
612            WasmMsg::Execute {
613                contract_addr,
614                msg,
615                funds,
616            } => {
617                let contract_addr = api.addr_validate(&contract_addr)?;
618                // first move the cash
619                self.send(
620                    api,
621                    storage,
622                    router,
623                    block,
624                    sender.clone(),
625                    contract_addr.clone().into(),
626                    &funds,
627                )?;
628
629                // then call the contract
630                let info = MessageInfo { sender, funds };
631                let response = self.call_execute(
632                    api,
633                    storage,
634                    contract_addr.clone(),
635                    router,
636                    block,
637                    info,
638                    msg.to_vec(),
639                )?;
640
641                let custom_event =
642                    Event::new("execute").add_attribute(CONTRACT_ATTR, &contract_addr);
643
644                let (sub_response, sub_messages) =
645                    self.build_app_response(&contract_addr, custom_event, response);
646
647                let mut app_response = self.process_response(
648                    api,
649                    router,
650                    storage,
651                    block,
652                    contract_addr,
653                    sub_response,
654                    sub_messages,
655                )?;
656                app_response.data = encode_response_data(app_response.data);
657                Ok(app_response)
658            }
659            WasmMsg::Instantiate {
660                admin,
661                code_id,
662                msg,
663                funds,
664                label,
665            } => self.process_wasm_msg_instantiate(
666                api, storage, router, block, sender, admin, code_id, msg, funds, label, None,
667            ),
668            #[cfg(feature = "cosmwasm_1_2")]
669            WasmMsg::Instantiate2 {
670                admin,
671                code_id,
672                msg,
673                funds,
674                label,
675                salt,
676            } => self.process_wasm_msg_instantiate(
677                api,
678                storage,
679                router,
680                block,
681                sender,
682                admin,
683                code_id,
684                msg,
685                funds,
686                label,
687                Some(salt),
688            ),
689            WasmMsg::Migrate {
690                contract_addr,
691                new_code_id,
692                msg,
693            } => {
694                let contract_addr = api.addr_validate(&contract_addr)?;
695
696                // check admin status and update the stored code_id
697                if new_code_id as usize > self.code_data.len() {
698                    bail!("Cannot migrate contract to unregistered code id");
699                }
700                let mut data = self.contract_data(storage, &contract_addr)?;
701                if data.admin != Some(sender) {
702                    bail!("Only admin can migrate contract: {:?}", data.admin);
703                }
704                data.code_id = new_code_id;
705                self.save_contract(storage, &contract_addr, &data)?;
706
707                // then call migrate
708                let res = self.call_migrate(
709                    contract_addr.clone(),
710                    api,
711                    storage,
712                    router,
713                    block,
714                    msg.to_vec(),
715                )?;
716
717                let custom_event = Event::new("migrate")
718                    .add_attribute(CONTRACT_ATTR, &contract_addr)
719                    .add_attribute("code_id", new_code_id.to_string());
720                let (res, msgs) = self.build_app_response(&contract_addr, custom_event, res);
721                let mut res =
722                    self.process_response(api, router, storage, block, contract_addr, res, msgs)?;
723                res.data = encode_response_data(res.data);
724                Ok(res)
725            }
726            WasmMsg::UpdateAdmin {
727                contract_addr,
728                admin,
729            } => self.update_admin(api, storage, sender, &contract_addr, Some(admin)),
730            WasmMsg::ClearAdmin { contract_addr } => {
731                self.update_admin(api, storage, sender, &contract_addr, None)
732            }
733            _ => unimplemented!("{}", Error::unsupported_wasm_message(msg)),
734        }
735    }
736
737    /// Processes WasmMsg::Instantiate and WasmMsg::Instantiate2 messages.
738    fn process_wasm_msg_instantiate(
739        &self,
740        api: &dyn Api,
741        storage: &mut dyn Storage,
742        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
743        block: &BlockInfo,
744        sender: Addr,
745        admin: Option<String>,
746        code_id: u64,
747        msg: Binary,
748        funds: Vec<Coin>,
749        label: String,
750        salt: Option<Binary>,
751    ) -> AnyResult<AppResponse> {
752        if label.is_empty() {
753            bail!("Label is required on all contracts");
754        }
755
756        let contract_addr = self.register_contract(
757            api,
758            storage,
759            code_id,
760            sender.clone(),
761            admin.map(Addr::unchecked),
762            label,
763            block.height,
764            salt,
765        )?;
766
767        // move the cash
768        self.send(
769            api,
770            storage,
771            router,
772            block,
773            sender.clone(),
774            contract_addr.clone().into(),
775            &funds,
776        )?;
777
778        // then call the contract
779        let info = MessageInfo { sender, funds };
780        let res = self.call_instantiate(
781            contract_addr.clone(),
782            api,
783            storage,
784            router,
785            block,
786            info,
787            msg.to_vec(),
788        )?;
789
790        let custom_event = Event::new("instantiate")
791            .add_attribute(CONTRACT_ATTR, &contract_addr)
792            .add_attribute("code_id", code_id.to_string());
793
794        let (res, msgs) = self.build_app_response(&contract_addr, custom_event, res);
795        let mut res = self.process_response(
796            api,
797            router,
798            storage,
799            block,
800            contract_addr.clone(),
801            res,
802            msgs,
803        )?;
804        res.data = Some(instantiate_response(res.data, &contract_addr));
805        Ok(res)
806    }
807
808    /// This will execute the given messages, making all changes to the local cache.
809    /// This *will* write some data to the cache if the message fails half-way through.
810    /// All sequential calls to RouterCache will be one atomic unit (all commit or all fail).
811    ///
812    /// For normal use cases, you can use Router::execute() or Router::execute_multi().
813    /// This is designed to be handled internally as part of larger process flows.
814    ///
815    /// The `data` on `AppResponse` is data returned from `reply` call, not from execution of
816    /// sub-message itself. In case if `reply` is not called, no `data` is set.
817    fn execute_submsg(
818        &self,
819        api: &dyn Api,
820        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
821        storage: &mut dyn Storage,
822        block: &BlockInfo,
823        contract: Addr,
824        msg: SubMsg<ExecC>,
825    ) -> AnyResult<AppResponse> {
826        let SubMsg {
827            msg,
828            id,
829            reply_on,
830            payload,
831            ..
832        } = msg;
833        // Prepare the message type URL, will be needed when calling `reply` entrypoint.
834        let type_url = Self::response_type_url(&msg);
835
836        // Execute the submessage in cache
837        let sub_message_result = transactional(storage, |write_cache, _| {
838            router.execute(api, write_cache, block, contract.clone(), msg)
839        });
840
841        // call reply if meaningful
842        if let Ok(mut r) = sub_message_result {
843            if matches!(reply_on, ReplyOn::Always | ReplyOn::Success) {
844                let reply = Reply {
845                    id,
846                    payload,
847                    gas_used: 0,
848                    result: SubMsgResult::Ok(
849                        #[allow(deprecated)]
850                        SubMsgResponse {
851                            events: r.events.clone(),
852                            data: r.data.clone(),
853                            msg_responses: vec![MsgResponse {
854                                type_url,
855                                value: r.data.unwrap_or_default(),
856                            }],
857                        },
858                    ),
859                };
860                // do reply and combine it with the original response
861                let reply_res = self.reply(api, router, storage, block, contract, reply)?;
862                // override data
863                r.data = reply_res.data;
864                // append the events
865                r.events.extend_from_slice(&reply_res.events);
866            } else {
867                // reply is not called, no data should be returned
868                r.data = None;
869            }
870            Ok(r)
871        } else if let Err(e) = sub_message_result {
872            if matches!(reply_on, ReplyOn::Always | ReplyOn::Error) {
873                let reply = Reply {
874                    id,
875                    payload,
876                    gas_used: 0,
877                    result: SubMsgResult::Err(format!("{:?}", e)),
878                };
879                self.reply(api, router, storage, block, contract, reply)
880            } else {
881                Err(e)
882            }
883        } else {
884            sub_message_result
885        }
886    }
887
888    fn reply(
889        &self,
890        api: &dyn Api,
891        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
892        storage: &mut dyn Storage,
893        block: &BlockInfo,
894        contract: Addr,
895        reply: Reply,
896    ) -> AnyResult<AppResponse> {
897        let ok_attr = if reply.result.is_ok() {
898            "handle_success"
899        } else {
900            "handle_failure"
901        };
902        let custom_event = Event::new("reply")
903            .add_attribute(CONTRACT_ATTR, &contract)
904            .add_attribute("mode", ok_attr);
905
906        let res = self.call_reply(contract.clone(), api, storage, router, block, reply)?;
907        let (res, msgs) = self.build_app_response(&contract, custom_event, res);
908        self.process_response(api, router, storage, block, contract, res, msgs)
909    }
910
911    /// Captures all the events, data and sub messages from the contract call.
912    ///
913    /// This function does not handle the messages.
914    fn build_app_response(
915        &self,
916        contract: &Addr,
917        custom_event: Event, // entry-point specific custom event added by x/wasm
918        response: Response<ExecC>,
919    ) -> (AppResponse, Vec<SubMsg<ExecC>>) {
920        let Response {
921            messages,
922            attributes,
923            events,
924            data,
925            ..
926        } = response;
927
928        // always add custom event
929        let mut app_events = Vec::with_capacity(2 + events.len());
930        app_events.push(custom_event);
931
932        // we only emit the `wasm` event if some attributes are specified
933        if !attributes.is_empty() {
934            // turn attributes into event and place it first
935            let wasm_event = Event::new("wasm")
936                .add_attribute(CONTRACT_ATTR, contract)
937                .add_attributes(attributes);
938            app_events.push(wasm_event);
939        }
940
941        // These need to get `wasm-` prefix to match the wasmd semantics (custom wasm messages cannot
942        // fake system level event types, like transfer from the bank module)
943        let wasm_events = events.into_iter().map(|mut ev| {
944            ev.ty = format!("wasm-{}", ev.ty);
945            ev.attributes
946                .insert(0, mock_wasmd_attr(CONTRACT_ATTR, contract));
947            ev
948        });
949        app_events.extend(wasm_events);
950
951        let app_response = AppResponse {
952            events: app_events,
953            data,
954        };
955        (app_response, messages)
956    }
957
958    fn process_response(
959        &self,
960        api: &dyn Api,
961        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
962        storage: &mut dyn Storage,
963        block: &BlockInfo,
964        contract: Addr,
965        response: AppResponse,
966        sub_messages: Vec<SubMsg<ExecC>>,
967    ) -> AnyResult<AppResponse> {
968        // Unpack the provided response.
969        let AppResponse {
970            mut events, data, ..
971        } = response;
972        // Recurse in all submessages.
973        let data = sub_messages
974            .into_iter()
975            .try_fold(data, |data, sub_message| {
976                // Execute the submessage.
977                let sub_response = self.execute_submsg(
978                    api,
979                    router,
980                    storage,
981                    block,
982                    contract.clone(),
983                    sub_message,
984                )?;
985                // COLLECT and append all events from the processed submessage.
986                events.extend_from_slice(&sub_response.events);
987                // REPLACE the data with value from the processes submessage (if not empty).
988                Ok::<_, AnyError>(sub_response.data.or(data))
989            })?;
990        // Return the response with updated data, events and message responses taken from
991        // all processed sub messages. Note that events and message responses are collected,
992        // but the data is replaced with the data from the last processes submessage.
993        Ok(AppResponse { events, data })
994    }
995
996    /// Creates a contract address and empty storage instance.
997    /// Returns the new contract address.
998    ///
999    /// You have to call init after this to set up the contract properly.
1000    /// These two steps are separated to have cleaner return values.
1001    pub fn register_contract(
1002        &self,
1003        api: &dyn Api,
1004        storage: &mut dyn Storage,
1005        code_id: u64,
1006        creator: Addr,
1007        admin: impl Into<Option<Addr>>,
1008        label: String,
1009        created: u64,
1010        salt: impl Into<Option<Binary>>,
1011    ) -> AnyResult<Addr> {
1012        // check if the contract's code with specified code_id exists
1013        if code_id as usize > self.code_data.len() {
1014            bail!("Cannot init contract with unregistered code id");
1015        }
1016
1017        // generate a new contract address
1018        let instance_id = self.instance_count(storage) as u64;
1019        let addr = if let Some(salt_binary) = salt.into() {
1020            // generate predictable contract address when salt is provided
1021            let code_data = self.code_data(code_id)?;
1022            let canonical_addr = &api.addr_canonicalize(creator.as_ref())?;
1023            self.address_generator.predictable_contract_address(
1024                api,
1025                storage,
1026                code_id,
1027                instance_id,
1028                code_data.checksum.as_slice(),
1029                canonical_addr,
1030                salt_binary.as_slice(),
1031            )?
1032        } else {
1033            // generate non-predictable contract address
1034            self.address_generator
1035                .contract_address(api, storage, code_id, instance_id)?
1036        };
1037
1038        // contract with the same address must not already exist
1039        if self.contract_data(storage, &addr).is_ok() {
1040            bail!(Error::duplicated_contract_address(addr));
1041        }
1042
1043        // prepare contract data and save new contract instance
1044        let info = ContractData {
1045            code_id,
1046            creator,
1047            admin: admin.into(),
1048            label,
1049            created,
1050        };
1051        self.save_contract(storage, &addr, &info)?;
1052        Ok(addr)
1053    }
1054
1055    /// Executes contract's `execute` entry-point.
1056    pub fn call_execute(
1057        &self,
1058        api: &dyn Api,
1059        storage: &mut dyn Storage,
1060        address: Addr,
1061        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
1062        block: &BlockInfo,
1063        info: MessageInfo,
1064        msg: Vec<u8>,
1065    ) -> AnyResult<Response<ExecC>> {
1066        Self::verify_response(self.with_storage(
1067            api,
1068            storage,
1069            router,
1070            block,
1071            address,
1072            |contract, deps, env| contract.execute(deps, env, info, msg),
1073        )?)
1074    }
1075
1076    /// Executes contract's `instantiate` entry-point.
1077    pub fn call_instantiate(
1078        &self,
1079        address: Addr,
1080        api: &dyn Api,
1081        storage: &mut dyn Storage,
1082        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
1083        block: &BlockInfo,
1084        info: MessageInfo,
1085        msg: Vec<u8>,
1086    ) -> AnyResult<Response<ExecC>> {
1087        Self::verify_response(self.with_storage(
1088            api,
1089            storage,
1090            router,
1091            block,
1092            address,
1093            |contract, deps, env| contract.instantiate(deps, env, info, msg),
1094        )?)
1095    }
1096
1097    /// Executes contract's `reply` entry-point.
1098    pub fn call_reply(
1099        &self,
1100        address: Addr,
1101        api: &dyn Api,
1102        storage: &mut dyn Storage,
1103        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
1104        block: &BlockInfo,
1105        reply: Reply,
1106    ) -> AnyResult<Response<ExecC>> {
1107        Self::verify_response(self.with_storage(
1108            api,
1109            storage,
1110            router,
1111            block,
1112            address,
1113            |contract, deps, env| contract.reply(deps, env, reply),
1114        )?)
1115    }
1116
1117    /// Executes contract's `sudo` entry-point.
1118    pub fn call_sudo(
1119        &self,
1120        address: Addr,
1121        api: &dyn Api,
1122        storage: &mut dyn Storage,
1123        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
1124        block: &BlockInfo,
1125        msg: Vec<u8>,
1126    ) -> AnyResult<Response<ExecC>> {
1127        Self::verify_response(self.with_storage(
1128            api,
1129            storage,
1130            router,
1131            block,
1132            address,
1133            |contract, deps, env| contract.sudo(deps, env, msg),
1134        )?)
1135    }
1136
1137    /// Executes contract's `migrate` entry-point.
1138    pub fn call_migrate(
1139        &self,
1140        address: Addr,
1141        api: &dyn Api,
1142        storage: &mut dyn Storage,
1143        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
1144        block: &BlockInfo,
1145        msg: Vec<u8>,
1146    ) -> AnyResult<Response<ExecC>> {
1147        Self::verify_response(self.with_storage(
1148            api,
1149            storage,
1150            router,
1151            block,
1152            address,
1153            |contract, deps, env| contract.migrate(deps, env, msg),
1154        )?)
1155    }
1156
1157    fn get_env<T: Into<Addr>>(&self, address: T, block: &BlockInfo) -> Env {
1158        Env {
1159            block: block.clone(),
1160            contract: ContractInfo {
1161                address: address.into(),
1162            },
1163            transaction: Some(TransactionInfo { index: 0 }),
1164        }
1165    }
1166
1167    fn with_storage_readonly<F, T>(
1168        &self,
1169        api: &dyn Api,
1170        storage: &dyn Storage,
1171        querier: &dyn Querier,
1172        block: &BlockInfo,
1173        address: Addr,
1174        action: F,
1175    ) -> AnyResult<T>
1176    where
1177        F: FnOnce(&dyn Contract<ExecC, QueryC>, Deps<QueryC>, Env) -> AnyResult<T>,
1178    {
1179        let contract = self.contract_data(storage, &address)?;
1180        let handler = self.contract_code(contract.code_id)?;
1181        let storage = self.contract_storage(storage, &address);
1182        let env = self.get_env(address, block);
1183
1184        let deps = Deps {
1185            storage: storage.as_ref(),
1186            api,
1187            querier: QuerierWrapper::new(querier),
1188        };
1189        action(handler, deps, env)
1190    }
1191
1192    fn with_storage<F, T>(
1193        &self,
1194        api: &dyn Api,
1195        storage: &mut dyn Storage,
1196        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
1197        block: &BlockInfo,
1198        address: Addr,
1199        action: F,
1200    ) -> AnyResult<T>
1201    where
1202        F: FnOnce(&dyn Contract<ExecC, QueryC>, DepsMut<QueryC>, Env) -> AnyResult<T>,
1203        ExecC: DeserializeOwned,
1204    {
1205        let contract = self.contract_data(storage, &address)?;
1206        let handler = self.contract_code(contract.code_id)?;
1207
1208        // We don't actually need a transaction here, as it is already embedded in a transactional.
1209        // execute_submsg or App.execute_multi.
1210        // However, we need to get write and read access to the same storage in two different objects,
1211        // and this is the only way I know how to do so.
1212        transactional(storage, |write_cache, read_store| {
1213            let mut contract_storage = self.contract_storage_mut(write_cache, &address);
1214            let querier = RouterQuerier::new(router, api, read_store, block);
1215            let env = self.get_env(address, block);
1216
1217            let deps = DepsMut {
1218                storage: contract_storage.as_mut(),
1219                api,
1220                querier: QuerierWrapper::new(&querier),
1221            };
1222            action(handler, deps, env)
1223        })
1224    }
1225
1226    /// Saves contract data in a storage under specified address.
1227    pub fn save_contract(
1228        &self,
1229        storage: &mut dyn Storage,
1230        address: &Addr,
1231        contract: &ContractData,
1232    ) -> AnyResult<()> {
1233        CONTRACTS
1234            .save(&mut prefixed(storage, NAMESPACE_WASM), address, contract)
1235            .map_err(Into::into)
1236    }
1237
1238    /// Returns the number of all contract instances.
1239    fn instance_count(&self, storage: &dyn Storage) -> usize {
1240        CONTRACTS
1241            .range_raw(
1242                &prefixed_read(storage, NAMESPACE_WASM),
1243                None,
1244                None,
1245                Order::Ascending,
1246            )
1247            .count()
1248    }
1249
1250    /// Returns the response type for specified message.
1251    fn response_type_url(msg: &CosmosMsg<ExecC>) -> String {
1252        const UNKNOWN: &str = "/unknown";
1253        match &msg {
1254            CosmosMsg::Bank(bank_msg) => match bank_msg {
1255                BankMsg::Send { .. } => "/cosmos.bank.v1beta1.MsgSendResponse",
1256                BankMsg::Burn { .. } => "/cosmos.bank.v1beta1.MsgBurnResponse",
1257                _ => UNKNOWN,
1258            },
1259            CosmosMsg::Custom(..) => UNKNOWN,
1260            #[cfg(feature = "staking")]
1261            CosmosMsg::Staking(staking_msg) => match staking_msg {
1262                StakingMsg::Delegate { .. } => "/cosmos.staking.v1beta1.MsgDelegateResponse",
1263                StakingMsg::Undelegate { .. } => "/cosmos.staking.v1beta1.MsgUndelegateResponse",
1264                StakingMsg::Redelegate { .. } => {
1265                    "/cosmos.staking.v1beta1.MsgBeginRedelegateResponse"
1266                }
1267                _ => UNKNOWN,
1268            },
1269            #[cfg(feature = "staking")]
1270            CosmosMsg::Distribution(distribution_msg) => match distribution_msg {
1271                #[cfg(feature = "cosmwasm_1_3")]
1272                DistributionMsg::FundCommunityPool { .. } => {
1273                    "/cosmos.distribution.v1beta1.MsgFundCommunityPoolResponse"
1274                }
1275                DistributionMsg::SetWithdrawAddress { .. } => {
1276                    "/cosmos.distribution.v1beta1.MsgSetWithdrawAddressResponse"
1277                }
1278                DistributionMsg::WithdrawDelegatorReward { .. } => {
1279                    "/cosmos.distribution.v1beta1.MsgWithdrawDelegatorRewardResponse"
1280                }
1281                _ => UNKNOWN,
1282            },
1283            #[cfg(feature = "stargate")]
1284            #[allow(deprecated)]
1285            CosmosMsg::Stargate { .. } => UNKNOWN,
1286            #[cfg(feature = "cosmwasm_2_0")]
1287            CosmosMsg::Any(..) => UNKNOWN,
1288            #[cfg(feature = "stargate")]
1289            CosmosMsg::Ibc(..) => UNKNOWN,
1290            CosmosMsg::Wasm(wasm_msg) => match wasm_msg {
1291                WasmMsg::Instantiate { .. } => "/cosmwasm.wasm.v1.MsgInstantiateContractResponse",
1292                #[cfg(feature = "cosmwasm_1_2")]
1293                WasmMsg::Instantiate2 { .. } => "/cosmwasm.wasm.v1.MsgInstantiateContract2Response",
1294                WasmMsg::Execute { .. } => "/cosmwasm.wasm.v1.MsgExecuteContractResponse",
1295                WasmMsg::Migrate { .. } => "/cosmwasm.wasm.v1.MsgMigrateContractResponse",
1296                WasmMsg::UpdateAdmin { .. } => "/cosmwasm.wasm.v1.MsgUpdateAdminResponse",
1297                WasmMsg::ClearAdmin { .. } => "/cosmwasm.wasm.v1.MsgClearAdminResponse",
1298                _ => UNKNOWN,
1299            },
1300            #[cfg(feature = "stargate")]
1301            CosmosMsg::Gov(gov_msg) => match gov_msg {
1302                GovMsg::Vote { .. } => "/cosmos.gov.v1beta1.MsgVoteResponse",
1303                #[cfg(feature = "cosmwasm_1_2")]
1304                GovMsg::VoteWeighted { .. } => "/cosmos.gov.v1beta1.MsgVoteWeightedResponse",
1305            },
1306            _ => UNKNOWN,
1307        }
1308        .to_string()
1309    }
1310}
1311
1312#[derive(Clone, PartialEq, Message)]
1313struct InstantiateResponse {
1314    #[prost(string, tag = "1")]
1315    pub address: String,
1316    #[prost(bytes, tag = "2")]
1317    pub data: Vec<u8>,
1318}
1319
1320fn instantiate_response(data: Option<Binary>, contact_address: &Addr) -> Binary {
1321    let data = data.unwrap_or_default().to_vec();
1322    let init_data = InstantiateResponse {
1323        address: contact_address.into(),
1324        data,
1325    };
1326    let mut new_data = Vec::<u8>::with_capacity(init_data.encoded_len());
1327    // the data must encode successfully
1328    init_data.encode(&mut new_data).unwrap();
1329    new_data.into()
1330}
1331
1332#[derive(Clone, PartialEq, Message)]
1333struct ExecuteResponse {
1334    #[prost(bytes, tag = "1")]
1335    pub data: Vec<u8>,
1336}
1337
1338/// Encodes the response data.
1339fn encode_response_data(data: Option<Binary>) -> Option<Binary> {
1340    data.map(|d| {
1341        let execute_response = ExecuteResponse { data: d.to_vec() };
1342        let mut encoded_data = Vec::<u8>::with_capacity(execute_response.encoded_len());
1343        execute_response.encode(&mut encoded_data).unwrap();
1344        encoded_data.into()
1345    })
1346}
1347
1348#[cfg(test)]
1349mod test {
1350    use super::*;
1351    use crate::app::Router;
1352    use crate::bank::BankKeeper;
1353    use crate::featured::staking::{DistributionKeeper, StakeKeeper};
1354    use crate::module::FailingModule;
1355    use crate::test_helpers::{caller, error, payout};
1356    use crate::transactions::StorageTransaction;
1357    use crate::{GovFailingModule, IbcFailingModule, StargateFailing};
1358    use cosmwasm_std::testing::{message_info, mock_env, MockApi, MockQuerier, MockStorage};
1359    #[cfg(feature = "cosmwasm_1_2")]
1360    use cosmwasm_std::CodeInfoResponse;
1361    use cosmwasm_std::{
1362        coin, from_json, to_json_vec, CanonicalAddr, CosmosMsg, Empty, HexBinary, StdError,
1363    };
1364
1365    /// Type alias for default build `Router` to make its reference in typical scenario
1366    type BasicRouter<ExecC = Empty, QueryC = Empty> = Router<
1367        BankKeeper,
1368        FailingModule<ExecC, QueryC, Empty>,
1369        WasmKeeper<ExecC, QueryC>,
1370        StakeKeeper,
1371        DistributionKeeper,
1372        IbcFailingModule,
1373        GovFailingModule,
1374        StargateFailing,
1375    >;
1376
1377    fn wasm_keeper() -> WasmKeeper<Empty, Empty> {
1378        WasmKeeper::new()
1379    }
1380
1381    fn mock_router() -> BasicRouter {
1382        Router {
1383            wasm: WasmKeeper::new(),
1384            bank: BankKeeper::new(),
1385            custom: FailingModule::new(),
1386            staking: StakeKeeper::new(),
1387            distribution: DistributionKeeper::new(),
1388            ibc: IbcFailingModule::new(),
1389            gov: GovFailingModule::new(),
1390            stargate: StargateFailing,
1391        }
1392    }
1393
1394    #[test]
1395    fn register_contract() {
1396        let api = MockApi::default();
1397
1398        // prepare user addresses
1399        let creator_addr = api.addr_make("creator");
1400        let user_addr = api.addr_make("foobar");
1401        let admin_addr = api.addr_make("admin");
1402        let unregistered_addr = api.addr_make("unregistered");
1403
1404        let mut wasm_storage = MockStorage::new();
1405        let mut wasm_keeper = wasm_keeper();
1406        let block = mock_env().block;
1407        let code_id = wasm_keeper.store_code(creator_addr, error::contract(false));
1408
1409        transactional(&mut wasm_storage, |cache, _| {
1410            // cannot register contract with unregistered codeId
1411            wasm_keeper.register_contract(
1412                &api,
1413                cache,
1414                code_id + 1,
1415                user_addr.clone(),
1416                admin_addr.clone(),
1417                "label".to_owned(),
1418                1000,
1419                None,
1420            )
1421        })
1422        .unwrap_err();
1423
1424        let contract_addr = transactional(&mut wasm_storage, |cache, _| {
1425            // we can register a new instance of this code
1426            wasm_keeper.register_contract(
1427                &api,
1428                cache,
1429                code_id,
1430                user_addr.clone(),
1431                admin_addr.clone(),
1432                "label".to_owned(),
1433                1000,
1434                None,
1435            )
1436        })
1437        .unwrap();
1438
1439        // verify contract data are as expected
1440        let contract_data = wasm_keeper
1441            .contract_data(&wasm_storage, &contract_addr)
1442            .unwrap();
1443
1444        assert_eq!(
1445            contract_data,
1446            ContractData {
1447                code_id,
1448                creator: user_addr.clone(),
1449                admin: admin_addr.into(),
1450                label: "label".to_owned(),
1451                created: 1000,
1452            }
1453        );
1454
1455        let err = transactional(&mut wasm_storage, |cache, _| {
1456            // now, we call this contract and see the error message from the contract
1457            let info = message_info(&user_addr, &[]);
1458            wasm_keeper.call_instantiate(
1459                contract_addr.clone(),
1460                &api,
1461                cache,
1462                &mock_router(),
1463                &block,
1464                info,
1465                b"{}".to_vec(),
1466            )
1467        })
1468        .unwrap_err();
1469
1470        // StdError from contract_error auto-converted to string
1471        assert_eq!(
1472            StdError::generic_err("Init failed"),
1473            err.downcast().unwrap()
1474        );
1475
1476        let err = transactional(&mut wasm_storage, |cache, _| {
1477            // and the error for calling an unregistered contract
1478            let info = message_info(&user_addr, &[]);
1479            wasm_keeper.call_instantiate(
1480                unregistered_addr,
1481                &api,
1482                cache,
1483                &mock_router(),
1484                &block,
1485                info,
1486                b"{}".to_vec(),
1487            )
1488        })
1489        .unwrap_err();
1490
1491        // Default error message from router when not found
1492        assert!(matches!(err.downcast().unwrap(), StdError::NotFound { .. }));
1493    }
1494
1495    #[test]
1496    fn query_contract_info() {
1497        let api = MockApi::default();
1498
1499        // prepare user addresses
1500        let creator_addr = api.addr_make("creator");
1501        let admin_addr = api.addr_make("admin");
1502
1503        let mut wasm_storage = MockStorage::new();
1504        let mut wasm_keeper = wasm_keeper();
1505        let block = mock_env().block;
1506        let code_id = wasm_keeper.store_code(creator_addr.clone(), payout::contract());
1507        assert_eq!(1, code_id);
1508
1509        let contract_addr = wasm_keeper
1510            .register_contract(
1511                &api,
1512                &mut wasm_storage,
1513                code_id,
1514                creator_addr.clone(),
1515                admin_addr.clone(),
1516                "label".to_owned(),
1517                1000,
1518                None,
1519            )
1520            .unwrap();
1521
1522        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1523        let query = WasmQuery::ContractInfo {
1524            contract_addr: contract_addr.into(),
1525        };
1526
1527        let contract_info = wasm_keeper
1528            .query(&api, &wasm_storage, &querier, &block, query)
1529            .unwrap();
1530
1531        let actual: ContractInfoResponse = from_json(contract_info).unwrap();
1532        let expected =
1533            ContractInfoResponse::new(code_id, creator_addr, admin_addr.into(), false, None);
1534        assert_eq!(expected, actual);
1535    }
1536
1537    #[test]
1538    #[cfg(feature = "cosmwasm_1_2")]
1539    fn query_code_info() {
1540        let api = MockApi::default();
1541        let wasm_storage = MockStorage::new();
1542        let mut wasm_keeper = wasm_keeper();
1543        let block = mock_env().block;
1544        let creator_addr = api.addr_make("creator");
1545        let code_id = wasm_keeper.store_code(creator_addr.clone(), payout::contract());
1546        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1547        let query = WasmQuery::CodeInfo { code_id };
1548        let code_info = wasm_keeper
1549            .query(&api, &wasm_storage, &querier, &block, query)
1550            .unwrap();
1551        let actual: CodeInfoResponse = from_json(code_info).unwrap();
1552        assert_eq!(code_id, actual.code_id);
1553        assert_eq!(creator_addr.as_str(), actual.creator.as_str());
1554        assert_eq!(32, actual.checksum.as_slice().len());
1555    }
1556
1557    #[test]
1558    #[cfg(feature = "cosmwasm_1_2")]
1559    fn different_contracts_must_have_different_checksum() {
1560        let api = MockApi::default();
1561        let creator_addr = api.addr_make("creator");
1562        let wasm_storage = MockStorage::new();
1563        let mut wasm_keeper = wasm_keeper();
1564        let block = mock_env().block;
1565        let code_id_payout = wasm_keeper.store_code(creator_addr.clone(), payout::contract());
1566        let code_id_caller = wasm_keeper.store_code(creator_addr, caller::contract());
1567        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1568        let query_payout = WasmQuery::CodeInfo {
1569            code_id: code_id_payout,
1570        };
1571        let query_caller = WasmQuery::CodeInfo {
1572            code_id: code_id_caller,
1573        };
1574        let code_info_payout = wasm_keeper
1575            .query(&api, &wasm_storage, &querier, &block, query_payout)
1576            .unwrap();
1577        let code_info_caller = wasm_keeper
1578            .query(&api, &wasm_storage, &querier, &block, query_caller)
1579            .unwrap();
1580        let info_payout: CodeInfoResponse = from_json(code_info_payout).unwrap();
1581        let info_caller: CodeInfoResponse = from_json(code_info_caller).unwrap();
1582        assert_eq!(code_id_payout, info_payout.code_id);
1583        assert_eq!(code_id_caller, info_caller.code_id);
1584        assert_ne!(info_caller.code_id, info_payout.code_id);
1585        assert_eq!(info_caller.creator, info_payout.creator);
1586        assert_ne!(info_caller.checksum, info_payout.checksum);
1587    }
1588
1589    #[test]
1590    #[cfg(feature = "cosmwasm_1_2")]
1591    fn querying_invalid_code_info_must_fail() {
1592        let api = MockApi::default();
1593        let wasm_storage = MockStorage::new();
1594        let wasm_keeper = wasm_keeper();
1595        let block = mock_env().block;
1596
1597        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1598        let query = WasmQuery::CodeInfo { code_id: 100 };
1599
1600        wasm_keeper
1601            .query(&api, &wasm_storage, &querier, &block, query)
1602            .unwrap_err();
1603    }
1604
1605    #[test]
1606    fn can_dump_raw_wasm_state() {
1607        let api = MockApi::default();
1608
1609        // prepare user addresses
1610        let creator_addr = api.addr_make("creator");
1611        let admin_addr = api.addr_make("admin");
1612        let user_addr = api.addr_make("foobar");
1613
1614        let mut wasm_keeper = wasm_keeper();
1615        let block = mock_env().block;
1616        let code_id = wasm_keeper.store_code(creator_addr, payout::contract());
1617
1618        let mut wasm_storage = MockStorage::new();
1619
1620        let contract_addr = wasm_keeper
1621            .register_contract(
1622                &api,
1623                &mut wasm_storage,
1624                code_id,
1625                user_addr.clone(),
1626                admin_addr,
1627                "label".to_owned(),
1628                1000,
1629                None,
1630            )
1631            .unwrap();
1632
1633        // make a contract with state
1634        let payout = coin(1500, "mlg");
1635        let msg = payout::InstantiateMessage {
1636            payout: payout.clone(),
1637        };
1638        wasm_keeper
1639            .call_instantiate(
1640                contract_addr.clone(),
1641                &api,
1642                &mut wasm_storage,
1643                &mock_router(),
1644                &block,
1645                message_info(&user_addr, &[]),
1646                to_json_vec(&msg).unwrap(),
1647            )
1648            .unwrap();
1649
1650        // dump state
1651        let state = wasm_keeper.dump_wasm_raw(&wasm_storage, &contract_addr);
1652        assert_eq!(state.len(), 2);
1653        // check contents
1654        let (k, v) = &state[0];
1655        assert_eq!(k.as_slice(), b"count");
1656        let count: u32 = from_json(v).unwrap();
1657        assert_eq!(count, 1);
1658        let (k, v) = &state[1];
1659        assert_eq!(k.as_slice(), b"payout");
1660        let stored_pay: payout::InstantiateMessage = from_json(v).unwrap();
1661        assert_eq!(stored_pay.payout, payout);
1662    }
1663
1664    #[test]
1665    fn contract_send_coins() {
1666        let api = MockApi::default();
1667
1668        // prepare user addresses
1669        let creator_addr = api.addr_make("creator");
1670        let user_addr = api.addr_make("foobar");
1671
1672        let mut wasm_keeper = wasm_keeper();
1673        let block = mock_env().block;
1674        let code_id = wasm_keeper.store_code(creator_addr, payout::contract());
1675
1676        let mut wasm_storage = MockStorage::new();
1677        let mut cache = StorageTransaction::new(&wasm_storage);
1678
1679        let contract_addr = wasm_keeper
1680            .register_contract(
1681                &api,
1682                &mut cache,
1683                code_id,
1684                user_addr.clone(),
1685                None,
1686                "label".to_owned(),
1687                1000,
1688                None,
1689            )
1690            .unwrap();
1691
1692        let payout = coin(100, "TGD");
1693
1694        // init the contract
1695        let info = message_info(&user_addr, &[]);
1696        let init_msg = to_json_vec(&payout::InstantiateMessage {
1697            payout: payout.clone(),
1698        })
1699        .unwrap();
1700        let res = wasm_keeper
1701            .call_instantiate(
1702                contract_addr.clone(),
1703                &api,
1704                &mut cache,
1705                &mock_router(),
1706                &block,
1707                info,
1708                init_msg,
1709            )
1710            .unwrap();
1711        assert_eq!(0, res.messages.len());
1712
1713        // execute the contract
1714        let info = message_info(&user_addr, &[]);
1715        let res = wasm_keeper
1716            .call_execute(
1717                &api,
1718                &mut cache,
1719                contract_addr.clone(),
1720                &mock_router(),
1721                &block,
1722                info,
1723                b"{}".to_vec(),
1724            )
1725            .unwrap();
1726        assert_eq!(1, res.messages.len());
1727        match &res.messages[0].msg {
1728            CosmosMsg::Bank(BankMsg::Send { to_address, amount }) => {
1729                assert_eq!(to_address.as_str(), user_addr.as_str());
1730                assert_eq!(amount.as_slice(), &[payout.clone()]);
1731            }
1732            m => panic!("Unexpected message {:?}", m),
1733        }
1734
1735        // and flush before query
1736        cache.prepare().commit(&mut wasm_storage);
1737
1738        // query the contract
1739        let query = to_json_vec(&payout::QueryMsg::Payout {}).unwrap();
1740        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1741        let data = wasm_keeper
1742            .query_smart(contract_addr, &api, &wasm_storage, &querier, &block, query)
1743            .unwrap();
1744        let res: payout::InstantiateMessage = from_json(data).unwrap();
1745        assert_eq!(res.payout, payout);
1746    }
1747
1748    fn assert_payout(
1749        router: &WasmKeeper<Empty, Empty>,
1750        storage: &mut dyn Storage,
1751        contract_addr: &Addr,
1752        payout: &Coin,
1753    ) {
1754        let api = MockApi::default();
1755        let user_addr = api.addr_make("silly");
1756        let info = message_info(&user_addr, &[]);
1757        let res = router
1758            .call_execute(
1759                &api,
1760                storage,
1761                contract_addr.clone(),
1762                &mock_router(),
1763                &mock_env().block,
1764                info,
1765                b"{}".to_vec(),
1766            )
1767            .unwrap();
1768        assert_eq!(1, res.messages.len());
1769        match &res.messages[0].msg {
1770            CosmosMsg::Bank(BankMsg::Send { to_address, amount }) => {
1771                assert_eq!(to_address.as_str(), user_addr.as_str());
1772                assert_eq!(amount.as_slice(), &[payout.clone()]);
1773            }
1774            m => panic!("Unexpected message {:?}", m),
1775        }
1776    }
1777
1778    fn assert_no_contract(storage: &dyn Storage, contract_addr: &Addr) {
1779        let contract = CONTRACTS.may_load(storage, contract_addr).unwrap();
1780        assert!(contract.is_none(), "{:?}", contract_addr);
1781    }
1782
1783    #[test]
1784    fn multi_level_wasm_cache() {
1785        let api = MockApi::default();
1786
1787        // prepare user addresses
1788        let creator_addr = api.addr_make("creator");
1789        let user_addr = api.addr_make("foobar");
1790        let user_addr_1 = api.addr_make("johnny");
1791
1792        let mut wasm_keeper = wasm_keeper();
1793        let block = mock_env().block;
1794        let code_id = wasm_keeper.store_code(creator_addr, payout::contract());
1795
1796        let mut wasm_storage = MockStorage::new();
1797
1798        let payout1 = coin(100, "TGD");
1799
1800        // set contract 1 and commit (on router)
1801        let contract1 = transactional(&mut wasm_storage, |cache, _| {
1802            let contract = wasm_keeper
1803                .register_contract(
1804                    &api,
1805                    cache,
1806                    code_id,
1807                    user_addr.clone(),
1808                    None,
1809                    "".to_string(),
1810                    1000,
1811                    None,
1812                )
1813                .unwrap();
1814            let info = message_info(&user_addr, &[]);
1815            let init_msg = to_json_vec(&payout::InstantiateMessage {
1816                payout: payout1.clone(),
1817            })
1818            .unwrap();
1819            wasm_keeper
1820                .call_instantiate(
1821                    contract.clone(),
1822                    &api,
1823                    cache,
1824                    &mock_router(),
1825                    &block,
1826                    info,
1827                    init_msg,
1828                )
1829                .unwrap();
1830
1831            Ok(contract)
1832        })
1833        .unwrap();
1834
1835        let payout2 = coin(50, "BTC");
1836        let payout3 = coin(1234, "ATOM");
1837
1838        // create a new cache and check we can use contract 1
1839        let (contract2, contract3) = transactional(&mut wasm_storage, |cache, wasm_reader| {
1840            assert_payout(&wasm_keeper, cache, &contract1, &payout1);
1841
1842            // create contract 2 and use it
1843            let contract2 = wasm_keeper
1844                .register_contract(
1845                    &api,
1846                    cache,
1847                    code_id,
1848                    user_addr.clone(),
1849                    None,
1850                    "".to_owned(),
1851                    1000,
1852                    None,
1853                )
1854                .unwrap();
1855            let info = message_info(&user_addr, &[]);
1856            let init_msg = to_json_vec(&payout::InstantiateMessage {
1857                payout: payout2.clone(),
1858            })
1859            .unwrap();
1860            let _res = wasm_keeper
1861                .call_instantiate(
1862                    contract2.clone(),
1863                    &api,
1864                    cache,
1865                    &mock_router(),
1866                    &block,
1867                    info,
1868                    init_msg,
1869                )
1870                .unwrap();
1871            assert_payout(&wasm_keeper, cache, &contract2, &payout2);
1872
1873            // create a level2 cache and check we can use contract 1 and contract 2
1874            let contract3 = transactional(cache, |cache2, read| {
1875                assert_payout(&wasm_keeper, cache2, &contract1, &payout1);
1876                assert_payout(&wasm_keeper, cache2, &contract2, &payout2);
1877
1878                // create a contract on level 2
1879                let contract3 = wasm_keeper
1880                    .register_contract(
1881                        &api,
1882                        cache2,
1883                        code_id,
1884                        user_addr,
1885                        None,
1886                        "".to_owned(),
1887                        1000,
1888                        None,
1889                    )
1890                    .unwrap();
1891                let info = message_info(&user_addr_1, &[]);
1892                let init_msg = to_json_vec(&payout::InstantiateMessage {
1893                    payout: payout3.clone(),
1894                })
1895                .unwrap();
1896                let _res = wasm_keeper
1897                    .call_instantiate(
1898                        contract3.clone(),
1899                        &api,
1900                        cache2,
1901                        &mock_router(),
1902                        &block,
1903                        info,
1904                        init_msg,
1905                    )
1906                    .unwrap();
1907                assert_payout(&wasm_keeper, cache2, &contract3, &payout3);
1908
1909                // ensure first cache still doesn't see this contract
1910                assert_no_contract(read, &contract3);
1911                Ok(contract3)
1912            })
1913            .unwrap();
1914
1915            // after applying transaction, all contracts present on cache
1916            assert_payout(&wasm_keeper, cache, &contract1, &payout1);
1917            assert_payout(&wasm_keeper, cache, &contract2, &payout2);
1918            assert_payout(&wasm_keeper, cache, &contract3, &payout3);
1919
1920            // but not yet the root router
1921            assert_no_contract(wasm_reader, &contract1);
1922            assert_no_contract(wasm_reader, &contract2);
1923            assert_no_contract(wasm_reader, &contract3);
1924
1925            Ok((contract2, contract3))
1926        })
1927        .unwrap();
1928
1929        // ensure that it is now applied to the router
1930        assert_payout(&wasm_keeper, &mut wasm_storage, &contract1, &payout1);
1931        assert_payout(&wasm_keeper, &mut wasm_storage, &contract2, &payout2);
1932        assert_payout(&wasm_keeper, &mut wasm_storage, &contract3, &payout3);
1933    }
1934
1935    fn assert_admin(
1936        storage: &dyn Storage,
1937        wasm_keeper: &WasmKeeper<Empty, Empty>,
1938        contract_addr: &impl ToString,
1939        admin: Option<Addr>,
1940    ) {
1941        let api = MockApi::default();
1942        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1943        // query
1944        let data = wasm_keeper
1945            .query(
1946                &api,
1947                storage,
1948                &querier,
1949                &mock_env().block,
1950                WasmQuery::ContractInfo {
1951                    contract_addr: contract_addr.to_string(),
1952                },
1953            )
1954            .unwrap();
1955        let res: ContractInfoResponse = from_json(data).unwrap();
1956        assert_eq!(res.admin, admin);
1957    }
1958
1959    #[test]
1960    fn update_clear_admin_works() {
1961        let api = MockApi::default();
1962        let mut wasm_keeper = wasm_keeper();
1963        let block = mock_env().block;
1964        let creator = api.addr_make("creator");
1965        let code_id = wasm_keeper.store_code(creator.clone(), caller::contract());
1966
1967        let mut wasm_storage = MockStorage::new();
1968
1969        let admin = api.addr_make("admin");
1970        let new_admin = api.addr_make("new_admin");
1971        let normal_user = api.addr_make("normal_user");
1972
1973        let contract_addr = wasm_keeper
1974            .register_contract(
1975                &api,
1976                &mut wasm_storage,
1977                code_id,
1978                creator,
1979                admin.clone(),
1980                "label".to_owned(),
1981                1000,
1982                None,
1983            )
1984            .unwrap();
1985
1986        // init the contract
1987        let info = message_info(&admin, &[]);
1988        let init_msg = to_json_vec(&Empty {}).unwrap();
1989        let res = wasm_keeper
1990            .call_instantiate(
1991                contract_addr.clone(),
1992                &api,
1993                &mut wasm_storage,
1994                &mock_router(),
1995                &block,
1996                info,
1997                init_msg,
1998            )
1999            .unwrap();
2000        assert_eq!(0, res.messages.len());
2001
2002        assert_admin(
2003            &wasm_storage,
2004            &wasm_keeper,
2005            &contract_addr,
2006            Some(admin.clone()),
2007        );
2008
2009        // non-admin should not be allowed to become admin on their own
2010        wasm_keeper
2011            .execute_wasm(
2012                &api,
2013                &mut wasm_storage,
2014                &mock_router(),
2015                &block,
2016                normal_user.clone(),
2017                WasmMsg::UpdateAdmin {
2018                    contract_addr: contract_addr.to_string(),
2019                    admin: normal_user.to_string(),
2020                },
2021            )
2022            .unwrap_err();
2023
2024        // should still be admin
2025        assert_admin(
2026            &wasm_storage,
2027            &wasm_keeper,
2028            &contract_addr,
2029            Some(admin.clone()),
2030        );
2031
2032        // admin should be allowed to transfer administration permissions
2033        let res = wasm_keeper
2034            .execute_wasm(
2035                &api,
2036                &mut wasm_storage,
2037                &mock_router(),
2038                &block,
2039                admin,
2040                WasmMsg::UpdateAdmin {
2041                    contract_addr: contract_addr.to_string(),
2042                    admin: new_admin.to_string(),
2043                },
2044            )
2045            .unwrap();
2046        assert_eq!(res.events.len(), 0);
2047
2048        // new_admin should now be admin
2049        assert_admin(
2050            &wasm_storage,
2051            &wasm_keeper,
2052            &contract_addr,
2053            Some(new_admin.clone()),
2054        );
2055
2056        // new_admin should now be able to clear to admin
2057        let res = wasm_keeper
2058            .execute_wasm(
2059                &api,
2060                &mut wasm_storage,
2061                &mock_router(),
2062                &block,
2063                new_admin,
2064                WasmMsg::ClearAdmin {
2065                    contract_addr: contract_addr.to_string(),
2066                },
2067            )
2068            .unwrap();
2069        assert_eq!(res.events.len(), 0);
2070
2071        // should have no admin now
2072        assert_admin(&wasm_storage, &wasm_keeper, &contract_addr, None);
2073    }
2074
2075    #[test]
2076    fn uses_simple_address_generator_by_default() {
2077        let api = MockApi::default();
2078        let mut wasm_keeper = wasm_keeper();
2079        let creator_addr = api.addr_make("creator");
2080        let code_id = wasm_keeper.store_code(creator_addr.clone(), payout::contract());
2081        assert_eq!(1, code_id);
2082
2083        let mut wasm_storage = MockStorage::new();
2084
2085        let admin = api.addr_make("admin");
2086        let contract_addr = wasm_keeper
2087            .register_contract(
2088                &api,
2089                &mut wasm_storage,
2090                code_id,
2091                creator_addr.clone(),
2092                admin.clone(),
2093                "label".to_owned(),
2094                1000,
2095                None,
2096            )
2097            .unwrap();
2098
2099        assert_eq!(
2100            contract_addr.as_str(),
2101            "cosmwasm1mzdhwvvh22wrt07w59wxyd58822qavwkx5lcej7aqfkpqqlhaqfsgn6fq2",
2102            "default address generator returned incorrect address"
2103        );
2104
2105        let salt = HexBinary::from_hex("c0ffee").unwrap();
2106
2107        let contract_addr = wasm_keeper
2108            .register_contract(
2109                &api,
2110                &mut wasm_storage,
2111                code_id,
2112                creator_addr.clone(),
2113                admin.clone(),
2114                "label".to_owned(),
2115                1000,
2116                Binary::from(salt.clone()),
2117            )
2118            .unwrap();
2119
2120        assert_eq!(
2121            contract_addr.as_str(),
2122            "cosmwasm1drhu6t78wacgm5qjzs4hvkv9fd9awa9henw7fh6vmzrhf7k2nkjsg3flns",
2123            "default address generator returned incorrect address"
2124        );
2125
2126        let code_id = wasm_keeper.store_code(creator_addr, payout::contract());
2127        assert_eq!(2, code_id);
2128
2129        let user_addr = api.addr_make("boobaz");
2130
2131        let contract_addr = wasm_keeper
2132            .register_contract(
2133                &api,
2134                &mut wasm_storage,
2135                code_id,
2136                user_addr,
2137                admin,
2138                "label".to_owned(),
2139                1000,
2140                Binary::from(salt),
2141            )
2142            .unwrap();
2143
2144        assert_eq!(
2145            contract_addr.as_str(),
2146            "cosmwasm13cfeertf2gny0rzp5jwqzst8crmfgvcd2lq5su0c9z66yxa45qdsdd0uxc",
2147            "default address generator returned incorrect address"
2148        );
2149    }
2150
2151    struct TestAddressGenerator {
2152        address: Addr,
2153        predictable_address: Addr,
2154    }
2155
2156    impl AddressGenerator for TestAddressGenerator {
2157        fn contract_address(
2158            &self,
2159            _api: &dyn Api,
2160            _storage: &mut dyn Storage,
2161            _code_id: u64,
2162            _instance_id: u64,
2163        ) -> AnyResult<Addr> {
2164            Ok(self.address.clone())
2165        }
2166
2167        fn predictable_contract_address(
2168            &self,
2169            _api: &dyn Api,
2170            _storage: &mut dyn Storage,
2171            _code_id: u64,
2172            _instance_id: u64,
2173            _checksum: &[u8],
2174            _creator: &CanonicalAddr,
2175            _salt: &[u8],
2176        ) -> AnyResult<Addr> {
2177            Ok(self.predictable_address.clone())
2178        }
2179    }
2180
2181    #[test]
2182    fn can_use_custom_address_generator() {
2183        let api = MockApi::default();
2184        let expected_addr = api.addr_make("address");
2185        let expected_predictable_addr = api.addr_make("predictable_address");
2186        let mut wasm_keeper: WasmKeeper<Empty, Empty> =
2187            WasmKeeper::new().with_address_generator(TestAddressGenerator {
2188                address: expected_addr.clone(),
2189                predictable_address: expected_predictable_addr.clone(),
2190            });
2191        let creator = api.addr_make("creator");
2192        let code_id = wasm_keeper.store_code(creator.clone(), payout::contract());
2193
2194        let mut wasm_storage = MockStorage::new();
2195
2196        let admin = api.addr_make("admin");
2197        let contract_addr = wasm_keeper
2198            .register_contract(
2199                &api,
2200                &mut wasm_storage,
2201                code_id,
2202                creator.clone(),
2203                admin.clone(),
2204                "label".to_owned(),
2205                1000,
2206                None,
2207            )
2208            .unwrap();
2209
2210        assert_eq!(
2211            contract_addr, expected_addr,
2212            "custom address generator returned incorrect address"
2213        );
2214
2215        let contract_addr = wasm_keeper
2216            .register_contract(
2217                &api,
2218                &mut wasm_storage,
2219                code_id,
2220                creator,
2221                admin,
2222                "label".to_owned(),
2223                1000,
2224                Binary::from(HexBinary::from_hex("23A74B8C").unwrap()),
2225            )
2226            .unwrap();
2227
2228        assert_eq!(
2229            contract_addr, expected_predictable_addr,
2230            "custom address generator returned incorrect address"
2231        );
2232    }
2233}