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
use crate::{ProofError, StorageProofChecker};
18
use cumulus_primitives_core::relay_chain;
19
use frame_support::assert_ok;
20
use parity_scale_codec::Encode;
21
use sp_core::H256;
22
use sp_runtime::traits::HashingFor;
23
use sp_std::collections::btree_map::BTreeMap;
24
use sp_trie::PrefixedMemoryDB;
25
use std::{fs::File, io::Write};
26

            
27
// Storage Root: 767caa877bcea0d34dd515a202b75efa41bffbc9f814ab59e2c1c96716d4c65d
28
pub const STORAGE_ROOT: &[u8] = &[
29
	118, 124, 170, 135, 123, 206, 160, 211, 77, 213, 21, 162, 2, 183, 94, 250, 65, 191, 251, 201,
30
	248, 20, 171, 89, 226, 193, 201, 103, 22, 212, 198, 93,
31
];
32

            
33
// Timestamp key: f0c365c3cf59d671eb72da0e7a4113c49f1f0515f462cdcf84e0f1d6045dfcbb
34
pub const TIMESTAMP_KEY: &[u8] = &[
35
	240, 195, 101, 195, 207, 89, 214, 113, 235, 114, 218, 14, 122, 65, 19, 196, 159, 31, 5, 21,
36
	244, 98, 205, 207, 132, 224, 241, 214, 4, 93, 252, 187,
37
];
38

            
39
// Total Issuance Key: c2261276cc9d1f8598ea4b6a74b15c2f57c875e4cff74148e4628f264b974c80
40
pub const TOTAL_ISSUANCE_KEY: &[u8] = &[
41
	194, 38, 18, 118, 204, 157, 31, 133, 152, 234, 75, 106, 116, 177, 92, 47, 87, 200, 117, 228,
42
	207, 247, 65, 72, 228, 98, 143, 38, 75, 151, 76, 128,
43
];
44

            
45
// Treasury Approval Key: 89d139e01a5eb2256f222e5fc5dbe6b33c9c1284130706f5aea0c8b3d4c54d89
46
pub const TREASURY_APPROVALS_KEY: &[u8] = &[
47
	137, 209, 57, 224, 26, 94, 178, 37, 111, 34, 46, 95, 197, 219, 230, 179, 60, 156, 18, 132, 19,
48
	7, 6, 245, 174, 160, 200, 179, 212, 197, 77, 137,
49
];
50

            
51
4
pub fn mock_proof() -> Vec<Vec<u8>> {
52
	use parity_scale_codec::Decode;
53
4
	Vec::decode(&mut &include_bytes!("../proof").to_vec()[..]).unwrap()
54
4
}
55

            
56
#[test]
57
1
fn test_storage_proof_checker_valid() {
58
1
	let proof = mock_proof();
59
1
	let storage_root = H256::from_slice(STORAGE_ROOT);
60
1

            
61
1
	assert_ok!(StorageProofChecker::<HashingFor<relay_chain::Block>>::new(
62
1
		storage_root,
63
1
		proof
64
1
	));
65
1
}
66

            
67
#[test]
68
1
fn test_storage_proof_checker_root_mismatch() {
69
1
	let proof = mock_proof();
70
1
	// A different storage root
71
1
	let storage_root = H256::from_slice(
72
1
		hex::decode("767caa877bcea0d34dd515a202b75efa41bffbc9f814ab59e2c1c96716d4c65e")
73
1
			.unwrap()
74
1
			.as_slice(),
75
1
	);
76
1

            
77
1
	assert_eq!(
78
1
		StorageProofChecker::<HashingFor<relay_chain::Block>>::new(storage_root, proof)
79
1
			.unwrap_err(),
80
1
		ProofError::RootMismatch
81
1
	);
82
1
}
83

            
84
#[test]
85
1
fn test_storage_proof_read_entries() {
86
1
	let proof = mock_proof();
87
1
	let storage_root = H256::from_slice(STORAGE_ROOT);
88
1
	let proof_checker =
89
1
		StorageProofChecker::<HashingFor<relay_chain::Block>>::new(storage_root, proof).unwrap();
90
1

            
91
1
	let timestamp = proof_checker.read_entry(TIMESTAMP_KEY).unwrap();
92
1
	let total_issuance = proof_checker.read_entry(TOTAL_ISSUANCE_KEY).unwrap();
93
1
	let approvals = proof_checker.read_entry(TREASURY_APPROVALS_KEY).unwrap();
94
1

            
95
1
	assert_eq!(
96
1
		timestamp,
97
1
		1_708_190_328_000u64.encode(),
98
		"Timestamp does not match the expected value"
99
	);
100
1
	assert_eq!(
101
1
		total_issuance,
102
1
		14_123_366_426_803_276_130u128.encode(),
103
		"Total issuance does not match the expected value"
104
	);
105
1
	assert_eq!(
106
1
		approvals,
107
1
		vec![607, 608, 609, 610, 611, 612, 613, 614, 615, 616, 617, 618, 619, 620, 621, 622, 623]
108
1
			.encode(),
109
		"Treasury Approvals does not match the expected value"
110
	);
111
1
}
112

            
113
#[test]
114
1
fn test_storage_proof_checker_absent() {
115
1
	let proof = mock_proof();
116
1
	let storage_root = H256::from_slice(STORAGE_ROOT);
117
1

            
118
1
	let proof_checker =
119
1
		StorageProofChecker::<HashingFor<relay_chain::Block>>::new(storage_root, proof).unwrap();
120
1

            
121
1
	// A key that is not present in the proof
122
1
	let key =
123
1
		hex::decode("89d139e01a5eb2256f222e5fc5dbe6b33c9c1284130706f5aea0c8b3d4c54d2c").unwrap();
124
1
	let value = proof_checker.read_entry(&key);
125
1
	assert_eq!(value.err(), Some(ProofError::Absent));
126
1
}
127

            
128
20
pub fn build_mocked_proof(
129
20
	entries: Vec<(Vec<u8>, Vec<u8>)>,
130
20
	keys: Vec<Vec<u8>>,
131
20
) -> (H256, Vec<Vec<u8>>) {
132
20
	let (db, root) = PrefixedMemoryDB::<HashingFor<relay_chain::Block>>::default_with_root();
133
20
	let state_version = Default::default();
134
20
	let mut backend = sp_state_machine::TrieBackendBuilder::new(db, root).build();
135
20

            
136
16805
	entries.into_iter().for_each(|(key, value)| {
137
16805
		backend.insert(vec![(None, vec![(key, Some(value))])], state_version);
138
16805
	});
139
20

            
140
20
	let root = *backend.root();
141
20
	let proof = sp_state_machine::prove_read(backend, keys).expect("prove read");
142
20

            
143
20
	(root, proof.into_iter_nodes().collect())
144
20
}
145

            
146
// Generate mocked proofs for the benchmarks. The proofs are generated for a set of
147
// keys and values, and then stored in a file. The proofs are then used in the benchmarks
148
// to simulate the proofs obtained from the relay chain.
149
#[test]
150
1
fn test_mocked_storage_proof() {
151
1
	// This set of entries generates proofs with number of nodes in proof increasing by 100 for
152
1
	// each entry (Number of Proof Node, Number of Entries)
153
1
	let entries: Vec<(u32, u32)> = vec![
154
1
		(100, 95),
155
1
		(200, 190),
156
1
		(300, 270),
157
1
		(400, 320),
158
1
		(500, 370),
159
1
		(600, 420),
160
1
		(700, 470),
161
1
		(800, 530),
162
1
		(900, 630),
163
1
		(1000, 730),
164
1
		(1100, 830),
165
1
		(1200, 930),
166
1
		(1300, 1030),
167
1
		(1400, 1130),
168
1
		(1500, 1230),
169
1
		(1600, 1330),
170
1
		(1700, 1430),
171
1
		(1800, 1530),
172
1
		(1900, 1630),
173
1
		(2000, 1730),
174
1
	];
175
1

            
176
1
	let mut proofs = BTreeMap::new();
177
20
	entries.into_iter().for_each(|(i, x)| {
178
16805
		let keys: Vec<Vec<u8>> = (1..x as u128).into_iter().map(|y| y.encode()).collect();
179
20
		let entries = keys
180
20
			.iter()
181
20
			.enumerate()
182
16805
			.map(|(i, key)| (key.clone(), (i as u128).encode()))
183
20
			.collect();
184
20

            
185
20
		let (state_root, proof) = build_mocked_proof(entries, keys);
186
20
		proofs.insert(i, (state_root, proof));
187
20
	});
188
1

            
189
1
	let mut file = File::create(format!("benchmark_proofs")).unwrap();
190
1
	file.write_all(&proofs.encode()).unwrap();
191
1
}