From 3e4f55d6b162515e7975fb363de43cee123abe32 Mon Sep 17 00:00:00 2001 From: quake Date: Tue, 10 Oct 2023 20:43:08 +0900 Subject: [PATCH] fix: should only return main chain data --- .../src/components/get_blocks_proof.rs | 55 +++---- .../src/components/get_last_state_proof.rs | 14 +- .../src/components/get_transactions_proof.rs | 152 ++++++++---------- 3 files changed, 91 insertions(+), 130 deletions(-) diff --git a/util/light-client-protocol-server/src/components/get_blocks_proof.rs b/util/light-client-protocol-server/src/components/get_blocks_proof.rs index f071177dd8..dd8e0ae174 100644 --- a/util/light-client-protocol-server/src/components/get_blocks_proof.rs +++ b/util/light-client-protocol-server/src/components/get_blocks_proof.rs @@ -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::(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 @@ -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::( self.peer, diff --git a/util/light-client-protocol-server/src/components/get_last_state_proof.rs b/util/light-client-protocol-server/src/components/get_last_state_proof.rs index b2352236d9..1129eea868 100644 --- a/util/light-client-protocol-server/src/components/get_last_state_proof.rs +++ b/util/light-client-protocol-server/src/components/get_last_state_proof.rs @@ -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::(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(); diff --git a/util/light-client-protocol-server/src/components/get_transactions_proof.rs b/util/light-client-protocol-server/src/components/get_transactions_proof.rs index f8b22c7a7a..4b77884cfe 100644 --- a/util/light-client-protocol-server/src/components/get_transactions_proof.rs +++ b/util/light-client-protocol-server/src/components/get_transactions_proof.rs @@ -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::(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::>() - }) - }) - .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::>(), - &txs_and_tx_indices - .iter() - .map(|(_, index)| *index as u32) - .collect::>(), - ) - .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::>(), + &txs_and_tx_indices + .iter() + .map(|(_, index)| *index as u32) + .collect::>(), + ) + .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::( self.peer,