Custom justification verification (#979)

* custom justification verification

* is_ancestor -> ensure_descendant

* PrecommitIsNotCommitDescendant

* Update primitives/header-chain/src/justification.rs

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

bridges/primitives/header-chain/src/justification.rs

@@ -20,103 +20,13 @@
 //! will ever be moved to the sp_finality_grandpa, we should reuse that implementation.
 
 use codec::{Decode, Encode};
-use finality_grandpa::{voter_set::VoterSet, Chain, Error as GrandpaError};
+use finality_grandpa::voter_set::VoterSet;
 use frame_support::RuntimeDebug;
 use sp_finality_grandpa::{AuthorityId, AuthoritySignature, SetId};
 use sp_runtime::traits::Header as HeaderT;
 use sp_std::collections::{btree_map::BTreeMap, btree_set::BTreeSet};
 use sp_std::prelude::*;
 
-/// Justification verification error.
-#[derive(RuntimeDebug, PartialEq)]
-pub enum Error {
-	/// Failed to decode justification.
-	JustificationDecode,
-	/// Justification is finalizing unexpected header.
-	InvalidJustificationTarget,
-	/// Invalid commit in justification.
-	InvalidJustificationCommit,
-	/// Justification has invalid authority singature.
-	InvalidAuthoritySignature,
-	/// The justification has precommit for the header that has no route from the target header.
-	InvalidPrecommitAncestryProof,
-	/// The justification has 'unused' headers in its precommit ancestries.
-	InvalidPrecommitAncestries,
-}
-
-/// Decode justification target.
-pub fn decode_justification_target<Header: HeaderT>(
-	raw_justification: &[u8],
-) -> Result<(Header::Hash, Header::Number), Error> {
-	GrandpaJustification::<Header>::decode(&mut &*raw_justification)
-		.map(|justification| (justification.commit.target_hash, justification.commit.target_number))
-		.map_err(|_| Error::JustificationDecode)
-}
-
-/// Verify that justification, that is generated by given authority set, finalizes given header.
-pub fn verify_justification<Header: HeaderT>(
-	finalized_target: (Header::Hash, Header::Number),
-	authorities_set_id: SetId,
-	authorities_set: &VoterSet<AuthorityId>,
-	justification: &GrandpaJustification<Header>,
-) -> Result<(), Error>
-where
-	Header::Number: finality_grandpa::BlockNumberOps,
-{
-	// Ensure that it is justification for the expected header
-	if (justification.commit.target_hash, justification.commit.target_number) != finalized_target {
-		return Err(Error::InvalidJustificationTarget);
-	}
-
-	// Validate commit of the justification. Note that `validate_commit()` assumes that all
-	// signatures are valid. We'll check the validity of the signatures later since they're more
-	// resource intensive to verify.
-	let ancestry_chain = AncestryChain::new(&justification.votes_ancestries);
-	match finality_grandpa::validate_commit(&justification.commit, authorities_set, &ancestry_chain) {
-		Ok(ref result) if result.ghost().is_some() => {}
-		_ => return Err(Error::InvalidJustificationCommit),
-	}
-
-	// Now that we know that the commit is correct, check authorities signatures
-	let mut buf = Vec::new();
-	let mut visited_hashes = BTreeSet::new();
-	for signed in &justification.commit.precommits {
-		if !sp_finality_grandpa::check_message_signature_with_buffer(
-			&finality_grandpa::Message::Precommit(signed.precommit.clone()),
-			&signed.id,
-			&signed.signature,
-			justification.round,
-			authorities_set_id,
-			&mut buf,
-		) {
-			return Err(Error::InvalidAuthoritySignature);
-		}
-
-		if justification.commit.target_hash == signed.precommit.target_hash {
-			continue;
-		}
-
-		match ancestry_chain.ancestry(justification.commit.target_hash, signed.precommit.target_hash) {
-			Ok(route) => {
-				// ancestry starts from parent hash but the precommit target hash has been visited
-				visited_hashes.insert(signed.precommit.target_hash);
-				visited_hashes.extend(route);
-			}
-			_ => {
-				// could this happen in practice? I don't think so, but original code has this check
-				return Err(Error::InvalidPrecommitAncestryProof);
-			}
-		}
-	}
-
-	// both iterators are `BTree*` iterators, so have the same order => safe to compare
-	if !visited_hashes.iter().eq(ancestry_chain.ancestry.keys()) {
-		return Err(Error::InvalidPrecommitAncestries);
-	}
-
-	Ok(())
-}
-
 /// A GRANDPA Justification is a proof that a given header was finalized
 /// at a certain height and with a certain set of authorities.
 ///
@@ -138,44 +48,172 @@ impl<H: HeaderT> crate::FinalityProof<H::Number> for GrandpaJustification<H> {
 	}
 }
 
-/// A utility trait implementing `finality_grandpa::Chain` using a given set of headers.
-#[derive(RuntimeDebug)]
-struct AncestryChain<Header: HeaderT> {
-	ancestry: BTreeMap<Header::Hash, Header::Hash>,
+/// Justification verification error.
+#[derive(RuntimeDebug, PartialEq)]
+pub enum Error {
+	/// Failed to decode justification.
+	JustificationDecode,
+	/// Justification is finalizing unexpected header.
+	InvalidJustificationTarget,
+	/// The authority has provided an invalid signature.
+	InvalidAuthoritySignature,
+	/// The justification contains precommit for header that is not a descendant of the commit header.
+	PrecommitIsNotCommitDescendant,
+	/// The cumulative weight of all votes in the justification is not enough to justify commit
+	/// header finalization.
+	TooLowCumulativeWeight,
+	/// The justification contains extra (unused) headers in its `votes_ancestries` field.
+	ExtraHeadersInVotesAncestries,
 }
 
-impl<Header: HeaderT> AncestryChain<Header> {
-	fn new(ancestry: &[Header]) -> AncestryChain<Header> {
-		AncestryChain {
-			ancestry: ancestry
-				.iter()
-				.map(|header| (header.hash(), *header.parent_hash()))
-				.collect(),
-		}
-	}
+/// Decode justification target.
+pub fn decode_justification_target<Header: HeaderT>(
+	raw_justification: &[u8],
+) -> Result<(Header::Hash, Header::Number), Error> {
+	GrandpaJustification::<Header>::decode(&mut &*raw_justification)
+		.map(|justification| (justification.commit.target_hash, justification.commit.target_number))
+		.map_err(|_| Error::JustificationDecode)
 }
 
-impl<Header: HeaderT> finality_grandpa::Chain<Header::Hash, Header::Number> for AncestryChain<Header>
+/// Verify that justification, that is generated by given authority set, finalizes given header.
+pub fn verify_justification<Header: HeaderT>(
+	finalized_target: (Header::Hash, Header::Number),
+	authorities_set_id: SetId,
+	authorities_set: &VoterSet<AuthorityId>,
+	justification: &GrandpaJustification<Header>,
+) -> Result<(), Error>
 where
 	Header::Number: finality_grandpa::BlockNumberOps,
 {
-	fn ancestry(&self, base: Header::Hash, block: Header::Hash) -> Result<Vec<Header::Hash>, GrandpaError> {
-		let mut route = Vec::new();
-		let mut current_hash = block;
-		loop {
-			if current_hash == base {
-				break;
-			}
-			match self.ancestry.get(&current_hash).cloned() {
-				Some(parent_hash) => {
-					current_hash = parent_hash;
-					route.push(current_hash);
-				}
-				_ => return Err(GrandpaError::NotDescendent),
-			}
-		}
-		route.pop(); // remove the base
+	// ensure that it is justification for the expected header
+	if (justification.commit.target_hash, justification.commit.target_number) != finalized_target {
+		return Err(Error::InvalidJustificationTarget);
+	}
 
-		Ok(route)
+	let mut chain = AncestryChain::new(&justification.votes_ancestries);
+	let mut signature_buffer = Vec::new();
+	let mut votes = BTreeSet::new();
+	let mut cumulative_weight = 0u64;
+	for signed in &justification.commit.precommits {
+		// authority must be in the set
+		let authority_info = match authorities_set.get(&signed.id) {
+			Some(authority_info) => authority_info,
+			None => {
+				// just ignore precommit from unknown authority as `finality_grandpa::import_precommit` does
+				continue;
+			}
+		};
+
+		// check if authority has already voted in the same round.
+		//
+		// there's a lot of code in `validate_commit` and `import_precommit` functions inside
+		// `finality-grandpa` crate (mostly related to reporing equivocations). But the only thing that we
+		// care about is that only first vote from the authority is accepted
+		if !votes.insert(signed.id.clone()) {
+			continue;
+		}
+
+		// everything below this line can't just `continue`, because state is already altered
+
+		// all precommits must be for block higher than the target
+		if signed.precommit.target_number < justification.commit.target_number {
+			return Err(Error::PrecommitIsNotCommitDescendant);
+		}
+		// all precommits must be for target block descendents
+		chain = chain.ensure_descendant(&justification.commit.target_hash, &signed.precommit.target_hash)?;
+		// since we know now that the precommit target is the descendant of the justification target,
+		// we may increase 'weight' of the justification target
+		//
+		// there's a lot of code in the `VoteGraph::insert` method inside `finality-grandpa` crate,
+		// but in the end it is only used to find GHOST, which we don't care about. The only thing
+		// that we care about is that the justification target has enough weight
+		cumulative_weight = cumulative_weight.checked_add(authority_info.weight().0.into()).expect(
+			"sum of weights of ALL authorities is expected not to overflow - this is guaranteed by\
+				existence of VoterSet;\
+				the order of loop conditions guarantees that we can account vote from same authority\
+				multiple times;\
+				thus we'll never overflow the u64::MAX;\
+				qed",
+		);
+		// verify authority signature
+		if !sp_finality_grandpa::check_message_signature_with_buffer(
+			&finality_grandpa::Message::Precommit(signed.precommit.clone()),
+			&signed.id,
+			&signed.signature,
+			justification.round,
+			authorities_set_id,
+			&mut signature_buffer,
+		) {
+			return Err(Error::InvalidAuthoritySignature);
+		}
+	}
+
+	// check that there are no extra headers in the justification
+	if !chain.unvisited.is_empty() {
+		return Err(Error::ExtraHeadersInVotesAncestries);
+	}
+
+	// check that the cumulative weight of validators voted for the justification target (or one
+	// of its descendents) is larger than required threshold.
+	let threshold = authorities_set.threshold().0.into();
+	if cumulative_weight >= threshold {
+		Ok(())
+	} else {
+		Err(Error::TooLowCumulativeWeight)
+	}
+}
+
+/// Votes ancestries with useful methods.
+#[derive(RuntimeDebug)]
+pub struct AncestryChain<Header: HeaderT> {
+	/// Header hash => parent header hash mapping.
+	pub parents: BTreeMap<Header::Hash, Header::Hash>,
+	/// Hashes of headers that weren't visited by `is_ancestor` method.
+	pub unvisited: BTreeSet<Header::Hash>,
+}
+
+impl<Header: HeaderT> AncestryChain<Header> {
+	/// Create new ancestry chain.
+	pub fn new(ancestry: &[Header]) -> AncestryChain<Header> {
+		let mut parents = BTreeMap::new();
+		let mut unvisited = BTreeSet::new();
+		for ancestor in ancestry {
+			let hash = ancestor.hash();
+			let parent_hash = *ancestor.parent_hash();
+			parents.insert(hash, parent_hash);
+			unvisited.insert(hash);
+		}
+		AncestryChain { parents, unvisited }
+	}
+
+	/// Returns `Err(_)` if `precommit_target` is a descendant of the `commit_target` block and `Ok(_)` otherwise.
+	pub fn ensure_descendant(
+		mut self,
+		commit_target: &Header::Hash,
+		precommit_target: &Header::Hash,
+	) -> Result<Self, Error> {
+		let mut current_hash = *precommit_target;
+		loop {
+			if current_hash == *commit_target {
+				break;
+			}
+
+			let is_visited_before = !self.unvisited.remove(&current_hash);
+			current_hash = match self.parents.get(&current_hash) {
+				Some(parent_hash) => {
+					if is_visited_before {
+						// `Some(parent_hash)` means that the `current_hash` is in the `parents` container
+						// `is_visited_before` means that it has been visited before in some of previous calls
+						// => since we assume that previous call has finished with `true`, this also will
+						//    be finished with `true`
+						return Ok(self);
+					}
+
+					*parent_hash
+				}
+				None => return Err(Error::PrecommitIsNotCommitDescendant),
+			};
+		}
+		Ok(self)
 	}
 }