Fix Babe revert when last finalized block is a leaf (paritytech#11500)

* Fix Babe revert when a leaf is the last finalized block

Without this fix the last finalized block weight data is wrongly removed
on revert scenario where the last finalized block is a leaf.

* Remove redundant check

* Added test to exercise the fix

* Rename test

* Give variables better names

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.

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);