Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpraynaud
jpraynaud committed Apr 30, 2024
1 parent d582f65 commit 09000ea
Show file tree
Hide file tree
Showing 2 changed files with 153 additions and 128 deletions.
3 changes: 2 additions & 1 deletion mithril-aggregator/src/dependency_injection/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1366,7 +1366,8 @@ impl DependenciesBuilder {
/// build Prover service
pub async fn build_prover_service(&mut self) -> Result<Arc<dyn ProverService>> {
let transaction_retriever = self.get_transaction_repository().await?;
let service = MithrilProverService::new(transaction_retriever);
let block_range_root_retriever = self.get_transaction_repository().await?;
let service = MithrilProverService::new(transaction_retriever, block_range_root_retriever);

Ok(Arc::new(service))
}
Expand Down
278 changes: 151 additions & 127 deletions mithril-aggregator/src/services/prover.rs
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
use std::{collections::BTreeMap, sync::Arc};
use std::{collections::HashMap, sync::Arc};

use anyhow::Context;
use async_trait::async_trait;

use mithril_common::{
crypto_helper::{MKMap, MKMapNode, MKTree, MKTreeNode},
crypto_helper::MKTree,
entities::{
BlockRange, CardanoDbBeacon, CardanoTransaction, CardanoTransactionsSetProof,
TransactionHash,
},
signable_builder::BlockRangeRootRetriever,
StdResult,
};

Expand Down Expand Up @@ -41,70 +41,38 @@ pub trait TransactionsRetriever: Sync + Send {
/// Mithril prover
pub struct MithrilProverService {
transaction_retriever: Arc<dyn TransactionsRetriever>,
block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
}

impl MithrilProverService {
/// Create a new Mithril prover
pub fn new(transaction_retriever: Arc<dyn TransactionsRetriever>) -> Self {
pub fn new(
transaction_retriever: Arc<dyn TransactionsRetriever>,
block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
) -> Self {
Self {
transaction_retriever,
block_range_root_retriever,
}
}

async fn get_transactions_by_hashes_with_block_range(
async fn get_transactions_grouped_by_block_range(
&self,
hashes: Vec<TransactionHash>,
) -> StdResult<Vec<(BlockRange, CardanoTransaction)>> {
let transactions = self.transaction_retriever.get_by_hashes(hashes).await?;
let transactions_with_block_range = transactions
.into_iter()
.map(|transaction| {
let block_range = BlockRange::from_block_number(transaction.block_number);
(block_range, transaction)
})
.collect::<Vec<_>>();
Ok(transactions_with_block_range)
}

fn compute_merkle_map_from_transactions(
&self,
transactions: Vec<CardanoTransaction>,
) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange>>> {
let mut transactions_by_block_ranges: BTreeMap<BlockRange, Vec<TransactionHash>> =
BTreeMap::new();
let mut last_transaction: Option<CardanoTransaction> = None;
hashes: &[TransactionHash],
) -> StdResult<HashMap<BlockRange, Vec<CardanoTransaction>>> {
let mut block_ranges_map = HashMap::new();
let transactions = self
.transaction_retriever
.get_by_hashes(hashes.to_vec())
.await?;
for transaction in transactions {
let block_range = BlockRange::from_block_number(transaction.block_number);
last_transaction = Some(transaction.clone());
transactions_by_block_ranges
.entry(block_range)
.or_default()
.push(transaction.transaction_hash);
}
let mut block_ranges = transactions_by_block_ranges.into_iter().try_fold(
vec![],
|mut acc, (block_range, transactions)| -> StdResult<Vec<(_, MKMapNode<_>)>> {
acc.push((block_range, MKTree::new(&transactions)?.into()));
Ok(acc)
},
)?;

// This is a temporary fix to avoid including an incomplete block ranges in the computation of the prover.
// This will be swiftly replaced by the new way of computing proof relying on the block range roots stored in database.
if let Some(transaction) = last_transaction {
if let Some((last_block_range, _)) = block_ranges.last() {
if transaction.block_number < last_block_range.end - 1 {
block_ranges.pop();
}
}
let block_range_transactions: &mut Vec<_> =
block_ranges_map.entry(block_range).or_insert(vec![]);
block_range_transactions.push(transaction)
}

let mk_hash_map = MKMap::new_from_iter(
block_ranges
)
.with_context(|| "ProverService failed to compute the merkelized structure that proves ownership of the transaction")?;

Ok(mk_hash_map)
Ok(block_ranges_map)
}
}

Expand All @@ -115,41 +83,41 @@ impl ProverService for MithrilProverService {
up_to: &CardanoDbBeacon,
transaction_hashes: &[TransactionHash],
) -> StdResult<Vec<CardanoTransactionsSetProof>> {
// 1 - Get transactions to prove per block range
let transactions_to_prove = self
.get_transactions_by_hashes_with_block_range(transaction_hashes.to_vec())
// 1 - Compute the set of block ranges with transactions to prove
let block_range_transactions = self
.get_transactions_grouped_by_block_range(transaction_hashes)
.await?;

// 2 - Compute Transactions Merkle Tree
let transactions = self.transaction_retriever.get_up_to(up_to).await?;
let mk_map = self.compute_merkle_map_from_transactions(transactions)?;

// 3 - Compute proof for each transaction to prove
let mut transaction_hashes_certified = vec![];
for (_block_range, transaction) in transactions_to_prove {
let mk_tree_node_transaction_hash: MKTreeNode =
transaction.transaction_hash.to_owned().into();
if mk_map
.compute_proof(&[mk_tree_node_transaction_hash])
.is_ok()
{
transaction_hashes_certified.push(transaction.transaction_hash);
}
// 2 - Compute block ranges sub Merkle trees
let mut mk_trees = HashMap::new();
for (block_range, transactions) in block_range_transactions {
let mk_tree = MKTree::new(
&transactions
.iter()
.map(|t| t.transaction_hash.clone())
.collect::<Vec<_>>(),
)?;
mk_trees.insert(block_range, mk_tree);
}

if !transaction_hashes_certified.is_empty() {
let mk_leaves: Vec<MKTreeNode> = transaction_hashes_certified
.iter()
.map(|h| h.to_owned().into())
.collect();
let mk_proof = mk_map.compute_proof(&mk_leaves)?;
let transactions_set_proof_batch =
CardanoTransactionsSetProof::new(transaction_hashes_certified, mk_proof);

Ok(vec![transactions_set_proof_batch])
} else {
Ok(vec![])
// 3 - Compute block range roots Merkle map
let mut mk_map = self
.block_range_root_retriever
.compute_merkle_map_from_block_range_roots(up_to.immutable_file_number)
.await?;

// 4 - Enrich the Merkle map with the block ranges Merkle trees
for (block_range, mk_tree) in mk_trees {
mk_map.insert(block_range, mk_tree.into())?;
}

// 5 - Compute the proof for each transaction
let mk_proof = mk_map.compute_proof(transaction_hashes)?;

Ok(vec![CardanoTransactionsSetProof::new(
transaction_hashes.to_vec(),
mk_proof,
)])
}
}

Expand All @@ -158,12 +126,31 @@ mod tests {
use std::cmp::max;

use anyhow::anyhow;
use mithril_common::entities::CardanoTransaction;
use mithril_common::crypto_helper::{MKMap, MKMapNode, MKTreeNode};
use mithril_common::entities::{CardanoTransaction, ImmutableFileNumber};
use mithril_common::test_utils::fake_data;
use mockall::mock;
use mockall::predicate::eq;

use super::*;

mock! {
pub BlockRangeRootRetrieverImpl { }

#[async_trait]
impl BlockRangeRootRetriever for BlockRangeRootRetrieverImpl {
async fn retrieve_block_range_roots(
&self,
up_to_beacon: ImmutableFileNumber,
) -> StdResult<Box<dyn Iterator<Item = (BlockRange, MKTreeNode)>>>;

async fn compute_merkle_map_from_block_range_roots(
&self,
up_to_beacon: ImmutableFileNumber,
) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange>>>;
}
}

fn generate_transactions(
total_transactions: usize,
) -> (Vec<TransactionHash>, Vec<CardanoTransaction>) {
Expand All @@ -185,30 +172,46 @@ mod tests {
(hashes, transactions)
}

fn build_prover<F>(retriever_mock_config: F) -> MithrilProverService
fn build_prover<F, G>(
transaction_retriever_mock_config: F,
block_range_root_retriever_mock_config: G,
) -> MithrilProverService
where
F: FnOnce(&mut MockTransactionsRetriever),
G: FnOnce(&mut MockBlockRangeRootRetrieverImpl),
{
let mut transaction_retriever = MockTransactionsRetriever::new();
retriever_mock_config(&mut transaction_retriever);
transaction_retriever_mock_config(&mut transaction_retriever);
let mut block_range_root_retriever = MockBlockRangeRootRetrieverImpl::new();
block_range_root_retriever_mock_config(&mut block_range_root_retriever);

MithrilProverService::new(Arc::new(transaction_retriever))
MithrilProverService::new(
Arc::new(transaction_retriever),
Arc::new(block_range_root_retriever),
)
}

#[tokio::test]
async fn compute_proof_for_one_set_with_multiple_transactions() {
let (transaction_hashes, transactions) = generate_transactions(3);
let prover = build_prover(|retriever_mock| {
let transactions_by_hashes_res = transactions.clone();
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions_by_hashes_res));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.return_once(move |_| Ok(transactions));
});
let prover = build_prover(
|retriever_mock| {
let transactions_by_hashes_res = transactions.clone();
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions_by_hashes_res));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.return_once(move |_| Ok(transactions));
},
|block_range_root_retriever_mock| {
block_range_root_retriever_mock
.expect_compute_merkle_map_from_block_range_roots()
.return_once(|_| MKMap::new(&[]));
},
);

let transactions_set_proof = prover
.compute_transactions_proofs(&fake_data::beacon(), &transaction_hashes)
Expand All @@ -225,16 +228,23 @@ mod tests {
#[tokio::test]
async fn cant_compute_proof_for_unknown_transaction() {
let (transaction_hashes, transactions) = generate_transactions(3);
let prover = build_prover(|retriever_mock| {
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.returning(|_| Ok(vec![]));
});
let prover = build_prover(
|retriever_mock| {
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.returning(|_| Ok(vec![]));
},
|block_range_root_retriever_mock| {
block_range_root_retriever_mock
.expect_compute_merkle_map_from_block_range_roots()
.return_once(|_| MKMap::new(&[]));
},
);

let transactions_set_proof = prover
.compute_transactions_proofs(&fake_data::beacon(), &transaction_hashes)
Expand All @@ -247,19 +257,26 @@ mod tests {
#[tokio::test]
async fn compute_proof_for_one_set_of_three_known_transactions_and_two_unknowns() {
let (transaction_hashes, transactions) = generate_transactions(5);
let prover = build_prover(|retriever_mock| {
// The last two are not in the "store"
let transactions = transactions[0..=2].to_vec();
let transactions_by_hashes_res = transactions.clone();
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions_by_hashes_res));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.return_once(move |_| Ok(transactions));
});
let prover = build_prover(
|retriever_mock| {
// The last two are not in the "store"
let transactions = transactions[0..=2].to_vec();
let transactions_by_hashes_res = transactions.clone();
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions_by_hashes_res));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.return_once(move |_| Ok(transactions));
},
|block_range_root_retriever_mock| {
block_range_root_retriever_mock
.expect_compute_merkle_map_from_block_range_roots()
.return_once(|_| MKMap::new(&[]));
},
);

let transactions_set_proof = prover
.compute_transactions_proofs(&fake_data::beacon(), &transaction_hashes)
Expand All @@ -277,12 +294,19 @@ mod tests {
#[tokio::test]
async fn cant_compute_proof_if_retriever_fail() {
let (transaction_hashes, _transactions) = generate_transactions(3);
let prover = build_prover(|retriever_mock| {
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.returning(|_| Err(anyhow!("Error")));
});
let prover = build_prover(
|retriever_mock| {
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.returning(|_| Err(anyhow!("Error")));
},
|block_range_root_retriever_mock| {
block_range_root_retriever_mock
.expect_compute_merkle_map_from_block_range_roots()
.return_once(|_| MKMap::new(&[]));
},
);

prover
.compute_transactions_proofs(&fake_data::beacon(), &transaction_hashes)
Expand Down

0 comments on commit 09000ea

Please sign in to comment.