Skip to content

Commit 385bc89

Browse files
quakequake
authored
Merge pull request #4182 from quake/quake/light-client-main-chain
fix: should only return main chain data
2 parents f45505f + 3e4f55d commit 385bc89

File tree

3 files changed

+91
-130
lines changed

3 files changed

+91
-130
lines changed

‎util/light-client-protocol-server/src/components/get_blocks_proof.rs

Lines changed: 22 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -40,14 +40,15 @@ impl<'a> GetBlocksProofProcess<'a> {
4040

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

43-
let last_hash = self.message.last_hash().to_entity();
44-
let last_block = if let Some(block) = snapshot.get_block(&last_hash) {
45-
block
46-
} else {
43+
let last_block_hash = self.message.last_hash().to_entity();
44+
if !snapshot.is_main_chain(&last_block_hash) {
4745
return self
4846
.protocol
4947
.reply_tip_state::<packed::SendBlocksProof>(self.peer, self.nc);
50-
};
48+
}
49+
let last_block = snapshot
50+
.get_block(&last_block_hash)
51+
.expect("block should be in store");
5152

5253
let block_hashes: Vec<_> = self
5354
.message
@@ -59,42 +60,30 @@ impl<'a> GetBlocksProofProcess<'a> {
5960
let mut uniq = HashSet::new();
6061
if !block_hashes
6162
.iter()
62-
.chain([last_hash.clone()].iter())
63+
.chain([last_block_hash].iter())
6364
.all(|hash| uniq.insert(hash))
6465
{
6566
return StatusCode::MalformedProtocolMessage
6667
.with_context("duplicate block hash exists");
6768
}
6869

69-
let (positions, block_headers, missing_blocks) = block_hashes
70+
let (found, missing): (Vec<_>, Vec<_>) = block_hashes
7071
.into_iter()
71-
.map(|block_hash| {
72-
snapshot
73-
.get_block_header(&block_hash)
74-
.map(|header| header.number())
75-
.filter(|number| *number != last_block.number())
76-
.and_then(|number| snapshot.get_ancestor(&last_hash, number))
77-
.filter(|header| header.hash() == block_hash)
78-
.ok_or(block_hash)
79-
})
80-
.fold(
81-
(Vec::new(), Vec::new(), Vec::new()),
82-
|(mut positions, mut block_headers, mut missing_blocks), result| {
83-
match result {
84-
Ok(header) => {
85-
positions.push(leaf_index_to_pos(header.number()));
86-
block_headers.push(header);
87-
}
88-
Err(block_hash) => {
89-
missing_blocks.push(block_hash);
90-
}
91-
}
92-
(positions, block_headers, missing_blocks)
93-
},
94-
);
72+
.partition(|block_hash| snapshot.is_main_chain(block_hash));
73+
74+
let mut positions = Vec::with_capacity(found.len());
75+
let mut block_headers = Vec::with_capacity(found.len());
76+
77+
for block_hash in found {
78+
let header = snapshot
79+
.get_block_header(&block_hash)
80+
.expect("header should be in store");
81+
positions.push(leaf_index_to_pos(header.number()));
82+
block_headers.push(header.data());
83+
}
9584

96-
let proved_items = block_headers.into_iter().map(|view| view.data()).pack();
97-
let missing_items = missing_blocks.pack();
85+
let proved_items = block_headers.pack();
86+
let missing_items = missing.pack();
9887

9988
self.protocol.reply_proof::<packed::SendBlocksProof>(
10089
self.peer,

‎util/light-client-protocol-server/src/components/get_last_state_proof.rs

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -204,18 +204,14 @@ impl<'a> GetLastStateProofProcess<'a> {
204204
let snapshot = self.protocol.shared.snapshot();
205205

206206
let last_block_hash = self.message.last_hash().to_entity();
207-
let last_block = if !snapshot.is_main_chain(&last_block_hash) {
207+
if !snapshot.is_main_chain(&last_block_hash) {
208208
return self
209209
.protocol
210210
.reply_tip_state::<packed::SendLastStateProof>(self.peer, self.nc);
211-
} else if let Some(block) = snapshot.get_block(&last_block_hash) {
212-
block
213-
} else {
214-
let errmsg = format!(
215-
"the block is in the main chain but not found, its hash is {last_block_hash:#x}"
216-
);
217-
return StatusCode::InternalError.with_context(errmsg);
218-
};
211+
}
212+
let last_block = snapshot
213+
.get_block(&last_block_hash)
214+
.expect("block should be in store");
219215

220216
let start_block_hash = self.message.start_hash().to_entity();
221217
let start_block_number: BlockNumber = self.message.start_number().unpack();

‎util/light-client-protocol-server/src/components/get_transactions_proof.rs

Lines changed: 64 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -40,107 +40,83 @@ impl<'a> GetTransactionsProofProcess<'a> {
4040

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

43-
let last_hash = self.message.last_hash().to_entity();
44-
let last_block = if let Some(block) = snapshot.get_block(&last_hash) {
45-
block
46-
} else {
43+
let last_block_hash = self.message.last_hash().to_entity();
44+
if !snapshot.is_main_chain(&last_block_hash) {
4745
return self
4846
.protocol
4947
.reply_tip_state::<packed::SendTransactionsProof>(self.peer, self.nc);
50-
};
48+
}
49+
let last_block = snapshot
50+
.get_block(&last_block_hash)
51+
.expect("block should be in store");
5152

52-
let (txs_in_blocks, missing_txs) = self
53+
let (found, missing): (Vec<_>, Vec<_>) = self
5354
.message
5455
.tx_hashes()
5556
.to_entity()
5657
.into_iter()
57-
.map(|tx_hash| {
58-
let tx_with_info = snapshot.get_transaction_with_info(&tx_hash);
59-
(tx_hash, tx_with_info)
60-
})
61-
.fold(
62-
(HashMap::new(), Vec::new()),
63-
|(mut found, mut missing_txs), (tx_hash, tx_with_info)| {
64-
if let Some((tx, tx_info)) = tx_with_info {
65-
found
66-
.entry(tx_info.block_hash)
67-
.or_insert_with(Vec::new)
68-
.push((tx, tx_info.index));
69-
} else {
70-
missing_txs.push(tx_hash);
71-
}
72-
(found, missing_txs)
73-
},
74-
);
75-
76-
let (positions, filtered_blocks, missing_txs) = txs_in_blocks
77-
.into_iter()
78-
.map(|(block_hash, txs_and_tx_indices)| {
58+
.partition(|tx_hash| {
7959
snapshot
80-
.get_block_header(&block_hash)
81-
.map(|header| header.number())
82-
.filter(|number| *number != last_block.number())
83-
.and_then(|number| snapshot.get_ancestor(&last_hash, number))
84-
.filter(|header| header.hash() == block_hash)
85-
.and_then(|_| snapshot.get_block(&block_hash))
86-
.map(|block| (block, txs_and_tx_indices.clone()))
87-
.ok_or_else(|| {
88-
txs_and_tx_indices
89-
.into_iter()
90-
.map(|(tx, _)| tx.hash())
91-
.collect::<Vec<_>>()
92-
})
93-
})
94-
.fold(
95-
(Vec::new(), Vec::new(), missing_txs),
96-
|(mut positions, mut filtered_blocks, mut missing_txs), result| {
97-
match result {
98-
Ok((block, txs_and_tx_indices)) => {
99-
let merkle_proof = CBMT::build_merkle_proof(
100-
&block
101-
.transactions()
102-
.iter()
103-
.map(|tx| tx.hash())
104-
.collect::<Vec<_>>(),
105-
&txs_and_tx_indices
106-
.iter()
107-
.map(|(_, index)| *index as u32)
108-
.collect::<Vec<_>>(),
109-
)
110-
.expect("build proof with verified inputs should be OK");
111-
112-
let txs: Vec<_> = txs_and_tx_indices
113-
.into_iter()
114-
.map(|(tx, _)| tx.data())
115-
.collect();
116-
117-
let filtered_block = packed::FilteredBlock::new_builder()
118-
.header(block.header().data())
119-
.witnesses_root(block.calc_witnesses_root())
120-
.transactions(txs.pack())
121-
.proof(
122-
packed::MerkleProof::new_builder()
123-
.indices(merkle_proof.indices().to_owned().pack())
124-
.lemmas(merkle_proof.lemmas().to_owned().pack())
125-
.build(),
126-
)
127-
.build();
128-
129-
positions.push(leaf_index_to_pos(block.number()));
130-
filtered_blocks.push(filtered_block);
131-
}
132-
Err(tx_hashes) => {
133-
missing_txs.extend(tx_hashes);
134-
}
135-
}
136-
(positions, filtered_blocks, missing_txs)
137-
},
138-
);
60+
.get_transaction_info(tx_hash)
61+
.map(|tx_info| snapshot.is_main_chain(&tx_info.block_hash))
62+
.unwrap_or_default()
63+
});
64+
65+
let mut txs_in_blocks = HashMap::new();
66+
for tx_hash in found {
67+
let (tx, tx_info) = snapshot
68+
.get_transaction_with_info(&tx_hash)
69+
.expect("tx exists");
70+
txs_in_blocks
71+
.entry(tx_info.block_hash)
72+
.or_insert_with(Vec::new)
73+
.push((tx, tx_info.index));
74+
}
75+
76+
let mut positions = Vec::with_capacity(txs_in_blocks.len());
77+
let mut filtered_blocks = Vec::with_capacity(txs_in_blocks.len());
78+
for (block_hash, txs_and_tx_indices) in txs_in_blocks.into_iter() {
79+
let block = snapshot
80+
.get_block(&block_hash)
81+
.expect("block should be in store");
82+
let merkle_proof = CBMT::build_merkle_proof(
83+
&block
84+
.transactions()
85+
.iter()
86+
.map(|tx| tx.hash())
87+
.collect::<Vec<_>>(),
88+
&txs_and_tx_indices
89+
.iter()
90+
.map(|(_, index)| *index as u32)
91+
.collect::<Vec<_>>(),
92+
)
93+
.expect("build proof with verified inputs should be OK");
94+
95+
let txs: Vec<_> = txs_and_tx_indices
96+
.into_iter()
97+
.map(|(tx, _)| tx.data())
98+
.collect();
99+
100+
let filtered_block = packed::FilteredBlock::new_builder()
101+
.header(block.header().data())
102+
.witnesses_root(block.calc_witnesses_root())
103+
.transactions(txs.pack())
104+
.proof(
105+
packed::MerkleProof::new_builder()
106+
.indices(merkle_proof.indices().to_owned().pack())
107+
.lemmas(merkle_proof.lemmas().to_owned().pack())
108+
.build(),
109+
)
110+
.build();
111+
112+
positions.push(leaf_index_to_pos(block.number()));
113+
filtered_blocks.push(filtered_block);
114+
}
139115

140116
let proved_items = packed::FilteredBlockVec::new_builder()
141117
.set(filtered_blocks)
142118
.build();
143-
let missing_items = missing_txs.pack();
119+
let missing_items = missing.pack();
144120

145121
self.protocol.reply_proof::<packed::SendTransactionsProof>(
146122
self.peer,

0 commit comments

Comments
 (0)