Typed Storage Keys (#1419)

* first iteration on storage multi keys * decoding values from concat style hashers * move util functions and remove comments * change codegen for storage keys and fix examples * trait bounds don't match scale value... * fix trait bounds and examples * reconstruct storage keys in iterations * build(deps): bump js-sys from 0.3.67 to 0.3.68 (#1428) Bumps [js-sys](https://github.com/rustwasm/wasm-bindgen) from 0.3.67 to 0.3.68. - [Release notes](https://github.com/rustwasm/wasm-bindgen/releases) - [Changelog](https://github.com/rustwasm/wasm-bindgen/blob/main/CHANGELOG.md) - [Commits](https://github.com/rustwasm/wasm-bindgen/commits) --- updated-dependencies: - dependency-name: js-sys dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * build(deps): bump clap from 4.4.18 to 4.5.0 (#1427) Bumps [clap](https://github.com/clap-rs/clap) from 4.4.18 to 4.5.0. - [Release notes](https://github.com/clap-rs/clap/releases) - [Changelog](https://github.com/clap-rs/clap/blob/master/CHANGELOG.md) - [Commits](https://github.com/clap-rs/clap/compare/v4.4.18...clap_complete-v4.5.0) --- updated-dependencies: - dependency-name: clap dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * build(deps): bump either from 1.9.0 to 1.10.0 (#1425) Bumps [either](https://github.com/rayon-rs/either) from 1.9.0 to 1.10.0. - [Commits](https://github.com/rayon-rs/either/compare/1.9.0...1.10.0) --- updated-dependencies: - dependency-name: either dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * build(deps): bump thiserror from 1.0.56 to 1.0.57 (#1424) Bumps [thiserror](https://github.com/dtolnay/thiserror) from 1.0.56 to 1.0.57. - [Release notes](https://github.com/dtolnay/thiserror/releases) - [Commits](https://github.com/dtolnay/thiserror/compare/1.0.56...1.0.57) --- updated-dependencies: - dependency-name: thiserror dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * build(deps): bump jsonrpsee from 0.21.0 to 0.22.0 (#1426) * build(deps): bump jsonrpsee from 0.21.0 to 0.22.0 Bumps [jsonrpsee](https://github.com/paritytech/jsonrpsee) from 0.21.0 to 0.22.0. - [Release notes](https://github.com/paritytech/jsonrpsee/releases) - [Changelog](https://github.com/paritytech/jsonrpsee/blob/master/CHANGELOG.md) - [Commits](https://github.com/paritytech/jsonrpsee/compare/v0.21.0...v0.22.0) --- updated-dependencies: - dependency-name: jsonrpsee dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> * Update Cargo.lock --------- Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: James Wilson <james@jsdw.me> Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com> * subxt: Derive `std::cmp` traits for subxt payloads and addresses (#1429) * subxt/tx: Derive std::cmp traits Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * subxt/runtime_api: Derive std::cmp traits Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * subxt/constants: Derive std::cmp traits Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * subxt/custom_values: Derive std::cmp traits Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * subxt/storage: Derive std::cmp traits Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * subxt: Fix non_canonical_partial_ord_impl clippy introduced in 1.73 Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * subxt: Add comment wrt derivative issue that triggers clippy warning Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * Update subxt/src/backend/mod.rs * Update subxt/src/constants/constant_address.rs --------- Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com> * fix clippy * add integration tests * fix doc tests * change hashing logic for hashers=1 * refactor * clippy and fmt * regenerate polkadot file which got changed by the automatic PR * nested design for storage keys * refactor codegen * codegen adjustments * fix storage hasher codegen test * Suggestions for storage value decoding (#1457) * Storage decode tweaks * doc tweak * more precise error when leftover or not enough bytes * integrate nits from PR * add fuzztest for storage keys, fix decoding bug * clippy and fmt * clippy * Niklas Suggestions * lifetime issues and iterator impls * fmt and clippy * regenerate polkadot.rs * fix storage key encoding for empty keys * rename trait methods for storage keys * fix hasher bug... * impl nits, add iterator struct seperate from `StorageHashers` * clippy fix * remove println --------- Signed-off-by: dependabot[bot] <support@github.com> Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: James Wilson <james@jsdw.me> Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com> Co-authored-by: Alexandru Vasile <60601340+lexnv@users.noreply.github.com>
2026-06-13 15:11:03 +00:00 · 2024-03-06 14:04:51 +01:00
parent 55e87c86ba
commit a2ee750365
16 changed files with 1960 additions and 1488 deletions
@@ -0,0 +1,466 @@
+use crate::{
+    error::{Error, MetadataError, StorageAddressError},
+    utils::{Encoded, Static},
+};
+use scale_decode::{visitor::IgnoreVisitor, DecodeAsType};
+use scale_encode::EncodeAsType;
+use scale_info::{PortableRegistry, TypeDef};
+use scale_value::Value;
+use subxt_metadata::{StorageEntryType, StorageHasher};
+
+use derivative::Derivative;
+
+use super::utils::hash_bytes;
+
+/// A collection of storage hashers paired with the type ids of the types they should hash.
+/// Can be created for each storage entry in the metadata via [`StorageHashers::new()`].
+#[derive(Debug)]
+pub(crate) struct StorageHashers {
+    hashers_and_ty_ids: Vec<(StorageHasher, u32)>,
+}
+
+impl StorageHashers {
+    /// Creates new [`StorageHashers`] from a storage entry. Looks at the [`StorageEntryType`] and
+    /// assigns a hasher to each type id that makes up the key.
+    pub fn new(storage_entry: &StorageEntryType, types: &PortableRegistry) -> Result<Self, Error> {
+        let mut hashers_and_ty_ids = vec![];
+        if let StorageEntryType::Map {
+            hashers, key_ty, ..
+        } = storage_entry
+        {
+            let ty = types
+                .resolve(*key_ty)
+                .ok_or(MetadataError::TypeNotFound(*key_ty))?;
+
+            if let TypeDef::Tuple(tuple) = &ty.type_def {
+                if hashers.len() == 1 {
+                    // use the same hasher for all fields, if only 1 hasher present:
+                    let hasher = hashers[0];
+                    for f in tuple.fields.iter() {
+                        hashers_and_ty_ids.push((hasher, f.id));
+                    }
+                } else if hashers.len() < tuple.fields.len() {
+                    return Err(StorageAddressError::WrongNumberOfHashers {
+                        hashers: hashers.len(),
+                        fields: tuple.fields.len(),
+                    }
+                    .into());
+                } else {
+                    for (i, f) in tuple.fields.iter().enumerate() {
+                        hashers_and_ty_ids.push((hashers[i], f.id));
+                    }
+                }
+            } else {
+                if hashers.len() != 1 {
+                    return Err(StorageAddressError::WrongNumberOfHashers {
+                        hashers: hashers.len(),
+                        fields: 1,
+                    }
+                    .into());
+                }
+                hashers_and_ty_ids.push((hashers[0], *key_ty));
+            };
+        }
+
+        Ok(Self { hashers_and_ty_ids })
+    }
+
+    /// Creates an iterator over the storage hashers and type ids.
+    pub fn iter(&self) -> StorageHashersIter<'_> {
+        StorageHashersIter {
+            hashers: self,
+            idx: 0,
+        }
+    }
+}
+
+/// An iterator over all type ids of the key and the respective hashers.
+/// See [`StorageHashers::iter()`].
+#[derive(Debug)]
+pub struct StorageHashersIter<'a> {
+    hashers: &'a StorageHashers,
+    idx: usize,
+}
+
+impl<'a> StorageHashersIter<'a> {
+    fn next_or_err(&mut self) -> Result<(StorageHasher, u32), Error> {
+        self.next().ok_or_else(|| {
+            StorageAddressError::TooManyKeys {
+                expected: self.hashers.hashers_and_ty_ids.len(),
+            }
+            .into()
+        })
+    }
+}
+
+impl<'a> Iterator for StorageHashersIter<'a> {
+    type Item = (StorageHasher, u32);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let item = self.hashers.hashers_and_ty_ids.get(self.idx).copied()?;
+        self.idx += 1;
+        Some(item)
+    }
+}
+
+impl<'a> ExactSizeIterator for StorageHashersIter<'a> {
+    fn len(&self) -> usize {
+        self.hashers.hashers_and_ty_ids.len() - self.idx
+    }
+}
+
+/// This trait should be implemented by anything that can be used as one or multiple storage keys.
+pub trait StorageKey {
+    /// Encodes the storage key into some bytes
+    fn encode_storage_key(
+        &self,
+        bytes: &mut Vec<u8>,
+        hashers: &mut StorageHashersIter,
+        types: &PortableRegistry,
+    ) -> Result<(), Error>;
+
+    /// Attempts to decode the StorageKey given some bytes and a set of hashers and type IDs that they are meant to represent.
+    /// The bytes passed to `decode` should start with:
+    /// - 1. some fixed size hash (for all hashers except `Identity`)
+    /// - 2. the plain key value itself (for `Identity`, `Blake2_128Concat` and `Twox64Concat` hashers)
+    fn decode_storage_key(
+        bytes: &mut &[u8],
+        hashers: &mut StorageHashersIter,
+        types: &PortableRegistry,
+    ) -> Result<Self, Error>
+    where
+        Self: Sized + 'static;
+}
+
+/// Implement `StorageKey` for `()` which can be used for keyless storage entries,
+/// or to otherwise just ignore some entry.
+impl StorageKey for () {
+    fn encode_storage_key(
+        &self,
+        _bytes: &mut Vec<u8>,
+        hashers: &mut StorageHashersIter,
+        _types: &PortableRegistry,
+    ) -> Result<(), Error> {
+        _ = hashers.next_or_err();
+        Ok(())
+    }
+
+    fn decode_storage_key(
+        bytes: &mut &[u8],
+        hashers: &mut StorageHashersIter,
+        types: &PortableRegistry,
+    ) -> Result<Self, Error> {
+        let (hasher, ty_id) = match hashers.next_or_err() {
+            Ok((hasher, ty_id)) => (hasher, ty_id),
+            Err(_) if bytes.is_empty() => return Ok(()),
+            Err(err) => return Err(err),
+        };
+        consume_hash_returning_key_bytes(bytes, hasher, ty_id, types)?;
+        Ok(())
+    }
+}
+
+/// A storage key for static encoded values.
+/// The original value is only present at construction, but can be decoded from the contained bytes.
+#[derive(Derivative)]
+#[derivative(Clone(bound = ""), Debug(bound = ""))]
+pub struct StaticStorageKey<K: ?Sized> {
+    bytes: Static<Encoded>,
+    _marker: std::marker::PhantomData<K>,
+}
+
+impl<K: codec::Encode + ?Sized> StaticStorageKey<K> {
+    /// Creates a new static storage key
+    pub fn new(key: &K) -> Self {
+        StaticStorageKey {
+            bytes: Static(Encoded(key.encode())),
+            _marker: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<K: codec::Decode + ?Sized> StaticStorageKey<K> {
+    /// Decodes the encoded inner bytes into the type `K`.
+    pub fn decoded(&self) -> Result<K, Error> {
+        let decoded = K::decode(&mut self.bytes())?;
+        Ok(decoded)
+    }
+}
+
+impl<K: ?Sized> StaticStorageKey<K> {
+    /// Returns the scale-encoded bytes that make up this key
+    pub fn bytes(&self) -> &[u8] {
+        &self.bytes.0 .0
+    }
+}
+
+// Note: The ?Sized bound is necessary to support e.g. `StorageKey<[u8]>`.
+impl<K: ?Sized> StorageKey for StaticStorageKey<K> {
+    fn encode_storage_key(
+        &self,
+        bytes: &mut Vec<u8>,
+        hashers: &mut StorageHashersIter,
+        types: &PortableRegistry,
+    ) -> Result<(), Error> {
+        let (hasher, ty_id) = hashers.next_or_err()?;
+        let encoded_value = self.bytes.encode_as_type(ty_id, types)?;
+        hash_bytes(&encoded_value, hasher, bytes);
+        Ok(())
+    }
+
+    fn decode_storage_key(
+        bytes: &mut &[u8],
+        hashers: &mut StorageHashersIter,
+        types: &PortableRegistry,
+    ) -> Result<Self, Error>
+    where
+        Self: Sized + 'static,
+    {
+        let (hasher, ty_id) = hashers.next_or_err()?;
+        let key_bytes = consume_hash_returning_key_bytes(bytes, hasher, ty_id, types)?;
+
+        // if the hasher had no key appended, we can't decode it into a `StaticStorageKey`.
+        let Some(key_bytes) = key_bytes else {
+            return Err(StorageAddressError::HasherCannotReconstructKey { ty_id, hasher }.into());
+        };
+
+        // Return the key bytes.
+        let key = StaticStorageKey {
+            bytes: Static(Encoded(key_bytes.to_vec())),
+            _marker: std::marker::PhantomData::<K>,
+        };
+        Ok(key)
+    }
+}
+
+impl StorageKey for Vec<scale_value::Value> {
+    fn encode_storage_key(
+        &self,
+        bytes: &mut Vec<u8>,
+        hashers: &mut StorageHashersIter,
+        types: &PortableRegistry,
+    ) -> Result<(), Error> {
+        for value in self.iter() {
+            let (hasher, ty_id) = hashers.next_or_err()?;
+            let encoded_value = value.encode_as_type(ty_id, types)?;
+            hash_bytes(&encoded_value, hasher, bytes);
+        }
+        Ok(())
+    }
+
+    fn decode_storage_key(
+        bytes: &mut &[u8],
+        hashers: &mut StorageHashersIter,
+        types: &PortableRegistry,
+    ) -> Result<Self, Error>
+    where
+        Self: Sized + 'static,
+    {
+        let mut result: Vec<scale_value::Value> = vec![];
+        for (hasher, ty_id) in hashers.by_ref() {
+            match consume_hash_returning_key_bytes(bytes, hasher, ty_id, types)? {
+                Some(value_bytes) => {
+                    let value = Value::decode_as_type(&mut &*value_bytes, ty_id, types)?;
+                    result.push(value.remove_context());
+                }
+                None => {
+                    result.push(Value::unnamed_composite([]));
+                }
+            }
+        }
+
+        // We've consumed all of the hashers, so we expect to also consume all of the bytes:
+        if !bytes.is_empty() {
+            return Err(StorageAddressError::TooManyBytes.into());
+        }
+
+        Ok(result)
+    }
+}
+
+// Skip over the hash bytes (including any key at the end), returning bytes
+// representing the key if one exists, or None if the hasher has no key appended.
+fn consume_hash_returning_key_bytes<'a>(
+    bytes: &mut &'a [u8],
+    hasher: StorageHasher,
+    ty_id: u32,
+    types: &PortableRegistry,
+) -> Result<Option<&'a [u8]>, Error> {
+    // Strip the bytes off for the actual hash, consuming them.
+    let bytes_to_strip = hasher.len_excluding_key();
+    if bytes.len() < bytes_to_strip {
+        return Err(StorageAddressError::NotEnoughBytes.into());
+    }
+    *bytes = &bytes[bytes_to_strip..];
+
+    // Now, find the bytes representing the key, consuming them.
+    let before_key = *bytes;
+    if hasher.ends_with_key() {
+        scale_decode::visitor::decode_with_visitor(bytes, ty_id, types, IgnoreVisitor)
+            .map_err(|err| Error::Decode(err.into()))?;
+        // Return the key bytes, having advanced the input cursor past them.
+        let key_bytes = &before_key[..before_key.len() - bytes.len()];
+
+        Ok(Some(key_bytes))
+    } else {
+        // There are no key bytes, so return None.
+        Ok(None)
+    }
+}
+
+/// Generates StorageKey implementations for tuples
+macro_rules! impl_tuples {
+    ($($ty:ident $n:tt),+) => {{
+        impl<$($ty: StorageKey),+> StorageKey for ($( $ty ),+) {
+            fn encode_storage_key(
+                &self,
+                bytes: &mut Vec<u8>,
+                hashers: &mut StorageHashersIter,
+                types: &PortableRegistry,
+            ) -> Result<(), Error> {
+                $(    self.$n.encode_storage_key(bytes, hashers, types)?;    )+
+                Ok(())
+            }
+
+            fn decode_storage_key(
+                bytes: &mut &[u8],
+                hashers: &mut StorageHashersIter,
+                types: &PortableRegistry,
+            ) -> Result<Self, Error>
+            where
+                Self: Sized + 'static,
+            {
+                Ok( ( $(    $ty::decode_storage_key(bytes, hashers, types)?,    )+ ) )
+            }
+        }
+    }};
+}
+
+#[rustfmt::skip]
+const _: () = {
+    impl_tuples!(A 0, B 1);
+    impl_tuples!(A 0, B 1, C 2);
+    impl_tuples!(A 0, B 1, C 2, D 3);
+    impl_tuples!(A 0, B 1, C 2, D 3, E 4);
+    impl_tuples!(A 0, B 1, C 2, D 3, E 4, F 5);
+    impl_tuples!(A 0, B 1, C 2, D 3, E 4, F 5, G 6);
+    impl_tuples!(A 0, B 1, C 2, D 3, E 4, F 5, G 6, H 7);
+};
+
+#[cfg(test)]
+mod tests {
+
+    use codec::Encode;
+    use scale_info::{meta_type, PortableRegistry, Registry, TypeInfo};
+    use subxt_metadata::StorageHasher;
+
+    use crate::utils::Era;
+
+    use super::{StaticStorageKey, StorageKey};
+
+    struct KeyBuilder {
+        registry: Registry,
+        bytes: Vec<u8>,
+        hashers_and_ty_ids: Vec<(StorageHasher, u32)>,
+    }
+
+    impl KeyBuilder {
+        fn new() -> KeyBuilder {
+            KeyBuilder {
+                registry: Registry::new(),
+                bytes: vec![],
+                hashers_and_ty_ids: vec![],
+            }
+        }
+
+        fn add<T: TypeInfo + Encode + 'static>(mut self, value: T, hasher: StorageHasher) -> Self {
+            let id = self.registry.register_type(&meta_type::<T>()).id;
+
+            self.hashers_and_ty_ids.push((hasher, id));
+            for _i in 0..hasher.len_excluding_key() {
+                self.bytes.push(0);
+            }
+            value.encode_to(&mut self.bytes);
+            self
+        }
+
+        fn build(self) -> (PortableRegistry, Vec<u8>, Vec<(StorageHasher, u32)>) {
+            (self.registry.into(), self.bytes, self.hashers_and_ty_ids)
+        }
+    }
+
+    #[test]
+    fn storage_key_decoding_fuzz() {
+        let hashers = [
+            StorageHasher::Blake2_128,
+            StorageHasher::Blake2_128Concat,
+            StorageHasher::Blake2_256,
+            StorageHasher::Identity,
+            StorageHasher::Twox128,
+            StorageHasher::Twox256,
+            StorageHasher::Twox64Concat,
+        ];
+
+        let key_preserving_hashers = [
+            StorageHasher::Blake2_128Concat,
+            StorageHasher::Identity,
+            StorageHasher::Twox64Concat,
+        ];
+
+        type T4A = (
+            (),
+            StaticStorageKey<u32>,
+            StaticStorageKey<String>,
+            StaticStorageKey<Era>,
+        );
+        type T4B = (
+            (),
+            (StaticStorageKey<u32>, StaticStorageKey<String>),
+            StaticStorageKey<Era>,
+        );
+        type T4C = (
+            ((), StaticStorageKey<u32>),
+            (StaticStorageKey<String>, StaticStorageKey<Era>),
+        );
+
+        let era = Era::Immortal;
+        for h0 in hashers {
+            for h1 in key_preserving_hashers {
+                for h2 in key_preserving_hashers {
+                    for h3 in key_preserving_hashers {
+                        let (types, bytes, hashers_and_ty_ids) = KeyBuilder::new()
+                            .add((), h0)
+                            .add(13u32, h1)
+                            .add("Hello", h2)
+                            .add(era, h3)
+                            .build();
+
+                        let hashers = super::StorageHashers { hashers_and_ty_ids };
+                        let keys_a =
+                            T4A::decode_storage_key(&mut &bytes[..], &mut hashers.iter(), &types)
+                                .unwrap();
+
+                        let keys_b =
+                            T4B::decode_storage_key(&mut &bytes[..], &mut hashers.iter(), &types)
+                                .unwrap();
+
+                        let keys_c =
+                            T4C::decode_storage_key(&mut &bytes[..], &mut hashers.iter(), &types)
+                                .unwrap();
+
+                        assert_eq!(keys_a.1.decoded().unwrap(), 13);
+                        assert_eq!(keys_b.1 .0.decoded().unwrap(), 13);
+                        assert_eq!(keys_c.0 .1.decoded().unwrap(), 13);
+
+                        assert_eq!(keys_a.2.decoded().unwrap(), "Hello");
+                        assert_eq!(keys_b.1 .1.decoded().unwrap(), "Hello");
+                        assert_eq!(keys_c.1 .0.decoded().unwrap(), "Hello");
+                        assert_eq!(keys_a.3.decoded().unwrap(), era);
+                        assert_eq!(keys_b.2.decoded().unwrap(), era);
+                        assert_eq!(keys_c.1 .1.decoded().unwrap(), era);
+                    }
+                }
+            }
+        }
+    }
+}