frame_decode/methods/
storage_encoder.rs

1// Copyright (C) 2022-2023 Parity Technologies (UK) Ltd. (admin@parity.io)
2// This file is a part of the scale-value crate.
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8//         http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use super::storage_type_info::{StorageHasher, StorageTypeInfo};
17use crate::methods::storage_type_info::StorageInfoError;
18use alloc::vec::Vec;
19use scale_type_resolver::TypeResolver;
20
21/// An error returned trying to encode storage keys.
22#[non_exhaustive]
23#[allow(missing_docs)]
24#[derive(Debug, thiserror::Error)]
25pub enum StorageKeyEncodeError {
26    #[error("Cannot get storage info: {0}")]
27    CannotGetInfo(StorageInfoError<'static>),
28    #[error("Failed to encode storage key: {0}")]
29    EncodeError(#[from] scale_encode::Error),
30    #[error("Too many keys provided: expected at most {max_keys_expected}")]
31    TooManyKeysProvided {
32        /// The maximum number of keys that were expected.
33        max_keys_expected: usize,
34    },
35}
36
37/// Encode a storage key prefix from a pallet name and storage entry name. This prefix
38/// is the first 32 bytes of any storage key which comes from a pallet, and is essentially
39/// `twox_128(pallet_name) + twox_128(storage_entry_name)`.
40pub fn encode_prefix(pallet_name: &str, storage_entry: &str) -> [u8; 32] {
41    let mut prefix = [0u8; 32];
42
43    let pallet_bytes = sp_crypto_hashing::twox_128(pallet_name.as_bytes());
44    let entry_bytes = sp_crypto_hashing::twox_128(storage_entry.as_bytes());
45
46    prefix[..16].copy_from_slice(&pallet_bytes);
47    prefix[16..].copy_from_slice(&entry_bytes);
48
49    prefix
50}
51
52/// Encode a complete storage key for a given pallet and storage entry and a set of keys that
53/// are each able to be encoded via [`scale_encode::EncodeAsType`].
54///
55/// This is the same as [`encode_storage_key_to`], but returns the encoded key as a `Vec<u8>`, rather
56/// than accepting a mutable output buffer.
57///
58/// # Example
59///
60/// ```rust
61/// use frame_decode::storage::encode_storage_key;
62/// use frame_metadata::RuntimeMetadata;
63/// use parity_scale_codec::Decode;
64///
65/// let metadata_bytes = std::fs::read("artifacts/metadata_10000000_9180.scale").unwrap();
66/// let RuntimeMetadata::V14(metadata) = RuntimeMetadata::decode(&mut &*metadata_bytes).unwrap() else { return };
67///
68/// let account_id = [0u8; 32];
69///
70/// // System.Account needs only one key to point at a specific value; the account ID.
71/// // We just fake an account ID for this example  by providing 32 0 bytes, but anything
72/// // which would `scale_encode::EncodeAsType` into 32 bytes would work.
73/// let encoded_key = encode_storage_key(
74///     "System",
75///     "Account",
76///     [account_id],
77///     &metadata,
78///     &metadata.types,
79/// ).unwrap();
80/// ```
81pub fn encode_storage_key<Info, Resolver, Keys>(
82    pallet_name: &str,
83    storage_entry: &str,
84    keys: Keys,
85    info: &Info,
86    type_resolver: &Resolver,
87) -> Result<Vec<u8>, StorageKeyEncodeError>
88where
89    Keys: IntoStorageKeys,
90    Info: StorageTypeInfo,
91    Info::TypeId: Clone + core::fmt::Debug,
92    Resolver: TypeResolver<TypeId = Info::TypeId>,
93{
94    // pre-allocate at least as many bytes as we need for the root/prefix.
95    let mut out = Vec::with_capacity(32);
96    encode_storage_key_to(
97        pallet_name,
98        storage_entry,
99        keys,
100        info,
101        type_resolver,
102        &mut out,
103    )?;
104    Ok(out)
105}
106
107/// Encode a complete storage key for a given pallet and storage entry and a set of keys that
108/// are each able to be encoded via [`scale_encode::EncodeAsType`].
109///
110/// # Example
111///
112/// ```rust
113/// use frame_decode::storage::encode_storage_key_to;
114/// use frame_metadata::RuntimeMetadata;
115/// use parity_scale_codec::Decode;
116///
117/// let metadata_bytes = std::fs::read("artifacts/metadata_10000000_9180.scale").unwrap();
118/// let RuntimeMetadata::V14(metadata) = RuntimeMetadata::decode(&mut &*metadata_bytes).unwrap() else { return };
119///
120/// let account_id = [0u8; 32];
121///
122/// // We'll encode the key to this.
123/// let mut encoded_key = Vec::new();
124///
125/// // System.Account needs only one key to point at a specific value; the account ID.
126/// // We just fake an account ID for this example  by providing 32 0 bytes, but anything
127/// // which would `scale_encode::EncodeAsType` into 32 bytes would work.
128/// encode_storage_key_to(
129///     "System",
130///     "Account",
131///     [account_id],
132///     &metadata,
133///     &metadata.types,
134///     &mut encoded_key,
135/// ).unwrap();
136/// ```
137pub fn encode_storage_key_to<Info, Resolver, Keys>(
138    pallet_name: &str,
139    storage_entry: &str,
140    keys: Keys,
141    info: &Info,
142    type_resolver: &Resolver,
143    out: &mut Vec<u8>,
144) -> Result<(), StorageKeyEncodeError>
145where
146    Keys: IntoStorageKeys,
147    Info: StorageTypeInfo,
148    Info::TypeId: Clone + core::fmt::Debug,
149    Resolver: TypeResolver<TypeId = Info::TypeId>,
150{
151    let storage_info = info
152        .get_storage_info(pallet_name, storage_entry)
153        .map_err(|e| StorageKeyEncodeError::CannotGetInfo(e.into_owned()))?;
154
155    encode_storage_key_with_info_to(
156        pallet_name,
157        storage_entry,
158        keys,
159        &storage_info,
160        type_resolver,
161        out,
162    )
163}
164
165/// Encode a complete storage key for a given pallet and storage entry and a set of keys that
166/// are each able to be encoded via [`scale_encode::EncodeAsType`].
167///
168/// Unlike [`encode_storage_key`], which obtains the storage info internally given the pallet and storage entry names,
169/// this function takes the storage info as an argument. This is useful if you already have the storage info available,
170/// for example if you are encoding multiple keys for the same storage entry.
171pub fn encode_storage_key_with_info_to<Resolver, Keys>(
172    pallet_name: &str,
173    storage_entry: &str,
174    keys: Keys,
175    storage_info: &crate::storage::StorageInfo<<Resolver as TypeResolver>::TypeId>,
176    type_resolver: &Resolver,
177    out: &mut Vec<u8>,
178) -> Result<(), StorageKeyEncodeError>
179where
180    Keys: IntoStorageKeys,
181    Resolver: TypeResolver,
182    <Resolver as TypeResolver>::TypeId: Clone + core::fmt::Debug,
183{
184    // If too many keys provided, bail early.
185    if storage_info.keys.len() < keys.num_keys() {
186        return Err(StorageKeyEncodeError::TooManyKeysProvided {
187            max_keys_expected: storage_info.keys.len(),
188        });
189    }
190
191    // Encode the prefix:
192    let prefix = encode_prefix(pallet_name, storage_entry);
193    out.extend_from_slice(&prefix);
194
195    // Encode the keys:
196    let mut keys = keys.into_storage_keys();
197    let mut temp = Vec::with_capacity(32);
198    for key_info in &storage_info.keys {
199        match keys.encode_next_key_to(key_info.key_id.clone(), type_resolver, &mut temp) {
200            None => break, // No more keys to encode.
201            Some(Err(e)) => return Err(StorageKeyEncodeError::EncodeError(e)),
202            Some(Ok(())) => { /* All ok */ }
203        };
204
205        match key_info.hasher {
206            StorageHasher::Blake2_128 => {
207                let hash = sp_crypto_hashing::blake2_128(&temp);
208                out.extend_from_slice(&hash);
209            }
210            StorageHasher::Blake2_256 => {
211                let hash = sp_crypto_hashing::blake2_256(&temp);
212                out.extend_from_slice(&hash);
213            }
214            StorageHasher::Blake2_128Concat => {
215                let hash = sp_crypto_hashing::blake2_128(&temp);
216                out.extend_from_slice(&hash);
217                out.extend_from_slice(&temp);
218            }
219            StorageHasher::Twox128 => {
220                let hash = sp_crypto_hashing::twox_128(&temp);
221                out.extend_from_slice(&hash);
222            }
223            StorageHasher::Twox256 => {
224                let hash = sp_crypto_hashing::twox_256(&temp);
225                out.extend_from_slice(&hash);
226            }
227            StorageHasher::Twox64Concat => {
228                let hash = sp_crypto_hashing::twox_64(&temp);
229                out.extend_from_slice(&hash);
230                out.extend_from_slice(&temp);
231            }
232            StorageHasher::Identity => {
233                out.extend_from_slice(&temp);
234            }
235        }
236
237        // Clear our temp space ready for the next key.
238        temp.clear();
239    }
240
241    Ok(())
242}
243
244/// This can be implemented for anything that can be converted into something implementing [`StorageKeys`].
245/// It is implemented by default for tuples up to length 10, vectors and arrays (where the values all implement
246/// [`scale_encode::EncodeAsType`]).
247pub trait IntoStorageKeys {
248    /// An implementation of [`StorageKeys`] that can be used to iterate through the keys.
249    type Keys: StorageKeys;
250    /// Return an implementation of [`StorageKeys`] for this type.
251    fn into_storage_keys(self) -> Self::Keys;
252    /// The number of storage keys this type represents.
253    fn num_keys(&self) -> usize;
254}
255
256/// Since [`scale_encode::EncodeAsType`] is not dyn safe, this trait is used to iterate through and encode a set of keys.
257pub trait StorageKeys {
258    /// Encode the next key, if there is one, into the provided output buffer.
259    fn encode_next_key_to<Resolver>(
260        &mut self,
261        type_id: Resolver::TypeId,
262        types: &Resolver,
263        out: &mut Vec<u8>,
264    ) -> Option<Result<(), scale_encode::Error>>
265    where
266        Resolver: TypeResolver;
267    /// Encode the next key, if there is one, and return the encoded bytes
268    fn encode_next_key<Resolver>(
269        &mut self,
270        type_id: Resolver::TypeId,
271        types: &Resolver,
272    ) -> Option<Result<Vec<u8>, scale_encode::Error>>
273    where
274        Resolver: TypeResolver,
275    {
276        let mut out = Vec::new();
277        self.encode_next_key_to(type_id, types, &mut out)
278            .map(|res| res.map(|_| out))
279    }
280}
281
282// Vecs of keys implement IntoStorageKeys.
283impl<K: scale_encode::EncodeAsType> IntoStorageKeys for Vec<K> {
284    type Keys = <Self as IntoIterator>::IntoIter;
285    fn num_keys(&self) -> usize {
286        self.len()
287    }
288    fn into_storage_keys(self) -> Self::Keys {
289        self.into_iter()
290    }
291}
292
293impl<K: scale_encode::EncodeAsType> StorageKeys for alloc::vec::IntoIter<K> {
294    fn encode_next_key_to<Resolver>(
295        &mut self,
296        type_id: Resolver::TypeId,
297        types: &Resolver,
298        out: &mut Vec<u8>,
299    ) -> Option<Result<(), scale_encode::Error>>
300    where
301        Resolver: TypeResolver,
302    {
303        let next_key = self.next()?;
304        if let Err(e) = next_key.encode_as_type_to(type_id, types, out) {
305            return Some(Err(e));
306        }
307        Some(Ok(()))
308    }
309}
310
311// As do arrays of keys.
312impl<K: scale_encode::EncodeAsType, const N: usize> IntoStorageKeys for [K; N] {
313    type Keys = <Self as IntoIterator>::IntoIter;
314    fn num_keys(&self) -> usize {
315        N
316    }
317    fn into_storage_keys(self) -> Self::Keys {
318        self.into_iter()
319    }
320}
321
322impl<K: scale_encode::EncodeAsType, const N: usize> StorageKeys for core::array::IntoIter<K, N> {
323    fn encode_next_key_to<Resolver>(
324        &mut self,
325        type_id: Resolver::TypeId,
326        types: &Resolver,
327        out: &mut Vec<u8>,
328    ) -> Option<Result<(), scale_encode::Error>>
329    where
330        Resolver: TypeResolver,
331    {
332        let next_key = self.next()?;
333        if let Err(e) = next_key.encode_as_type_to(type_id, types, out) {
334            return Some(Err(e));
335        }
336        Some(Ok(()))
337    }
338}
339
340// Empty tuples can be used as a placeholder for no storage keys.
341impl IntoStorageKeys for () {
342    type Keys = ();
343    fn num_keys(&self) -> usize {
344        0
345    }
346    fn into_storage_keys(self) -> Self::Keys {}
347}
348
349impl StorageKeys for () {
350    fn encode_next_key_to<Resolver>(
351        &mut self,
352        _type_id: Resolver::TypeId,
353        _types: &Resolver,
354        _out: &mut Vec<u8>,
355    ) -> Option<Result<(), scale_encode::Error>>
356    where
357        Resolver: TypeResolver,
358    {
359        None
360    }
361}
362
363// Tuples of different lengths can be used as storage keys, too.
364macro_rules! impl_tuple_storage_keys {
365    ($($ty:ident $number:tt),*) => {
366        const _: () = {
367            const TUPLE_LEN: usize = 0 $(+ $number - $number + 1)*;
368
369            impl <$($ty: scale_encode::EncodeAsType),*> IntoStorageKeys for ($($ty,)*) {
370                type Keys = TupleIter<$($ty),*>;
371                fn num_keys(&self) -> usize {
372                    TUPLE_LEN
373                }
374                fn into_storage_keys(self) -> Self::Keys {
375                    TupleIter {
376                        idx: 0,
377                        items: self,
378                    }
379                }
380            }
381
382            pub struct TupleIter<$($ty),*> {
383                idx: usize,
384                items: ($($ty,)*)
385            }
386
387            impl <$($ty: scale_encode::EncodeAsType),*> StorageKeys for TupleIter<$($ty),*> {
388                fn encode_next_key_to<Resolver>(&mut self, type_id: Resolver::TypeId, types: &Resolver, out: &mut Vec<u8>) -> Option<Result<(), scale_encode::Error>>
389                where
390                    Resolver: TypeResolver,
391                {
392                    $(
393                        if self.idx == $number {
394                            let item = &self.items.$number;
395                            if let Err(e) = item.encode_as_type_to(type_id, types, out) {
396                                return Some(Err(e));
397                            }
398                            self.idx += 1;
399                            return Some(Ok(()));
400                        }
401                    )*
402                    None
403                }
404            }
405        };
406    };
407}
408
409impl_tuple_storage_keys!(A 0);
410impl_tuple_storage_keys!(A 0, B 1);
411impl_tuple_storage_keys!(A 0, B 1, C 2);
412impl_tuple_storage_keys!(A 0, B 1, C 2, D 3);
413impl_tuple_storage_keys!(A 0, B 1, C 2, D 3, E 4);
414impl_tuple_storage_keys!(A 0, B 1, C 2, D 3, E 4, F 5);
415impl_tuple_storage_keys!(A 0, B 1, C 2, D 3, E 4, F 5, G 6);
416impl_tuple_storage_keys!(A 0, B 1, C 2, D 3, E 4, F 5, G 6, H 7);
417impl_tuple_storage_keys!(A 0, B 1, C 2, D 3, E 4, F 5, G 6, H 7, I 8);
418impl_tuple_storage_keys!(A 0, B 1, C 2, D 3, E 4, F 5, G 6, H 7, I 8, J 9);
419
420#[cfg(test)]
421mod test {
422    use super::*;
423
424    #[test]
425    fn test_tuple_storage_keys() {
426        use parity_scale_codec::Encode;
427        use scale_info_legacy::LookupName;
428
429        // We just need some basic types to test with.
430        let types = crate::legacy_types::polkadot::relay_chain();
431        let types = types.for_spec_version(0);
432
433        let keys = (123u16, true, "hello");
434        assert_eq!(keys.num_keys(), 3);
435        let mut storage_keys = keys.into_storage_keys();
436
437        let val = storage_keys
438            .encode_next_key(LookupName::parse("u64").unwrap(), &types)
439            .unwrap()
440            .unwrap();
441        assert_eq!(val, 123u64.encode());
442
443        let val = storage_keys
444            .encode_next_key(LookupName::parse("bool").unwrap(), &types)
445            .unwrap()
446            .unwrap();
447        assert_eq!(val, true.encode());
448
449        let val = storage_keys
450            .encode_next_key(LookupName::parse("String").unwrap(), &types)
451            .unwrap()
452            .unwrap();
453        assert_eq!(val, "hello".encode());
454
455        assert!(storage_keys
456            .encode_next_key(LookupName::parse("foo").unwrap(), &types)
457            .is_none());
458        assert!(storage_keys
459            .encode_next_key(LookupName::parse("foo").unwrap(), &types)
460            .is_none());
461    }
462
463    #[test]
464    fn test_vec_storage_keys() {
465        use parity_scale_codec::Encode;
466        use scale_info_legacy::LookupName;
467
468        // We just need some basic types to test with.
469        let types = crate::legacy_types::polkadot::relay_chain();
470        let types = types.for_spec_version(0);
471
472        let keys = vec![123u16, 456u16, 789u16];
473        assert_eq!(keys.num_keys(), 3);
474        let mut storage_keys = keys.into_storage_keys();
475
476        let val = storage_keys
477            .encode_next_key(LookupName::parse("u64").unwrap(), &types)
478            .unwrap()
479            .unwrap();
480        assert_eq!(val, 123u64.encode());
481
482        let val = storage_keys
483            .encode_next_key(LookupName::parse("u16").unwrap(), &types)
484            .unwrap()
485            .unwrap();
486        assert_eq!(val, 456u16.encode());
487
488        let val = storage_keys
489            .encode_next_key(LookupName::parse("u32").unwrap(), &types)
490            .unwrap()
491            .unwrap();
492        assert_eq!(val, 789u32.encode());
493
494        assert!(storage_keys
495            .encode_next_key(LookupName::parse("foo").unwrap(), &types)
496            .is_none());
497        assert!(storage_keys
498            .encode_next_key(LookupName::parse("foo").unwrap(), &types)
499            .is_none());
500    }
501
502    #[test]
503    fn test_array_storage_keys() {
504        use parity_scale_codec::Encode;
505        use scale_info_legacy::LookupName;
506
507        // We just need some basic types to test with.
508        let types = crate::legacy_types::polkadot::relay_chain();
509        let types = types.for_spec_version(0);
510
511        let keys: [u16; 3] = [123, 456, 789];
512        assert_eq!(keys.num_keys(), 3);
513        let mut storage_keys = keys.into_storage_keys();
514
515        let val = storage_keys
516            .encode_next_key(LookupName::parse("u64").unwrap(), &types)
517            .unwrap()
518            .unwrap();
519        assert_eq!(val, 123u64.encode());
520
521        let val = storage_keys
522            .encode_next_key(LookupName::parse("u16").unwrap(), &types)
523            .unwrap()
524            .unwrap();
525        assert_eq!(val, 456u16.encode());
526
527        let val = storage_keys
528            .encode_next_key(LookupName::parse("u32").unwrap(), &types)
529            .unwrap()
530            .unwrap();
531        assert_eq!(val, 789u32.encode());
532
533        assert!(storage_keys
534            .encode_next_key(LookupName::parse("foo").unwrap(), &types)
535            .is_none());
536        assert!(storage_keys
537            .encode_next_key(LookupName::parse("foo").unwrap(), &types)
538            .is_none());
539    }
540
541    #[test]
542    fn test_encode_storage_key() {
543        use crate::storage::encode_storage_key;
544        use frame_metadata::RuntimeMetadata;
545        use parity_scale_codec::Decode;
546
547        let metadata_bytes = std::fs::read("artifacts/metadata_10000000_9180.scale").unwrap();
548        let RuntimeMetadata::V14(metadata) =
549            RuntimeMetadata::decode(&mut &*metadata_bytes).unwrap()
550        else {
551            return;
552        };
553
554        let account_id = [0u8; 32];
555
556        // System.Account needs only one key to point at a specific value; the account ID.
557        // We just fake an account ID for this example by providing 32 0 bytes, but anything
558        // which would `scale_encode::EncodeAsType` into 32 bytes should work.
559        encode_storage_key(
560            "System",
561            "Account",
562            [account_id],
563            &metadata,
564            &metadata.types,
565        )
566        .expect("Encoding should work");
567        encode_storage_key(
568            "System",
569            "Account",
570            (account_id,),
571            &metadata,
572            &metadata.types,
573        )
574        .expect("Encoding should work");
575
576        // We provide no additional keys, so we should get the prefix only.
577        let out = encode_storage_key("System", "Account", (), &metadata, &metadata.types)
578            .expect("Encoding should work");
579        assert_eq!(&out, &encode_prefix("System", "Account"));
580
581        // We provide too many additional keys so should get an error.
582        let err = encode_storage_key(
583            "System",
584            "Account",
585            (account_id, 123u16),
586            &metadata,
587            &metadata.types,
588        );
589        assert!(matches!(
590            err,
591            Err(StorageKeyEncodeError::TooManyKeysProvided {
592                max_keys_expected: 1
593            })
594        ));
595    }
596}
frame_decode/methods/storage_encoder.rs

frame_decode/methods/
storage_encoder.rs
