diff --git a/bridges/bin/runtime-common/src/messages.rs b/bridges/bin/runtime-common/src/messages.rs index 0fe9935dbdb6..03801d5279d0 100644 --- a/bridges/bin/runtime-common/src/messages.rs +++ b/bridges/bin/runtime-common/src/messages.rs @@ -395,6 +395,7 @@ mod tests { use codec::Encode; use sp_core::H256; use sp_runtime::traits::Header as _; + use sp_trie::accessed_nodes_tracker::Error as AccessedNodesTrackerError; #[test] fn verify_chain_message_rejects_message_with_too_large_declared_weight() { @@ -541,7 +542,7 @@ mod tests { target::verify_messages_proof::(proof, 10) },), Err(VerificationError::HeaderChain(HeaderChainError::StorageProof( - StorageProofError::DuplicateNodesInProof + StorageProofError::StorageProof(sp_trie::StorageProofError::DuplicateNodes.into()) ))), ); } @@ -553,7 +554,9 @@ mod tests { proof.storage_proof.push(vec![42]); target::verify_messages_proof::(proof, 10) },), - Err(VerificationError::StorageProof(StorageProofError::UnusedNodesInTheProof)), + Err(VerificationError::StorageProof(StorageProofError::AccessedNodesTracker( + AccessedNodesTrackerError::UnusedNodes.into() + ))), ); } diff --git a/bridges/primitives/runtime/src/storage_proof.rs b/bridges/primitives/runtime/src/storage_proof.rs index 1b706aa66c16..251ceec5a9ee 100644 --- a/bridges/primitives/runtime/src/storage_proof.rs +++ b/bridges/primitives/runtime/src/storage_proof.rs @@ -21,15 +21,16 @@ use codec::{Decode, Encode}; use frame_support::PalletError; use hash_db::{HashDB, Hasher, EMPTY_PREFIX}; use scale_info::TypeInfo; -use sp_std::{boxed::Box, collections::btree_set::BTreeSet, vec::Vec}; +use sp_std::{boxed::Box, vec::Vec}; +pub use sp_trie::RawStorageProof; use sp_trie::{ - read_trie_value, LayoutV1, MemoryDB, Recorder, StorageProof, Trie, TrieConfiguration, + accessed_nodes_tracker::{AccessedNodesTracker, Error as AccessedNodesTrackerError}, + read_trie_value, + recorder_ext::RecorderExt, + LayoutV1, MemoryDB, Recorder, StorageProof, StorageProofError, Trie, TrieConfiguration, TrieDBBuilder, TrieError, TrieHash, }; -/// Raw storage proof type (just raw trie nodes). -pub type RawStorageProof = Vec>; - /// Storage proof size requirements. /// /// This is currently used by benchmarks when generating storage proofs. @@ -51,10 +52,9 @@ pub struct StorageProofChecker where H: Hasher, { - proof_nodes_count: usize, root: H::Out, db: MemoryDB, - recorder: Recorder>, + accessed_nodes_tracker: AccessedNodesTracker, } impl StorageProofChecker @@ -65,52 +65,39 @@ where /// /// This returns an error if the given proof is invalid with respect to the given root. pub fn new(root: H::Out, proof: RawStorageProof) -> Result { - // 1. we don't want extra items in the storage proof - // 2. `StorageProof` is storing all trie nodes in the `BTreeSet` - // - // => someone could simply add duplicate items to the proof and we won't be - // able to detect that by just using `StorageProof` - // - // => let's check it when we are converting our "raw proof" into `StorageProof` - let proof_nodes_count = proof.len(); - let proof = StorageProof::new(proof); - if proof_nodes_count != proof.iter_nodes().count() { - return Err(Error::DuplicateNodesInProof) - } + let proof = StorageProof::new_with_duplicate_nodes_check(proof) + .map_err(|e| Error::StorageProof(e.into()))?; + + let recorder = AccessedNodesTracker::new(proof.len()); let db = proof.into_memory_db(); if !db.contains(&root, EMPTY_PREFIX) { return Err(Error::StorageRootMismatch) } - let recorder = Recorder::default(); - let checker = StorageProofChecker { proof_nodes_count, root, db, recorder }; - Ok(checker) + Ok(StorageProofChecker { root, db, accessed_nodes_tracker: recorder }) } /// Returns error if the proof has some nodes that are left intact by previous `read_value` /// calls. - pub fn ensure_no_unused_nodes(mut self) -> Result<(), Error> { - let visited_nodes = self - .recorder - .drain() - .into_iter() - .map(|record| record.data) - .collect::>(); - let visited_nodes_count = visited_nodes.len(); - if self.proof_nodes_count == visited_nodes_count { - Ok(()) - } else { - Err(Error::UnusedNodesInTheProof) - } + pub fn ensure_no_unused_nodes(self) -> Result<(), Error> { + self.accessed_nodes_tracker + .ensure_no_unused_nodes() + .map_err(|e| Error::AccessedNodesTracker(e.into())) } /// Reads a value from the available subset of storage. If the value cannot be read due to an /// incomplete or otherwise invalid proof, this function returns an error. pub fn read_value(&mut self, key: &[u8]) -> Result>, Error> { // LayoutV1 or LayoutV0 is identical for proof that only read values. - read_trie_value::, _>(&self.db, &self.root, key, Some(&mut self.recorder), None) - .map_err(|_| Error::StorageValueUnavailable) + read_trie_value::, _>( + &self.db, + &self.root, + key, + Some(&mut self.accessed_nodes_tracker), + None, + ) + .map_err(|_| Error::StorageValueUnavailable) } /// Reads and decodes a value from the available subset of storage. If the value cannot be read @@ -145,10 +132,10 @@ where /// Storage proof related errors. #[derive(Encode, Decode, Clone, Eq, PartialEq, PalletError, Debug, TypeInfo)] pub enum Error { - /// Duplicate trie nodes are found in the proof. - DuplicateNodesInProof, - /// Unused trie nodes are found in the proof. - UnusedNodesInTheProof, + /// Error generated by the `AccessedNodesTrackerError`. + AccessedNodesTracker(StrippableError), + /// Error originating in the `storage_proof` module. + StorageProof(StrippableError), /// Expected storage root is missing from the proof. StorageRootMismatch, /// Unable to reach expected storage value using provided trie nodes. @@ -202,15 +189,7 @@ where trie.get(&key)?; } - // recorder may record the same trie node multiple times and we don't want duplicate nodes - // in our proofs => let's deduplicate it by collecting to the BTreeSet first - Ok(recorder - .drain() - .into_iter() - .map(|n| n.data.to_vec()) - .collect::>() - .into_iter() - .collect()) + Ok(recorder.into_raw_storage_proof()) } #[cfg(test)] @@ -243,30 +222,22 @@ pub mod tests { ); } - #[test] - fn proof_with_duplicate_items_is_rejected() { - let (root, mut proof) = craft_valid_storage_proof(); - proof.push(proof.first().unwrap().clone()); - - assert_eq!( - StorageProofChecker::::new(root, proof).map(drop), - Err(Error::DuplicateNodesInProof), - ); - } - #[test] fn proof_with_unused_items_is_rejected() { let (root, proof) = craft_valid_storage_proof(); let mut checker = StorageProofChecker::::new(root, proof.clone()).unwrap(); - checker.read_value(b"key1").unwrap(); + checker.read_value(b"key1").unwrap().unwrap(); checker.read_value(b"key2").unwrap(); checker.read_value(b"key4").unwrap(); checker.read_value(b"key22").unwrap(); assert_eq!(checker.ensure_no_unused_nodes(), Ok(())); let checker = StorageProofChecker::::new(root, proof).unwrap(); - assert_eq!(checker.ensure_no_unused_nodes(), Err(Error::UnusedNodesInTheProof)); + assert_eq!( + checker.ensure_no_unused_nodes(), + Err(Error::AccessedNodesTracker(AccessedNodesTrackerError::UnusedNodes.into())) + ); } } diff --git a/substrate/frame/session/src/historical/mod.rs b/substrate/frame/session/src/historical/mod.rs index b9cecea1a7f7..618497e3d54d 100644 --- a/substrate/frame/session/src/historical/mod.rs +++ b/substrate/frame/session/src/historical/mod.rs @@ -37,21 +37,22 @@ use sp_runtime::{ }; use sp_session::{MembershipProof, ValidatorCount}; use sp_staking::SessionIndex; -use sp_std::prelude::*; +use sp_std::{fmt::Debug, prelude::*}; use sp_trie::{ trie_types::{TrieDBBuilder, TrieDBMutBuilderV0}, - LayoutV0, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX, + LayoutV0, MemoryDB, Recorder, StorageProof, Trie, TrieMut, TrieRecorder, }; use frame_support::{ print, traits::{KeyOwnerProofSystem, ValidatorSet, ValidatorSetWithIdentification}, - Parameter, + Parameter, LOG_TARGET, }; use crate::{self as pallet_session, Pallet as Session}; pub use pallet::*; +use sp_trie::{accessed_nodes_tracker::AccessedNodesTracker, recorder_ext::RecorderExt}; #[frame_support::pallet] pub mod pallet { @@ -118,6 +119,16 @@ impl Pallet { } }) } + + fn full_id_validators() -> Vec<(T::ValidatorId, T::FullIdentification)> { + >::validators() + .into_iter() + .filter_map(|validator| { + T::FullIdentificationOf::convert(validator.clone()) + .map(|full_id| (validator, full_id)) + }) + .collect::>() + } } impl ValidatorSet for Pallet { @@ -264,35 +275,16 @@ impl ProvingTrie { Ok(ProvingTrie { db, root }) } - fn from_nodes(root: T::Hash, nodes: &[Vec]) -> Self { - use sp_trie::HashDBT; - - let mut memory_db = MemoryDB::default(); - for node in nodes { - HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]); - } - - ProvingTrie { db: memory_db, root } + fn from_proof(root: T::Hash, proof: StorageProof) -> Self { + ProvingTrie { db: proof.into_memory_db(), root } } /// Prove the full verification data for a given key and key ID. pub fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option>> { let mut recorder = Recorder::>::new(); - { - let trie = - TrieDBBuilder::new(&self.db, &self.root).with_recorder(&mut recorder).build(); - let val_idx = (key_id, key_data).using_encoded(|s| { - trie.get(s).ok()?.and_then(|raw| u32::decode(&mut &*raw).ok()) - })?; - - val_idx.using_encoded(|s| { - trie.get(s) - .ok()? - .and_then(|raw| >::decode(&mut &*raw).ok()) - })?; - } + self.query(key_id, key_data, Some(&mut recorder)); - Some(recorder.drain().into_iter().map(|r| r.data).collect()) + Some(recorder.into_raw_storage_proof()) } /// Access the underlying trie root. @@ -300,10 +292,17 @@ impl ProvingTrie { &self.root } - // Check a proof contained within the current memory-db. Returns `None` if the - // nodes within the current `MemoryDB` are insufficient to query the item. - fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option> { - let trie = TrieDBBuilder::new(&self.db, &self.root).build(); + /// Search for a key inside the proof. + fn query( + &self, + key_id: KeyTypeId, + key_data: &[u8], + recorder: Option<&mut dyn TrieRecorder>, + ) -> Option> { + let trie = TrieDBBuilder::new(&self.db, &self.root) + .with_optional_recorder(recorder) + .build(); + let val_idx = (key_id, key_data) .using_encoded(|s| trie.get(s)) .ok()? @@ -322,13 +321,7 @@ impl> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet Option { let session = >::current_index(); - let validators = >::validators() - .into_iter() - .filter_map(|validator| { - T::FullIdentificationOf::convert(validator.clone()) - .map(|full_id| (validator, full_id)) - }) - .collect::>(); + let validators = Self::full_id_validators(); let count = validators.len() as ValidatorCount; @@ -343,30 +336,35 @@ impl> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet Option> { - let (id, data) = key; - - if proof.session == >::current_index() { - >::key_owner(id, data.as_ref()).and_then(|owner| { - T::FullIdentificationOf::convert(owner.clone()).and_then(move |id| { - let count = >::validators().len() as ValidatorCount; - - if count != proof.validator_count { - return None - } + fn print_error(e: E) { + log::error!( + target: LOG_TARGET, + "Rejecting equivocation report because of key ownership proof error: {:?}", e + ); + } - Some((owner, id)) - }) - }) + let (id, data) = key; + let (root, count) = if proof.session == >::current_index() { + let validators = Self::full_id_validators(); + let count = validators.len() as ValidatorCount; + let trie = ProvingTrie::::generate_for(validators).ok()?; + (trie.root, count) } else { - let (root, count) = >::get(&proof.session)?; - - if count != proof.validator_count { - return None - } + >::get(&proof.session)? + }; - let trie = ProvingTrie::::from_nodes(root, &proof.trie_nodes); - trie.query(id, data.as_ref()) + if count != proof.validator_count { + return None } + + let proof = StorageProof::new_with_duplicate_nodes_check(proof.trie_nodes) + .map_err(print_error) + .ok()?; + let mut accessed_nodes_tracker = AccessedNodesTracker::::new(proof.len()); + let trie = ProvingTrie::::from_proof(root, proof); + let res = trie.query(id, data.as_ref(), Some(&mut accessed_nodes_tracker))?; + accessed_nodes_tracker.ensure_no_unused_nodes().map_err(print_error).ok()?; + Some(res) } } diff --git a/substrate/primitives/trie/src/accessed_nodes_tracker.rs b/substrate/primitives/trie/src/accessed_nodes_tracker.rs new file mode 100644 index 000000000000..378e3c2812c0 --- /dev/null +++ b/substrate/primitives/trie/src/accessed_nodes_tracker.rs @@ -0,0 +1,119 @@ +// This file is part of Substrate. + +// Copyright (C) Parity Technologies (UK) Ltd. +// SPDX-License-Identifier: Apache-2.0 + +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Helpers for checking for duplicate nodes. + +use alloc::collections::BTreeSet; +use core::hash::Hash; +use scale_info::TypeInfo; +use sp_core::{Decode, Encode}; +use trie_db::{RecordedForKey, TrieAccess, TrieRecorder}; + +/// Error associated with the `AccessedNodesTracker` module. +#[derive(Encode, Decode, Clone, Eq, PartialEq, Debug, TypeInfo)] +pub enum Error { + /// The proof contains unused nodes. + UnusedNodes, +} + +/// Helper struct used to ensure that a storage proof doesn't contain duplicate or unused nodes. +/// +/// The struct needs to be used as a `TrieRecorder` and `ensure_no_unused_nodes()` has to be called +/// to actually perform the check. +pub struct AccessedNodesTracker { + proof_nodes_count: usize, + recorder: BTreeSet, +} + +impl AccessedNodesTracker { + /// Create a new instance of `RedundantNodesChecker`, starting from a `RawStorageProof`. + pub fn new(proof_nodes_count: usize) -> Self { + Self { proof_nodes_count, recorder: BTreeSet::new() } + } + + /// Ensure that all the nodes in the proof have been accessed. + pub fn ensure_no_unused_nodes(self) -> Result<(), Error> { + if self.proof_nodes_count != self.recorder.len() { + return Err(Error::UnusedNodes) + } + + Ok(()) + } +} + +impl TrieRecorder for AccessedNodesTracker { + fn record(&mut self, access: TrieAccess) { + match access { + TrieAccess::NodeOwned { hash, .. } | + TrieAccess::EncodedNode { hash, .. } | + TrieAccess::Value { hash, .. } => { + self.recorder.insert(hash); + }, + _ => {}, + } + } + + fn trie_nodes_recorded_for_key(&self, _key: &[u8]) -> RecordedForKey { + RecordedForKey::None + } +} + +#[cfg(test)] +pub mod tests { + use super::*; + use crate::{tests::create_storage_proof, StorageProof}; + use hash_db::Hasher; + use trie_db::{Trie, TrieDBBuilder}; + + type Hash = ::Out; + type Layout = crate::LayoutV1; + + const TEST_DATA: &[(&[u8], &[u8])] = + &[(b"key1", &[1; 64]), (b"key2", &[2; 64]), (b"key3", &[3; 64])]; + + #[test] + fn proof_with_unused_nodes_is_rejected() { + let (raw_proof, root) = create_storage_proof::(TEST_DATA); + let proof = StorageProof::new(raw_proof.clone()); + let proof_nodes_count = proof.len(); + + let mut accessed_nodes_tracker = AccessedNodesTracker::::new(proof_nodes_count); + { + let db = proof.clone().into_memory_db(); + let trie = TrieDBBuilder::::new(&db, &root) + .with_recorder(&mut accessed_nodes_tracker) + .build(); + + trie.get(b"key1").unwrap().unwrap(); + trie.get(b"key2").unwrap().unwrap(); + trie.get(b"key3").unwrap().unwrap(); + } + assert_eq!(accessed_nodes_tracker.ensure_no_unused_nodes(), Ok(())); + + let mut accessed_nodes_tracker = AccessedNodesTracker::::new(proof_nodes_count); + { + let db = proof.into_memory_db(); + let trie = TrieDBBuilder::::new(&db, &root) + .with_recorder(&mut accessed_nodes_tracker) + .build(); + + trie.get(b"key1").unwrap().unwrap(); + trie.get(b"key2").unwrap().unwrap(); + } + assert_eq!(accessed_nodes_tracker.ensure_no_unused_nodes(), Err(Error::UnusedNodes)); + } +} diff --git a/substrate/primitives/trie/src/lib.rs b/substrate/primitives/trie/src/lib.rs index 0c14e3af196d..ef6b6a5743c2 100644 --- a/substrate/primitives/trie/src/lib.rs +++ b/substrate/primitives/trie/src/lib.rs @@ -21,6 +21,7 @@ extern crate alloc; +pub mod accessed_nodes_tracker; #[cfg(feature = "std")] pub mod cache; mod error; @@ -28,6 +29,7 @@ mod node_codec; mod node_header; #[cfg(feature = "std")] pub mod recorder; +pub mod recorder_ext; mod storage_proof; mod trie_codec; mod trie_stream; @@ -46,7 +48,7 @@ use hash_db::{Hasher, Prefix}; pub use memory_db::{prefixed_key, HashKey, KeyFunction, PrefixedKey}; /// The Substrate format implementation of `NodeCodec`. pub use node_codec::NodeCodec; -pub use storage_proof::{CompactProof, StorageProof}; +pub use storage_proof::{CompactProof, StorageProof, StorageProofError}; /// Trie codec reexport, mainly child trie support /// for trie compact proof. pub use trie_codec::{decode_compact, encode_compact, Error as CompactProofError}; @@ -64,6 +66,9 @@ pub use trie_db::{proof::VerifyError, MerkleValue}; /// The Substrate format implementation of `TrieStream`. pub use trie_stream::TrieStream; +/// Raw storage proof type (just raw trie nodes). +pub type RawStorageProof = Vec>; + /// substrate trie layout pub struct LayoutV0(PhantomData); @@ -616,6 +621,50 @@ mod tests { type MemoryDBMeta = memory_db::MemoryDB, trie_db::DBValue>; + pub fn create_trie( + data: &[(&[u8], &[u8])], + ) -> (MemoryDB, trie_db::TrieHash) { + let mut db = MemoryDB::default(); + let mut root = Default::default(); + + { + let mut trie = trie_db::TrieDBMutBuilder::::new(&mut db, &mut root).build(); + for (k, v) in data { + trie.insert(k, v).expect("Inserts data"); + } + } + + let mut recorder = Recorder::::new(); + { + let trie = trie_db::TrieDBBuilder::::new(&mut db, &mut root) + .with_recorder(&mut recorder) + .build(); + for (k, _v) in data { + trie.get(k).unwrap(); + } + } + + (db, root) + } + + pub fn create_storage_proof( + data: &[(&[u8], &[u8])], + ) -> (RawStorageProof, trie_db::TrieHash) { + let (db, root) = create_trie::(data); + + let mut recorder = Recorder::::new(); + { + let trie = trie_db::TrieDBBuilder::::new(&db, &root) + .with_recorder(&mut recorder) + .build(); + for (k, _v) in data { + trie.get(k).unwrap(); + } + } + + (recorder.drain().into_iter().map(|record| record.data).collect(), root) + } + fn hashed_null_node() -> TrieHash { ::hashed_null_node() } diff --git a/substrate/primitives/trie/src/recorder.rs b/substrate/primitives/trie/src/recorder.rs index 22a22b33b370..2886577eddc6 100644 --- a/substrate/primitives/trie/src/recorder.rs +++ b/substrate/primitives/trie/src/recorder.rs @@ -145,7 +145,7 @@ impl Recorder { /// Convert the recording to a [`StorageProof`]. /// - /// In contrast to [`Self::drain_storage_proof`] this doesn't consumes and doesn't clears the + /// In contrast to [`Self::drain_storage_proof`] this doesn't consume and doesn't clear the /// recordings. /// /// Returns the [`StorageProof`]. @@ -429,7 +429,8 @@ impl<'a, H: Hasher> trie_db::TrieRecorder for TrieRecorder<'a, H> { #[cfg(test)] mod tests { use super::*; - use trie_db::{Trie, TrieDBBuilder, TrieDBMutBuilder, TrieHash, TrieMut, TrieRecorder}; + use crate::tests::create_trie; + use trie_db::{Trie, TrieDBBuilder, TrieRecorder}; type MemoryDB = crate::MemoryDB; type Layout = crate::LayoutV1; @@ -438,23 +439,9 @@ mod tests { const TEST_DATA: &[(&[u8], &[u8])] = &[(b"key1", &[1; 64]), (b"key2", &[2; 64]), (b"key3", &[3; 64]), (b"key4", &[4; 64])]; - fn create_trie() -> (MemoryDB, TrieHash) { - let mut db = MemoryDB::default(); - let mut root = Default::default(); - - { - let mut trie = TrieDBMutBuilder::::new(&mut db, &mut root).build(); - for (k, v) in TEST_DATA { - trie.insert(k, v).expect("Inserts data"); - } - } - - (db, root) - } - #[test] fn recorder_works() { - let (db, root) = create_trie(); + let (db, root) = create_trie::(TEST_DATA); let recorder = Recorder::default(); @@ -498,7 +485,7 @@ mod tests { #[test] fn recorder_transactions_rollback_work() { - let (db, root) = create_trie(); + let (db, root) = create_trie::(TEST_DATA); let recorder = Recorder::default(); let mut stats = vec![RecorderStats::default()]; @@ -547,7 +534,7 @@ mod tests { #[test] fn recorder_transactions_commit_work() { - let (db, root) = create_trie(); + let (db, root) = create_trie::(TEST_DATA); let recorder = Recorder::default(); @@ -586,7 +573,7 @@ mod tests { #[test] fn recorder_transactions_commit_and_rollback_work() { - let (db, root) = create_trie(); + let (db, root) = create_trie::(TEST_DATA); let recorder = Recorder::default(); @@ -645,7 +632,7 @@ mod tests { #[test] fn recorder_transaction_accessed_keys_works() { let key = TEST_DATA[0].0; - let (db, root) = create_trie(); + let (db, root) = create_trie::(TEST_DATA); let recorder = Recorder::default(); diff --git a/substrate/primitives/trie/src/recorder_ext.rs b/substrate/primitives/trie/src/recorder_ext.rs new file mode 100644 index 000000000000..866d5b72c5d6 --- /dev/null +++ b/substrate/primitives/trie/src/recorder_ext.rs @@ -0,0 +1,47 @@ +// This file is part of Substrate. + +// Copyright (C) Parity Technologies (UK) Ltd. +// SPDX-License-Identifier: Apache-2.0 + +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Extension for the default recorder. + +use crate::RawStorageProof; +use alloc::{collections::BTreeSet, vec::Vec}; +use trie_db::{Recorder, TrieLayout}; + +/// Convenience extension for the `Recorder` struct. +/// +/// Used to deduplicate some logic. +pub trait RecorderExt +where + Self: Sized, +{ + /// Convert the recorder into a `BTreeSet`. + fn into_set(self) -> BTreeSet>; + + /// Convert the recorder into a `RawStorageProof`, avoiding duplicate nodes. + fn into_raw_storage_proof(self) -> RawStorageProof { + // The recorder may record the same trie node multiple times, + // and we don't want duplicate nodes in our proofs + // => let's deduplicate it by collecting to a BTreeSet first + self.into_set().into_iter().collect() + } +} + +impl RecorderExt for Recorder { + fn into_set(mut self) -> BTreeSet> { + self.drain().into_iter().map(|record| record.data).collect::>() + } +} diff --git a/substrate/primitives/trie/src/storage_proof.rs b/substrate/primitives/trie/src/storage_proof.rs index e46c49be19cb..a9f6298742f6 100644 --- a/substrate/primitives/trie/src/storage_proof.rs +++ b/substrate/primitives/trie/src/storage_proof.rs @@ -25,6 +25,13 @@ use scale_info::TypeInfo; // with `LayoutV0`. use crate::LayoutV1 as Layout; +/// Error associated with the `storage_proof` module. +#[derive(Encode, Decode, Clone, Eq, PartialEq, Debug, TypeInfo)] +pub enum StorageProofError { + /// The proof contains duplicate nodes. + DuplicateNodes, +} + /// A proof that some set of key-value pairs are included in the storage trie. The proof contains /// the storage values so that the partial storage backend can be reconstructed by a verifier that /// does not already have access to the key-value pairs. @@ -43,6 +50,22 @@ impl StorageProof { StorageProof { trie_nodes: BTreeSet::from_iter(trie_nodes) } } + /// Constructs a storage proof from a subset of encoded trie nodes in a storage backend. + /// + /// Returns an error if the provided subset of encoded trie nodes contains duplicates. + pub fn new_with_duplicate_nodes_check( + trie_nodes: impl IntoIterator>, + ) -> Result { + let mut trie_nodes_set = BTreeSet::new(); + for node in trie_nodes { + if !trie_nodes_set.insert(node) { + return Err(StorageProofError::DuplicateNodes); + } + } + + Ok(StorageProof { trie_nodes: trie_nodes_set }) + } + /// Returns a new empty proof. /// /// An empty proof is capable of only proving trivial statements (ie. that an empty set of @@ -56,6 +79,11 @@ impl StorageProof { self.trie_nodes.is_empty() } + /// Returns the number of nodes in the proof. + pub fn len(&self) -> usize { + self.trie_nodes.len() + } + /// Convert into an iterator over encoded trie nodes in lexicographical order constructed /// from the proof. pub fn into_iter_nodes(self) -> impl Sized + DoubleEndedIterator> { @@ -198,3 +226,23 @@ impl CompactProof { Ok((db, root)) } } + +#[cfg(test)] +pub mod tests { + use super::*; + use crate::{tests::create_storage_proof, StorageProof}; + + type Layout = crate::LayoutV1; + + const TEST_DATA: &[(&[u8], &[u8])] = + &[(b"key1", &[1; 64]), (b"key2", &[2; 64]), (b"key3", &[3; 64]), (b"key11", &[4; 64])]; + + #[test] + fn proof_with_duplicate_nodes_is_rejected() { + let (raw_proof, _root) = create_storage_proof::(TEST_DATA); + assert!(matches!( + StorageProof::new_with_duplicate_nodes_check(raw_proof), + Err(StorageProofError::DuplicateNodes) + )); + } +}