grandpa-rpc: allow proving finality of blocks from latest authority set (#8585)

* grandpa: use new latest stored justification in prove_finality

* grandpa: include end in range in FinalityProof::unknown_headers

* grandpa: typo in comment

* grandpa: remove ProvableJustification

* grandpa: revert unnecessary changes

* grandpa: extend AuthoritySetChangeId and cleanup get_set_id

* grandpa: move check_finality_proof to the test module

* grandpa: warn on missing authority set changes data

* grandpa: add missing use statement

* grandpa: simplify finality_proof tests

* grandpa: additional tests for finality_proof

Co-authored-by: André Silva <andrerfosilva@gmail.com>

substrate/client/finality-grandpa/src/authorities.rs

@@ -18,6 +18,10 @@
 
 //! Utilities for dealing with authorities, authority sets, and handoffs.
 
+use std::cmp::Ord;
+use std::fmt::Debug;
+use std::ops::Add;
+
 use fork_tree::ForkTree;
 use parking_lot::MappedMutexGuard;
 use finality_grandpa::voter_set::VoterSet;
@@ -27,9 +31,7 @@ use sc_telemetry::{telemetry, TelemetryHandle, CONSENSUS_INFO};
 use sp_finality_grandpa::{AuthorityId, AuthorityList};
 use sc_consensus::shared_data::{SharedData, SharedDataLocked};
 
-use std::cmp::Ord;
-use std::fmt::Debug;
-use std::ops::Add;
+use crate::SetId;
 
 /// Error type returned on operations on the `AuthoritySet`.
 #[derive(Debug, derive_more::Display)]
@@ -684,6 +686,20 @@ impl<H, N: Add<Output=N> + Clone> PendingChange<H, N> {
 #[derive(Debug, Encode, Decode, Clone, PartialEq)]
 pub struct AuthoritySetChanges<N>(Vec<(u64, N)>);
 
+/// The response when querying for a set id for a specific block. Either we get a set id
+/// together with a block number for the last block in the set, or that the requested block is in the
+/// latest set, or that we don't know what set id the given block belongs to.
+#[derive(Debug, PartialEq)]
+pub enum AuthoritySetChangeId<N> {
+	/// The requested block is in the latest set.
+	Latest,
+	/// Tuple containing the set id and the last block number of that set.
+	Set(SetId, N),
+	/// We don't know which set id the request block belongs to (this can only happen due to missing
+	/// data).
+	Unknown,
+}
+
 impl<N> From<Vec<(u64, N)>> for AuthoritySetChanges<N> {
 	fn from(changes: Vec<(u64, N)>) -> AuthoritySetChanges<N> {
 		AuthoritySetChanges(changes)
@@ -699,7 +715,15 @@ impl<N: Ord + Clone> AuthoritySetChanges<N> {
 		self.0.push((set_id, block_number));
 	}
 
-	pub(crate) fn get_set_id(&self, block_number: N) -> Option<(u64, N)> {
+	pub(crate) fn get_set_id(&self, block_number: N) -> AuthoritySetChangeId<N> {
+		if self.0
+			.last()
+			.map(|last_auth_change| last_auth_change.1 < block_number)
+			.unwrap_or(false)
+		{
+			return AuthoritySetChangeId::Latest;
+		}
+
 		let idx = self.0
 			.binary_search_by_key(&block_number, |(_, n)| n.clone())
 			.unwrap_or_else(|b| b);
@@ -711,16 +735,16 @@ impl<N: Ord + Clone> AuthoritySetChanges<N> {
 				let (prev_set_id, _) = self.0[idx - 1usize];
 				if set_id != prev_set_id + 1u64 {
 					// Without the preceding set_id we don't have a well-defined start.
-					return None;
+					return AuthoritySetChangeId::Unknown;
 				}
 			} else if set_id != 0 {
 				// If this is the first index, yet not the first set id then it's not well-defined
 				// that we are in the right set id.
-				return None;
+				return AuthoritySetChangeId::Unknown;
 			}
-			Some((set_id, block_number))
+			AuthoritySetChangeId::Set(set_id, block_number)
 		} else {
-			None
+			AuthoritySetChangeId::Unknown
 		}
 	}
 
@@ -1660,11 +1684,11 @@ mod tests {
 		authority_set_changes.append(1, 81);
 		authority_set_changes.append(2, 121);
 
-		assert_eq!(authority_set_changes.get_set_id(20), Some((0, 41)));
-		assert_eq!(authority_set_changes.get_set_id(40), Some((0, 41)));
-		assert_eq!(authority_set_changes.get_set_id(41), Some((0, 41)));
-		assert_eq!(authority_set_changes.get_set_id(42), Some((1, 81)));
-		assert_eq!(authority_set_changes.get_set_id(141), None);
+		assert_eq!(authority_set_changes.get_set_id(20), AuthoritySetChangeId::Set(0, 41));
+		assert_eq!(authority_set_changes.get_set_id(40), AuthoritySetChangeId::Set(0, 41));
+		assert_eq!(authority_set_changes.get_set_id(41), AuthoritySetChangeId::Set(0, 41));
+		assert_eq!(authority_set_changes.get_set_id(42), AuthoritySetChangeId::Set(1, 81));
+		assert_eq!(authority_set_changes.get_set_id(141), AuthoritySetChangeId::Latest);
 	}
 
 	#[test]
@@ -1674,11 +1698,11 @@ mod tests {
 		authority_set_changes.append(3, 81);
 		authority_set_changes.append(4, 121);
 
-		assert_eq!(authority_set_changes.get_set_id(20), None);
-		assert_eq!(authority_set_changes.get_set_id(40), None);
-		assert_eq!(authority_set_changes.get_set_id(41), None);
-		assert_eq!(authority_set_changes.get_set_id(42), Some((3, 81)));
-		assert_eq!(authority_set_changes.get_set_id(141), None);
+		assert_eq!(authority_set_changes.get_set_id(20), AuthoritySetChangeId::Unknown);
+		assert_eq!(authority_set_changes.get_set_id(40), AuthoritySetChangeId::Unknown);
+		assert_eq!(authority_set_changes.get_set_id(41), AuthoritySetChangeId::Unknown);
+		assert_eq!(authority_set_changes.get_set_id(42), AuthoritySetChangeId::Set(3, 81));
+		assert_eq!(authority_set_changes.get_set_id(141), AuthoritySetChangeId::Latest);
 	}
 
 	#[test]

substrate/client/finality-grandpa/src/finality_proof.rs

@@ -36,23 +36,23 @@
 //! finality proof (that finalizes some block C that is ancestor of the B and descendant
 //! of the U) could be returned.
 
-use log::trace;
+use log::{trace, warn};
 use std::sync::Arc;
 
-use finality_grandpa::BlockNumberOps;
 use parity_scale_codec::{Encode, Decode};
-use sp_blockchain::{Backend as BlockchainBackend, Error as ClientError, Result as ClientResult};
+use sp_blockchain::{Backend as BlockchainBackend, HeaderBackend};
+use sp_finality_grandpa::GRANDPA_ENGINE_ID;
 use sp_runtime::{
-	EncodedJustification, generic::BlockId,
+	generic::BlockId,
 	traits::{NumberFor, Block as BlockT, Header as HeaderT, One},
 };
 use sc_client_api::backend::Backend;
-use sp_finality_grandpa::{AuthorityId, GRANDPA_ENGINE_ID};
 
-use crate::authorities::AuthoritySetChanges;
-use crate::justification::GrandpaJustification;
-use crate::SharedAuthoritySet;
-use crate::VoterSet;
+use crate::{
+	SharedAuthoritySet, best_justification,
+	authorities::{AuthoritySetChangeId, AuthoritySetChanges},
+	justification::GrandpaJustification,
+};
 
 const MAX_UNKNOWN_HEADERS: usize = 100_000;
 
@@ -97,14 +97,13 @@ where
 impl<B, Block> FinalityProofProvider<B, Block>
 where
 	Block: BlockT,
-	NumberFor<Block>: BlockNumberOps,
 	B: Backend<Block> + Send + Sync + 'static,
 {
 	/// Prove finality for the given block number by returning a Justification for the last block of
 	/// the authority set.
 	pub fn prove_finality(
 		&self,
-		block: NumberFor<Block>
+		block: NumberFor<Block>,
 	) -> Result<Option<Vec<u8>>, FinalityProofError> {
 		let authority_set_changes = if let Some(changes) = self
 			.shared_authority_set
@@ -116,8 +115,8 @@ where
 			return Ok(None);
 		};
 
-		prove_finality::<_, _, GrandpaJustification<Block>>(
-			&*self.backend.blockchain(),
+		prove_finality(
+			&*self.backend,
 			authority_set_changes,
 			block,
 		)
@@ -151,19 +150,19 @@ pub enum FinalityProofError {
 	Client(sp_blockchain::Error),
 }
 
-fn prove_finality<Block, B, J>(
-	blockchain: &B,
+fn prove_finality<Block, B>(
+	backend: &B,
 	authority_set_changes: AuthoritySetChanges<NumberFor<Block>>,
 	block: NumberFor<Block>,
 ) -> Result<Option<Vec<u8>>, FinalityProofError>
 where
 	Block: BlockT,
-	B: BlockchainBackend<Block>,
-	J: ProvableJustification<Block::Header>,
+	B: Backend<Block>,
 {
-	// Early-return if we sure that there are no blocks finalized AFTER begin block
-	let info = blockchain.info();
-	if info.finalized_number <= block {
+	// Early-return if we are sure that there are no blocks finalized that cover the requested
+	// block.
+	let info = backend.blockchain().info();
+	if info.finalized_number < block {
 		let err = format!(
 			"Requested finality proof for descendant of #{} while we only have finalized #{}.",
 			block,
@@ -173,45 +172,60 @@ where
 		return Err(FinalityProofError::BlockNotYetFinalized);
 	}
 
-	// Get set_id the block belongs to, and the last block of the set which should contain a
-	// Justification we can use to prove the requested block.
-	let (_, last_block_for_set) = if let Some(id) = authority_set_changes.get_set_id(block) {
-		id
-	} else {
-		trace!(
-			target: "afg",
-			"AuthoritySetChanges does not cover the requested block #{}. \
-			Maybe the subscription API is more appropriate.",
-			block,
-		);
-		return Err(FinalityProofError::BlockNotInAuthoritySetChanges);
-	};
-
-	// Get the Justification stored at the last block of the set
-	let last_block_for_set_id = BlockId::Number(last_block_for_set);
-	let justification =
-		if let Some(grandpa_justification) = blockchain.justifications(last_block_for_set_id)?
-			.and_then(|justifications| justifications.into_justification(GRANDPA_ENGINE_ID))
-		{
-			grandpa_justification
-		} else {
-			trace!(
+	let (justification, just_block) = match authority_set_changes.get_set_id(block) {
+		AuthoritySetChangeId::Latest => {
+			if let Some(justification) = best_justification(backend)?
+				.map(|j: GrandpaJustification<Block>| (j.encode(), j.target().0))
+			{
+				justification
+			} else {
+				trace!(
+					target: "afg",
+					"No justification found for the latest finalized block. \
+					Returning empty proof.",
+				);
+				return Ok(None);
+			}
+		}
+		AuthoritySetChangeId::Set(_, last_block_for_set) => {
+			let last_block_for_set_id = BlockId::Number(last_block_for_set);
+			let justification = if let Some(grandpa_justification) = backend
+				.blockchain()
+				.justifications(last_block_for_set_id)?
+				.and_then(|justifications| justifications.into_justification(GRANDPA_ENGINE_ID))
+			{
+				grandpa_justification
+			} else {
+				trace!(
+					target: "afg",
+					"No justification found when making finality proof for {}. \
+					Returning empty proof.",
+					block,
+				);
+				return Ok(None);
+			};
+			(justification, last_block_for_set)
+		}
+		AuthoritySetChangeId::Unknown => {
+			warn!(
 				target: "afg",
-				"No justification found when making finality proof for {}. Returning empty proof.",
+				"AuthoritySetChanges does not cover the requested block #{} due to missing data. \
+				 You need to resync to populate AuthoritySetChanges properly.",
 				block,
 			);
-			return Ok(None);
-		};
+			return Err(FinalityProofError::BlockNotInAuthoritySetChanges);
+		}
+	};
 
 	// Collect all headers from the requested block until the last block of the set
 	let unknown_headers = {
 		let mut headers = Vec::new();
 		let mut current = block + One::one();
 		loop {
-			if current >= last_block_for_set || headers.len() >= MAX_UNKNOWN_HEADERS {
+			if current > just_block || headers.len() >= MAX_UNKNOWN_HEADERS {
 				break;
 			}
-			headers.push(blockchain.expect_header(BlockId::Number(current))?);
+			headers.push(backend.blockchain().expect_header(BlockId::Number(current))?);
 			current += One::one();
 		}
 		headers
@@ -219,7 +233,7 @@ where
 
 	Ok(Some(
 		FinalityProof {
-			block: blockchain.expect_block_hash_from_id(&last_block_for_set_id)?,
+			block: backend.blockchain().expect_block_hash_from_id(&BlockId::Number(just_block))?,
 			justification,
 			unknown_headers,
 		}
@@ -227,97 +241,49 @@ where
 	))
 }
 
-/// Check GRANDPA proof-of-finality for the given block.
-///
-/// Returns the vector of headers that MUST be validated + imported
-/// AND if at least one of those headers is invalid, all other MUST be considered invalid.
-///
-/// This is currently not used, and exists primarily as an example of how to check finality proofs.
-#[cfg(test)]
-fn check_finality_proof<Header: HeaderT, J>(
-	current_set_id: u64,
-	current_authorities: sp_finality_grandpa::AuthorityList,
-	remote_proof: Vec<u8>,
-) -> ClientResult<FinalityProof<Header>>
-where
-	J: ProvableJustification<Header>,
-{
-	let proof = FinalityProof::<Header>::decode(&mut &remote_proof[..])
-		.map_err(|_| ClientError::BadJustification("failed to decode finality proof".into()))?;
-
-	let justification: J = Decode::decode(&mut &proof.justification[..])
-		.map_err(|_| ClientError::JustificationDecode)?;
-	justification.verify(current_set_id, &current_authorities)?;
-
-	Ok(proof)
-}
-
-/// Justification used to prove block finality.
-pub trait ProvableJustification<Header: HeaderT>: Encode + Decode {
-	/// Verify justification with respect to authorities set and authorities set id.
-	fn verify(&self, set_id: u64, authorities: &[(AuthorityId, u64)]) -> ClientResult<()>;
-
-	/// Decode and verify justification.
-	fn decode_and_verify(
-		justification: &EncodedJustification,
-		set_id: u64,
-		authorities: &[(AuthorityId, u64)],
-	) -> ClientResult<Self> {
-		let justification =
-			Self::decode(&mut &**justification).map_err(|_| ClientError::JustificationDecode)?;
-		justification.verify(set_id, authorities)?;
-		Ok(justification)
-	}
-}
-
-impl<Block: BlockT> ProvableJustification<Block::Header> for GrandpaJustification<Block>
-where
-	NumberFor<Block>: BlockNumberOps,
-{
-	fn verify(&self, set_id: u64, authorities: &[(AuthorityId, u64)]) -> ClientResult<()> {
-		let authorities = VoterSet::new(authorities.iter().cloned()).ok_or(
-			ClientError::Consensus(sp_consensus::Error::InvalidAuthoritiesSet),
-		)?;
-
-		GrandpaJustification::verify_with_voter_set(self, set_id, &authorities)
-	}
-}
-
 #[cfg(test)]
 pub(crate) mod tests {
 	use super::*;
-	use crate::authorities::AuthoritySetChanges;
+	use crate::{
+		authorities::AuthoritySetChanges, BlockNumberOps, ClientError, SetId,
+	};
+	use futures::executor::block_on;
+	use sc_block_builder::BlockBuilderProvider;
+	use sc_client_api::{apply_aux, LockImportRun};
+	use sp_consensus::BlockOrigin;
 	use sp_core::crypto::Public;
-	use sp_runtime::Justifications;
-	use sp_finality_grandpa::AuthorityList;
-	use sc_client_api::NewBlockState;
-	use sc_client_api::in_mem::Blockchain as InMemoryBlockchain;
-	use substrate_test_runtime_client::runtime::{Block, Header, H256};
+	use sp_finality_grandpa::{AuthorityId, GRANDPA_ENGINE_ID as ID};
+	use sp_keyring::Ed25519Keyring;
+	use substrate_test_runtime_client::{
+		runtime::{Block, Header, H256},
+		Backend as TestBackend, ClientBlockImportExt, ClientExt, DefaultTestClientBuilderExt,
+		TestClient, TestClientBuilder, TestClientBuilderExt,
+	};
+
+	/// Check GRANDPA proof-of-finality for the given block.
+	///
+	/// Returns the vector of headers that MUST be validated + imported
+	/// AND if at least one of those headers is invalid, all other MUST be considered invalid.
+	fn check_finality_proof<Block: BlockT>(
+		current_set_id: SetId,
+		current_authorities: sp_finality_grandpa::AuthorityList,
+		remote_proof: Vec<u8>,
+	) -> sp_blockchain::Result<super::FinalityProof<Block::Header>>
+	where
+		NumberFor<Block>: BlockNumberOps,
+	{
+		let proof = super::FinalityProof::<Block::Header>::decode(&mut &remote_proof[..])
+			.map_err(|_| ClientError::BadJustification("failed to decode finality proof".into()))?;
+
+		let justification: GrandpaJustification<Block> = Decode::decode(&mut &proof.justification[..])
+			.map_err(|_| ClientError::JustificationDecode)?;
+		justification.verify(current_set_id, &current_authorities)?;
+
+		Ok(proof)
+	}
 
 	pub(crate) type FinalityProof = super::FinalityProof<Header>;
 
-	#[derive(Debug, PartialEq, Encode, Decode)]
-	pub struct TestJustification(pub (u64, AuthorityList), pub Vec<u8>);
-
-	impl ProvableJustification<Header> for TestJustification {
-		fn verify(&self, set_id: u64, authorities: &[(AuthorityId, u64)]) -> ClientResult<()> {
-			if (self.0).0 != set_id || (self.0).1 != authorities {
-				return Err(ClientError::BadJustification("test".into()));
-			}
-
-			Ok(())
-		}
-	}
-
-	#[derive(Debug, PartialEq, Encode, Decode)]
-	pub struct TestBlockJustification(TestJustification, u64, H256);
-
-	impl ProvableJustification<Header> for TestBlockJustification {
-		fn verify(&self, set_id: u64, authorities: &[(AuthorityId, u64)]) -> ClientResult<()> {
-			self.0.verify(set_id, authorities)
-		}
-	}
-
 	fn header(number: u64) -> Header {
 		let parent_hash = match number {
 			0 => Default::default(),
@@ -332,57 +298,64 @@ pub(crate) mod tests {
 		)
 	}
 
-	fn test_blockchain() -> InMemoryBlockchain<Block> {
-		use sp_finality_grandpa::GRANDPA_ENGINE_ID as ID;
-		let blockchain = InMemoryBlockchain::<Block>::new();
-		let just0 = Some(Justifications::from((ID, vec![0])));
-		let just1 = Some(Justifications::from((ID, vec![1])));
-		let just2 = None;
-		let just3 = Some(Justifications::from((ID, vec![3])));
-		blockchain.insert(header(0).hash(), header(0), just0, None, NewBlockState::Final).unwrap();
-		blockchain.insert(header(1).hash(), header(1), just1, None, NewBlockState::Final).unwrap();
-		blockchain.insert(header(2).hash(), header(2), just2, None, NewBlockState::Best).unwrap();
-		blockchain.insert(header(3).hash(), header(3), just3, None, NewBlockState::Final).unwrap();
-		blockchain
+	fn test_blockchain(
+		number_of_blocks: u64,
+		to_finalize: &[u64],
+	) -> (Arc<TestClient>, Arc<TestBackend>, Vec<Block>) {
+		let builder = TestClientBuilder::new();
+		let backend = builder.backend();
+		let mut client = Arc::new(builder.build());
+
+		let mut blocks = Vec::new();
+		for _ in 0..number_of_blocks {
+			let block = client.new_block(Default::default()).unwrap().build().unwrap().block;
+			block_on(client.import(BlockOrigin::Own, block.clone())).unwrap();
+			blocks.push(block);
+		}
+
+		for block in to_finalize {
+			client.finalize_block(BlockId::Number(*block), None).unwrap();
+		}
+		(client, backend, blocks)
+	}
+
+	fn store_best_justification(client: &TestClient, just: &GrandpaJustification<Block>) {
+		client.lock_import_and_run(|import_op| {
+			crate::aux_schema::update_best_justification(
+				just,
+				|insert| apply_aux(import_op, insert, &[]),
+			)
+		})
+		.unwrap();
 	}
 
 	#[test]
 	fn finality_proof_fails_if_no_more_last_finalized_blocks() {
-		use sp_finality_grandpa::GRANDPA_ENGINE_ID as ID;
-		let blockchain = test_blockchain();
-		let just1 = Some(Justifications::from((ID, vec![1])));
-		let just2 = Some(Justifications::from((ID, vec![2])));
-		blockchain.insert(header(4).hash(), header(4), just1, None, NewBlockState::Best).unwrap();
-		blockchain.insert(header(5).hash(), header(5), just2, None, NewBlockState::Best).unwrap();
+		let (_, backend, _) = test_blockchain(6, &[4]);
+		let authority_set_changes = AuthoritySetChanges::empty();
 
-		let mut authority_set_changes = AuthoritySetChanges::empty();
-		authority_set_changes.append(0, 5);
-
-		// The last finalized block is 3, so we cannot provide further justifications.
-		let proof_of_4 = prove_finality::<_, _, TestJustification>(
-			&blockchain,
+		// The last finalized block is 4, so we cannot provide further justifications.
+		let proof_of_5 = prove_finality(
+			&*backend,
 			authority_set_changes,
-			*header(4).number(),
+			5,
 		);
-		assert!(matches!(proof_of_4, Err(FinalityProofError::BlockNotYetFinalized)));
+		assert!(matches!(proof_of_5, Err(FinalityProofError::BlockNotYetFinalized)));
 	}
 
 	#[test]
 	fn finality_proof_is_none_if_no_justification_known() {
-		let blockchain = test_blockchain();
-		blockchain
-			.insert(header(4).hash(), header(4), None, None, NewBlockState::Final)
-			.unwrap();
+		let (_, backend, _) = test_blockchain(6, &[4]);
 
 		let mut authority_set_changes = AuthoritySetChanges::empty();
 		authority_set_changes.append(0, 4);
 
 		// Block 4 is finalized without justification
 		// => we can't prove finality of 3
-		let proof_of_3 = prove_finality::<_, _, TestJustification>(
-			&blockchain,
+		let proof_of_3 = prove_finality(
+			&*backend,
 			authority_set_changes,
-			*header(3).number(),
+			3,
 		)
 		.unwrap();
 		assert_eq!(proof_of_3, None);
@@ -391,7 +364,7 @@ pub(crate) mod tests {
 	#[test]
 	fn finality_proof_check_fails_when_proof_decode_fails() {
 		// When we can't decode proof from Vec<u8>
-		check_finality_proof::<_, TestJustification>(
+		check_finality_proof::<Block>(
 			1,
 			vec![(AuthorityId::from_slice(&[3u8; 32]), 1u64)],
 			vec![42],
@@ -402,92 +375,208 @@ pub(crate) mod tests {
 	#[test]
 	fn finality_proof_check_fails_when_proof_is_empty() {
 		// When decoded proof has zero length
-		check_finality_proof::<_, TestJustification>(
+		check_finality_proof::<Block>(
 			1,
 			vec![(AuthorityId::from_slice(&[3u8; 32]), 1u64)],
-			Vec::<TestJustification>::new().encode(),
+			Vec::<GrandpaJustification<Block>>::new().encode(),
 		)
 		.unwrap_err();
 	}
 
 	#[test]
-	fn finality_proof_check_works() {
-		let auth = vec![(AuthorityId::from_slice(&[3u8; 32]), 1u64)];
+	fn finality_proof_check_fails_with_incomplete_justification() {
+		let (client, _, blocks) = test_blockchain(8, &[4, 5, 8]);
+
+		// Create a commit without precommits
+		let commit = finality_grandpa::Commit {
+			target_hash: blocks[7].hash(),
+			target_number: *blocks[7].header().number(),
+			precommits: Vec::new(),
+		};
+		let grandpa_just = GrandpaJustification::from_commit(&client, 8, commit).unwrap();
+
 		let finality_proof = FinalityProof {
 			block: header(2).hash(),
-			justification: TestJustification((1, auth.clone()), vec![7]).encode(),
+			justification: grandpa_just.encode(),
 			unknown_headers: Vec::new(),
 		};
-		let proof = check_finality_proof::<_, TestJustification>(
+
+		check_finality_proof::<Block>(
 			1,
-			auth.clone(),
+			vec![(AuthorityId::from_slice(&[3u8; 32]), 1u64)],
 			finality_proof.encode(),
-		)
-		.unwrap();
-		assert_eq!(proof, finality_proof);
+		).unwrap_err();
+	}
+
+	fn create_commit<S, Id>(
+		block: Block,
+		round: u64,
+		set_id: SetId,
+		auth: &[Ed25519Keyring]
+	) -> finality_grandpa::Commit<H256, u64, S, Id>
+	where
+		Id: From<sp_core::ed25519::Public>,
+		S: From<sp_core::ed25519::Signature>,
+	{
+		let mut precommits = Vec::new();
+
+		for voter in auth {
+			let precommit = finality_grandpa::Precommit {
+				target_hash: block.hash(),
+				target_number: *block.header().number(),
+			};
+
+			let msg = finality_grandpa::Message::Precommit(precommit.clone());
+			let encoded = sp_finality_grandpa::localized_payload(round, set_id, &msg);
+			let signature = voter.sign(&encoded[..]).into();
+
+			let signed_precommit = finality_grandpa::SignedPrecommit {
+				precommit,
+				signature,
+				id: voter.public().into(),
+			};
+			precommits.push(signed_precommit);
+		}
+
+		finality_grandpa::Commit {
+			target_hash: block.hash(),
+			target_number: *block.header().number(),
+			precommits,
+		}
+	}
+
+	#[test]
+	fn finality_proof_check_works_with_correct_justification() {
+		let (client, _, blocks) = test_blockchain(8, &[4, 5, 8]);
+
+		let alice = Ed25519Keyring::Alice;
+		let set_id = 1;
+		let round = 8;
+		let commit = create_commit(blocks[7].clone(), round, set_id, &[alice]);
+		let grandpa_just = GrandpaJustification::from_commit(&client, round, commit).unwrap();
+
+		let finality_proof = FinalityProof {
+			block: header(2).hash(),
+			justification: grandpa_just.encode(),
+			unknown_headers: Vec::new(),
+		};
+		assert_eq!(
+			finality_proof,
+			check_finality_proof::<Block>(
+				set_id,
+				vec![(alice.public().into(), 1u64)],
+				finality_proof.encode(),
+			)
+			.unwrap(),
+		);
 	}
 
 	#[test]
 	fn finality_proof_using_authority_set_changes_fails_with_undefined_start() {
-		use sp_finality_grandpa::GRANDPA_ENGINE_ID as ID;
-		let blockchain = test_blockchain();
-		let auth = vec![(AuthorityId::from_slice(&[1u8; 32]), 1u64)];
-		let grandpa_just4 = TestJustification((0, auth.clone()), vec![4]).encode();
-		let grandpa_just7 = TestJustification((1, auth.clone()), vec![7]).encode();
-		let just4 = Some(Justifications::from((ID, grandpa_just4)));
-		let just7 = Some(Justifications::from((ID, grandpa_just7)));
-		blockchain.insert(header(4).hash(), header(4), just4, None, NewBlockState::Final).unwrap();
-		blockchain.insert(header(5).hash(), header(5), None, None, NewBlockState::Final).unwrap();
-		blockchain.insert(header(6).hash(), header(6), None, None, NewBlockState::Final).unwrap();
-		blockchain.insert(header(7).hash(), header(7), just7, None, NewBlockState::Final).unwrap();
+		let (_, backend, _) = test_blockchain(8, &[4, 5, 8]);
 
 		// We have stored the correct block number for the relevant set, but as we are missing the
 		// block for the preceding set the start is not well-defined.
 		let mut authority_set_changes = AuthoritySetChanges::empty();
-		authority_set_changes.append(1, 7);
+		authority_set_changes.append(1, 8);
 
-		let proof_of_5 = prove_finality::<_, _, TestJustification>(
-			&blockchain,
+		let proof_of_6 = prove_finality(
+			&*backend,
 			authority_set_changes,
-			*header(5).number(),
+			6,
 		);
-		assert!(matches!(proof_of_5, Err(FinalityProofError::BlockNotInAuthoritySetChanges)));
+		assert!(matches!(proof_of_6, Err(FinalityProofError::BlockNotInAuthoritySetChanges)));
 	}
 
 	#[test]
 	fn finality_proof_using_authority_set_changes_works() {
-		use sp_finality_grandpa::GRANDPA_ENGINE_ID as ID;
-		let blockchain = test_blockchain();
-		let auth = vec![(AuthorityId::from_slice(&[1u8; 32]), 1u64)];
-		let grandpa_just4 = TestJustification((0, auth.clone()), vec![4]).encode();
-		let grandpa_just7 = TestJustification((1, auth.clone()), vec![7]).encode();
-		let just4 = Some(Justifications::from((ID, grandpa_just4)));
-		let just7 = Some(Justifications::from((ID, grandpa_just7.clone())));
-		blockchain.insert(header(4).hash(), header(4), just4, None, NewBlockState::Final) .unwrap();
-		blockchain.insert(header(5).hash(), header(5), None, None, NewBlockState::Final) .unwrap();
-		blockchain.insert(header(6).hash(), header(6), None, None, NewBlockState::Final).unwrap();
-		blockchain.insert(header(7).hash(), header(7), just7, None, NewBlockState::Final).unwrap();
+		let (client, backend, blocks) = test_blockchain(8, &[4, 5]);
+		let block7 = &blocks[6];
+		let block8 = &blocks[7];
 
+		let round = 8;
+		let commit = create_commit(block8.clone(), round, 1, &[Ed25519Keyring::Alice]);
+		let grandpa_just8 = GrandpaJustification::from_commit(&client, round, commit).unwrap();
+
+		client.finalize_block(
+			BlockId::Number(8),
+			Some((ID, grandpa_just8.encode().clone()))
+		)
+		.unwrap();
+
+		// Authority set change at block 8, so the justification stored there will be used in the
+		// FinalityProof for block 6
 		let mut authority_set_changes = AuthoritySetChanges::empty();
-		authority_set_changes.append(0, 4);
-		authority_set_changes.append(1, 7);
+		authority_set_changes.append(0, 5);
+		authority_set_changes.append(1, 8);
 
-		let proof_of_5: FinalityProof = Decode::decode(
-			&mut &prove_finality::<_, _, TestJustification>(
-				&blockchain,
-				authority_set_changes,
-				*header(5).number(),
+		let proof_of_6: FinalityProof = Decode::decode(
+			&mut &prove_finality(
+				&*backend,
+				authority_set_changes.clone(),
+				6,
 			)
 			.unwrap()
 			.unwrap()[..],
 		)
 		.unwrap();
 		assert_eq!(
-			proof_of_5,
+			proof_of_6,
 			FinalityProof {
-				block: header(7).hash(),
-				justification: grandpa_just7,
-				unknown_headers: vec![header(6)],
+				block: block8.hash(),
+				justification: grandpa_just8.encode(),
+				unknown_headers: vec![block7.header().clone(), block8.header().clone()],
+			},
+		);
+	}
+
+	#[test]
+	fn finality_proof_in_last_set_fails_without_latest() {
+		let (_, backend, _) = test_blockchain(8, &[4, 5, 8]);
+
+		// No recent authority set change, so we are in the latest set, and we will try to pickup
+		// the best stored justification, for which there is none in this case.
+		let mut authority_set_changes = AuthoritySetChanges::empty();
+		authority_set_changes.append(0, 5);
+
+		assert!(matches!(
+			prove_finality(&*backend, authority_set_changes, 6),
+			Ok(None),
+		));
+	}
+
+	#[test]
+	fn finality_proof_in_last_set_using_latest_justification_works() {
+		let (client, backend, blocks) = test_blockchain(8, &[4, 5, 8]);
+		let block7 = &blocks[6];
+		let block8 = &blocks[7];
+
+		let round = 8;
+		let commit = create_commit(block8.clone(), round, 1, &[Ed25519Keyring::Alice]);
+		let grandpa_just8 = GrandpaJustification::from_commit(&client, round, commit).unwrap();
+		store_best_justification(&client, &grandpa_just8);
+
+		// No recent authority set change, so we are in the latest set, and will pickup the best
+		// stored justification
+		let mut authority_set_changes = AuthoritySetChanges::empty();
+		authority_set_changes.append(0, 5);
+
+		let proof_of_6: FinalityProof = Decode::decode(
+			&mut &prove_finality(
+				&*backend,
+				authority_set_changes,
+				6,
+			)
+			.unwrap()
+			.unwrap()[..],
+		)
+		.unwrap();
+		assert_eq!(
+			proof_of_6,
+			FinalityProof {
+				block: block8.hash(),
+				justification: grandpa_just8.encode(),
+				unknown_headers: vec![block7.header().clone(), block8.header().clone()],
 			}
 		);
 	}