1use async_trait::async_trait;
2use rayon::prelude::*;
3use slog::{Logger, debug, info};
4use std::{
5 collections::{BTreeMap, BTreeSet, HashMap},
6 sync::Arc,
7 time::Duration,
8};
9
10use mithril_common::{
11 StdResult,
12 crypto_helper::{MKMap, MKMapNode, MKTree, MKTreeStorer},
13 entities::{
14 BlockNumber, BlockRange, CardanoTransaction, CardanoTransactionsSetProof, TransactionHash,
15 },
16 logging::LoggerExtensions,
17 signable_builder::BlockRangeRootRetriever,
18};
19use mithril_resource_pool::ResourcePool;
20
21#[cfg_attr(test, mockall::automock)]
23#[async_trait]
24pub trait ProverService: Sync + Send {
25 async fn compute_transactions_proofs(
27 &self,
28 up_to: BlockNumber,
29 transaction_hashes: &[TransactionHash],
30 ) -> StdResult<Vec<CardanoTransactionsSetProof>>;
31
32 async fn compute_cache(&self, up_to: BlockNumber) -> StdResult<()>;
34}
35
36#[cfg_attr(test, mockall::automock)]
38#[async_trait]
39pub trait TransactionsRetriever: Sync + Send {
40 async fn get_by_hashes(
42 &self,
43 hashes: Vec<TransactionHash>,
44 up_to: BlockNumber,
45 ) -> StdResult<Vec<CardanoTransaction>>;
46
47 async fn get_by_block_ranges(
49 &self,
50 block_ranges: Vec<BlockRange>,
51 ) -> StdResult<Vec<CardanoTransaction>>;
52}
53
54pub struct MithrilProverService<S: MKTreeStorer> {
56 transaction_retriever: Arc<dyn TransactionsRetriever>,
57 block_range_root_retriever: Arc<dyn BlockRangeRootRetriever<S>>,
58 mk_map_pool: ResourcePool<MKMap<BlockRange, MKMapNode<BlockRange, S>, S>>,
59 logger: Logger,
60}
61
62impl<S: MKTreeStorer> MithrilProverService<S> {
63 pub fn new(
65 transaction_retriever: Arc<dyn TransactionsRetriever>,
66 block_range_root_retriever: Arc<dyn BlockRangeRootRetriever<S>>,
67 mk_map_pool_size: usize,
68 logger: Logger,
69 ) -> Self {
70 Self {
71 transaction_retriever,
72 block_range_root_retriever,
73 mk_map_pool: ResourcePool::new(mk_map_pool_size, vec![]),
74 logger: logger.new_with_component_name::<Self>(),
75 }
76 }
77
78 async fn get_block_ranges(
79 &self,
80 transaction_hashes: &[TransactionHash],
81 up_to: BlockNumber,
82 ) -> StdResult<Vec<BlockRange>> {
83 let transactions = self
84 .transaction_retriever
85 .get_by_hashes(transaction_hashes.to_vec(), up_to)
86 .await?;
87 let block_ranges = transactions
88 .iter()
89 .map(|t| BlockRange::from_block_number(t.block_number))
90 .collect::<BTreeSet<_>>();
91
92 Ok(block_ranges.into_iter().collect::<Vec<_>>())
93 }
94
95 async fn get_all_transactions_for_block_ranges(
97 &self,
98 block_ranges: &[BlockRange],
99 ) -> StdResult<HashMap<BlockRange, Vec<CardanoTransaction>>> {
100 let mut block_ranges_map = HashMap::new();
101 let transactions = self
102 .transaction_retriever
103 .get_by_block_ranges(block_ranges.to_vec())
104 .await?;
105 for transaction in transactions {
106 let block_range = BlockRange::from_block_number(transaction.block_number);
107 let block_range_transactions: &mut Vec<_> =
108 block_ranges_map.entry(block_range).or_insert(vec![]);
109 block_range_transactions.push(transaction)
110 }
111
112 Ok(block_ranges_map)
113 }
114}
115
116#[async_trait]
117impl<S: MKTreeStorer> ProverService for MithrilProverService<S> {
118 async fn compute_transactions_proofs(
119 &self,
120 up_to: BlockNumber,
121 transaction_hashes: &[TransactionHash],
122 ) -> StdResult<Vec<CardanoTransactionsSetProof>> {
123 let block_ranges_transactions = self.get_block_ranges(transaction_hashes, up_to).await?;
125 let block_range_transactions = self
126 .get_all_transactions_for_block_ranges(&block_ranges_transactions)
127 .await?;
128
129 let mk_trees: StdResult<Vec<(BlockRange, MKTree<S>)>> = block_range_transactions
131 .into_iter()
132 .map(|(block_range, transactions)| {
133 let mk_tree = MKTree::new(&transactions)?;
134 Ok((block_range, mk_tree))
135 })
136 .collect();
137 let mk_trees = BTreeMap::from_iter(mk_trees?);
138
139 let acquire_timeout = Duration::from_millis(1000);
141 let mut mk_map = self.mk_map_pool.acquire_resource(acquire_timeout)?;
142
143 for (block_range, mk_tree) in mk_trees {
145 mk_map.replace(block_range, mk_tree.into())?;
146 }
147
148 match mk_map.compute_proof(transaction_hashes) {
150 Ok(mk_proof) => {
151 self.mk_map_pool.give_back_resource_pool_item(mk_map)?;
152 let mk_proof_leaves = mk_proof.leaves();
153 let transaction_hashes_certified: Vec<TransactionHash> = transaction_hashes
154 .iter()
155 .filter(|hash| mk_proof_leaves.contains(&hash.as_str().into()))
156 .cloned()
157 .collect();
158
159 Ok(vec![CardanoTransactionsSetProof::new(
160 transaction_hashes_certified,
161 mk_proof,
162 )])
163 }
164 _ => Ok(vec![]),
165 }
166 }
167
168 async fn compute_cache(&self, up_to: BlockNumber) -> StdResult<()> {
169 let pool_size = self.mk_map_pool.size();
170 info!(
171 self.logger, "Starts computing the Merkle map pool resource of size {pool_size}";
172 "up_to_block_number" => *up_to,
173 );
174 let mk_map_cache = self
175 .block_range_root_retriever
176 .compute_merkle_map_from_block_range_roots(up_to)
177 .await?;
178 let mk_maps_new = (1..=pool_size)
179 .into_par_iter()
180 .map(|i| {
181 debug!(
182 self.logger,
183 "Computing the Merkle map pool resource {i}/{pool_size}"
184 );
185 mk_map_cache.clone()
186 })
187 .collect::<Vec<MKMap<_, _, _>>>();
188 debug!(self.logger, "Draining the Merkle map pool");
189 let discriminant_new = self.mk_map_pool.discriminant()? + 1;
190 self.mk_map_pool.set_discriminant(discriminant_new)?;
191 self.mk_map_pool.clear();
192 debug!(
193 self.logger,
194 "Giving back new resources to the Merkle map pool"
195 );
196 mk_maps_new
197 .into_iter()
198 .map(|mk_map| self.mk_map_pool.give_back_resource(mk_map, discriminant_new))
199 .collect::<StdResult<Vec<_>>>()?;
200 info!(
201 self.logger,
202 "Completed computing the Merkle map pool resource of size {pool_size}"
203 );
204
205 Ok(())
206 }
207}
208
209#[cfg(test)]
210mod tests {
211 use anyhow::anyhow;
212 use mithril_common::crypto_helper::{
213 MKMap, MKMapNode, MKTreeNode, MKTreeStoreInMemory, MKTreeStorer,
214 };
215 use mithril_common::entities::CardanoTransaction;
216 use mithril_common::test_utils::CardanoTransactionsBuilder;
217 use mockall::mock;
218 use mockall::predicate::eq;
219
220 use crate::test_tools::TestLogger;
221
222 use super::*;
223
224 mock! {
225 pub BlockRangeRootRetrieverImpl<S: MKTreeStorer> { }
226
227 #[async_trait]
228 impl<S: MKTreeStorer> BlockRangeRootRetriever<S> for BlockRangeRootRetrieverImpl<S> {
229 async fn retrieve_block_range_roots<'a>(
230 &'a self,
231 up_to_beacon: BlockNumber,
232 ) -> StdResult<Box<dyn Iterator<Item = (BlockRange, MKTreeNode)> + 'a>>;
233
234 async fn compute_merkle_map_from_block_range_roots(
235 &self,
236 up_to_beacon: BlockNumber,
237 ) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange, S>, S>>;
238 }
239 }
240
241 mod test_data {
242 use mithril_common::crypto_helper::MKTreeStoreInMemory;
243
244 use super::*;
245
246 pub fn filter_transactions_for_indices(
247 indices: &[usize],
248 transactions: &[CardanoTransaction],
249 ) -> Vec<CardanoTransaction> {
250 transactions
251 .iter()
252 .enumerate()
253 .filter(|(i, _)| indices.contains(i))
254 .map(|(_, t)| t.to_owned())
255 .collect()
256 }
257
258 pub fn map_to_transaction_hashes(
259 transactions: &[CardanoTransaction],
260 ) -> Vec<TransactionHash> {
261 transactions.iter().map(|t| t.transaction_hash.clone()).collect()
262 }
263
264 pub fn transactions_group_by_block_range(
265 transactions: &[CardanoTransaction],
266 ) -> BTreeMap<BlockRange, Vec<CardanoTransaction>> {
267 let mut block_ranges_map = BTreeMap::new();
268 for transaction in transactions {
269 let block_range = BlockRange::from_block_number(transaction.block_number);
270 let block_range_transactions: &mut Vec<_> =
271 block_ranges_map.entry(block_range).or_insert(vec![]);
272 block_range_transactions.push(transaction.to_owned())
273 }
274
275 block_ranges_map
276 }
277
278 pub fn filter_transactions_for_block_ranges(
279 block_ranges: &[BlockRange],
280 transactions: &[CardanoTransaction],
281 ) -> Vec<CardanoTransaction> {
282 transactions
283 .iter()
284 .filter(|t| block_ranges.contains(&BlockRange::from_block_number(t.block_number)))
285 .map(|t| t.to_owned())
286 .collect()
287 }
288
289 pub fn compute_mk_map_from_block_ranges_map(
290 block_ranges_map: BTreeMap<BlockRange, Vec<CardanoTransaction>>,
291 ) -> MKMap<BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>, MKTreeStoreInMemory>
292 {
293 MKMap::new_from_iter(
294 block_ranges_map.into_iter().map(|(block_range, transactions)| {
295 (
296 block_range,
297 MKMapNode::TreeNode(
298 MKTree::<MKTreeStoreInMemory>::new(&transactions)
299 .unwrap()
300 .compute_root()
301 .unwrap()
302 .clone(),
303 ),
304 )
305 }),
306 )
307 .unwrap()
308 }
309
310 pub fn compute_beacon_from_transactions(
311 transactions: &[CardanoTransaction],
312 ) -> BlockNumber {
313 let max_transaction = transactions.iter().max_by_key(|t| t.block_number).unwrap();
314 max_transaction.block_number
315 }
316
317 pub struct TestData {
318 pub transaction_hashes_to_prove: Vec<TransactionHash>,
319 pub block_ranges_map: BTreeMap<BlockRange, Vec<CardanoTransaction>>,
320 pub block_ranges_to_prove: Vec<BlockRange>,
321 pub all_transactions_in_block_ranges_to_prove: Vec<CardanoTransaction>,
322 pub beacon: BlockNumber,
323 }
324
325 pub fn build_test_data(
326 transactions_to_prove: &[CardanoTransaction],
327 transactions: &[CardanoTransaction],
328 ) -> TestData {
329 let transaction_hashes_to_prove = map_to_transaction_hashes(transactions_to_prove);
330 let block_ranges_map = transactions_group_by_block_range(transactions);
331 let block_ranges_map_to_prove =
332 transactions_group_by_block_range(transactions_to_prove);
333 let block_ranges_to_prove =
334 block_ranges_map_to_prove.keys().cloned().collect::<Vec<_>>();
335 let all_transactions_in_block_ranges_to_prove =
336 filter_transactions_for_block_ranges(&block_ranges_to_prove, transactions);
337 let beacon = compute_beacon_from_transactions(transactions);
338
339 TestData {
340 transaction_hashes_to_prove,
341 block_ranges_map,
342 block_ranges_to_prove,
343 all_transactions_in_block_ranges_to_prove,
344 beacon,
345 }
346 }
347 }
348
349 fn build_prover<F, G, S: MKTreeStorer + 'static>(
350 transaction_retriever_mock_config: F,
351 block_range_root_retriever_mock_config: G,
352 ) -> MithrilProverService<S>
353 where
354 F: FnOnce(&mut MockTransactionsRetriever),
355 G: FnOnce(&mut MockBlockRangeRootRetrieverImpl<S>),
356 {
357 let mut transaction_retriever = MockTransactionsRetriever::new();
358 transaction_retriever_mock_config(&mut transaction_retriever);
359 let mut block_range_root_retriever = MockBlockRangeRootRetrieverImpl::new();
360 block_range_root_retriever_mock_config(&mut block_range_root_retriever);
361 let mk_map_pool_size = 1;
362
363 MithrilProverService::new(
364 Arc::new(transaction_retriever),
365 Arc::new(block_range_root_retriever),
366 mk_map_pool_size,
367 TestLogger::stdout(),
368 )
369 }
370
371 #[tokio::test]
372 async fn compute_proof_for_one_set_of_three_certified_transactions() {
373 let transactions = CardanoTransactionsBuilder::new()
374 .max_transactions_per_block(1)
375 .blocks_per_block_range(3)
376 .build_block_ranges(5);
377 let transactions_to_prove =
378 test_data::filter_transactions_for_indices(&[1, 2, 4], &transactions);
379 let test_data = test_data::build_test_data(&transactions_to_prove, &transactions);
380 let prover = build_prover(
381 |transaction_retriever_mock| {
382 let transaction_hashes_to_prove = test_data.transaction_hashes_to_prove.clone();
383 let transactions_to_prove = transactions_to_prove.clone();
384 transaction_retriever_mock
385 .expect_get_by_hashes()
386 .with(eq(transaction_hashes_to_prove), eq(test_data.beacon))
387 .return_once(move |_, _| Ok(transactions_to_prove));
388
389 let block_ranges_to_prove = test_data.block_ranges_to_prove.clone();
390 let all_transactions_in_block_ranges_to_prove =
391 test_data.all_transactions_in_block_ranges_to_prove.clone();
392 transaction_retriever_mock
393 .expect_get_by_block_ranges()
394 .with(eq(block_ranges_to_prove))
395 .return_once(move |_| Ok(all_transactions_in_block_ranges_to_prove));
396 },
397 |block_range_root_retriever_mock| {
398 let block_ranges_map = test_data.block_ranges_map.clone();
399 block_range_root_retriever_mock
400 .expect_compute_merkle_map_from_block_range_roots()
401 .return_once(|_| {
402 Ok(test_data::compute_mk_map_from_block_ranges_map(
403 block_ranges_map,
404 ))
405 });
406 },
407 );
408 prover.compute_cache(test_data.beacon).await.unwrap();
409
410 let transactions_set_proof = prover
411 .compute_transactions_proofs(test_data.beacon, &test_data.transaction_hashes_to_prove)
412 .await
413 .unwrap();
414
415 assert_eq!(transactions_set_proof.len(), 1);
416 assert_eq!(
417 transactions_set_proof[0].transactions_hashes(),
418 test_data.transaction_hashes_to_prove
419 );
420 transactions_set_proof[0].verify().unwrap();
421 }
422
423 #[tokio::test]
424 async fn cant_compute_proof_for_not_yet_certified_transaction() {
425 let transactions = CardanoTransactionsBuilder::new()
426 .max_transactions_per_block(1)
427 .blocks_per_block_range(3)
428 .build_block_ranges(5);
429 let transactions_to_prove =
430 test_data::filter_transactions_for_indices(&[1, 2, 4], &transactions);
431 let test_data = test_data::build_test_data(&transactions_to_prove, &transactions);
432 let prover = build_prover(
433 |transaction_retriever_mock| {
434 let transaction_hashes_to_prove = test_data.transaction_hashes_to_prove.clone();
435 transaction_retriever_mock
436 .expect_get_by_hashes()
437 .with(eq(transaction_hashes_to_prove), eq(test_data.beacon))
438 .return_once(move |_, _| Ok(vec![]));
439 transaction_retriever_mock
440 .expect_get_by_block_ranges()
441 .with(eq(vec![]))
442 .return_once(move |_| Ok(vec![]));
443 },
444 |block_range_root_retriever_mock| {
445 let block_ranges_map = test_data.block_ranges_map.clone();
446 block_range_root_retriever_mock
447 .expect_compute_merkle_map_from_block_range_roots()
448 .return_once(|_| {
449 Ok(test_data::compute_mk_map_from_block_ranges_map(
450 block_ranges_map,
451 ))
452 });
453 },
454 );
455 prover.compute_cache(test_data.beacon).await.unwrap();
456
457 let transactions_set_proof = prover
458 .compute_transactions_proofs(test_data.beacon, &test_data.transaction_hashes_to_prove)
459 .await
460 .unwrap();
461
462 assert_eq!(transactions_set_proof.len(), 0);
463 }
464
465 #[tokio::test]
466 async fn cant_compute_proof_for_unknown_transaction() {
467 let transactions = CardanoTransactionsBuilder::new()
468 .max_transactions_per_block(1)
469 .blocks_per_block_range(3)
470 .build_block_ranges(5);
471 let transactions_to_prove = test_data::filter_transactions_for_indices(&[], &transactions);
472 let mut test_data = test_data::build_test_data(&transactions_to_prove, &transactions);
473 test_data.transaction_hashes_to_prove = vec!["tx-unknown-123".to_string()];
474 let prover = build_prover(
475 |transaction_retriever_mock| {
476 let transaction_hashes_to_prove = test_data.transaction_hashes_to_prove.clone();
477 let transactions_to_prove = transactions_to_prove.clone();
478 transaction_retriever_mock
479 .expect_get_by_hashes()
480 .with(eq(transaction_hashes_to_prove), eq(test_data.beacon))
481 .return_once(move |_, _| Ok(transactions_to_prove));
482
483 let block_ranges_to_prove = test_data.block_ranges_to_prove.clone();
484 let all_transactions_in_block_ranges_to_prove =
485 test_data.all_transactions_in_block_ranges_to_prove.clone();
486 transaction_retriever_mock
487 .expect_get_by_block_ranges()
488 .with(eq(block_ranges_to_prove))
489 .return_once(move |_| Ok(all_transactions_in_block_ranges_to_prove));
490 },
491 |block_range_root_retriever_mock| {
492 let block_ranges_map = test_data.block_ranges_map.clone();
493 block_range_root_retriever_mock
494 .expect_compute_merkle_map_from_block_range_roots()
495 .return_once(|_| {
496 Ok(test_data::compute_mk_map_from_block_ranges_map(
497 block_ranges_map,
498 ))
499 });
500 },
501 );
502 prover.compute_cache(test_data.beacon).await.unwrap();
503
504 let transactions_set_proof = prover
505 .compute_transactions_proofs(test_data.beacon, &test_data.transaction_hashes_to_prove)
506 .await
507 .unwrap();
508
509 assert_eq!(transactions_set_proof.len(), 0);
510 }
511
512 #[tokio::test]
513 async fn compute_proof_for_one_set_of_three_certified_transactions_and_two_unknowns() {
514 let transactions = CardanoTransactionsBuilder::new()
515 .max_transactions_per_block(1)
516 .blocks_per_block_range(3)
517 .build_block_ranges(5);
518 let transactions_to_prove =
519 test_data::filter_transactions_for_indices(&[1, 2, 4], &transactions);
520 let transaction_hashes_unknown =
521 vec!["tx-unknown-123".to_string(), "tx-unknown-456".to_string()];
522 let mut test_data = test_data::build_test_data(&transactions_to_prove, &transactions);
523 let transaction_hashes_known = test_data.transaction_hashes_to_prove.clone();
524 test_data.transaction_hashes_to_prove = [
525 test_data.transaction_hashes_to_prove.clone(),
526 transaction_hashes_unknown,
527 ]
528 .concat();
529 let prover = build_prover(
530 |transaction_retriever_mock| {
531 let transaction_hashes_to_prove = test_data.transaction_hashes_to_prove.clone();
532 let transactions_to_prove = transactions_to_prove.clone();
533 transaction_retriever_mock
534 .expect_get_by_hashes()
535 .with(eq(transaction_hashes_to_prove), eq(test_data.beacon))
536 .return_once(move |_, _| Ok(transactions_to_prove));
537
538 let block_ranges_to_prove = test_data.block_ranges_to_prove.clone();
539 let all_transactions_in_block_ranges_to_prove =
540 test_data.all_transactions_in_block_ranges_to_prove.clone();
541 transaction_retriever_mock
542 .expect_get_by_block_ranges()
543 .with(eq(block_ranges_to_prove))
544 .return_once(move |_| Ok(all_transactions_in_block_ranges_to_prove));
545 },
546 |block_range_root_retriever_mock| {
547 let block_ranges_map = test_data.block_ranges_map.clone();
548 block_range_root_retriever_mock
549 .expect_compute_merkle_map_from_block_range_roots()
550 .return_once(|_| {
551 Ok(test_data::compute_mk_map_from_block_ranges_map(
552 block_ranges_map,
553 ))
554 });
555 },
556 );
557 prover.compute_cache(test_data.beacon).await.unwrap();
558
559 let transactions_set_proof = prover
560 .compute_transactions_proofs(test_data.beacon, &test_data.transaction_hashes_to_prove)
561 .await
562 .unwrap();
563
564 assert_eq!(transactions_set_proof.len(), 1);
565 assert_eq!(
566 transactions_set_proof[0].transactions_hashes(),
567 transaction_hashes_known
568 );
569 transactions_set_proof[0].verify().unwrap();
570 }
571
572 #[tokio::test]
573 async fn cant_compute_proof_if_transaction_retriever_fails() {
574 let transactions = CardanoTransactionsBuilder::new()
575 .max_transactions_per_block(1)
576 .blocks_per_block_range(3)
577 .build_block_ranges(5);
578 let transactions_to_prove =
579 test_data::filter_transactions_for_indices(&[1, 2, 4], &transactions);
580 let test_data = test_data::build_test_data(&transactions_to_prove, &transactions);
581 let prover = build_prover::<_, _, MKTreeStoreInMemory>(
582 |transaction_retriever_mock| {
583 transaction_retriever_mock
584 .expect_get_by_hashes()
585 .returning(|_, _| Err(anyhow!("Error")));
586 },
587 |block_range_root_retriever_mock| {
588 block_range_root_retriever_mock
589 .expect_compute_merkle_map_from_block_range_roots()
590 .return_once(|_| MKMap::new(&[]));
591 },
592 );
593 prover.compute_cache(test_data.beacon).await.unwrap();
594
595 prover
596 .compute_transactions_proofs(test_data.beacon, &test_data.transaction_hashes_to_prove)
597 .await
598 .expect_err("Should have failed because of transaction retriever failure");
599 }
600
601 #[tokio::test]
602 async fn cant_compute_proof_if_block_range_root_retriever_fails() {
603 let transactions = CardanoTransactionsBuilder::new()
604 .max_transactions_per_block(1)
605 .blocks_per_block_range(3)
606 .build_block_ranges(5);
607 let transactions_to_prove =
608 test_data::filter_transactions_for_indices(&[1, 2, 4], &transactions);
609 let test_data = test_data::build_test_data(&transactions_to_prove, &transactions);
610 let prover = build_prover::<_, _, MKTreeStoreInMemory>(
611 |transaction_retriever_mock| {
612 let transactions_to_prove = transactions_to_prove.clone();
613 transaction_retriever_mock
614 .expect_get_by_hashes()
615 .return_once(move |_, _| Ok(transactions_to_prove));
616
617 let all_transactions_in_block_ranges_to_prove =
618 test_data.all_transactions_in_block_ranges_to_prove.clone();
619 transaction_retriever_mock
620 .expect_get_by_block_ranges()
621 .return_once(move |_| Ok(all_transactions_in_block_ranges_to_prove));
622 },
623 |block_range_root_retriever_mock| {
624 block_range_root_retriever_mock
625 .expect_compute_merkle_map_from_block_range_roots()
626 .return_once(|_| Err(anyhow!("Error")));
627 },
628 );
629
630 prover
631 .compute_transactions_proofs(test_data.beacon, &test_data.transaction_hashes_to_prove)
632 .await
633 .expect_err("Should have failed because of block range root retriever failure");
634 }
635}