1use super::StorageInstance;
21use crate::{
22 storage::types::{
23 CountedStorageMapInstance, CountedStorageNMapInstance, Counter, KeyGenerator,
24 QueryKindTrait,
25 },
26 traits::{PartialStorageInfoTrait, StorageInfo},
27 StorageHasher,
28};
29use alloc::{vec, vec::Vec};
30use codec::{Decode, DecodeAll, FullCodec};
31use impl_trait_for_tuples::impl_for_tuples;
32use sp_core::Get;
33
34pub trait TryDecodeEntireStorage {
42 fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>>;
44}
45
46#[cfg_attr(all(not(feature = "tuples-96"), not(feature = "tuples-128")), impl_for_tuples(64))]
47#[cfg_attr(all(feature = "tuples-96", not(feature = "tuples-128")), impl_for_tuples(96))]
48#[cfg_attr(feature = "tuples-128", impl_for_tuples(128))]
49impl TryDecodeEntireStorage for Tuple {
50 fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
51 let mut errors = Vec::new();
52 let mut len = 0usize;
53
54 for_tuples!(#(
55 match Tuple::try_decode_entire_state() {
56 Ok(bytes) => len += bytes,
57 Err(errs) => errors.extend(errs),
58 }
59 )*);
60
61 if errors.is_empty() {
62 Ok(len)
63 } else {
64 Err(errors)
65 }
66 }
67}
68
69#[derive(Clone, PartialEq, Eq)]
71pub struct TryDecodeEntireStorageError {
72 pub key: Vec<u8>,
74 pub raw: Option<Vec<u8>>,
76 pub info: StorageInfo,
78}
79
80impl core::fmt::Display for TryDecodeEntireStorageError {
81 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
82 write!(
83 f,
84 "`{}::{}` key `{}` is undecodable",
85 &alloc::str::from_utf8(&self.info.pallet_name).unwrap_or("<invalid>"),
86 &alloc::str::from_utf8(&self.info.storage_name).unwrap_or("<invalid>"),
87 array_bytes::bytes2hex("0x", &self.key)
88 )
89 }
90}
91
92impl core::fmt::Debug for TryDecodeEntireStorageError {
93 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
94 write!(
95 f,
96 "key: {} value: {} info: {:?}",
97 array_bytes::bytes2hex("0x", &self.key),
98 array_bytes::bytes2hex("0x", self.raw.clone().unwrap_or_default()),
99 self.info
100 )
101 }
102}
103
104fn decode_storage_info<V: Decode>(
110 info: StorageInfo,
111) -> Result<usize, Vec<TryDecodeEntireStorageError>> {
112 let mut decoded = 0;
113
114 let decode_key = |key: &[u8]| match sp_io::storage::get(key) {
115 None => Ok(0),
116 Some(bytes) => {
117 let len = bytes.len();
118 <V as DecodeAll>::decode_all(&mut bytes.as_ref()).map_err(|_| {
119 TryDecodeEntireStorageError {
120 key: key.to_vec(),
121 raw: Some(bytes.to_vec()),
122 info: info.clone(),
123 }
124 })?;
125
126 Ok::<usize, _>(len)
127 },
128 };
129
130 let mut errors = vec![];
131 let mut next_key = Some(info.prefix.clone());
132 loop {
133 match next_key {
134 Some(key) if key.starts_with(&info.prefix) => {
135 match decode_key(&key) {
136 Ok(bytes) => {
137 decoded += bytes;
138 },
139 Err(e) => errors.push(e),
140 };
141 next_key = sp_io::storage::next_key(&key);
142 },
143 _ => break,
144 }
145 }
146
147 if errors.is_empty() {
148 Ok(decoded)
149 } else {
150 Err(errors)
151 }
152}
153
154impl<Prefix, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
155 for crate::storage::types::StorageValue<Prefix, Value, QueryKind, OnEmpty>
156where
157 Prefix: StorageInstance,
158 Value: FullCodec,
159 QueryKind: QueryKindTrait<Value, OnEmpty>,
160 OnEmpty: Get<QueryKind::Query> + 'static,
161{
162 fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
163 let info = Self::partial_storage_info()
164 .first()
165 .cloned()
166 .expect("Value has only one storage info; qed");
167 decode_storage_info::<Value>(info)
168 }
169}
170
171impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
172 for crate::storage::types::StorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues>
173where
174 Prefix: StorageInstance,
175 Hasher: StorageHasher,
176 Key: FullCodec,
177 Value: FullCodec,
178 QueryKind: QueryKindTrait<Value, OnEmpty>,
179 OnEmpty: Get<QueryKind::Query> + 'static,
180 MaxValues: Get<Option<u32>>,
181{
182 fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
183 let info = Self::partial_storage_info()
184 .first()
185 .cloned()
186 .expect("Map has only one storage info; qed");
187 decode_storage_info::<Value>(info)
188 }
189}
190
191impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
192 for crate::storage::types::CountedStorageMap<
193 Prefix,
194 Hasher,
195 Key,
196 Value,
197 QueryKind,
198 OnEmpty,
199 MaxValues,
200 >
201where
202 Prefix: CountedStorageMapInstance,
203 Hasher: StorageHasher,
204 Key: FullCodec,
205 Value: FullCodec,
206 QueryKind: QueryKindTrait<Value, OnEmpty>,
207 OnEmpty: Get<QueryKind::Query> + 'static,
208 MaxValues: Get<Option<u32>>,
209{
210 fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
211 let (map_info, counter_info) = match &Self::partial_storage_info()[..] {
212 [a, b] => (a.clone(), b.clone()),
213 _ => panic!("Counted map has two storage info items; qed"),
214 };
215 let mut decoded = decode_storage_info::<Counter>(counter_info)?;
216 decoded += decode_storage_info::<Value>(map_info)?;
217 Ok(decoded)
218 }
219}
220
221impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty, MaxValues>
222 TryDecodeEntireStorage
223 for crate::storage::types::StorageDoubleMap<
224 Prefix,
225 Hasher1,
226 Key1,
227 Hasher2,
228 Key2,
229 Value,
230 QueryKind,
231 OnEmpty,
232 MaxValues,
233 >
234where
235 Prefix: StorageInstance,
236 Hasher1: StorageHasher,
237 Key1: FullCodec,
238 Hasher2: StorageHasher,
239 Key2: FullCodec,
240 Value: FullCodec,
241 QueryKind: QueryKindTrait<Value, OnEmpty>,
242 OnEmpty: Get<QueryKind::Query> + 'static,
243 MaxValues: Get<Option<u32>>,
244{
245 fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
246 let info = Self::partial_storage_info()
247 .first()
248 .cloned()
249 .expect("Double-map has only one storage info; qed");
250 decode_storage_info::<Value>(info)
251 }
252}
253
254impl<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
255 for crate::storage::types::StorageNMap<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues>
256where
257 Prefix: StorageInstance,
258 Key: KeyGenerator,
259 Value: FullCodec,
260 QueryKind: QueryKindTrait<Value, OnEmpty>,
261 OnEmpty: Get<QueryKind::Query> + 'static,
262 MaxValues: Get<Option<u32>>,
263{
264 fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
265 let info = Self::partial_storage_info()
266 .first()
267 .cloned()
268 .expect("N-map has only one storage info; qed");
269 decode_storage_info::<Value>(info)
270 }
271}
272
273impl<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
274 for crate::storage::types::CountedStorageNMap<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues>
275where
276 Prefix: CountedStorageNMapInstance,
277 Key: KeyGenerator,
278 Value: FullCodec,
279 QueryKind: QueryKindTrait<Value, OnEmpty>,
280 OnEmpty: Get<QueryKind::Query> + 'static,
281 MaxValues: Get<Option<u32>>,
282{
283 fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
284 let (map_info, counter_info) = match &Self::partial_storage_info()[..] {
285 [a, b] => (a.clone(), b.clone()),
286 _ => panic!("Counted NMap has two storage info items; qed"),
287 };
288
289 let mut decoded = decode_storage_info::<Counter>(counter_info)?;
290 decoded += decode_storage_info::<Value>(map_info)?;
291 Ok(decoded)
292 }
293}
294
295#[cfg(test)]
296mod tests {
297 use super::*;
298 use crate::{
299 storage::types::{self, CountedStorageMapInstance, CountedStorageNMapInstance, Key},
300 Blake2_128Concat,
301 };
302
303 type H = Blake2_128Concat;
304
305 macro_rules! build_prefix {
306 ($name:ident) => {
307 struct $name;
308 impl StorageInstance for $name {
309 fn pallet_prefix() -> &'static str {
310 "test_pallet"
311 }
312 const STORAGE_PREFIX: &'static str = stringify!($name);
313 }
314 };
315 }
316
317 build_prefix!(ValuePrefix);
318 type Value = types::StorageValue<ValuePrefix, u32>;
319
320 build_prefix!(MapPrefix);
321 type Map = types::StorageMap<MapPrefix, H, u32, u32>;
322
323 build_prefix!(CMapCounterPrefix);
324 build_prefix!(CMapPrefix);
325 impl CountedStorageMapInstance for CMapPrefix {
326 type CounterPrefix = CMapCounterPrefix;
327 }
328 type CMap = types::CountedStorageMap<CMapPrefix, H, u8, u16>;
329
330 build_prefix!(DMapPrefix);
331 type DMap = types::StorageDoubleMap<DMapPrefix, H, u32, H, u32, u32>;
332
333 build_prefix!(NMapPrefix);
334 type NMap = types::StorageNMap<NMapPrefix, (Key<H, u8>, Key<H, u8>), u128>;
335
336 build_prefix!(CountedNMapCounterPrefix);
337 build_prefix!(CountedNMapPrefix);
338 impl CountedStorageNMapInstance for CountedNMapPrefix {
339 type CounterPrefix = CountedNMapCounterPrefix;
340 }
341 type CNMap = types::CountedStorageNMap<CountedNMapPrefix, (Key<H, u8>, Key<H, u8>), u128>;
342
343 #[test]
344 fn try_decode_entire_state_value_works() {
345 sp_io::TestExternalities::new_empty().execute_with(|| {
346 assert_eq!(Value::try_decode_entire_state(), Ok(0));
347
348 Value::put(42);
349 assert_eq!(Value::try_decode_entire_state(), Ok(4));
350
351 Value::kill();
352 assert_eq!(Value::try_decode_entire_state(), Ok(0));
353
354 sp_io::storage::set(&Value::hashed_key(), &[0u8, 1]);
356 assert!(Value::try_decode_entire_state().is_err());
357 })
358 }
359
360 #[test]
361 fn try_decode_entire_state_map_works() {
362 sp_io::TestExternalities::new_empty().execute_with(|| {
363 assert_eq!(Map::try_decode_entire_state(), Ok(0));
364
365 Map::insert(0, 42);
366 assert_eq!(Map::try_decode_entire_state(), Ok(4));
367
368 Map::insert(0, 42);
369 assert_eq!(Map::try_decode_entire_state(), Ok(4));
370
371 Map::insert(1, 42);
372 assert_eq!(Map::try_decode_entire_state(), Ok(8));
373
374 Map::remove(0);
375 assert_eq!(Map::try_decode_entire_state(), Ok(4));
376
377 sp_io::storage::set(&Map::hashed_key_for(2), &[0u8, 1]);
379 assert!(Map::try_decode_entire_state().is_err());
380 assert_eq!(Map::try_decode_entire_state().unwrap_err().len(), 1);
381
382 sp_io::storage::set(&Map::hashed_key_for(3), &[0u8, 1]);
384 assert!(Map::try_decode_entire_state().is_err());
385 assert_eq!(Map::try_decode_entire_state().unwrap_err().len(), 2);
386 })
387 }
388
389 #[test]
390 fn try_decode_entire_state_counted_map_works() {
391 sp_io::TestExternalities::new_empty().execute_with(|| {
392 assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
394
395 let counter = 4;
396 let value_size = std::mem::size_of::<u16>();
397
398 CMap::insert(0, 42);
399 assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
400
401 CMap::insert(0, 42);
402 assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
403
404 CMap::insert(1, 42);
405 assert_eq!(CMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
406
407 CMap::remove(0);
408 assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
409
410 let _ = CMap::clear(u32::MAX, None);
412 assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
413
414 sp_io::storage::set(&CMap::hashed_key_for(2), &[0u8]);
416 assert!(CMap::try_decode_entire_state().is_err());
417 })
418 }
419
420 #[test]
421 fn try_decode_entire_state_double_works() {
422 sp_io::TestExternalities::new_empty().execute_with(|| {
423 assert_eq!(DMap::try_decode_entire_state(), Ok(0));
424
425 DMap::insert(0, 0, 42);
426 assert_eq!(DMap::try_decode_entire_state(), Ok(4));
427
428 DMap::insert(0, 0, 42);
429 assert_eq!(DMap::try_decode_entire_state(), Ok(4));
430
431 DMap::insert(0, 1, 42);
432 assert_eq!(DMap::try_decode_entire_state(), Ok(8));
433
434 DMap::insert(1, 0, 42);
435 assert_eq!(DMap::try_decode_entire_state(), Ok(12));
436
437 DMap::remove(0, 0);
438 assert_eq!(DMap::try_decode_entire_state(), Ok(8));
439
440 sp_io::storage::set(&DMap::hashed_key_for(1, 1), &[0u8, 1]);
442 assert!(DMap::try_decode_entire_state().is_err());
443 })
444 }
445
446 #[test]
447 fn try_decode_entire_state_n_map_works() {
448 sp_io::TestExternalities::new_empty().execute_with(|| {
449 assert_eq!(NMap::try_decode_entire_state(), Ok(0));
450
451 let value_size = std::mem::size_of::<u128>();
452
453 NMap::insert((0u8, 0), 42);
454 assert_eq!(NMap::try_decode_entire_state(), Ok(value_size));
455
456 NMap::insert((0, 0), 42);
457 assert_eq!(NMap::try_decode_entire_state(), Ok(value_size));
458
459 NMap::insert((0, 1), 42);
460 assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 2));
461
462 NMap::insert((1, 0), 42);
463 assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 3));
464
465 NMap::remove((0, 0));
466 assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 2));
467
468 sp_io::storage::set(&NMap::hashed_key_for((1, 1)), &[0u8, 1]);
470 assert!(NMap::try_decode_entire_state().is_err());
471 })
472 }
473
474 #[test]
475 fn try_decode_entire_state_counted_n_map_works() {
476 sp_io::TestExternalities::new_empty().execute_with(|| {
477 sp_io::TestExternalities::new_empty().execute_with(|| {
478 assert_eq!(NMap::try_decode_entire_state(), Ok(0));
479
480 let value_size = std::mem::size_of::<u128>();
481 let counter = 4;
482
483 CNMap::insert((0u8, 0), 42);
484 assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size + counter));
485
486 CNMap::insert((0, 0), 42);
487 assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size + counter));
488
489 CNMap::insert((0, 1), 42);
490 assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
491
492 CNMap::insert((1, 0), 42);
493 assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 3 + counter));
494
495 CNMap::remove((0, 0));
496 assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
497
498 sp_io::storage::set(&CNMap::hashed_key_for((1, 1)), &[0u8, 1]);
500 assert!(CNMap::try_decode_entire_state().is_err());
501 })
502 })
503 }
504
505 #[test]
506 fn extra_bytes_are_rejected() {
507 sp_io::TestExternalities::new_empty().execute_with(|| {
508 assert_eq!(Map::try_decode_entire_state(), Ok(0));
509
510 sp_io::storage::set(&Map::hashed_key_for(2), &[0u8, 1, 3, 4, 5, 6]);
512 assert!(Map::try_decode_entire_state().is_err());
513 })
514 }
515
516 #[test]
517 fn try_decode_entire_state_tuple_of_storage_works() {
518 sp_io::TestExternalities::new_empty().execute_with(|| {
519 assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(0));
520
521 Value::put(42);
522 assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(4));
523
524 Map::insert(0, 42);
525 assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(8));
526 });
527 }
528}