clone_cw_multi_test/wasm_emulation/query/

wasm.rs

use std::marker::PhantomData;

use crate::addons::MockApiBech32;
use crate::prefixed_storage::get_full_contract_storage_namespace;
use crate::queries::wasm::WasmRemoteQuerier;
use crate::wasm_emulation::query::gas::{
    GAS_COST_ALL_QUERIES, GAS_COST_CONTRACT_INFO, GAS_COST_RAW_COSMWASM_QUERY,
};
use crate::wasm_emulation::query::mock_querier::QueryResultWithGas;
use crate::wasm_emulation::query::MockQuerier;
use crate::Contract;

use crate::wasm_emulation::contract::WasmContract;
use cosmwasm_std::testing::MockStorage;
use cosmwasm_vm::GasInfo;

use cosmwasm_std::{
    to_json_binary, Addr, ContractInfoResponse, CustomMsg, CustomQuery, OwnedDeps, Storage,
    SystemError, SystemResult,
};
use cosmwasm_std::{ContractInfo, ContractResult};

use cosmwasm_std::WasmQuery;
use cw_orch::daemon::queriers::CosmWasm;
use serde::de::DeserializeOwned;

use crate::wasm_emulation::channel::RemoteChannel;

use super::mock_querier::ForkState;

pub struct WasmQuerier<
    ExecC: CustomMsg + DeserializeOwned + 'static,
    QueryC: CustomQuery + DeserializeOwned + 'static,
> {
    fork_state: ForkState<ExecC, QueryC>,
}

impl<
        ExecC: CustomMsg + DeserializeOwned + 'static,
        QueryC: CustomQuery + DeserializeOwned + 'static,
    > WasmQuerier<ExecC, QueryC>
{
    pub fn new(fork_state: ForkState<ExecC, QueryC>) -> Self {
        Self { fork_state }
    }

    pub fn query(&self, remote: RemoteChannel, request: &WasmQuery) -> QueryResultWithGas {
        match request {
            WasmQuery::ContractInfo { contract_addr } => {
                let addr = Addr::unchecked(contract_addr);
                let data = if let Some(local_contract) = self
                    .fork_state
                    .querier_storage
                    .wasm
                    .contracts
                    .get(contract_addr)
                {
                    local_contract.clone()
                } else {
                    let maybe_distant_contract = WasmRemoteQuerier::load_distant_contract(
                        self.fork_state.remote.clone(),
                        &addr,
                    );

                    if let Ok(contract) = maybe_distant_contract {
                        contract
                    } else {
                        return (
                            SystemResult::Err(SystemError::NoSuchContract {
                                addr: contract_addr.to_string(),
                            }),
                            GasInfo::with_externally_used(GAS_COST_CONTRACT_INFO),
                        );
                    }
                };
                let response =
                    ContractInfoResponse::new(data.code_id, data.creator, data.admin, false, None);
                (
                    SystemResult::Ok(to_json_binary(&response).into()),
                    GasInfo::with_externally_used(GAS_COST_CONTRACT_INFO),
                )
            }
            WasmQuery::Raw { contract_addr, key } => {
                // We first try to load that information locally
                let mut total_key =
                    get_full_contract_storage_namespace(&Addr::unchecked(contract_addr)).to_vec();
                total_key.extend_from_slice(key);

                if let Some(value) = self
                    .fork_state
                    .querier_storage
                    .wasm
                    .storage
                    .iter()
                    .find(|e| e.0 == total_key)
                {
                    (
                        SystemResult::Ok(ContractResult::Ok(value.1.clone().into())),
                        GasInfo::with_externally_used(GAS_COST_RAW_COSMWASM_QUERY),
                    )
                } else {
                    (
                        SystemResult::Ok(
                            WasmRemoteQuerier::raw_query(
                                remote,
                                &Addr::unchecked(contract_addr),
                                key.clone(),
                            )
                            .map(Into::into)
                            .into(),
                        ),
                        GasInfo::with_externally_used(GAS_COST_RAW_COSMWASM_QUERY),
                    )
                }
            }
            WasmQuery::Smart { contract_addr, msg } => {
                let addr = Addr::unchecked(contract_addr);

                let mut storage = MockStorage::default();
                // Set the storage
                for (key, value) in self
                    .fork_state
                    .querier_storage
                    .wasm
                    .get_contract_storage(&addr)
                {
                    storage.set(&key, &value);
                }

                let deps = OwnedDeps {
                    storage,
                    api: MockApiBech32::new(&self.fork_state.remote.pub_address_prefix),
                    querier: MockQuerier::new(self.fork_state.clone()),
                    custom_query_type: PhantomData::<QueryC>,
                };
                let mut env = self.fork_state.local_state.env.clone();
                env.contract = ContractInfo {
                    address: addr.clone(),
                };

                // First we get the code id corresponding to the contract
                let code_id = if let Some(local_contract) = self
                    .fork_state
                    .querier_storage
                    .wasm
                    .contracts
                    .get(contract_addr)
                {
                    local_contract.code_id
                } else {
                    let wasm_querier = CosmWasm::new_sync(remote.channel.clone(), &remote.rt);

                    let code_info = remote
                        .rt
                        .block_on(wasm_querier._contract_info(&addr))
                        .unwrap();

                    code_info.code_id
                };

                // Then, we get the corresponding wasm contract

                // If the contract data is already defined in our storage, we load it from there
                let result = if let Some(code) = self
                    .fork_state
                    .querier_storage
                    .wasm
                    .codes
                    .get(&(code_id as usize))
                {
                    // Local Wasm Contract case
                    <WasmContract as Contract<ExecC, QueryC>>::query(
                        code,
                        deps.as_ref(),
                        env,
                        msg.to_vec(),
                        self.fork_state.clone(),
                    )
                } else if let Some(local_contract) = self
                    .fork_state
                    .local_state
                    .contracts
                    .get(&(code_id as usize))
                {
                    // Local Rust Contract case
                    unsafe {
                        local_contract.as_ref().unwrap().query(
                            deps.as_ref(),
                            env,
                            msg.to_vec(),
                            self.fork_state.clone(),
                        )
                    }
                } else {
                    // Distant Registered Contract case
                    // TODO, this should be part of the cache as well
                    // However, it's not really possible to register that data inside the App, because this is deep in the execution layer
                    <WasmContract as Contract<ExecC, QueryC>>::query(
                        &WasmContract::new_distant_code_id(code_id, remote.clone()),
                        deps.as_ref(),
                        env,
                        msg.to_vec(),
                        self.fork_state.clone(),
                    )
                };

                let result = match result {
                    Err(e) => {
                        return (
                            SystemResult::Err(SystemError::InvalidRequest {
                                error: format!("Error querying a contract: {}", e),
                                request: msg.clone(),
                            }),
                            GasInfo::with_externally_used(0),
                        )
                    }
                    Ok(result) => result,
                };

                (
                    SystemResult::Ok(ContractResult::Ok(result)),
                    GasInfo::with_externally_used(GAS_COST_ALL_QUERIES),
                )
            }
            WasmQuery::CodeInfo { code_id } => {
                let code_data = self
                    .fork_state
                    .querier_storage
                    .wasm
                    .code_data
                    .get(&(*code_id as usize));
                let res = if let Some(code_data) = code_data {
                    cosmwasm_std::CodeInfoResponse::new(
                        *code_id,
                        code_data.creator.clone(),
                        code_data.checksum,
                    )
                } else {
                    let maybe_code_info =
                        WasmRemoteQuerier::code_info(self.fork_state.remote.clone(), *code_id);

                    if let Ok(code_info) = maybe_code_info {
                        code_info
                    } else {
                        return (
                            SystemResult::Err(SystemError::NoSuchCode { code_id: *code_id }),
                            GasInfo::with_externally_used(GAS_COST_CONTRACT_INFO),
                        );
                    }
                };
                (
                    SystemResult::Ok(to_json_binary(&res).into()),
                    GasInfo::with_externally_used(GAS_COST_CONTRACT_INFO),
                )
            }
            _ => unimplemented!(),
        }
    }
}