diff --git a/src/protocols/light_client/components/send_last_state_proof.rs b/src/protocols/light_client/components/send_last_state_proof.rs index 40d1794..3ee706a 100644 --- a/src/protocols/light_client/components/send_last_state_proof.rs +++ b/src/protocols/light_client/components/send_last_state_proof.rs @@ -150,10 +150,10 @@ impl<'a> SendLastStateProofProcess<'a> { // Check if headers are continuous. if reorg_count != 0 { - return_if_failed!(check_continuous_headers(&headers[..reorg_count - 1])); + return_if_failed!(check_continuous_headers(&headers[..reorg_count])); } return_if_failed!(check_continuous_headers( - &headers[reorg_count + sampled_count..] + &headers[(reorg_count + sampled_count)..] )); // Verify MMR proof @@ -249,6 +249,30 @@ impl<'a> SendLastStateProofProcess<'a> { } } else if reorg_count == 0 { new_last_headers + } else if sampled_count == 0 + && check_continuous_headers(&headers[(reorg_count - 1)..=reorg_count]) + .is_ok() + { + // At the above part: + // - `reorg_count != 0`, `headers[..reorg_count]` are checked. + // - `headers[(reorg_count + sampled_count)..]` are always checked. + // So, only check `reorg_count - 1` and `reorg_count` is enough to prove + // all header are continuous. + let required_count = last_n_blocks - last_n_count; + let old_last_headers = &headers[..reorg_count]; + let old_last_headers_len = old_last_headers.len(); + let tmp_headers = old_last_headers + .iter() + .skip(old_last_headers_len.saturating_sub(required_count)) + .map(ToOwned::to_owned) + .chain(new_last_headers) + .collect::>(); + debug!("peer {}: reorg but no enough blocks, so all headers should be continuous, \ + reorg: {reorg_count}, sampled: {sampled_count}, last_n_calculate: {last_n_count}, \ + last_n_real: {}, last_n_param: {last_n_blocks}, original_request: {original_request}", + self.peer_index, tmp_headers.len() + ); + tmp_headers } else { // If this branch is reached, the follow conditions must be satisfied: // - No previous prove state. diff --git a/src/protocols/light_client/peers.rs b/src/protocols/light_client/peers.rs index dd4677e..ff04ecc 100644 --- a/src/protocols/light_client/peers.rs +++ b/src/protocols/light_client/peers.rs @@ -1285,6 +1285,13 @@ impl Peers { self.inner.get(index).map(|peer| peer.clone()) } + #[cfg(test)] + pub(crate) fn mock_initialized(&self, index: PeerIndex) { + if let Some(mut peer) = self.inner.get_mut(&index) { + _ = peer.state.take(); + } + } + #[cfg(test)] pub(crate) fn mock_prove_request( &self, diff --git a/src/tests/protocols/light_client/send_last_state_proof.rs b/src/tests/protocols/light_client/send_last_state_proof.rs index e7778cb..59a71a6 100644 --- a/src/tests/protocols/light_client/send_last_state_proof.rs +++ b/src/tests/protocols/light_client/send_last_state_proof.rs @@ -1647,6 +1647,18 @@ async fn reorg_rollback_5_blocks() { test_with_reorg_blocks(param).await; } +#[tokio::test(flavor = "multi_thread")] +async fn reorg_rollback_after_restart_and_last_n_blocks_is_not_enough() { + let param = ReorgTestParameter { + last_number: 30, + rollback_blocks_count: 3, + last_n_blocks: 20, + restart: true, + ..Default::default() + }; + test_with_reorg_blocks(param).await; +} + #[tokio::test(flavor = "multi_thread")] async fn reorg_detect_long_fork_turn_1() { let param = ReorgTestParameter { @@ -1682,6 +1694,8 @@ struct ReorgTestParameter { long_fork_detected: bool, expected_last_headers_count_opt: Option, result: StatusCode, + // Mock "restart" state: after restart, the first received "last state" is on a forked chain. + restart: bool, } async fn test_with_reorg_blocks(param: ReorgTestParameter) { @@ -1834,6 +1848,11 @@ async fn test_with_reorg_blocks(param: ReorgTestParameter) { assert_eq!(chain.shared().snapshot().tip_number(), last_number); } + if param.restart { + protocol.peers().mock_initialized(peer_index); + protocol.peers().request_last_state(peer_index).unwrap(); + } + // Run the test. { let mut prove_request = chain.build_prove_request(