Skip to content

Commit

Permalink
Merge pull request #4182 from quake/quake/light-client-main-chain
Browse files Browse the repository at this point in the history 
fix: should only return main chain data
  • Loading branch information
@quake
quake authored Oct 10, 2023
2 parents f45505f + 3e4f55d commit 385bc89
Show file tree
Hide file tree
Showing 3 changed files with 91 additions and 130 deletions.
55 changes: 22 additions & 33 deletions util/light-client-protocol-server/src/components/get_blocks_proof.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,14 +40,15 @@ impl<'a> GetBlocksProofProcess<'a> {

let snapshot = self.protocol.shared.snapshot();

let last_hash = self.message.last_hash().to_entity();
let last_block = if let Some(block) = snapshot.get_block(&last_hash) {
block
} else {
let last_block_hash = self.message.last_hash().to_entity();
if !snapshot.is_main_chain(&last_block_hash) {
return self
.protocol
.reply_tip_state::<packed::SendBlocksProof>(self.peer, self.nc);
};
}
let last_block = snapshot
.get_block(&last_block_hash)
.expect("block should be in store");

let block_hashes: Vec<_> = self
.message
Expand All @@ -59,42 +60,30 @@ impl<'a> GetBlocksProofProcess<'a> {
let mut uniq = HashSet::new();
if !block_hashes
.iter()
.chain([last_hash.clone()].iter())
.chain([last_block_hash].iter())
.all(|hash| uniq.insert(hash))
{
return StatusCode::MalformedProtocolMessage
.with_context("duplicate block hash exists");
}

let (positions, block_headers, missing_blocks) = block_hashes
let (found, missing): (Vec<_>, Vec<_>) = block_hashes
.into_iter()
.map(|block_hash| {
snapshot
.get_block_header(&block_hash)
.map(|header| header.number())
.filter(|number| *number != last_block.number())
.and_then(|number| snapshot.get_ancestor(&last_hash, number))
.filter(|header| header.hash() == block_hash)
.ok_or(block_hash)
})
.fold(
(Vec::new(), Vec::new(), Vec::new()),
|(mut positions, mut block_headers, mut missing_blocks), result| {
match result {
Ok(header) => {
positions.push(leaf_index_to_pos(header.number()));
block_headers.push(header);
}
Err(block_hash) => {
missing_blocks.push(block_hash);
}
}
(positions, block_headers, missing_blocks)
},
);
.partition(|block_hash| snapshot.is_main_chain(block_hash));

let mut positions = Vec::with_capacity(found.len());
let mut block_headers = Vec::with_capacity(found.len());

for block_hash in found {
let header = snapshot
.get_block_header(&block_hash)
.expect("header should be in store");
positions.push(leaf_index_to_pos(header.number()));
block_headers.push(header.data());
}

let proved_items = block_headers.into_iter().map(|view| view.data()).pack();
let missing_items = missing_blocks.pack();
let proved_items = block_headers.pack();
let missing_items = missing.pack();

self.protocol.reply_proof::<packed::SendBlocksProof>(
self.peer,
Expand Down
14 changes: 5 additions & 9 deletions util/light-client-protocol-server/src/components/get_last_state_proof.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -204,18 +204,14 @@ impl<'a> GetLastStateProofProcess<'a> {
let snapshot = self.protocol.shared.snapshot();

let last_block_hash = self.message.last_hash().to_entity();
let last_block = if !snapshot.is_main_chain(&last_block_hash) {
if !snapshot.is_main_chain(&last_block_hash) {
return self
.protocol
.reply_tip_state::<packed::SendLastStateProof>(self.peer, self.nc);
} else if let Some(block) = snapshot.get_block(&last_block_hash) {
block
} else {
let errmsg = format!(
"the block is in the main chain but not found, its hash is {last_block_hash:#x}"
);
return StatusCode::InternalError.with_context(errmsg);
};
}
let last_block = snapshot
.get_block(&last_block_hash)
.expect("block should be in store");

let start_block_hash = self.message.start_hash().to_entity();
let start_block_number: BlockNumber = self.message.start_number().unpack();
Expand Down
152 changes: 64 additions & 88 deletions util/light-client-protocol-server/src/components/get_transactions_proof.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,107 +40,83 @@ impl<'a> GetTransactionsProofProcess<'a> {

let snapshot = self.protocol.shared.snapshot();

let last_hash = self.message.last_hash().to_entity();
let last_block = if let Some(block) = snapshot.get_block(&last_hash) {
block
} else {
let last_block_hash = self.message.last_hash().to_entity();
if !snapshot.is_main_chain(&last_block_hash) {
return self
.protocol
.reply_tip_state::<packed::SendTransactionsProof>(self.peer, self.nc);
};
}
let last_block = snapshot
.get_block(&last_block_hash)
.expect("block should be in store");

let (txs_in_blocks, missing_txs) = self
let (found, missing): (Vec<_>, Vec<_>) = self
.message
.tx_hashes()
.to_entity()
.into_iter()
.map(|tx_hash| {
let tx_with_info = snapshot.get_transaction_with_info(&tx_hash);
(tx_hash, tx_with_info)
})
.fold(
(HashMap::new(), Vec::new()),
|(mut found, mut missing_txs), (tx_hash, tx_with_info)| {
if let Some((tx, tx_info)) = tx_with_info {
found
.entry(tx_info.block_hash)
.or_insert_with(Vec::new)
.push((tx, tx_info.index));
} else {
missing_txs.push(tx_hash);
}
(found, missing_txs)
},
);

let (positions, filtered_blocks, missing_txs) = txs_in_blocks
.into_iter()
.map(|(block_hash, txs_and_tx_indices)| {
.partition(|tx_hash| {
snapshot
.get_block_header(&block_hash)
.map(|header| header.number())
.filter(|number| *number != last_block.number())
.and_then(|number| snapshot.get_ancestor(&last_hash, number))
.filter(|header| header.hash() == block_hash)
.and_then(|_| snapshot.get_block(&block_hash))
.map(|block| (block, txs_and_tx_indices.clone()))
.ok_or_else(|| {
txs_and_tx_indices
.into_iter()
.map(|(tx, _)| tx.hash())
.collect::<Vec<_>>()
})
})
.fold(
(Vec::new(), Vec::new(), missing_txs),
|(mut positions, mut filtered_blocks, mut missing_txs), result| {
match result {
Ok((block, txs_and_tx_indices)) => {
let merkle_proof = CBMT::build_merkle_proof(
&block
.transactions()
.iter()
.map(|tx| tx.hash())
.collect::<Vec<_>>(),
&txs_and_tx_indices
.iter()
.map(|(_, index)| *index as u32)
.collect::<Vec<_>>(),
)
.expect("build proof with verified inputs should be OK");

let txs: Vec<_> = txs_and_tx_indices
.into_iter()
.map(|(tx, _)| tx.data())
.collect();

let filtered_block = packed::FilteredBlock::new_builder()
.header(block.header().data())
.witnesses_root(block.calc_witnesses_root())
.transactions(txs.pack())
.proof(
packed::MerkleProof::new_builder()
.indices(merkle_proof.indices().to_owned().pack())
.lemmas(merkle_proof.lemmas().to_owned().pack())
.build(),
)
.build();

positions.push(leaf_index_to_pos(block.number()));
filtered_blocks.push(filtered_block);
}
Err(tx_hashes) => {
missing_txs.extend(tx_hashes);
}
}
(positions, filtered_blocks, missing_txs)
},
);
.get_transaction_info(tx_hash)
.map(|tx_info| snapshot.is_main_chain(&tx_info.block_hash))
.unwrap_or_default()
});

let mut txs_in_blocks = HashMap::new();
for tx_hash in found {
let (tx, tx_info) = snapshot
.get_transaction_with_info(&tx_hash)
.expect("tx exists");
txs_in_blocks
.entry(tx_info.block_hash)
.or_insert_with(Vec::new)
.push((tx, tx_info.index));
}

let mut positions = Vec::with_capacity(txs_in_blocks.len());
let mut filtered_blocks = Vec::with_capacity(txs_in_blocks.len());
for (block_hash, txs_and_tx_indices) in txs_in_blocks.into_iter() {
let block = snapshot
.get_block(&block_hash)
.expect("block should be in store");
let merkle_proof = CBMT::build_merkle_proof(
&block
.transactions()
.iter()
.map(|tx| tx.hash())
.collect::<Vec<_>>(),
&txs_and_tx_indices
.iter()
.map(|(_, index)| *index as u32)
.collect::<Vec<_>>(),
)
.expect("build proof with verified inputs should be OK");

let txs: Vec<_> = txs_and_tx_indices
.into_iter()
.map(|(tx, _)| tx.data())
.collect();

let filtered_block = packed::FilteredBlock::new_builder()
.header(block.header().data())
.witnesses_root(block.calc_witnesses_root())
.transactions(txs.pack())
.proof(
packed::MerkleProof::new_builder()
.indices(merkle_proof.indices().to_owned().pack())
.lemmas(merkle_proof.lemmas().to_owned().pack())
.build(),
)
.build();

positions.push(leaf_index_to_pos(block.number()));
filtered_blocks.push(filtered_block);
}

let proved_items = packed::FilteredBlockVec::new_builder()
.set(filtered_blocks)
.build();
let missing_items = missing_txs.pack();
let missing_items = missing.pack();

self.protocol.reply_proof::<packed::SendTransactionsProof>(
self.peer,
Expand Down

0 comments on commit 385bc89

Please sign in to comment.