1
// Copyright 2025 Moonbeam Foundation.
2
// This file is part of Moonbeam.
3

            
4
// Moonbeam is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Moonbeam is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
16

            
17
#![cfg_attr(not(feature = "std"), no_std)]
18

            
19
use cumulus_primitives_core::relay_chain;
20
use sp_core::H256;
21
use sp_runtime::traits::{Hash, HashingFor};
22
use sp_state_machine::{Backend, TrieBackend, TrieBackendBuilder};
23
use sp_std::vec::Vec;
24
use sp_trie::{HashDBT, MemoryDB, StorageProof, EMPTY_PREFIX};
25

            
26
#[cfg(test)]
27
mod tests;
28

            
29
#[derive(Debug, PartialEq)]
30
pub enum ProofError {
31
	// The storage root in the proof does not match the expected storage root.
32
	RootMismatch,
33
	// The proof is invalid.
34
	Proof,
35
	// The key is not present in the proof.
36
	Absent,
37
	// Block number is not present
38
	BlockNumberNotPresent,
39
}
40

            
41
pub type RawStorageProof = Vec<Vec<u8>>;
42

            
43
/// A storage proof checker. It is used to verify a storage proof against a well-known storage root,
44
/// and return the value of the storage item if the proof is valid.
45
#[derive(Debug)]
46
pub struct StorageProofChecker<H>
47
where
48
	H: Hash,
49
{
50
	trie_backend: TrieBackend<MemoryDB<H>, H>,
51
}
52

            
53
impl<H: Hash> StorageProofChecker<H> {
54
	/// Create a new storage proof checker. Returns an error if the given `storage_root` is not
55
	/// present in the proof.
56
10
	pub fn new(
57
10
		storage_root: H::Out,
58
10
		raw_proof: impl IntoIterator<Item = Vec<u8>>,
59
10
	) -> Result<Self, ProofError> {
60
10
		let storage_proof = StorageProof::new(raw_proof);
61
10
		let db = storage_proof.into_memory_db::<H>();
62
10

            
63
10
		if !db.contains(&storage_root, EMPTY_PREFIX) {
64
3
			return Err(ProofError::RootMismatch);
65
7
		}
66
7
		let trie_backend = TrieBackendBuilder::new(db, storage_root).build();
67
7

            
68
7
		Ok(Self { trie_backend })
69
10
	}
70

            
71
	/// Returns the value of the storage given the key, if the proof is valid.
72
	/// Returns `Err` if the proof is invalid, or if the value specified by the key according to the
73
	/// proof is not present.
74
13
	pub fn read_entry(&self, key: &[u8]) -> Result<Vec<u8>, ProofError> {
75
13
		self.trie_backend
76
13
			.storage(key)
77
13
			.map_err(|_| ProofError::Proof)?
78
13
			.ok_or(ProofError::Absent)
79
13
	}
80

            
81
2
	pub fn read_entries(&self, keys: &[&[u8]]) -> Result<Vec<Vec<u8>>, ProofError> {
82
2
		let mut values = Vec::new();
83
8
		for key in keys {
84
7
			let value = self.read_entry(key)?;
85
6
			values.push(value);
86
		}
87
1
		Ok(values)
88
2
	}
89
}
90

            
91
3
pub fn verify_entry(
92
3
	expected_root: H256,
93
3
	proof: impl IntoIterator<Item = Vec<u8>>,
94
3
	key: &[u8],
95
3
) -> Result<Vec<u8>, ProofError> {
96
2
	let proof_checker =
97
3
		StorageProofChecker::<HashingFor<relay_chain::Block>>::new(expected_root, proof)?;
98

            
99
2
	proof_checker.read_entry(key)
100
3
}
101

            
102
3
pub fn verify_entries(
103
3
	expected_root: H256,
104
3
	proof: impl IntoIterator<Item = Vec<u8>>,
105
3
	keys: &[&[u8]],
106
3
) -> Result<Vec<Vec<u8>>, ProofError> {
107
2
	let proof_checker =
108
3
		StorageProofChecker::<HashingFor<relay_chain::Block>>::new(expected_root, proof)?;
109

            
110
2
	proof_checker.read_entries(keys)
111
3
}