Skip to content

Commit

Permalink
feat: add 'BlockRangeRootRetriever' trait in Cardano transaction sign…
Browse files Browse the repository at this point in the history 
…able builder
  • Loading branch information
@jpraynaud
jpraynaud committed Apr 29, 2024
1 parent b914203 commit 1bd52fa
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 162 deletions.
25 changes: 21 additions & 4 deletions mithril-common/src/messages/cardano_transactions_proof.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -314,9 +314,11 @@ mod tests {

#[cfg(feature = "fs")]
mod fs_only {
use crate::entities::{CardanoDbBeacon, CardanoTransaction};
use crate::crypto_helper::{MKMap, MKMapNode};
use crate::entities::{BlockRange, CardanoDbBeacon, CardanoTransaction};
use crate::signable_builder::{
CardanoTransactionsSignableBuilder, MockTransactionsImporter, SignableBuilder,
CardanoTransactionsSignableBuilder, MockBlockRangeRootRetriever,
MockTransactionsImporter, SignableBuilder,
};
use slog::Logger;
use std::sync::Arc;
Expand Down Expand Up @@ -376,13 +378,28 @@ mod tests {
transactions: &[CardanoTransaction],
immutable_file_number: u64,
) -> ProtocolMessage {
let transactions = transactions.to_vec();
let mut transaction_importer = MockTransactionsImporter::new();
transaction_importer
.expect_import()
.return_once(move |_| Ok(transactions));
.return_once(move |_| Ok(()));
let mut block_range_root_retriever = MockBlockRangeRootRetriever::new();

let transactions_imported = transactions.to_vec();
block_range_root_retriever
.expect_compute_merkle_map_from_block_range_roots()
.return_once(move |_| {
MKMap::<BlockRange, MKMapNode<BlockRange>>::new_from_iter(
transactions_imported.into_iter().map(|tx| {
(
BlockRange::from_block_number(tx.block_number),
MKMapNode::TreeNode(tx.transaction_hash.clone().into()),
)
}),
)
});
let cardano_transaction_signable_builder = CardanoTransactionsSignableBuilder::new(
Arc::new(transaction_importer),
Arc::new(block_range_root_retriever),
Logger::root(slog::Discard, slog::o!()),
);
cardano_transaction_signable_builder
Expand Down
231 changes: 73 additions & 158 deletions mithril-common/src/signable_builder/cardano_transactions.rs
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,12 @@
use std::{collections::HashMap, sync::Arc};
use std::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use slog::{debug, Logger};

use crate::{
crypto_helper::{MKMap, MKMapNode, MKTree, MKTreeNode},
entities::{
BlockRange, CardanoDbBeacon, CardanoTransaction, ProtocolMessage, ProtocolMessagePartKey,
TransactionHash,
},
crypto_helper::{MKMap, MKMapNode, MKTreeNode},
entities::{BlockRange, CardanoDbBeacon, ProtocolMessage, ProtocolMessagePartKey},
signable_builder::SignableBuilder,
StdResult,
};
Expand All @@ -23,65 +20,55 @@ use mockall::automock;
#[async_trait]
pub trait TransactionsImporter: Send + Sync {
/// Returns all transactions up to the given beacon
async fn import(&self, up_to_beacon: ImmutableFileNumber)
-> StdResult<Vec<CardanoTransaction>>;
async fn import(&self, up_to_beacon: ImmutableFileNumber) -> StdResult<()>;
}

/// Block Range Merkle roots retriever
#[cfg_attr(test, automock)]
#[async_trait]
pub trait BlockRangeRootRetriever: Send + Sync {
/// Returns a Merkle map of the block ranges roots up to a given beacon
async fn retrieve_block_range_roots(
&self,
up_to_beacon: ImmutableFileNumber,
) -> StdResult<Box<dyn Iterator<Item = (BlockRange, MKTreeNode)>>>;

/// Returns a Merkle map of the block ranges roots up to a given beacon
async fn compute_merkle_map_from_block_range_roots(
&self,
up_to_beacon: ImmutableFileNumber,
) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange>>> {
let block_range_roots_iterator = self
.retrieve_block_range_roots(up_to_beacon)
.await?
.map(|(block_range, root)| (block_range, root.into()));
let mk_hash_map = MKMap::new_from_iter(block_range_roots_iterator)
.with_context(|| "ProverService failed to compute the merkelized structure that proves ownership of the transaction")?;

Ok(mk_hash_map)
}
}

/// A [CardanoTransactionsSignableBuilder] builder
pub struct CardanoTransactionsSignableBuilder {
transaction_importer: Arc<dyn TransactionsImporter>,
block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
logger: Logger,
}

impl CardanoTransactionsSignableBuilder {
/// Constructor
pub fn new(transaction_importer: Arc<dyn TransactionsImporter>, logger: Logger) -> Self {
pub fn new(
transaction_importer: Arc<dyn TransactionsImporter>,
block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
logger: Logger,
) -> Self {
Self {
transaction_importer,
block_range_root_retriever,
logger,
}
}

// Note: Code duplicated from aggregator Prover service as is.
// This will be not be the case when we use the cached intermediate merkle roots.
fn compute_merkle_map_from_transactions(
&self,
transactions: Vec<CardanoTransaction>,
) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange>>> {
let mut transactions_by_block_ranges: HashMap<BlockRange, Vec<TransactionHash>> =
HashMap::new();
for transaction in transactions {
let block_range = BlockRange::from_block_number(transaction.block_number);
transactions_by_block_ranges
.entry(block_range)
.or_default()
.push(transaction.transaction_hash);
}
let mk_hash_map = MKMap::new_from_iter(
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)
},
)?,
)
.with_context(|| "ProverService failed to compute the merkelized structure that proves ownership of the transaction")?;

Ok(mk_hash_map)
}

fn compute_merkle_root(&self, transactions: Vec<CardanoTransaction>) -> StdResult<MKTreeNode> {
let mk_map = self.compute_merkle_map_from_transactions(transactions)?;

let mk_root = mk_map.compute_root().with_context(|| {
"CardanoTransactionsSignableBuilder failed to compute MKHashMap root"
})?;

Ok(mk_root)
}
}

#[async_trait]
Expand All @@ -95,11 +82,15 @@ impl SignableBuilder<CardanoDbBeacon> for CardanoTransactionsSignableBuilder {
"Compute protocol message for CardanoTransactions at beacon: {beacon}"
);

let transactions = self
.transaction_importer
self.transaction_importer
.import(beacon.immutable_file_number)
.await?;
let mk_root = self.compute_merkle_root(transactions)?;

let mk_root = self
.block_range_root_retriever
.compute_merkle_map_from_block_range_roots(beacon.immutable_file_number)
.await?
.compute_root()?;

let mut protocol_message = ProtocolMessage::new();
protocol_message.set_message_part(
Expand All @@ -117,104 +108,21 @@ impl SignableBuilder<CardanoDbBeacon> for CardanoTransactionsSignableBuilder {

#[cfg(test)]
mod tests {
use crate::test_utils::TestLogger;

use super::*;

#[tokio::test]
async fn test_compute_merkle_root_in_same_block_range() {
let transaction_1 = CardanoTransaction::new("tx-hash-123", 1, 10, "block_hash", 1);
let transaction_2 = CardanoTransaction::new("tx-hash-456", 2, 20, "block_hash", 1);
let transaction_3 = CardanoTransaction::new("tx-hash-789", 3, 30, "block_hash", 1);
let transaction_4 = CardanoTransaction::new("tx-hash-abc", 4, 40, "block_hash", 1);

for tx in [
&transaction_1,
&transaction_2,
&transaction_3,
&transaction_4,
] {
assert!(
tx.block_number < BlockRange::LENGTH,
"all manipulated transactions should be in the same block range"
);
}

let cardano_transaction_signable_builder = CardanoTransactionsSignableBuilder::new(
Arc::new(MockTransactionsImporter::new()),
TestLogger::stdout(),
);

let merkle_root_reference = cardano_transaction_signable_builder
.compute_merkle_root(vec![
transaction_1.clone(),
transaction_2.clone(),
transaction_3.clone(),
])
.unwrap();
use crate::{entities::CardanoTransaction, test_utils::TestLogger};

{
let transactions_set = vec![transaction_1.clone()];
let mk_root = cardano_transaction_signable_builder
.compute_merkle_root(transactions_set)
.unwrap();
assert_ne!(merkle_root_reference, mk_root);
}
{
let transactions_set = vec![transaction_1.clone(), transaction_2.clone()];
let mk_root = cardano_transaction_signable_builder
.compute_merkle_root(transactions_set)
.unwrap();
assert_ne!(merkle_root_reference, mk_root);
}
{
let transactions_set = vec![
transaction_1.clone(),
transaction_2.clone(),
transaction_3.clone(),
transaction_4.clone(),
];
let mk_root = cardano_transaction_signable_builder
.compute_merkle_root(transactions_set)
.unwrap();
assert_ne!(merkle_root_reference, mk_root);
}

{
// In a same block range Transactions in a different order return a different merkle root.
let transactions_set = vec![
transaction_1.clone(),
transaction_3.clone(),
transaction_2.clone(),
];
let mk_root = cardano_transaction_signable_builder
.compute_merkle_root(transactions_set)
.unwrap();
assert_ne!(merkle_root_reference, mk_root);
}
}

#[tokio::test]
async fn test_compute_merkle_root_order_of_block_range_does_not_matter() {
let transaction_1 =
CardanoTransaction::new("tx-hash-123", BlockRange::LENGTH - 1, 10, "block_hash", 1);
let transaction_2 =
CardanoTransaction::new("tx-hash-456", BlockRange::LENGTH + 1, 20, "block_hash", 1);

let cardano_transaction_signable_builder = CardanoTransactionsSignableBuilder::new(
Arc::new(MockTransactionsImporter::new()),
TestLogger::stdout(),
);

let merkle_root_reference = cardano_transaction_signable_builder
.compute_merkle_root(vec![transaction_1.clone(), transaction_2.clone()])
.unwrap();

let mk_root = cardano_transaction_signable_builder
.compute_merkle_root(vec![transaction_2.clone(), transaction_1.clone()])
.unwrap();
use super::*;

assert_eq!(merkle_root_reference, mk_root);
fn compute_mk_map_from_transactions(
transactions: Vec<CardanoTransaction>,
) -> MKMap<BlockRange, MKMapNode<BlockRange>> {
MKMap::new_from_iter(transactions.iter().map(|tx| {
(
BlockRange::from_block_number(tx.block_number),
MKMapNode::TreeNode(tx.transaction_hash.clone().into()),
)
}))
.unwrap()
}

#[tokio::test]
Expand All @@ -229,13 +137,20 @@ mod tests {
CardanoTransaction::new("tx-hash-456", 20, 2, "block_hash", 12),
CardanoTransaction::new("tx-hash-789", 30, 3, "block_hash", 13),
];
let imported_transactions = transactions.clone();
let mk_map = compute_mk_map_from_transactions(transactions.clone());
let mut transaction_importer = MockTransactionsImporter::new();
transaction_importer
.expect_import()
.return_once(move |_| Ok(imported_transactions));
.return_once(move |_| Ok(()));
let retrieved_transactions = transactions.clone();
let mut block_range_root_retriever = MockBlockRangeRootRetriever::new();
block_range_root_retriever
.expect_compute_merkle_map_from_block_range_roots()
.return_once(move |_| Ok(compute_mk_map_from_transactions(retrieved_transactions)));

let cardano_transactions_signable_builder = CardanoTransactionsSignableBuilder::new(
Arc::new(transaction_importer),
Arc::new(block_range_root_retriever),
TestLogger::stdout(),
);

Expand All @@ -246,13 +161,10 @@ mod tests {
.unwrap();

// Assert
let mk_root = cardano_transactions_signable_builder
.compute_merkle_root(transactions)
.unwrap();
let mut signable_expected = ProtocolMessage::new();
signable_expected.set_message_part(
ProtocolMessagePartKey::CardanoTransactionsMerkleRoot,
mk_root.to_hex(),
mk_map.compute_root().unwrap().to_hex(),
);
signable_expected.set_message_part(
ProtocolMessagePartKey::LatestImmutableFileNumber,
Expand All @@ -262,14 +174,17 @@ mod tests {
}

#[tokio::test]
async fn test_compute_signable_with_no_transaction_return_error() {
async fn test_compute_signable_with_no_block_range_root_return_error() {
let beacon = CardanoDbBeacon::default();
let mut transaction_importer = MockTransactionsImporter::new();
transaction_importer
.expect_import()
.return_once(|_| Ok(vec![]));
transaction_importer.expect_import().return_once(|_| Ok(()));
let mut block_range_root_retriever = MockBlockRangeRootRetriever::new();
block_range_root_retriever
.expect_compute_merkle_map_from_block_range_roots()
.return_once(move |_| Ok(compute_mk_map_from_transactions(vec![])));
let cardano_transactions_signable_builder = CardanoTransactionsSignableBuilder::new(
Arc::new(transaction_importer),
Arc::new(block_range_root_retriever),
TestLogger::stdout(),
);

Expand Down

0 comments on commit 1bd52fa

Please sign in to comment.