1use anyhow::{
18 anyhow,
19 Result,
20};
21use contract_transcode::{
22 ContractMessageTranscoder,
23 Value,
24};
25use ink_env::Environment;
26use ink_metadata::layout::{
27 Layout,
28 StructLayout,
29};
30use itertools::Itertools;
31use scale::{
32 Decode,
33 Encode,
34};
35use scale_info::{
36 form::PortableForm,
37 Type,
38};
39use serde::{
40 Serialize,
41 Serializer,
42};
43use sp_core::{
44 hexdisplay::AsBytesRef,
45 storage::ChildInfo,
46};
47use std::{
48 collections::BTreeMap,
49 fmt::{
50 self,
51 Display,
52 Formatter,
53 },
54 marker::PhantomData,
55};
56use subxt::{
57 backend::{
58 legacy::{
59 rpc_methods::Bytes,
60 LegacyRpcMethods,
61 },
62 rpc::{
63 rpc_params,
64 RpcClient,
65 },
66 },
67 ext::scale_decode::IntoVisitor,
68 Config,
69 OnlineClient,
70};
71
72use super::{
73 fetch_contract_info,
74 url_to_string,
75 ContractInfo,
76 TrieId,
77};
78
79pub struct ContractStorage<C: Config, E: Environment> {
80 rpc: ContractStorageRpc<C>,
81 _phantom: PhantomData<fn() -> E>,
82}
83
84impl<C: Config, E: Environment> ContractStorage<C, E>
85where
86 C::AccountId: AsRef<[u8]> + Display + IntoVisitor,
87 C::Hash: IntoVisitor,
88 E::Balance: IntoVisitor + Serialize,
89{
90 pub fn new(rpc: ContractStorageRpc<C>) -> Self {
91 Self {
92 rpc,
93 _phantom: Default::default(),
94 }
95 }
96
97 pub async fn version(&self) -> Result<u16> {
101 self.rpc
102 .client
103 .storage()
104 .at_latest()
105 .await?
106 .storage_version("Contracts")
107 .await
108 .map_err(|e| {
109 anyhow!("The storage version for the contracts pallet could not be determined: {e}")
110 })
111 }
112
113 pub async fn load_contract_storage_data(
115 &self,
116 contract_account: &C::AccountId,
117 ) -> Result<ContractStorageData> {
118 let contract_info = self.rpc.fetch_contract_info::<E>(contract_account).await?;
119 let trie_id = contract_info.trie_id();
120
121 let mut storage_keys = Vec::new();
122 let mut storage_values = Vec::new();
123 const KEYS_COUNT: u32 = 1000;
124 loop {
125 let mut keys = self
126 .rpc
127 .fetch_storage_keys_paged(
128 trie_id,
129 None,
130 KEYS_COUNT,
131 storage_keys.last().map(|k: &Bytes| k.as_bytes_ref()),
132 None,
133 )
134 .await?;
135 let keys_count = keys.len();
136 let mut values = self.rpc.fetch_storage_entries(trie_id, &keys, None).await?;
137 assert_eq!(
138 keys_count,
139 values.len(),
140 "storage keys and values must be the same length"
141 );
142 storage_keys.append(&mut keys);
143 storage_values.append(&mut values);
144
145 if (keys_count as u32) < KEYS_COUNT {
146 break
147 }
148 }
149
150 let storage = storage_keys
151 .into_iter()
152 .zip(storage_values.into_iter())
153 .filter_map(|(key, value)| value.map(|v| (key, v)))
154 .collect();
155
156 let contract_storage = ContractStorageData(storage);
157 Ok(contract_storage)
158 }
159
160 pub async fn load_contract_storage_with_layout(
161 &self,
162 contract_account: &C::AccountId,
163 decoder: &ContractMessageTranscoder,
164 ) -> Result<ContractStorageLayout> {
165 let data = self.load_contract_storage_data(contract_account).await?;
166 ContractStorageLayout::new(data, decoder)
167 }
168}
169
170#[derive(Serialize, Debug)]
172pub struct ContractStorageData(BTreeMap<Bytes, Bytes>);
173
174impl ContractStorageData {
175 pub fn new(data: BTreeMap<Bytes, Bytes>) -> Self {
177 Self(data)
178 }
179}
180
181#[derive(Serialize, Debug)]
183pub struct RootKeyEntry {
184 #[serde(serialize_with = "RootKeyEntry::key_as_hex")]
185 pub root_key: u32,
186 pub path: Vec<String>,
187 pub type_id: u32,
188}
189
190impl RootKeyEntry {
191 fn key_as_hex<S>(key: &u32, serializer: S) -> Result<S::Ok, S::Error>
192 where
193 S: Serializer,
194 {
195 serializer.serialize_str(format!("0x{}", hex::encode(key.encode())).as_str())
196 }
197}
198
199#[derive(Serialize, Debug)]
200pub struct Mapping {
201 #[serde(flatten)]
202 root: RootKeyEntry,
203 map: Vec<(Value, Value)>,
204}
205
206impl Mapping {
207 pub fn new(root: RootKeyEntry, value: Vec<(Value, Value)>) -> Mapping {
209 Mapping { root, map: value }
210 }
211
212 pub fn root(&self) -> &RootKeyEntry {
214 &self.root
215 }
216
217 pub fn iter(&self) -> impl DoubleEndedIterator<Item = &(Value, Value)> {
219 self.map.iter()
220 }
221}
222
223impl Display for Mapping {
224 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
225 let len = self.map.len();
226 for (i, e) in self.map.iter().enumerate() {
227 write!(f, "Mapping {{ {} => {} }}", e.0, e.1)?;
228 if i + 1 < len {
229 writeln!(f)?;
230 }
231 }
232 Ok(())
233 }
234}
235
236#[derive(Serialize, Debug)]
237pub struct Lazy {
238 #[serde(flatten)]
239 root: RootKeyEntry,
240 value: Value,
241}
242
243impl Lazy {
244 pub fn new(root: RootKeyEntry, value: Value) -> Lazy {
246 Lazy { root, value }
247 }
248
249 pub fn root(&self) -> &RootKeyEntry {
251 &self.root
252 }
253
254 pub fn value(&self) -> &Value {
256 &self.value
257 }
258}
259
260impl Display for Lazy {
261 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
262 write!(f, "Lazy {{ {} }}", self.value)
263 }
264}
265
266#[derive(Serialize, Debug)]
267pub struct StorageVec {
268 #[serde(flatten)]
269 root: RootKeyEntry,
270 len: u32,
271 vec: Vec<Value>,
272}
273
274impl StorageVec {
275 pub fn new(root: RootKeyEntry, len: u32, value: Vec<Value>) -> StorageVec {
277 StorageVec {
278 root,
279 len,
280 vec: value,
281 }
282 }
283
284 pub fn root(&self) -> &RootKeyEntry {
286 &self.root
287 }
288
289 pub fn len(&self) -> u32 {
291 self.len
292 }
293
294 pub fn values(&self) -> impl Iterator<Item = &Value> {
296 self.vec.iter()
297 }
298}
299
300impl Display for StorageVec {
301 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
302 for (i, v) in self.vec.iter().enumerate() {
303 write!(f, "StorageVec [{}] {{ [{}] => {} }}", self.len, i, v)?;
304 if i + 1 < self.len as usize {
305 writeln!(f)?;
306 }
307 }
308 Ok(())
309 }
310}
311
312#[derive(Serialize, Debug)]
313pub struct Packed {
314 #[serde(flatten)]
315 root: RootKeyEntry,
316 value: Value,
317}
318
319impl Packed {
320 pub fn new(root: RootKeyEntry, value: Value) -> Packed {
322 Packed { root, value }
323 }
324
325 pub fn root(&self) -> &RootKeyEntry {
327 &self.root
328 }
329
330 pub fn value(&self) -> &Value {
332 &self.value
333 }
334}
335
336impl Display for Packed {
337 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
338 write!(f, "{}", self.value)
339 }
340}
341
342#[derive(Serialize, Debug)]
344pub enum ContractStorageCell {
345 Mapping(Mapping),
346 Lazy(Lazy),
347 StorageVec(StorageVec),
348 Packed(Packed),
349}
350
351impl ContractStorageCell {
352 fn root(&self) -> &RootKeyEntry {
353 match self {
354 Self::Mapping(mapping) => mapping.root(),
355 Self::Lazy(lazy) => lazy.root(),
356 Self::StorageVec(storage_vec) => storage_vec.root(),
357 Self::Packed(packed) => packed.root(),
358 }
359 }
360
361 pub fn path(&self) -> String {
363 self.root().path.join("::")
364 }
365
366 pub fn parent(&self) -> String {
368 self.root().path.last().cloned().unwrap_or_default()
369 }
370
371 pub fn root_key(&self) -> String {
373 hex::encode(self.root().root_key.encode())
374 }
375}
376
377impl Display for ContractStorageCell {
378 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
379 match self {
380 Self::Mapping(mapping) => mapping.fmt(f),
381 Self::Lazy(lazy) => lazy.fmt(f),
382 Self::StorageVec(storage_vec) => storage_vec.fmt(f),
383 Self::Packed(value) => value.fmt(f),
384 }
385 }
386}
387
388#[derive(Serialize, Debug)]
390pub struct ContractStorageLayout {
391 cells: Vec<ContractStorageCell>,
392}
393
394impl ContractStorageLayout {
395 pub fn new(
398 data: ContractStorageData,
399 decoder: &ContractMessageTranscoder,
400 ) -> Result<Self> {
401 let layout = decoder.metadata().layout();
402 let registry = decoder.metadata().registry();
403 let mut path_stack = vec!["root".to_string()];
404 let mut root_key_entries: Vec<RootKeyEntry> = Vec::new();
405 Self::collect_root_key_entries(layout, &mut path_stack, &mut root_key_entries);
406
407 let mut cells = data
408 .0
409 .into_iter()
410 .map(|(key, value)| {
411 let (root_key, mapping_key) = Self::key_parts(&key);
412 (root_key, (mapping_key, value))
413 })
414 .into_group_map()
415 .into_iter()
416 .map(|(root_key, mut data)| {
417 let root_key_entry = root_key_entries
418 .iter()
419 .find(|e| e.root_key == root_key)
420 .ok_or(anyhow!(
421 "Root key {} not found for the RootLayout",
422 root_key
423 ))?;
424 let type_def = registry.resolve(root_key_entry.type_id).ok_or(
425 anyhow!("Type {} not found in the registry", root_key_entry.type_id),
426 )?;
427 let root = RootKeyEntry {
428 path: root_key_entry.path.clone(),
429 type_id: root_key_entry.type_id,
430 root_key,
431 };
432 match type_def.path.to_string().as_str() {
433 "ink_storage::lazy::mapping::Mapping" => {
434 let key_type_id = Self::param_type_id(type_def, "K")
435 .ok_or(anyhow!("Param `K` not found in type registry"))?;
436 let value_type_id = Self::param_type_id(type_def, "V")
437 .ok_or(anyhow!("Param `V` not found in type registry"))?;
438 let value = Self::decode_to_mapping(
439 data,
440 key_type_id,
441 value_type_id,
442 decoder,
443 )?;
444 Ok(ContractStorageCell::Mapping(Mapping::new(root, value)))
445 }
446 "ink_storage::lazy::vec::StorageVec" => {
447 data.sort_by(|a, b| a.0.cmp(&b.0));
449 let raw_len = data
451 .first()
452 .ok_or(anyhow!("Length of the StorageVec not found"))?
453 .1
454 .clone();
455 let len = u32::decode(&mut raw_len.as_bytes_ref())?;
456 let value_type_id = Self::param_type_id(type_def, "V")
457 .ok_or(anyhow!("Param `V` not found in type registry"))?;
458 let value =
459 Self::decode_to_vec(&data[1..], value_type_id, decoder)?;
460 Ok(ContractStorageCell::StorageVec(StorageVec::new(
461 root, len, value,
462 )))
463 }
464 "ink_storage::lazy::Lazy" => {
465 let value_type_id = Self::param_type_id(type_def, "V")
466 .ok_or(anyhow!("Param `V` not found in type registry"))?;
467 let raw_value =
468 data.first().ok_or(anyhow!("Empty storage cell"))?.1.clone();
469 let value = decoder
470 .decode(value_type_id, &mut raw_value.as_bytes_ref())?;
471 Ok(ContractStorageCell::Lazy(Lazy::new(root, value)))
472 }
473 _ => {
474 let raw_value =
475 data.first().ok_or(anyhow!("Empty storage cell"))?.1.clone();
476 let value = decoder
477 .decode(root.type_id, &mut raw_value.as_bytes_ref())?;
478 Ok(ContractStorageCell::Packed(Packed::new(root, value)))
479 }
480 }
481 })
482 .collect::<Result<Vec<_>>>()?;
483
484 cells.sort_by_key(|k| k.path());
485
486 Ok(Self { cells })
487 }
488
489 pub fn iter(&self) -> impl Iterator<Item = &ContractStorageCell> {
491 self.cells.iter()
492 }
493
494 fn decode_to_mapping(
495 data: Vec<(Option<Bytes>, Bytes)>,
496 key_type_id: u32,
497 value_type_id: u32,
498 decoder: &ContractMessageTranscoder,
499 ) -> Result<Vec<(Value, Value)>> {
500 data.into_iter()
501 .map(|(k, v)| {
502 let k = k.ok_or(anyhow!("The Mapping key is missing in the map"))?;
503 let key = decoder.decode(key_type_id, &mut k.as_bytes_ref())?;
504 let value = decoder.decode(value_type_id, &mut v.as_bytes_ref())?;
505 Ok((key, value))
506 })
507 .collect()
508 }
509
510 fn decode_to_vec(
511 data: &[(Option<Bytes>, Bytes)],
512 value_type_id: u32,
513 decoder: &ContractMessageTranscoder,
514 ) -> Result<Vec<Value>> {
515 data.iter()
516 .map(|(_, v)| {
517 let value = decoder.decode(value_type_id, &mut v.as_bytes_ref())?;
518 Ok(value)
519 })
520 .collect()
521 }
522
523 fn collect_root_key_entries(
524 layout: &Layout<PortableForm>,
525 path: &mut Vec<String>,
526 entries: &mut Vec<RootKeyEntry>,
527 ) {
528 match layout {
529 Layout::Root(root) => {
530 entries.push(RootKeyEntry {
531 root_key: *root.root_key().key(),
532 path: path.clone(),
533 type_id: root.ty().id,
534 });
535 Self::collect_root_key_entries(root.layout(), path, entries);
536 }
537 Layout::Struct(struct_layout) => {
538 Self::struct_entries(struct_layout, path, entries)
539 }
540 Layout::Enum(enum_layout) => {
541 path.push(enum_layout.name().to_string());
542 for (variant, struct_layout) in enum_layout.variants() {
543 path.push(variant.value().to_string());
544 Self::struct_entries(struct_layout, path, entries);
545 path.pop();
546 }
547 path.pop();
548 }
549 Layout::Hash(_) => {
550 unimplemented!("Layout::Hash is not currently be constructed")
551 }
552 Layout::Array(_) | Layout::Leaf(_) => {}
553 }
554 }
555
556 fn struct_entries(
557 struct_layout: &StructLayout<PortableForm>,
558 path: &mut Vec<String>,
559 entries: &mut Vec<RootKeyEntry>,
560 ) {
561 let struct_label = struct_layout.name().to_string();
562 path.push(struct_label);
563 for field in struct_layout.fields() {
564 path.push(field.name().to_string());
565 Self::collect_root_key_entries(field.layout(), path, entries);
566 path.pop();
567 }
568 path.pop();
569 }
570
571 fn key_parts(key: &Bytes) -> (u32, Option<Bytes>) {
577 assert!(key.0.len() >= 20, "key must be at least 20 bytes");
578 let mut root_key_bytes = [0u8; 4];
579 root_key_bytes.copy_from_slice(&key.0[16..20]);
580
581 let root_key = <u32 as scale::Decode>::decode(&mut &root_key_bytes[..])
583 .expect("root key is 4 bytes, it always decodes successfully to a u32; qed");
584
585 let mapping_key = if key.0.len() > 20 {
586 Some(Bytes::from(key.0[20..].to_vec()))
587 } else {
588 None
589 };
590
591 (root_key, mapping_key)
592 }
593
594 fn param_type_id(type_def: &Type<PortableForm>, param_name: &str) -> Option<u32> {
596 Some(
597 type_def
598 .type_params
599 .iter()
600 .find(|&e| e.name == param_name)?
601 .ty?
602 .id,
603 )
604 }
605}
606
607pub struct ContractStorageRpc<C: Config> {
609 rpc_client: RpcClient,
610 rpc_methods: LegacyRpcMethods<C>,
611 client: OnlineClient<C>,
612}
613
614impl<C: Config> ContractStorageRpc<C>
615where
616 C::AccountId: AsRef<[u8]> + Display + IntoVisitor,
617 C::Hash: IntoVisitor,
618{
619 pub async fn new(url: &url::Url) -> Result<Self> {
621 let rpc_client = RpcClient::from_url(url_to_string(url)).await?;
622 let client = OnlineClient::from_rpc_client(rpc_client.clone()).await?;
623 let rpc_methods = LegacyRpcMethods::new(rpc_client.clone());
624
625 Ok(Self {
626 rpc_client,
627 rpc_methods,
628 client,
629 })
630 }
631
632 pub async fn fetch_contract_info<E: Environment>(
634 &self,
635 contract: &C::AccountId,
636 ) -> Result<ContractInfo<C::Hash, E::Balance>>
637 where
638 E::Balance: IntoVisitor,
639 {
640 fetch_contract_info::<C, E>(contract, &self.rpc_methods, &self.client).await
641 }
642
643 pub async fn fetch_contract_storage(
647 &self,
648 trie_id: &TrieId,
649 key: &Bytes,
650 block_hash: Option<C::Hash>,
651 ) -> Result<Option<Bytes>> {
652 let child_storage_key =
653 ChildInfo::new_default(trie_id.as_ref()).into_prefixed_storage_key();
654 let params = rpc_params![child_storage_key, key, block_hash];
655 let data: Option<Bytes> = self
656 .rpc_client
657 .request("childstate_getStorage", params)
658 .await?;
659 Ok(data)
660 }
661
662 pub async fn fetch_storage_keys_paged(
664 &self,
665 trie_id: &TrieId,
666 prefix: Option<&[u8]>,
667 count: u32,
668 start_key: Option<&[u8]>,
669 block_hash: Option<C::Hash>,
670 ) -> Result<Vec<Bytes>> {
671 let child_storage_key =
672 ChildInfo::new_default(trie_id.as_ref()).into_prefixed_storage_key();
673 let prefix_hex = prefix.map(|p| format!("0x{}", hex::encode(p)));
674 let start_key_hex = start_key.map(|k| format!("0x{}", hex::encode(k)));
675 let params = rpc_params![
676 child_storage_key,
677 prefix_hex,
678 count,
679 start_key_hex,
680 block_hash
681 ];
682 let data: Vec<Bytes> = self
683 .rpc_client
684 .request("childstate_getKeysPaged", params)
685 .await?;
686 Ok(data)
687 }
688
689 pub async fn fetch_storage_entries(
691 &self,
692 trie_id: &TrieId,
693 keys: &[Bytes],
694 block_hash: Option<C::Hash>,
695 ) -> Result<Vec<Option<Bytes>>> {
696 let child_storage_key =
697 ChildInfo::new_default(trie_id.as_ref()).into_prefixed_storage_key();
698 let params = rpc_params![child_storage_key, keys, block_hash];
699 let data: Vec<Option<Bytes>> = self
700 .rpc_client
701 .request("childstate_getStorageEntries", params)
702 .await?;
703 Ok(data)
704 }
705}