diff --git a/substrate/client/consensus/babe/src/lib.rs b/substrate/client/consensus/babe/src/lib.rs index 683df9ddacd62ef92e4c869d09cc7c8f237b8655..84a3c618b38039b306bceb0f060efea3e74f91d4 100644 --- a/substrate/client/consensus/babe/src/lib.rs +++ b/substrate/client/consensus/babe/src/lib.rs @@ -126,7 +126,7 @@ use sp_inherents::{CreateInherentDataProviders, InherentData, InherentDataProvid use sp_keystore::{SyncCryptoStore, SyncCryptoStorePtr}; use sp_runtime::{ generic::{BlockId, OpaqueDigestItemId}, - traits::{Block as BlockT, Header, NumberFor, One, SaturatedConversion, Saturating, Zero}, + traits::{Block as BlockT, Header, NumberFor, SaturatedConversion, Saturating, Zero}, DigestItem, }; @@ -1875,13 +1875,10 @@ where let finalized = client.info().finalized_number; let revertible = blocks.min(best_number - finalized); - let number = best_number - revertible; - let hash = client - .block_hash_from_id(&BlockId::Number(number))? - .ok_or(ClientError::Backend(format!( - "Unexpected hash lookup failure for block number: {}", - number - )))?; + let revert_up_to_number = best_number - revertible; + let revert_up_to_hash = client.hash(revert_up_to_number)?.ok_or(ClientError::Backend( + format!("Unexpected hash lookup failure for block number: {}", revert_up_to_number), + ))?; // Revert epoch changes tree. @@ -1890,34 +1887,37 @@ where aux_schema::load_epoch_changes::<Block, Client>(&*client, config.genesis_config())?; let mut epoch_changes = epoch_changes.shared_data(); - if number == Zero::zero() { + if revert_up_to_number == Zero::zero() { // Special case, no epoch changes data were present on genesis. *epoch_changes = EpochChangesFor::<Block, Epoch>::default(); } else { - epoch_changes.revert(descendent_query(&*client), hash, number); + epoch_changes.revert(descendent_query(&*client), revert_up_to_hash, revert_up_to_number); } // Remove block weights added after the revert point. let mut weight_keys = HashSet::with_capacity(revertible.saturated_into()); + let leaves = backend.blockchain().leaves()?.into_iter().filter(|&leaf| { - sp_blockchain::tree_route(&*client, hash, leaf) + sp_blockchain::tree_route(&*client, revert_up_to_hash, leaf) .map(|route| route.retracted().is_empty()) .unwrap_or_default() }); + for leaf in leaves { let mut hash = leaf; - // Insert parent after parent until we don't hit an already processed - // branch or we reach a direct child of the rollback point. - while weight_keys.insert(aux_schema::block_weight_key(hash)) { + loop { let meta = client.header_metadata(hash)?; - if meta.number <= number + One::one() { - // We've reached a child of the revert point, stop here. + if meta.number <= revert_up_to_number || + !weight_keys.insert(aux_schema::block_weight_key(hash)) + { + // We've reached the revert point or an already processed branch, stop here. break } - hash = client.header_metadata(hash)?.parent; + hash = meta.parent; } } + let weight_keys: Vec<_> = weight_keys.iter().map(|val| val.as_slice()).collect(); // Write epoch changes and remove weights in one shot. diff --git a/substrate/client/consensus/babe/src/tests.rs b/substrate/client/consensus/babe/src/tests.rs index 9875ff00673f59e84c4fd3b0557c1e01421eb3bf..e0590fc0cd86efd5ebe684c238809ce5bdb6534e 100644 --- a/substrate/client/consensus/babe/src/tests.rs +++ b/substrate/client/consensus/babe/src/tests.rs @@ -814,6 +814,44 @@ fn revert_prunes_epoch_changes_and_removes_weights() { assert!(weight_data_check(&fork3, false)); } +#[test] +fn revert_not_allowed_for_finalized() { + let mut net = BabeTestNet::new(1); + + let peer = net.peer(0); + let data = peer.data.as_ref().expect("babe link set up during initialization"); + + let client = peer.client().as_client(); + let backend = peer.client().as_backend(); + let mut block_import = data.block_import.lock().take().expect("import set up during init"); + + let mut proposer_factory = DummyFactory { + client: client.clone(), + config: data.link.config.clone(), + epoch_changes: data.link.epoch_changes.clone(), + mutator: Arc::new(|_, _| ()), + }; + + let mut propose_and_import_blocks_wrap = |parent_id, n| { + propose_and_import_blocks(&client, &mut proposer_factory, &mut block_import, parent_id, n) + }; + + let canon = propose_and_import_blocks_wrap(BlockId::Number(0), 3); + + // Finalize best block + client.finalize_block(BlockId::Hash(canon[2]), None, false).unwrap(); + + // Revert canon chain to last finalized block + revert(client.clone(), backend, 100).expect("revert should work for baked test scenario"); + + let weight_data_check = |hashes: &[Hash], expected: bool| { + hashes.iter().all(|hash| { + aux_schema::load_block_weight(&*client, hash).unwrap().is_some() == expected + }) + }; + assert!(weight_data_check(&canon, true)); +} + #[test] fn importing_epoch_change_block_prunes_tree() { let mut net = BabeTestNet::new(1);