Ensure key ownership proof is optimal

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::recorder_ext::Error as RecorderExtError;
 
 	#[test]
 	fn verify_chain_message_rejects_message_with_too_large_declared_weight() {
@@ -541,7 +542,7 @@ mod tests {
 				target::verify_messages_proof::<OnThisChainBridge>(proof, 10)
 			},),
 			Err(VerificationError::HeaderChain(HeaderChainError::StorageProof(
-				StorageProofError::DuplicateNodesInProof
+				StorageProofError::RecorderExt(RecorderExtError::DuplicateNodes.into())
 			))),
 		);
 	}
@@ -553,7 +554,9 @@ mod tests {
 				proof.storage_proof.push(vec![42]);
 				target::verify_messages_proof::<OnThisChainBridge>(proof, 10)
 			},),
-			Err(VerificationError::StorageProof(StorageProofError::UnusedNodesInTheProof)),
+			Err(VerificationError::StorageProof(StorageProofError::RecorderExt(
+				RecorderExtError::UnusedNodes.into()
+			))),
 		);
 	}
 

bridges/primitives/runtime/src/storage_proof.rs

@@ -21,15 +21,14 @@ 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,
-	TrieDBBuilder, TrieError, TrieHash,
+	read_trie_value,
+	recorder_ext::{Error as RecorderExtError, RecorderExt, RedundantNodesChecker},
+	LayoutV1, MemoryDB, Recorder, Trie, TrieConfiguration, TrieDBBuilder, TrieError, TrieHash,
 };
 
-/// Raw storage proof type (just raw trie nodes).
-pub type RawStorageProof = Vec<Vec<u8>>;
-
 /// Storage proof size requirements.
 ///
 /// This is currently used by benchmarks when generating storage proofs.
@@ -51,10 +50,9 @@ pub struct StorageProofChecker<H>
 where
 	H: Hasher,
 {
-	proof_nodes_count: usize,
 	root: H::Out,
 	db: MemoryDB<H>,
-	recorder: Recorder<LayoutV1<H>>,
+	redundant_nodes_checker: RedundantNodesChecker<H::Out>,
 }
 
 impl<H> StorageProofChecker<H>
@@ -65,52 +63,37 @@ 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<Self, Error> {
-		// 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 (recorder, proof) =
+			RedundantNodesChecker::new(proof).map_err(|e| Error::RecorderExt(e.into()))?;
 
 		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, redundant_nodes_checker: 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::<BTreeSet<_>>();
-		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.redundant_nodes_checker
+			.ensure_no_unused_nodes()
+			.map_err(|e| Error::RecorderExt(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<Option<Vec<u8>>, Error> {
 		// LayoutV1 or LayoutV0 is identical for proof that only read values.
-		read_trie_value::<LayoutV1<H>, _>(&self.db, &self.root, key, Some(&mut self.recorder), None)
-			.map_err(|_| Error::StorageValueUnavailable)
+		read_trie_value::<LayoutV1<H>, _>(
+			&self.db,
+			&self.root,
+			key,
+			Some(&mut self.redundant_nodes_checker),
+			None,
+		)
+		.map_err(|_| Error::StorageValueUnavailable)
 	}
 
 	/// Reads and decodes a value from the available subset of storage. If the value cannot be read
@@ -145,10 +128,8 @@ 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 originating in the `recorder_ext` module.
+	RecorderExt(StrippableError<RecorderExtError>),
 	/// Expected storage root is missing from the proof.
 	StorageRootMismatch,
 	/// Unable to reach expected storage value using provided trie nodes.
@@ -202,15 +183,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::<BTreeSet<_>>()
-		.into_iter()
-		.collect())
+	Ok(recorder.into_raw_storage_proof())
 }
 
 #[cfg(test)]
@@ -250,7 +223,7 @@ pub mod tests {
 
 		assert_eq!(
 			StorageProofChecker::<sp_core::Blake2Hasher>::new(root, proof).map(drop),
-			Err(Error::DuplicateNodesInProof),
+			Err(Error::RecorderExt(RecorderExtError::DuplicateNodes.into())),
 		);
 	}
 
@@ -267,6 +240,9 @@ pub mod tests {
 		assert_eq!(checker.ensure_no_unused_nodes(), Ok(()));
 
 		let checker = StorageProofChecker::<sp_core::Blake2Hasher>::new(root, proof).unwrap();
-		assert_eq!(checker.ensure_no_unused_nodes(), Err(Error::UnusedNodesInTheProof));
+		assert_eq!(
+			checker.ensure_no_unused_nodes(),
+			Err(Error::RecorderExt(RecorderExtError::UnusedNodes.into()))
+		);
 	}
 }

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::recorder_ext::{RecorderExt, RedundantNodesChecker};
 
 #[frame_support::pallet]
 pub mod pallet {
@@ -118,6 +119,16 @@ impl<T: Config> Pallet<T> {
 			}
 		})
 	}
+
+	fn full_id_validators() -> Vec<(T::ValidatorId, T::FullIdentification)> {
+		<Session<T>>::validators()
+			.into_iter()
+			.filter_map(|validator| {
+				T::FullIdentificationOf::convert(validator.clone())
+					.map(|full_id| (validator, full_id))
+			})
+			.collect::<Vec<_>>()
+	}
 }
 
 impl<T: Config> ValidatorSet<T::AccountId> for Pallet<T> {
@@ -264,46 +275,34 @@ impl<T: Config> ProvingTrie<T> {
 		Ok(ProvingTrie { db, root })
 	}
 
-	fn from_nodes(root: T::Hash, nodes: &[Vec<u8>]) -> 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<Vec<Vec<u8>>> {
 		let mut recorder = Recorder::<LayoutV0<T::Hashing>>::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| <IdentificationTuple<T>>::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.
 	pub fn root(&self) -> &T::Hash {
 		&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<IdentificationTuple<T>> {
-		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<T::Hash>>,
+	) -> Option<IdentificationTuple<T>> {
+		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<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T
 
 	fn prove(key: (KeyTypeId, D)) -> Option<Self::Proof> {
 		let session = <Session<T>>::current_index();
-		let validators = <Session<T>>::validators()
-			.into_iter()
-			.filter_map(|validator| {
-				T::FullIdentificationOf::convert(validator.clone())
-					.map(|full_id| (validator, full_id))
-			})
-			.collect::<Vec<_>>();
+		let validators = Self::full_id_validators();
 
 		let count = validators.len() as ValidatorCount;
 
@@ -343,30 +336,35 @@ impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T
 	}
 
 	fn check_proof(key: (KeyTypeId, D), proof: Self::Proof) -> Option<IdentificationTuple<T>> {
-		let (id, data) = key;
-
-		if proof.session == <Session<T>>::current_index() {
-			<Session<T>>::key_owner(id, data.as_ref()).and_then(|owner| {
-				T::FullIdentificationOf::convert(owner.clone()).and_then(move |id| {
-					let count = <Session<T>>::validators().len() as ValidatorCount;
-
-					if count != proof.validator_count {
-						return None
-					}
+		fn print_error<E: Debug>(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 == <Session<T>>::current_index() {
+			let validators = Self::full_id_validators();
+			let count = validators.len() as ValidatorCount;
+			let trie = ProvingTrie::<T>::generate_for(validators).ok()?;
+			(trie.root, count)
 		} else {
-			let (root, count) = <HistoricalSessions<T>>::get(&proof.session)?;
-
-			if count != proof.validator_count {
-				return None
-			}
+			<HistoricalSessions<T>>::get(&proof.session)?
+		};
 
-			let trie = ProvingTrie::<T>::from_nodes(root, &proof.trie_nodes);
-			trie.query(id, data.as_ref())
+		if count != proof.validator_count {
+			return None
 		}
+
+		let (mut redundant_nodes_checker, proof) =
+			RedundantNodesChecker::<T::Hash>::new(proof.trie_nodes)
+				.map_err(print_error)
+				.ok()?;
+		let trie = ProvingTrie::<T>::from_proof(root, proof);
+		let res = trie.query(id, data.as_ref(), Some(&mut redundant_nodes_checker))?;
+		redundant_nodes_checker.ensure_no_unused_nodes().map_err(print_error).ok()?;
+		Some(res)
 	}
 }
 

substrate/primitives/trie/src/lib.rs

@@ -28,6 +28,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;
@@ -64,6 +65,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<Vec<u8>>;
+
 /// substrate trie layout
 pub struct LayoutV0<H>(PhantomData<H>);
 

substrate/primitives/trie/src/recorder.rs

@@ -145,7 +145,7 @@ impl<H: Hasher> Recorder<H> {
 
 	/// 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 clears the
 	/// recordings.
 	///
 	/// Returns the [`StorageProof`].