1use frame_support::PalletError;
20use sp_core::RuntimeDebug;
21use sp_std::vec::Vec;
22use sp_trie::{
23 accessed_nodes_tracker::AccessedNodesTracker, read_trie_value, LayoutV1, MemoryDB, StorageProof,
24};
25
26use codec::{Decode, Encode};
27use hash_db::{HashDB, Hasher, EMPTY_PREFIX};
28use scale_info::TypeInfo;
29#[cfg(feature = "test-helpers")]
30use sp_trie::{recorder_ext::RecorderExt, Recorder, TrieDBBuilder, TrieError, TrieHash};
31#[cfg(feature = "test-helpers")]
32use trie_db::{Trie, TrieConfiguration, TrieDBMut};
33
34#[derive(Clone, Encode, Decode, RuntimeDebug, PartialEq, Eq, PalletError, TypeInfo)]
36pub enum StorageProofError {
37 UnableToGenerateTrieProof,
39 InvalidProof,
41 UnsortedEntries,
43 UnavailableKey,
45 EmptyVal,
47 DecodeError,
49 UnusedKey,
51
52 StorageRootMismatch,
54 StorageValueUnavailable,
56 DuplicateNodes,
58}
59
60impl From<sp_trie::StorageProofError> for StorageProofError {
61 fn from(e: sp_trie::StorageProofError) -> Self {
62 match e {
63 sp_trie::StorageProofError::DuplicateNodes => StorageProofError::DuplicateNodes,
64 }
65 }
66}
67
68impl From<sp_trie::accessed_nodes_tracker::Error> for StorageProofError {
69 fn from(e: sp_trie::accessed_nodes_tracker::Error) -> Self {
70 match e {
71 sp_trie::accessed_nodes_tracker::Error::UnusedNodes => StorageProofError::UnusedKey,
72 }
73 }
74}
75
76pub type RawStorageProof = sp_trie::RawStorageProof;
78
79pub fn raw_storage_proof_size(raw_storage_proof: &RawStorageProof) -> usize {
81 raw_storage_proof
82 .iter()
83 .fold(0usize, |sum, node| sum.saturating_add(node.len()))
84}
85
86#[cfg(feature = "test-helpers")]
90#[derive(Clone, Copy, Debug, Default)]
91pub struct UnverifiedStorageProofParams {
92 pub db_size: Option<u32>,
94}
95
96#[cfg(feature = "test-helpers")]
97impl UnverifiedStorageProofParams {
98 pub fn from_db_size(db_size: u32) -> Self {
100 Self { db_size: Some(db_size) }
101 }
102}
103
104pub struct StorageProofChecker<H>
109where
110 H: Hasher,
111{
112 root: H::Out,
113 db: MemoryDB<H>,
114 accessed_nodes_tracker: AccessedNodesTracker<H::Out>,
115}
116
117impl<H> StorageProofChecker<H>
118where
119 H: Hasher,
120{
121 pub fn new(root: H::Out, proof: RawStorageProof) -> Result<Self, StorageProofError> {
125 let proof = StorageProof::new_with_duplicate_nodes_check(proof)?;
126
127 let recorder = AccessedNodesTracker::new(proof.len());
128
129 let db = proof.into_memory_db();
130 if !db.contains(&root, EMPTY_PREFIX) {
131 return Err(StorageProofError::StorageRootMismatch)
132 }
133
134 Ok(StorageProofChecker { root, db, accessed_nodes_tracker: recorder })
135 }
136
137 pub fn ensure_no_unused_nodes(self) -> Result<(), StorageProofError> {
140 self.accessed_nodes_tracker.ensure_no_unused_nodes().map_err(Into::into)
141 }
142
143 pub fn read_value(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageProofError> {
146 read_trie_value::<LayoutV1<H>, _>(
148 &self.db,
149 &self.root,
150 key,
151 Some(&mut self.accessed_nodes_tracker),
152 None,
153 )
154 .map_err(|_| StorageProofError::StorageValueUnavailable)
155 }
156
157 pub fn read_and_decode_value<T: Decode>(
161 &mut self,
162 key: &[u8],
163 ) -> Result<Option<T>, StorageProofError> {
164 self.read_value(key).and_then(|v| {
165 v.map(|v| {
166 T::decode(&mut &v[..]).map_err(|e| {
167 log::warn!(target: "bridge-storage-proofs", "read_and_decode_value error: {e:?}");
168 StorageProofError::DecodeError
169 })
170 })
171 .transpose()
172 })
173 }
174
175 pub fn read_and_decode_mandatory_value<T: Decode>(
179 &mut self,
180 key: &[u8],
181 ) -> Result<T, StorageProofError> {
182 self.read_and_decode_value(key)?.ok_or(StorageProofError::EmptyVal)
183 }
184
185 pub fn read_and_decode_opt_value<T: Decode>(
189 &mut self,
190 key: &[u8],
191 ) -> Result<Option<T>, StorageProofError> {
192 match self.read_and_decode_value(key) {
193 Ok(outbound_lane_data) => Ok(outbound_lane_data),
194 Err(StorageProofError::StorageValueUnavailable) => Ok(None),
195 Err(e) => Err(e),
196 }
197 }
198}
199
200#[cfg(feature = "test-helpers")]
202pub fn grow_storage_value(mut value: Vec<u8>, params: &UnverifiedStorageProofParams) -> Vec<u8> {
203 if let Some(db_size) = params.db_size {
204 if db_size as usize > value.len() {
205 value.extend(sp_std::iter::repeat(42u8).take(db_size as usize - value.len()));
206 }
207 }
208 value
209}
210
211#[cfg(feature = "test-helpers")]
217pub fn grow_storage_proof<L: TrieConfiguration>(
218 trie: &mut TrieDBMut<L>,
219 prefix: Vec<u8>,
220 num_extra_nodes: usize,
221) {
222 use sp_trie::TrieMut;
223
224 let mut added_nodes = 0;
225 for i in 0..prefix.len() {
226 let mut prefix = prefix[0..=i].to_vec();
227 let first_nibble = (prefix[i] & 0xf0) >> 4;
229 let second_nibble = prefix[i] & 0x0f;
230
231 for branch in 1..=15 {
233 if added_nodes >= num_extra_nodes {
234 return
235 }
236
237 prefix[i] = (first_nibble.wrapping_add(branch) % 16) << 4;
239 trie.insert(&prefix, &[0; 32])
240 .map_err(|_| "TrieMut::insert has failed")
241 .expect("TrieMut::insert should not fail in benchmarks");
242 added_nodes += 1;
243 }
244
245 for branch in 1..=15 {
247 if added_nodes >= num_extra_nodes {
248 return
249 }
250
251 prefix[i] = (first_nibble << 4) | (second_nibble.wrapping_add(branch) % 16);
252 trie.insert(&prefix, &[0; 32])
253 .map_err(|_| "TrieMut::insert has failed")
254 .expect("TrieMut::insert should not fail in benchmarks");
255 added_nodes += 1;
256 }
257 }
258
259 assert_eq!(added_nodes, num_extra_nodes)
260}
261
262#[cfg(feature = "test-helpers")]
264pub fn record_all_keys<L: TrieConfiguration, DB>(
265 db: &DB,
266 root: &TrieHash<L>,
267) -> Result<RawStorageProof, sp_std::boxed::Box<TrieError<L>>>
268where
269 DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>,
270{
271 let mut recorder = Recorder::<L>::new();
272 let trie = TrieDBBuilder::<L>::new(db, root).with_recorder(&mut recorder).build();
273 for x in trie.iter()? {
274 let (key, _) = x?;
275 trie.get(&key)?;
276 }
277
278 Ok(recorder.into_raw_storage_proof())
279}
280
281#[cfg(feature = "std")]
285pub fn craft_valid_storage_proof() -> (sp_core::H256, RawStorageProof) {
286 use sp_state_machine::{backend::Backend, prove_read, InMemoryBackend};
287
288 let state_version = sp_runtime::StateVersion::default();
289
290 let backend = <InMemoryBackend<sp_core::Blake2Hasher>>::from((
292 sp_std::vec![
293 (None, vec![(b"key1".to_vec(), Some(b"value1".to_vec()))]),
294 (None, vec![(b"key2".to_vec(), Some(b"value2".to_vec()))]),
295 (None, vec![(b"key3".to_vec(), Some(b"value3".to_vec()))]),
296 (None, vec![(b"key4".to_vec(), Some((42u64, 42u32, 42u16, 42u8).encode()))]),
297 (None, vec![(b"key11".to_vec(), Some(vec![0u8; 32]))]),
299 ],
300 state_version,
301 ));
302 let root = backend.storage_root(sp_std::iter::empty(), state_version).0;
303 let proof =
304 prove_read(backend, &[&b"key1"[..], &b"key2"[..], &b"key4"[..], &b"key22"[..]]).unwrap();
305
306 (root, proof.into_nodes().into_iter().collect())
307}
308
309#[cfg(test)]
310pub mod tests_for_storage_proof_checker {
311 use super::*;
312 use codec::Encode;
313
314 #[test]
315 fn storage_proof_check() {
316 let (root, proof) = craft_valid_storage_proof();
317
318 let mut checker =
320 <StorageProofChecker<sp_core::Blake2Hasher>>::new(root, proof.clone()).unwrap();
321 assert_eq!(checker.read_value(b"key1"), Ok(Some(b"value1".to_vec())));
322 assert_eq!(checker.read_value(b"key2"), Ok(Some(b"value2".to_vec())));
323 assert_eq!(checker.read_value(b"key4"), Ok(Some((42u64, 42u32, 42u16, 42u8).encode())));
324 assert_eq!(
325 checker.read_value(b"key11111"),
326 Err(StorageProofError::StorageValueUnavailable)
327 );
328 assert_eq!(checker.read_value(b"key22"), Ok(None));
329 assert_eq!(checker.read_and_decode_value(b"key4"), Ok(Some((42u64, 42u32, 42u16, 42u8))),);
330 assert!(matches!(
331 checker.read_and_decode_value::<[u8; 64]>(b"key4"),
332 Err(StorageProofError::DecodeError),
333 ));
334
335 assert_eq!(
337 <StorageProofChecker<sp_core::Blake2Hasher>>::new(sp_core::H256::random(), proof).err(),
338 Some(StorageProofError::StorageRootMismatch)
339 );
340 }
341
342 #[test]
343 fn proof_with_unused_items_is_rejected() {
344 let (root, proof) = craft_valid_storage_proof();
345
346 let mut checker =
347 StorageProofChecker::<sp_core::Blake2Hasher>::new(root, proof.clone()).unwrap();
348 checker.read_value(b"key1").unwrap().unwrap();
349 checker.read_value(b"key2").unwrap();
350 checker.read_value(b"key4").unwrap();
351 checker.read_value(b"key22").unwrap();
352 assert_eq!(checker.ensure_no_unused_nodes(), Ok(()));
353
354 let checker = StorageProofChecker::<sp_core::Blake2Hasher>::new(root, proof).unwrap();
355 assert_eq!(checker.ensure_no_unused_nodes(), Err(StorageProofError::UnusedKey));
356 }
357}