clone_cw_multi_test/wasm_emulation/storage/
analyzer.rs

1use crate::{
2    prefixed_storage::{decode_length, to_length_prefixed, CONTRACT_STORAGE_PREFIX},
3    wasm_emulation::channel::RemoteChannel,
4    BankKeeper, Distribution, Gov, Ibc, Module, Staking, WasmKeeper,
5};
6use cosmwasm_std::{Addr, Api, Coin, CustomMsg, CustomQuery, Storage};
7use cw_orch::prelude::BankQuerier;
8use cw_utils::NativeBalance;
9use rustc_serialize::json::Json;
10use serde::de::DeserializeOwned;
11use serde::Serialize;
12use treediff::diff;
13use treediff::tools::Recorder;
14
15use crate::wasm::NAMESPACE_WASM;
16
17use crate::{wasm_emulation::input::QuerierStorage, App};
18
19use anyhow::Result as AnyResult;
20
21#[derive(Serialize)]
22pub struct SerializableCoin {
23    amount: String,
24    denom: String,
25}
26
27pub struct StorageAnalyzer {
28    pub storage: QuerierStorage,
29    pub remote: RemoteChannel,
30}
31
32pub type CustomWasmKeeper<CustomT> =
33    WasmKeeper<<CustomT as Module>::ExecT, <CustomT as Module>::QueryT>;
34
35impl StorageAnalyzer {
36    pub fn new<ApiT, StorageT, CustomT, StakingT, DistrT, IbcT, GovT>(
37        app: &App<
38            BankKeeper,
39            ApiT,
40            StorageT,
41            CustomT,
42            CustomWasmKeeper<CustomT>,
43            StakingT,
44            DistrT,
45            IbcT,
46            GovT,
47        >,
48    ) -> AnyResult<Self>
49    where
50        CustomT::ExecT: CustomMsg + DeserializeOwned + 'static,
51        CustomT::QueryT: CustomQuery + DeserializeOwned + 'static,
52
53        ApiT: Api,
54        StorageT: Storage,
55        CustomT: Module,
56        StakingT: Staking,
57        DistrT: Distribution,
58        IbcT: Ibc,
59        GovT: Gov,
60    {
61        Ok(Self {
62            storage: app.get_querier_storage()?,
63            remote: app.remote.clone(),
64        })
65    }
66
67    pub fn get_contract_storage(
68        &self,
69        contract_addr: impl Into<String>,
70    ) -> Vec<(Vec<u8>, Vec<u8>)> {
71        self.storage
72            .wasm
73            .get_contract_storage(&Addr::unchecked(contract_addr.into()))
74    }
75
76    pub fn readable_storage(&self, contract_addr: impl Into<String>) -> Vec<(String, String)> {
77        self.storage
78            .wasm
79            .get_contract_storage(&Addr::unchecked(contract_addr.into()))
80            .into_iter()
81            .map(|(key, value)| {
82                (
83                    String::from_utf8_lossy(&key).to_string(),
84                    String::from_utf8_lossy(&value).to_string(),
85                )
86            })
87            .collect()
88    }
89
90    /// We leverage the data structure we introduced for contracts to get their storage easily
91    pub fn all_contract_storage(&self) -> Vec<(String, Vec<u8>, Vec<u8>)> {
92        // In all wasm storage keys, we look for the `contract_addr/(...)/` pattern in the key
93        self.storage
94            .wasm
95            .storage
96            .iter()
97            .filter(|(key, _)| {
98                // The key must contain the NAMESPACE_WASM prefix
99                let prefix = to_length_prefixed(NAMESPACE_WASM);
100                key.len() >= prefix.len() && key[..prefix.len()] == prefix
101            })
102            .filter_map(|(key, value)| {
103                // Now we need to get the contract addr from the namespace
104
105                let prefix_len = to_length_prefixed(NAMESPACE_WASM).len();
106                let resulting_key = &key[prefix_len..];
107                let addr_len: usize = decode_length([resulting_key[0], resulting_key[1]])
108                    .try_into()
109                    .unwrap();
110                let contract_addr_addr =
111                    String::from_utf8_lossy(&resulting_key[2..(addr_len + 2)]).to_string();
112
113                let split: Vec<_> = contract_addr_addr.split('/').collect();
114                if split.len() != 2 || format!("{}/", split[0]) != CONTRACT_STORAGE_PREFIX {
115                    return None;
116                }
117
118                Some((
119                    split[1].to_string(),
120                    resulting_key[addr_len + 2..].to_vec(),
121                    value.clone(),
122                ))
123            })
124            .collect()
125    }
126
127    pub fn all_readable_contract_storage(&self) -> Vec<(String, String, String)> {
128        self.all_contract_storage()
129            .into_iter()
130            .map(|(contract, key, value)| {
131                (
132                    contract,
133                    String::from_utf8_lossy(&key).to_string(),
134                    String::from_utf8_lossy(&value).to_string(),
135                )
136            })
137            .collect()
138    }
139
140    pub fn compare_all_readable_contract_storage(&self) {
141        let wasm_querier = cw_orch::daemon::queriers::CosmWasm::new_sync(
142            self.remote.channel.clone(),
143            &self.remote.rt,
144        );
145        self.all_contract_storage()
146            .into_iter()
147            .for_each(|(contract_addr, key, value)| {
148                // We look for the data at that key on the contract
149                let distant_data = self.remote.rt.block_on(
150                    wasm_querier
151                        ._contract_raw_state(&Addr::unchecked(contract_addr.clone()), key.clone()),
152                );
153
154                if let Ok(data) = distant_data {
155                    let local_json: Json =
156                        if let Ok(v) = String::from_utf8_lossy(&value).to_string().parse() {
157                            v
158                        } else {
159                            log::info!(
160                                "Storage at {}, and key {}, was : {:x?}, now {:x?}",
161                                contract_addr,
162                                String::from_utf8_lossy(&key).to_string(),
163                                data.data,
164                                value
165                            );
166                            return;
167                        };
168                    let distant_json: Json =
169                        if let Ok(v) = String::from_utf8_lossy(&data.data).to_string().parse() {
170                            v
171                        } else {
172                            log::info!(
173                                "Storage at {}, and key {}, was : {:x?}, now {:x?}",
174                                contract_addr,
175                                String::from_utf8_lossy(&key).to_string(),
176                                data.data,
177                                value
178                            );
179                            return;
180                        };
181
182                    let mut d = Recorder::default();
183                    diff(&distant_json, &local_json, &mut d);
184
185                    let changes: Vec<_> = d
186                        .calls
187                        .iter()
188                        .filter(|change| {
189                            !matches!(change, treediff::tools::ChangeType::Unchanged(..))
190                        })
191                        .collect();
192
193                    log::info!(
194                        "Storage at {}, and key {}, changed like so : {:?}",
195                        contract_addr,
196                        String::from_utf8_lossy(&key).to_string(),
197                        changes
198                    );
199                } else if let Ok(v) = String::from_utf8_lossy(&value).to_string().parse::<Json>() {
200                    log::info!(
201                        "Storage at {}, and key {}, is new : {}",
202                        contract_addr,
203                        String::from_utf8_lossy(&key).to_string(),
204                        v
205                    );
206                } else {
207                    log::info!(
208                        "Storage at {}, and key {}, is new : {:?}",
209                        contract_addr,
210                        String::from_utf8_lossy(&key).to_string(),
211                        value
212                    );
213                }
214            });
215    }
216
217    pub fn get_balance(&self, addr: impl Into<String>) -> Vec<Coin> {
218        let addr: String = addr.into();
219        self.storage
220            .bank
221            .storage
222            .iter()
223            .find(|(a, _)| a.as_str() == addr)
224            .map(|(_, b)| b.0.clone())
225            .unwrap_or(vec![])
226    }
227
228    pub fn compare_all_balances(&self) {
229        let bank_querier = cw_orch::daemon::queriers::Bank {
230            channel: self.remote.channel.clone(),
231            rt_handle: Some(self.remote.rt.clone()),
232        };
233        self.get_all_local_balances()
234            .into_iter()
235            .for_each(|(addr, balances)| {
236                // We look for the data at that key on the contract
237                let distant_data = bank_querier.balance(&addr, None);
238
239                if let Ok(distant_coins) = distant_data {
240                    let distant_coins = serde_json::to_string(&distant_coins).unwrap();
241                    let distant_coins: Json = distant_coins.parse().unwrap();
242
243                    let local_coins = serde_json::to_string(&balances.0).unwrap();
244                    let local_coins: Json = local_coins.parse().unwrap();
245
246                    let mut d = Recorder::default();
247                    diff(&distant_coins, &local_coins, &mut d);
248
249                    let changes: Vec<_> = d
250                        .calls
251                        .iter()
252                        .filter(|change| {
253                            !matches!(change, treediff::tools::ChangeType::Unchanged(..))
254                        })
255                        .collect();
256
257                    log::info!("Bank balance for {} changed like so : {:?}", addr, changes);
258                }
259            });
260    }
261
262    pub fn get_all_local_balances(&self) -> Vec<(Addr, NativeBalance)> {
263        self.storage.bank.storage.clone()
264    }
265}
clone_cw_multi_test/wasm_emulation/storage/analyzer.rs

clone_cw_multi_test/wasm_emulation/storage/
analyzer.rs
