1#[cfg(test)]
4mod tests;
5
6use parking_lot::Mutex;
7use schnellru::{ByLength, LruMap};
8use sp_consensus_slots::Slot;
9use sp_consensus_subspace::{PotNextSlotInput, PotParametersChange};
10use std::num::NonZeroU32;
11use std::sync::Arc;
12use subspace_core_primitives::pot::{PotCheckpoints, PotOutput, PotSeed};
13
14#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
15struct CacheKey {
16 seed: PotSeed,
17 slot_iterations: NonZeroU32,
18}
19
20#[derive(Debug, Clone)]
21struct CacheValue {
22 checkpoints: Arc<Mutex<Option<PotCheckpoints>>>,
23}
24
25#[derive(Debug, Clone)]
27pub struct PotVerifier {
28 genesis_seed: PotSeed,
29 cache: Arc<Mutex<LruMap<CacheKey, CacheValue>>>,
30}
31
32impl PotVerifier {
33 pub fn new(genesis_seed: PotSeed, cache_size: u32) -> Self {
34 Self {
35 genesis_seed,
36 cache: Arc::new(Mutex::new(LruMap::new(ByLength::new(cache_size)))),
37 }
38 }
39
40 pub fn inject_verified_checkpoints(
42 &self,
43 seed: PotSeed,
44 slot_iterations: NonZeroU32,
45 checkpoints: PotCheckpoints,
46 ) {
47 self.cache.lock().insert(
48 CacheKey {
49 seed,
50 slot_iterations,
51 },
52 CacheValue {
53 checkpoints: Arc::new(Mutex::new(Some(checkpoints))),
54 },
55 );
56 }
57
58 pub fn genesis_seed(&self) -> PotSeed {
60 self.genesis_seed
61 }
62
63 pub fn get_checkpoints(
66 &self,
67 slot_iterations: NonZeroU32,
68 seed: PotSeed,
69 ) -> Option<PotCheckpoints> {
70 self.calculate_checkpoints(slot_iterations, seed, true)
71 }
72
73 pub fn try_get_checkpoints(
75 &self,
76 slot_iterations: NonZeroU32,
77 seed: PotSeed,
78 ) -> Option<PotCheckpoints> {
79 let cache_key = CacheKey {
80 seed,
81 slot_iterations,
82 };
83
84 self.cache
85 .lock()
86 .get(&cache_key)
87 .and_then(|value| value.checkpoints.try_lock()?.as_ref().copied())
88 }
89
90 pub fn is_output_valid(
100 &self,
101 input: PotNextSlotInput,
102 slots: Slot,
103 output: PotOutput,
104 maybe_parameters_change: Option<PotParametersChange>,
105 ) -> bool {
106 self.is_output_valid_internal(input, slots, output, maybe_parameters_change, true)
107 }
108
109 pub fn try_is_output_valid(
113 &self,
114 input: PotNextSlotInput,
115 slots: Slot,
116 output: PotOutput,
117 maybe_parameters_change: Option<PotParametersChange>,
118 ) -> bool {
119 self.is_output_valid_internal(input, slots, output, maybe_parameters_change, false)
120 }
121
122 fn is_output_valid_internal(
123 &self,
124 mut input: PotNextSlotInput,
125 slots: Slot,
126 output: PotOutput,
127 maybe_parameters_change: Option<PotParametersChange>,
128 do_proving_if_necessary: bool,
129 ) -> bool {
130 let mut slots = u64::from(slots);
131
132 loop {
133 let maybe_calculated_checkpoints = self.calculate_checkpoints(
134 input.slot_iterations,
135 input.seed,
136 do_proving_if_necessary,
137 );
138
139 let Some(calculated_checkpoints) = maybe_calculated_checkpoints else {
140 return false;
141 };
142 let calculated_output = calculated_checkpoints.output();
143
144 slots -= 1;
145
146 if slots == 0 {
147 return output == calculated_output;
148 }
149
150 input = PotNextSlotInput::derive(
151 input.slot_iterations,
152 input.slot,
153 calculated_output,
154 &maybe_parameters_change,
155 );
156 }
157 }
158
159 fn calculate_checkpoints(
161 &self,
162 slot_iterations: NonZeroU32,
163 seed: PotSeed,
164 do_proving_if_necessary: bool,
165 ) -> Option<PotCheckpoints> {
166 let cache_key = CacheKey {
167 seed,
168 slot_iterations,
169 };
170
171 loop {
172 let mut cache = self.cache.lock();
173 let maybe_cache_value = cache.get(&cache_key).cloned();
174 if let Some(cache_value) = maybe_cache_value {
175 drop(cache);
176 let correct_checkpoints = cache_value.checkpoints.lock();
177 if let Some(correct_checkpoints) = *correct_checkpoints {
178 return Some(correct_checkpoints);
179 }
180
181 continue;
183 }
184
185 if !do_proving_if_necessary {
186 return None;
188 }
189
190 let cache_value = CacheValue {
191 checkpoints: Arc::default(),
192 };
193 let checkpoints = Arc::clone(&cache_value.checkpoints);
194 let mut checkpoints = checkpoints
196 .try_lock()
197 .expect("No one can access this mutex yet; qed");
198 cache.insert(cache_key, cache_value);
200 drop(cache);
202
203 let proving_result = subspace_proof_of_time::prove(seed, slot_iterations);
204
205 let Ok(generated_checkpoints) = proving_result else {
206 drop(checkpoints);
208
209 let maybe_removed_cache_value = self.cache.lock().remove(&cache_key);
211 if let Some(removed_cache_value) = maybe_removed_cache_value {
212 let removed_verified_value = removed_cache_value.checkpoints.lock().is_some();
215 if removed_verified_value {
216 self.cache.lock().insert(cache_key, removed_cache_value);
217 }
218 }
219 return None;
220 };
221
222 checkpoints.replace(generated_checkpoints);
223 return Some(generated_checkpoints);
224 }
225 }
226
227 pub fn verify_checkpoints(
229 &self,
230 seed: PotSeed,
231 slot_iterations: NonZeroU32,
232 checkpoints: &PotCheckpoints,
233 ) -> bool {
234 self.verify_checkpoints_internal(seed, slot_iterations, checkpoints)
235 }
236
237 fn verify_checkpoints_internal(
238 &self,
239 seed: PotSeed,
240 slot_iterations: NonZeroU32,
241 checkpoints: &PotCheckpoints,
242 ) -> bool {
243 let cache_key = CacheKey {
244 seed,
245 slot_iterations,
246 };
247
248 loop {
249 let mut cache = self.cache.lock();
250 if let Some(cache_value) = cache.get(&cache_key).cloned() {
251 drop(cache);
252 let correct_checkpoints = cache_value.checkpoints.lock();
253 if let Some(correct_checkpoints) = correct_checkpoints.as_ref() {
254 return checkpoints == correct_checkpoints;
255 }
256
257 continue;
259 }
260
261 let cache_value = CacheValue {
262 checkpoints: Arc::default(),
263 };
264 let correct_checkpoints = Arc::clone(&cache_value.checkpoints);
265 let mut correct_checkpoints = correct_checkpoints
267 .try_lock()
268 .expect("No one can access this mutex yet; qed");
269 cache.insert(cache_key, cache_value);
271 drop(cache);
273
274 let verified_successfully =
275 subspace_proof_of_time::verify(seed, slot_iterations, checkpoints.as_slice())
276 .unwrap_or_default();
277
278 if !verified_successfully {
279 drop(correct_checkpoints);
281
282 let maybe_removed_cache_value = self.cache.lock().remove(&cache_key);
284 if let Some(removed_cache_value) = maybe_removed_cache_value {
285 let removed_verified_value = removed_cache_value.checkpoints.lock().is_some();
288 if removed_verified_value {
289 self.cache.lock().insert(cache_key, removed_cache_value);
290 }
291 }
292 return false;
293 }
294
295 correct_checkpoints.replace(*checkpoints);
297
298 return true;
299 }
300 }
301}