Support Tracking Forks in Substrate Pallet (#409)

* Support multiple "best headers" from different forks

* Update the name of a test

* Add note about multiple scheduled changes

* Disallow multiple scheduled authority set changes

* Return multiple best headers from Runtime APIs

* Remove invalid test write-up

* Add some sketch-ups of test scenarios

* Clean up test scenarios

* Add module for testing forks

* Write headers to storage

* Add way to check expected outcome for header imports

* Add support for importing finality proofs

* Support importing headers which schedule changes

* Write out test scenario using new framework

* Map authority set changes across forks

Gets all the tests in the `forks` module passing

* Remove basic tests

These were used when working on the initial test helper

* Prevent multiple pending set changes on the same fork

* Remove old test which allowed imports past unfinalized header

* Ignore failing test (for now)

* Rewrite `if` comparison using `match` and `cmp`

Fixes Clippy warning: `comparison_chain`

* Add helper for writing test headers with default characteristics

* Fix test that checked authority set updates

* Make note about importing headers on different unfinalized fork

* Perform some cleanup on the fork module

* Fix Clippy complaints

* Provide list of unfinalized headers to Runtime APIs

* Add proofs to expect() calls

* Make tests the focus of the forks module

* Allow specific errors to be checked in fork tests

* Remove unused method

* Replace unreachable() statement with expect()

* Rename storage `unfinalized_headers() `to make its purpose more clear

* Update Runtime API name in relayer to match pallet

* Commit `unfinalized_headers` changes I forgot to add

* Rename ChainTipHeight to BestHeight

* Make schedule_next_set_change require a mutable reference

* Remove check for key when enacting authority set

We only expect to take the happy-path in the pallet anyways, so this check
to save ourselves the time spent decoding the entry isn't really used.

* Clear justification set when writing headers to storage

* Clarify why we only allow one set change per fork

* Change best_headers() to return HeaderIDs

Prevents us from returning full headers (which are more expensive to
get from storage) since we only care about header IDs (number, hash)
anyways.

* Fix Clippy complaint

* Make note about equivocations

* Use HeaderIds when returning incomplete headers

This change stops returning full headers which are more expensive
to get from storage than header Ids (number, hash) are. Clients likely
don't need the full header anyways which is why this change is fine.

* Introduce HeaderId type to reduce type complexity

* Add signal hash to storage during genesis config

* Return error instead of expect()-ing

* Fix Clippy lint about `ok_or` fn call

* Rename `forks` module to indicate that it's for testing

* Use `const` for `expect()` proofs

* Remove check that key exists before `kill`-ing value

Co-authored-by: Svyatoslav Nikolsky <svyatonik@gmail.com>

bridges/bin/millau/runtime/src/lib.rs

@@ -516,9 +516,8 @@ impl_runtime_apis! {
 	}
 
 	impl bp_rialto::RialtoHeaderApi<Block> for Runtime {
-		fn best_block() -> (bp_rialto::BlockNumber, bp_rialto::Hash) {
-			let header = BridgeRialto::best_header();
-			(header.number, header.hash())
+		fn best_blocks() -> Vec<(bp_rialto::BlockNumber, bp_rialto::Hash)> {
+			BridgeRialto::best_headers()
 		}
 
 		fn finalized_block() -> (bp_rialto::BlockNumber, bp_rialto::Hash) {
@@ -527,13 +526,7 @@ impl_runtime_apis! {
 		}
 
 		fn incomplete_headers() -> Vec<(bp_rialto::BlockNumber, bp_rialto::Hash)> {
-			// Since the pallet doesn't accept multiple scheduled changes right now
-			// we can only have one header requiring a justification at any time.
-			if let Some(header) = BridgeRialto::requires_justification() {
-				vec![(header.number, header.hash())]
-			} else {
-				vec![]
-			}
+			BridgeRialto::require_justifications()
 		}
 
 		fn is_known_block(hash: bp_rialto::Hash) -> bool {

bridges/bin/rialto/runtime/src/lib.rs

@@ -599,9 +599,8 @@ impl_runtime_apis! {
 	}
 
 	impl bp_millau::MillauHeaderApi<Block> for Runtime {
-		fn best_block() -> (bp_millau::BlockNumber, bp_millau::Hash) {
-			let header = BridgeMillau::best_header();
-			(header.number, header.hash())
+		fn best_blocks() -> Vec<(bp_millau::BlockNumber, bp_millau::Hash)> {
+			BridgeMillau::best_headers()
 		}
 
 		fn finalized_block() -> (bp_millau::BlockNumber, bp_millau::Hash) {
@@ -610,13 +609,7 @@ impl_runtime_apis! {
 		}
 
 		fn incomplete_headers() -> Vec<(bp_millau::BlockNumber, bp_millau::Hash)> {
-			// Since the pallet doesn't accept multiple scheduled changes right now
-			// we can only have one header requiring a justification at any time.
-			if let Some(header) = BridgeMillau::requires_justification() {
-				vec![(header.number, header.hash())]
-			} else {
-				vec![]
-			}
+			BridgeMillau::require_justifications()
 		}
 
 		fn is_known_block(hash: bp_millau::Hash) -> bool {

bridges/modules/substrate/src/fork_tests.rs

@@ -0,0 +1,511 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// This file is part of Parity Bridges Common.
+
+// Parity Bridges Common is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Parity Bridges Common is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Parity Bridges Common.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Tests for checking that behaviour of importing headers and finality proofs works correctly.
+//!
+//! The tests are built around the idea that we will be importing headers on different forks and we
+//! should be able to check that we're correctly importing headers, scheduling changes, and
+//! finalizing headers across different forks.
+//!
+//! Each test is depicted using beautiful ASCII art. The symbols used in the tests are the
+//! following:
+//!
+//! - S|N: Schedules change in N blocks
+//! - E: Enacts change
+//! - F: Finalized
+//! - FN: Finality proof imported for header N
+//!
+//! Each diagram also comes with an import order. This is important since we expect things to fail
+//! when headers or proofs are imported in a certain order.
+//!
+//! Tests can be read as follows:
+//!
+//! ## Example Import 1
+//!
+//! (Type::Header(2, 1, None, None), Ok(()))
+//!
+//! Import header 2 on fork 1. This does not create a fork, or schedule an authority set change. We
+//! expect this header import to be succesful.
+//!
+//! ## Example Import 2
+//!
+//! (Type::Header(4, 2, Some((3, 1)), Some(0)), Ok(()))
+//!
+//! Import header 4 on fork 2. This header starts a new fork from header 3 on fork 1. It also
+//! schedules a change with a delay of 0 blocks. It should be succesfully imported.
+//!
+//! ## Example Import 3
+//!
+//! (Type::Finality(2, 1), Err(FinalizationError::OldHeader.into()))
+//!
+//! Import a finality proof for header 2 on fork 1. This finalty proof should fail to be imported
+//! because the header is an old header.
+
+use crate::justification::tests::*;
+use crate::mock::{helpers::*, *};
+use crate::storage::{AuthoritySet, ImportedHeader};
+use crate::verifier::*;
+use crate::{BestFinalized, BestHeight, BridgeStorage, NextScheduledChange, PalletStorage};
+use codec::Encode;
+use frame_support::{IterableStorageMap, StorageValue};
+use sp_finality_grandpa::{ConsensusLog, GRANDPA_ENGINE_ID};
+use sp_runtime::{Digest, DigestItem};
+use std::collections::BTreeMap;
+
+type ForkId = u64;
+type Delay = u64;
+
+// Indicates when to start a new fork. The first item in the tuple
+// will be the parent header of the header starting this fork.
+type ForksAt = Option<(TestNumber, ForkId)>;
+type ScheduledChangeAt = Option<Delay>;
+
+#[derive(Debug)]
+enum Type {
+	Header(TestNumber, ForkId, ForksAt, ScheduledChangeAt),
+	Finality(TestNumber, ForkId),
+}
+
+// Order: 1, 2, 2', 3, 3''
+//
+//          / [3'']
+//    / [2']
+// [1] <- [2] <- [3]
+#[test]
+fn fork_can_import_headers_on_different_forks() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, None), Ok(())),
+			(Type::Header(2, 2, Some((1, 1)), None), Ok(())),
+			(Type::Header(3, 1, None, None), Ok(())),
+			(Type::Header(3, 3, Some((2, 2)), None), Ok(())),
+		];
+
+		create_chain(&mut storage, &mut chain);
+
+		let best_headers = storage.best_headers();
+		assert_eq!(best_headers.len(), 2);
+		assert_eq!(<BestHeight<TestRuntime>>::get(), 3);
+	})
+}
+
+// Order: 1, 2, 2', F2, F2'
+//
+// [1] <- [2: F]
+//   \ [2']
+//
+// Not allowed to finalize 2'
+#[test]
+fn fork_does_not_allow_competing_finality_proofs() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, None), Ok(())),
+			(Type::Header(2, 2, Some((1, 1)), None), Ok(())),
+			(Type::Finality(2, 1), Ok(())),
+			(Type::Finality(2, 2), Err(FinalizationError::OldHeader.into())),
+		];
+
+		create_chain(&mut storage, &mut chain);
+	})
+}
+
+// Order: 1, 2, 3, F2, 3
+//
+// [1] <- [2: S|0] <- [3]
+//
+// Not allowed to import 3 until we get F2
+//
+// Note: Grandpa would technically allow 3 to be imported as long as it didn't try and enact an
+// authority set change. However, since we expect finality proofs to be imported quickly we've
+// decided to simplify our import process and disallow header imports until we get a finality proof.
+#[test]
+fn fork_waits_for_finality_proof_before_importing_header_past_one_which_enacts_a_change() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, Some(0)), Ok(())),
+			(
+				Type::Header(3, 1, None, None),
+				Err(ImportError::AwaitingFinalityProof.into()),
+			),
+			(Type::Finality(2, 1), Ok(())),
+			(Type::Header(3, 1, None, None), Ok(())),
+		];
+
+		create_chain(&mut storage, &mut chain);
+	})
+}
+
+// Order: 1, 2, F2, 3
+//
+// [1] <- [2: S|1] <- [3: S|0]
+//
+// Grandpa can have multiple authority set changes pending on the same fork. However, we've decided
+// to introduce a limit of _one_ pending authority set change per fork in order to simplify pallet
+// logic and to prevent DoS attacks if Grandpa finality were to temporarily stall for a long time
+// (we'd have to perform a lot of expensive ancestry checks to catch back up).
+#[test]
+fn fork_does_not_allow_multiple_scheduled_changes_on_the_same_fork() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, Some(1)), Ok(())),
+			(
+				Type::Header(3, 1, None, Some(0)),
+				Err(ImportError::PendingAuthoritySetChange.into()),
+			),
+			(Type::Finality(2, 1), Ok(())),
+			(Type::Header(3, 1, None, Some(0)), Ok(())),
+		];
+
+		create_chain(&mut storage, &mut chain);
+	})
+}
+
+// Order: 1, 2, 2'
+//
+//   / [2': S|0]
+// [1] <- [2: S|0]
+//
+// Both 2 and 2' should be marked as needing justifications since they enact changes.
+#[test]
+fn fork_correctly_tracks_which_headers_require_finality_proofs() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, Some(0)), Ok(())),
+			(Type::Header(2, 2, Some((1, 1)), Some(0)), Ok(())),
+		];
+
+		create_chain(&mut storage, &mut chain);
+
+		let header_ids = storage.missing_justifications();
+		assert_eq!(header_ids.len(), 2);
+		assert!(header_ids[0].hash != header_ids[1].hash);
+		assert_eq!(header_ids[0].number, 2);
+		assert_eq!(header_ids[1].number, 2);
+	})
+}
+
+// Order: 1, 2, 2', 3', F2, 3, 4'
+//
+//   / [2': S|1] <- [3'] <- [4']
+// [1] <- [2: S|0] <- [3]
+//
+//
+// Not allowed to import 3 or 4'
+// Can only import 3 after we get the finality proof for 2
+#[test]
+fn fork_does_not_allow_importing_past_header_that_enacts_changes_on_forks() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, Some(0)), Ok(())),
+			(Type::Header(2, 2, Some((1, 1)), Some(1)), Ok(())),
+			(
+				Type::Header(3, 1, None, None),
+				Err(ImportError::AwaitingFinalityProof.into()),
+			),
+			(Type::Header(3, 2, None, None), Ok(())),
+			(Type::Finality(2, 1), Ok(())),
+			(Type::Header(3, 1, None, None), Ok(())),
+			(
+				Type::Header(4, 2, None, None),
+				Err(ImportError::AwaitingFinalityProof.into()),
+			),
+		];
+
+		create_chain(&mut storage, &mut chain);
+
+		// Since we can't query the map directly to check if we applied the right authority set
+		// change (we don't know the header hash of 2) we need to get a little clever.
+		let mut next_change = <NextScheduledChange<TestRuntime>>::iter();
+		let (_, scheduled_change_on_fork) = next_change.next().unwrap();
+		assert_eq!(scheduled_change_on_fork.height, 3);
+
+		// Sanity check to make sure we enacted the change on the canonical change
+		assert_eq!(next_change.next(), None);
+	})
+}
+
+// Order: 1, 2, 3, 2', 3'
+//
+//   / [2'] <- [3']
+// [1] <- [2: S|0] <- [3]
+//
+// Not allowed to import 3
+// Fine to import 2' and 3'
+#[test]
+fn fork_allows_importing_on_different_fork_while_waiting_for_finality_proof() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, Some(0)), Ok(())),
+			(
+				Type::Header(3, 1, None, None),
+				Err(ImportError::AwaitingFinalityProof.into()),
+			),
+			(Type::Header(2, 2, Some((1, 1)), None), Ok(())),
+			(Type::Header(3, 2, None, None), Ok(())),
+		];
+
+		create_chain(&mut storage, &mut chain);
+	})
+}
+
+// Order: 1, 2, 2', F2, 3, 3'
+//
+//   / [2'] <- [3']
+// [1] <- [2: F] <- [3]
+//
+// In our current implementation we're allowed to keep building on fork 2 for as long as our hearts'
+// content. However, we'll never be able to finalize anything on that fork. We'd have to check for
+// ancestry with `best_finalized` on every import which will get expensive.
+//
+// I think this is fine as long as we run pruning every so often to clean up these dead forks.
+#[test]
+fn fork_allows_importing_on_different_fork_past_finalized_header() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, Some(0)), Ok(())),
+			(Type::Header(2, 2, Some((1, 1)), None), Ok(())),
+			(Type::Finality(2, 1), Ok(())),
+			(Type::Header(3, 1, None, None), Ok(())),
+			(Type::Header(3, 2, None, None), Ok(())),
+		];
+
+		create_chain(&mut storage, &mut chain);
+	})
+}
+
+// Order: 1, 2, 3, 4, 3', 4'
+//
+//                  / [3': E] <- [4']
+// [1] <- [2: S|1] <- [3: E] <- [4]
+//
+// Not allowed to import {4|4'}
+#[test]
+fn fork_can_track_scheduled_changes_across_forks() {
+	run_test(|| {
+		let mut storage = PalletStorage::<TestRuntime>::new();
+
+		let mut chain = vec![
+			(Type::Header(1, 1, None, None), Ok(())),
+			(Type::Header(2, 1, None, Some(1)), Ok(())),
+			(Type::Header(3, 1, None, None), Ok(())),
+			(
+				Type::Header(4, 1, None, None),
+				Err(ImportError::AwaitingFinalityProof.into()),
+			),
+			(Type::Header(3, 2, Some((2, 1)), None), Ok(())),
+			(
+				Type::Header(4, 2, None, None),
+				Err(ImportError::AwaitingFinalityProof.into()),
+			),
+		];
+
+		create_chain(&mut storage, &mut chain);
+	})
+}
+
+#[derive(Debug, PartialEq)]
+enum TestError {
+	Import(ImportError),
+	Finality(FinalizationError),
+}
+
+impl From<ImportError> for TestError {
+	fn from(e: ImportError) -> Self {
+		TestError::Import(e)
+	}
+}
+
+impl From<FinalizationError> for TestError {
+	fn from(e: FinalizationError) -> Self {
+		TestError::Finality(e)
+	}
+}
+
+// Builds a fork-aware representation of a blockchain given a list of headers.
+//
+// Takes a list of headers and finality proof operations which will be applied in order. The
+// expected outcome for each operation is also required.
+//
+// The first header in the list will be used as the genesis header and will be manually imported
+// into storage.
+fn create_chain<S>(storage: &mut S, chain: &mut Vec<(Type, Result<(), TestError>)>)
+where
+	S: BridgeStorage<Header = TestHeader> + Clone,
+{
+	let mut map = BTreeMap::new();
+	let mut verifier = Verifier {
+		storage: storage.clone(),
+	};
+	initialize_genesis(storage, &mut map, chain.remove(0).0);
+
+	for h in chain {
+		match h {
+			(Type::Header(num, fork_id, does_fork, schedules_change), expected_result) => {
+				// If we've never seen this fork before
+				if !map.contains_key(&fork_id) {
+					// Let's get the info about where to start the fork
+					if let Some((parent_num, forked_from_id)) = does_fork {
+						let fork = &*map.get(&forked_from_id).unwrap();
+						let parent = fork
+							.iter()
+							.find(|h| h.number == *parent_num)
+							.expect("Trying to fork on a parent which doesn't exist");
+
+						let mut header = test_header(*num);
+						header.parent_hash = parent.hash();
+						header.state_root = [*fork_id as u8; 32].into();
+
+						if let Some(delay) = schedules_change {
+							header.digest = change_log(*delay);
+						}
+
+						// Try and import into storage
+						let res = verifier.import_header(header.clone()).map_err(TestError::Import);
+						assert_eq!(
+							res, *expected_result,
+							"Expected {:?} while importing header ({}, {}), got {:?}",
+							*expected_result, *num, *fork_id, res,
+						);
+
+						// Let's mark the header down in a new fork
+						if res.is_ok() {
+							map.insert(*fork_id, vec![header]);
+						}
+					}
+				} else {
+					// We've seen this fork before so let's append our new header to it
+					let parent_hash = {
+						let fork = &*map.get(&fork_id).unwrap();
+						fork.last().unwrap().hash()
+					};
+
+					let mut header = test_header(*num);
+					header.parent_hash = parent_hash;
+
+					// Doing this to make sure headers at the same height but on
+					// different forks have different hashes
+					header.state_root = [*fork_id as u8; 32].into();
+
+					if let Some(delay) = schedules_change {
+						header.digest = change_log(*delay);
+					}
+
+					let res = verifier.import_header(header.clone()).map_err(TestError::Import);
+					assert_eq!(
+						res, *expected_result,
+						"Expected {:?} while importing header ({}, {}), got {:?}",
+						*expected_result, *num, *fork_id, res,
+					);
+
+					if res.is_ok() {
+						map.get_mut(&fork_id).unwrap().push(header);
+					}
+				}
+			}
+			(Type::Finality(num, fork_id), expected_result) => {
+				let header = map[fork_id]
+					.iter()
+					.find(|h| h.number == *num)
+					.expect("Trying to finalize block that doesn't exist");
+
+				// This is technically equivocating (accepting the same justification on the same
+				// `grandpa_round`).
+				//
+				// See for more: https://github.com/paritytech/parity-bridges-common/issues/430
+				let grandpa_round = 1;
+				let set_id = 1;
+				let authorities = authority_list();
+				let justification =
+					make_justification_for_header(&header, grandpa_round, set_id, &authorities).encode();
+
+				let res = verifier
+					.import_finality_proof(header.hash(), justification.into())
+					.map_err(TestError::Finality);
+				assert_eq!(
+					res, *expected_result,
+					"Expected {:?} while importing finality proof for header ({}, {}), got {:?}",
+					*expected_result, *num, *fork_id, res,
+				);
+			}
+		}
+	}
+
+	for (key, value) in map.iter() {
+		println!("{}: {:#?}", key, value);
+	}
+}
+
+fn initialize_genesis<S>(storage: &mut S, map: &mut BTreeMap<TestNumber, Vec<TestHeader>>, genesis: Type)
+where
+	S: BridgeStorage<Header = TestHeader>,
+{
+	if let Type::Header(num, fork_id, None, None) = genesis {
+		let genesis = test_header(num);
+		map.insert(fork_id, vec![genesis.clone()]);
+
+		let genesis = ImportedHeader {
+			header: genesis,
+			requires_justification: false,
+			is_finalized: true,
+			signal_hash: None,
+		};
+
+		<BestFinalized<TestRuntime>>::put(genesis.hash());
+		storage.write_header(&genesis);
+	} else {
+		panic!("Unexpected genesis block format {:#?}", genesis)
+	}
+
+	let set_id = 1;
+	let authorities = authority_list();
+	let authority_set = AuthoritySet::new(authorities, set_id);
+	storage.update_current_authority_set(authority_set);
+}
+
+fn change_log(delay: u64) -> Digest<TestHash> {
+	let consensus_log = ConsensusLog::<TestNumber>::ScheduledChange(sp_finality_grandpa::ScheduledChange {
+		next_authorities: vec![(alice(), 1), (bob(), 1)],
+		delay,
+	});
+
+	Digest::<TestHash> {
+		logs: vec![DigestItem::Consensus(GRANDPA_ENGINE_ID, consensus_log.encode())],
+	}
+}

bridges/modules/substrate/src/lib.rs

@@ -36,6 +36,7 @@ use bp_runtime::{BlockNumberOf, Chain, HashOf, HeaderOf};
 use frame_support::{decl_error, decl_module, decl_storage, dispatch::DispatchResult};
 use frame_system::ensure_signed;
 use sp_runtime::traits::Header as HeaderT;
+use sp_runtime::RuntimeDebug;
 use sp_std::{marker::PhantomData, prelude::*};
 
 // Re-export since the node uses these when configuring genesis
@@ -51,10 +52,8 @@ mod verifier;
 #[cfg(test)]
 mod mock;
 
-pub trait Trait: frame_system::Trait {
-	/// Chain that we are bridging here.
-	type BridgedChain: Chain;
-}
+#[cfg(test)]
+mod fork_tests;
 
 /// Block number of the bridged chain.
 pub(crate) type BridgedBlockNumber<T> = BlockNumberOf<<T as Trait>::BridgedChain>;
@@ -63,25 +62,45 @@ pub(crate) type BridgedBlockHash<T> = HashOf<<T as Trait>::BridgedChain>;
 /// Header of the bridged chain.
 pub(crate) type BridgedHeader<T> = HeaderOf<<T as Trait>::BridgedChain>;
 
+/// A convenience type identifying headers.
+#[derive(RuntimeDebug, PartialEq)]
+pub struct HeaderId<H: HeaderT> {
+	/// The block number of the header.
+	pub number: H::Number,
+	/// The hash of the header.
+	pub hash: H::Hash,
+}
+
+pub trait Trait: frame_system::Trait {
+	/// Chain that we are bridging here.
+	type BridgedChain: Chain;
+}
+
 decl_storage! {
 	trait Store for Module<T: Trait> as SubstrateBridge {
+		/// The number of the highest block(s) we know of.
+		BestHeight: BridgedBlockNumber<T>;
 		/// Hash of the header at the highest known height.
-		BestHeader: BridgedBlockHash<T>;
+		///
+		/// If there are multiple headers at the same "best" height
+		/// this will contain all of their hashes.
+		BestHeaders: Vec<BridgedBlockHash<T>>;
 		/// Hash of the best finalized header.
 		BestFinalized: BridgedBlockHash<T>;
-		/// A header which enacts an authority set change and therefore
-		/// requires a Grandpa justification.
-		// Since we won't always have an authority set change scheduled we
-		// won't always have a header which needs a justification.
-		RequiresJustification: Option<BridgedBlockHash<T>>;
+		/// The set of header IDs (number, hash) which enact an authority set change and therefore
+		/// require a Grandpa justification.
+		RequiresJustification: map hasher(identity) BridgedBlockHash<T> => BridgedBlockNumber<T>;
 		/// Headers which have been imported into the pallet.
 		ImportedHeaders: map hasher(identity) BridgedBlockHash<T> => Option<ImportedHeader<BridgedHeader<T>>>;
 		/// The current Grandpa Authority set.
 		CurrentAuthoritySet: AuthoritySet;
-		/// The next scheduled authority set change.
+		/// The next scheduled authority set change for a given fork.
+		///
+		/// The fork is indicated by the header which _signals_ the change (key in the mapping).
+		/// Note that this is different than a header which _enacts_ a change.
 		// Grandpa doesn't require there to always be a pending change. In fact, most of the time
 		// there will be no pending change available.
-		NextScheduledChange: Option<ScheduledChange<BridgedBlockNumber<T>>>;
+		NextScheduledChange: map hasher(identity) BridgedBlockHash<T> => Option<ScheduledChange<BridgedBlockNumber<T>>>;
 	}
 	add_extra_genesis {
 		config(initial_header): Option<BridgedHeader<T>>;
@@ -98,25 +117,37 @@ decl_storage! {
 				.initial_header
 				.clone()
 				.expect("An initial header is needed");
+			let initial_hash = initial_header.hash();
 
-			<BestHeader<T>>::put(initial_header.hash());
-			<BestFinalized<T>>::put(initial_header.hash());
-			<ImportedHeaders<T>>::insert(
-				initial_header.hash(),
-				ImportedHeader {
-					header: initial_header,
-					requires_justification: false,
-					is_finalized: true,
-				},
-			);
+			<BestHeight<T>>::put(initial_header.number());
+			<BestHeaders<T>>::put(vec![initial_hash]);
+			<BestFinalized<T>>::put(initial_hash);
 
 			let authority_set =
 				AuthoritySet::new(config.initial_authority_list.clone(), config.initial_set_id);
 			CurrentAuthoritySet::put(authority_set);
 
+			let mut signal_hash = None;
 			if let Some(ref change) = config.first_scheduled_change {
-				<NextScheduledChange<T>>::put(change);
+				assert!(
+					change.height > *initial_header.number(),
+					"Changes must be scheduled past initial header."
+				);
+
+				signal_hash = Some(initial_hash);
+				<NextScheduledChange<T>>::insert(initial_hash, change);
 			};
+
+			<ImportedHeaders<T>>::insert(
+				initial_hash,
+				ImportedHeader {
+					header: initial_header,
+					requires_justification: false,
+					is_finalized: true,
+					signal_hash,
+				},
+			);
+
 		})
 	}
 }
@@ -188,11 +219,13 @@ decl_module! {
 }
 
 impl<T: Trait> Module<T> {
-	/// Get the highest header that the pallet knows of.
-	// In a future where we support forks this could be a Vec of headers
-	// since we may have multiple headers at the same height.
-	pub fn best_header() -> BridgedHeader<T> {
-		PalletStorage::<T>::new().best_header().header
+	/// Get the highest header(s) that the pallet knows of.
+	pub fn best_headers() -> Vec<(BridgedBlockNumber<T>, BridgedBlockHash<T>)> {
+		PalletStorage::<T>::new()
+			.best_headers()
+			.iter()
+			.map(|id| (id.number, id.hash))
+			.collect()
 	}
 
 	/// Get the best finalized header the pallet knows of.
@@ -224,18 +257,15 @@ impl<T: Trait> Module<T> {
 		}
 	}
 
-	/// Return the latest header which enacts an authority set change
-	/// and still needs a finality proof.
+	/// Returns a list of headers which require finality proofs.
 	///
-	/// Will return None if there are no headers which are missing finality proofs.
-	pub fn requires_justification() -> Option<BridgedHeader<T>> {
-		let storage = PalletStorage::<T>::new();
-		let hash = storage.unfinalized_header()?;
-		let imported_header = storage.header_by_hash(hash).expect(
-			"We write a header to storage before marking it as unfinalized, therefore
-			this must always exist if we got an unfinalized header hash.",
-		);
-		Some(imported_header.header)
+	/// These headers require proofs because they enact authority set changes.
+	pub fn require_justifications() -> Vec<(BridgedBlockNumber<T>, BridgedBlockHash<T>)> {
+		PalletStorage::<T>::new()
+			.missing_justifications()
+			.iter()
+			.map(|id| (id.number, id.hash))
+			.collect()
 	}
 }
 
@@ -248,11 +278,8 @@ pub trait BridgeStorage {
 	/// Write a header to storage.
 	fn write_header(&mut self, header: &ImportedHeader<Self::Header>);
 
-	/// Get the header at the highest known height.
-	fn best_header(&self) -> ImportedHeader<Self::Header>;
-
-	/// Update the header at the highest height.
-	fn update_best_header(&mut self, hash: <Self::Header as HeaderT>::Hash);
+	/// Get the header(s) at the highest known height.
+	fn best_headers(&self) -> Vec<HeaderId<Self::Header>>;
 
 	/// Get the best finalized header the pallet knows of.
 	fn best_finalized_header(&self) -> ImportedHeader<Self::Header>;
@@ -263,14 +290,10 @@ pub trait BridgeStorage {
 	/// Check if a particular header is known to the pallet.
 	fn header_exists(&self, hash: <Self::Header as HeaderT>::Hash) -> bool;
 
-	/// Return a header which requires a justification. A header will require
-	/// a justification when it enacts an new authority set.
-	fn unfinalized_header(&self) -> Option<<Self::Header as HeaderT>::Hash>;
-
-	/// Mark a header as eventually requiring a justification.
+	/// Returns a list of headers which require justifications.
 	///
-	/// If None is passed the storage item is cleared.
-	fn update_unfinalized_header(&mut self, hash: Option<<Self::Header as HeaderT>::Hash>);
+	/// A header will require a justification if it enacts a new authority set.
+	fn missing_justifications(&self) -> Vec<HeaderId<Self::Header>>;
 
 	/// Get a specific header by its hash.
 	///
@@ -288,13 +311,22 @@ pub trait BridgeStorage {
 	/// Replace the current authority set with the next scheduled set.
 	///
 	/// Returns an error if there is no scheduled authority set to enact.
-	fn enact_authority_set(&mut self) -> Result<(), ()>;
+	fn enact_authority_set(&mut self, signal_hash: <Self::Header as HeaderT>::Hash) -> Result<(), ()>;
 
 	/// Get the next scheduled Grandpa authority set change.
-	fn scheduled_set_change(&self) -> Option<ScheduledChange<<Self::Header as HeaderT>::Number>>;
+	fn scheduled_set_change(
+		&self,
+		signal_hash: <Self::Header as HeaderT>::Hash,
+	) -> Option<ScheduledChange<<Self::Header as HeaderT>::Number>>;
 
 	/// Schedule a Grandpa authority set change in the future.
-	fn schedule_next_set_change(&self, next_change: ScheduledChange<<Self::Header as HeaderT>::Number>);
+	///
+	/// Takes the hash of the header which scheduled this particular change.
+	fn schedule_next_set_change(
+		&mut self,
+		signal_hash: <Self::Header as HeaderT>::Hash,
+		next_change: ScheduledChange<<Self::Header as HeaderT>::Number>,
+	);
 }
 
 /// Used to interact with the pallet storage in a more abstract way.
@@ -311,18 +343,43 @@ impl<T: Trait> BridgeStorage for PalletStorage<T> {
 	type Header = BridgedHeader<T>;
 
 	fn write_header(&mut self, header: &ImportedHeader<BridgedHeader<T>>) {
-		let hash = header.header.hash();
+		use core::cmp::Ordering;
+
+		let hash = header.hash();
+		let current_height = header.number();
+		let best_height = <BestHeight<T>>::get();
+
+		match current_height.cmp(&best_height) {
+			Ordering::Equal => {
+				<BestHeaders<T>>::append(hash);
+			}
+			Ordering::Greater => {
+				<BestHeaders<T>>::kill();
+				<BestHeaders<T>>::append(hash);
+				<BestHeight<T>>::put(current_height);
+			}
+			Ordering::Less => {
+				// This is fine. We can still have a valid header, but it might just be on a
+				// different fork and at a lower height than the "best" overall header.
+			}
+		}
+
+		if header.requires_justification {
+			<RequiresJustification<T>>::insert(hash, current_height);
+		} else {
+			// If the key doesn't exist this is a no-op, so it's fine to call it often
+			<RequiresJustification<T>>::remove(hash);
+		}
+
 		<ImportedHeaders<T>>::insert(hash, header);
 	}
 
-	fn best_header(&self) -> ImportedHeader<BridgedHeader<T>> {
-		let hash = <BestHeader<T>>::get();
-		self.header_by_hash(hash)
-			.expect("A header must have been written at genesis, therefore this must always exist")
-	}
-
-	fn update_best_header(&mut self, hash: BridgedBlockHash<T>) {
-		<BestHeader<T>>::put(hash)
+	fn best_headers(&self) -> Vec<HeaderId<BridgedHeader<T>>> {
+		let number = <BestHeight<T>>::get();
+		<BestHeaders<T>>::get()
+			.iter()
+			.map(|hash| HeaderId { number, hash: *hash })
+			.collect()
 	}
 
 	fn best_finalized_header(&self) -> ImportedHeader<BridgedHeader<T>> {
@@ -332,7 +389,7 @@ impl<T: Trait> BridgeStorage for PalletStorage<T> {
 	}
 
 	fn update_best_finalized(&self, hash: BridgedBlockHash<T>) {
-		<BestFinalized<T>>::put(hash)
+		<BestFinalized<T>>::put(hash);
 	}
 
 	fn header_exists(&self, hash: BridgedBlockHash<T>) -> bool {
@@ -343,16 +400,10 @@ impl<T: Trait> BridgeStorage for PalletStorage<T> {
 		<ImportedHeaders<T>>::get(hash)
 	}
 
-	fn unfinalized_header(&self) -> Option<BridgedBlockHash<T>> {
-		<RequiresJustification<T>>::get()
-	}
-
-	fn update_unfinalized_header(&mut self, hash: Option<<Self::Header as HeaderT>::Hash>) {
-		if let Some(hash) = hash {
-			<RequiresJustification<T>>::put(hash);
-		} else {
-			<RequiresJustification<T>>::kill();
-		}
+	fn missing_justifications(&self) -> Vec<HeaderId<BridgedHeader<T>>> {
+		<RequiresJustification<T>>::iter()
+			.map(|(hash, number)| HeaderId { number, hash })
+			.collect()
 	}
 
 	fn current_authority_set(&self) -> AuthoritySet {
@@ -363,24 +414,22 @@ impl<T: Trait> BridgeStorage for PalletStorage<T> {
 		CurrentAuthoritySet::put(new_set)
 	}
 
-	fn enact_authority_set(&mut self) -> Result<(), ()> {
-		if <NextScheduledChange<T>>::exists() {
-			let new_set = <NextScheduledChange<T>>::take()
-				.expect("Ensured that entry existed in storage")
-				.authority_set;
+	fn enact_authority_set(&mut self, signal_hash: BridgedBlockHash<T>) -> Result<(), ()> {
+		let new_set = <NextScheduledChange<T>>::take(signal_hash).ok_or(())?.authority_set;
 		self.update_current_authority_set(new_set);
 
 		Ok(())
-		} else {
-			Err(())
-		}
 	}
 
-	fn scheduled_set_change(&self) -> Option<ScheduledChange<BridgedBlockNumber<T>>> {
-		<NextScheduledChange<T>>::get()
+	fn scheduled_set_change(&self, signal_hash: BridgedBlockHash<T>) -> Option<ScheduledChange<BridgedBlockNumber<T>>> {
+		<NextScheduledChange<T>>::get(signal_hash)
 	}
 
-	fn schedule_next_set_change(&self, next_change: ScheduledChange<BridgedBlockNumber<T>>) {
-		<NextScheduledChange<T>>::put(next_change)
+	fn schedule_next_set_change(
+		&mut self,
+		signal_hash: BridgedBlockHash<T>,
+		next_change: ScheduledChange<BridgedBlockNumber<T>>,
+	) {
+		<NextScheduledChange<T>>::insert(signal_hash, next_change)
 	}
 }

bridges/modules/substrate/src/storage.rs

@@ -21,6 +21,7 @@ use core::default::Default;
 #[cfg(feature = "std")]
 use serde::{Deserialize, Serialize};
 use sp_finality_grandpa::{AuthorityList, SetId};
+use sp_runtime::traits::Header as HeaderT;
 use sp_runtime::RuntimeDebug;
 
 /// A Grandpa Authority List and ID.
@@ -54,7 +55,7 @@ pub struct ScheduledChange<N> {
 
 /// A more useful representation of a header for storage purposes.
 #[derive(Default, Encode, Decode, Clone, RuntimeDebug, PartialEq)]
-pub struct ImportedHeader<H> {
+pub struct ImportedHeader<H: HeaderT> {
 	/// A plain Substrate header.
 	pub header: H,
 	/// Does this header enact a new authority set change. If it does
@@ -63,9 +64,12 @@ pub struct ImportedHeader<H> {
 	/// Has this header been finalized, either explicitly via a justification,
 	/// or implicitly via one of its children getting finalized.
 	pub is_finalized: bool,
+	/// The hash of the header which scheduled a change on this fork. If there are currently
+	/// not pending changes on this fork this will be empty.
+	pub signal_hash: Option<H::Hash>,
 }
 
-impl<H> core::ops::Deref for ImportedHeader<H> {
+impl<H: HeaderT> core::ops::Deref for ImportedHeader<H> {
 	type Target = H;
 
 	fn deref(&self) -> &H {

bridges/modules/substrate/src/verifier.rs

@@ -54,7 +54,7 @@ impl From<Vec<u8>> for FinalityProof {
 /// Errors which can happen while importing a header.
 #[derive(RuntimeDebug, PartialEq)]
 pub enum ImportError {
-	/// This header is older than our latest finalized block, thus not useful.
+	/// This header is at the same height or older than our latest finalized block, thus not useful.
 	OldHeader,
 	/// This header has already been imported by the pallet.
 	HeaderAlreadyExists,
@@ -68,6 +68,13 @@ pub enum ImportError {
 	/// the authority weights being empty or overflowing the `AuthorityWeight`
 	/// type.
 	InvalidAuthoritySet,
+	/// This header is not allowed to be imported since an ancestor requires a finality proof.
+	///
+	/// This can happen if an ancestor is supposed to enact an authority set change.
+	AwaitingFinalityProof,
+	/// This header schedules an authority set change even though we're still waiting
+	/// for an old authority set change to be enacted on this fork.
+	PendingAuthoritySetChange,
 }
 
 /// Errors which can happen while verifying a headers finality.
@@ -79,7 +86,7 @@ pub enum FinalizationError {
 	PrematureJustification,
 	/// We failed to verify this header's ancestry.
 	AncestryCheckFailed,
-	/// This header is older than our latest finalized block, thus not useful.
+	/// This header is at the same height or older than our latest finalized block, thus not useful.
 	OldHeader,
 	/// The given justification was not able to finalize the given header.
 	///
@@ -127,22 +134,37 @@ where
 			return Err(ImportError::InvalidChildNumber);
 		}
 
-		// This header requires a justification since it enacts an authority set change. We don't
-		// need to act on it right away (we'll update the set once this header gets finalized), but
+		// A header requires a justification if it enacts an authority set change. We don't
+		// need to act on it right away (we'll update the set once the header gets finalized), but
 		// we need to make a note of it.
 		//
-		// TODO: This assumes that we can only have one authority set change pending at a time.
-		// This is not strictly true as Grandpa may schedule multiple changes on a given chain
-		// if the "next next" change is scheduled after the "delay" period of the "next" change
-		let requires_justification = if let Some(change) = self.storage.scheduled_set_change() {
-			change.height == *header.number()
+		// Note: This assumes that we can only have one authority set change pending per fork at a
+		// time. While this is not strictly true of Grandpa (it can have multiple pending changes,
+		// even across forks), this assumption simplifies our tracking of authority set changes.
+		let mut signal_hash = parent_header.signal_hash;
+		let scheduled_change = find_scheduled_change(&header);
+
+		// Check if our fork is expecting an authority set change
+		let requires_justification = if let Some(hash) = signal_hash {
+			const PROOF: &str = "If the header has a signal hash it means there's an accompanying set
+							change in storage, therefore this must always be valid.";
+			let pending_change = self.storage.scheduled_set_change(hash).expect(PROOF);
+
+			if scheduled_change.is_some() {
+				return Err(ImportError::PendingAuthoritySetChange);
+			}
+
+			if *header.number() > pending_change.height {
+				return Err(ImportError::AwaitingFinalityProof);
+			}
+
+			pending_change.height == *header.number()
 		} else {
 			// Since we don't currently have a pending authority set change let's check if the header
 			// contains a log indicating when the next change should be.
-			if let Some(change) = find_scheduled_change(&header) {
+			if let Some(change) = scheduled_change {
 				let mut total_weight = 0u64;
 
-				// We need to make sure that we don't overflow the `AuthorityWeight` type.
 				for (_id, weight) in &change.next_authorities {
 					total_weight = total_weight
 						.checked_add(*weight)
@@ -163,12 +185,16 @@ where
 				let height = (*header.number())
 					.checked_add(&change.delay)
 					.ok_or(ImportError::ScheduledHeightOverflow)?;
+
 				let scheduled_change = ScheduledChange {
 					authority_set: next_set,
 					height,
 				};
 
-				self.storage.schedule_next_set_change(scheduled_change);
+				// Note: It's important that the signal hash is updated if a header schedules a
+				// change or else we end up with inconsistencies in other places.
+				signal_hash = Some(hash);
+				self.storage.schedule_next_set_change(hash, scheduled_change);
 
 				// If the delay is 0 this header will enact the change it signaled
 				height == *header.number()
@@ -181,16 +207,9 @@ where
 			header,
 			requires_justification,
 			is_finalized: false,
+			signal_hash,
 		});
 
-		if requires_justification {
-			self.storage.update_unfinalized_header(Some(hash));
-		}
-
-		// Since we're not dealing with forks at the moment we know that
-		// the header we just got will be the one at the best height
-		self.storage.update_best_header(hash);
-
 		Ok(())
 	}
 
@@ -258,19 +277,25 @@ where
 		// new authority set change. When we finalize the header we need to update the current
 		// authority set.
 		if header.requires_justification {
+			const SIGNAL_HASH_PROOF: &str = "When we import a header we only mark it as
+			`requires_justification` if we have checked that it contains a signal hash. Therefore
+			this must always be valid.";
+
+			const ENACT_SET_PROOF: &str =
+				"Headers must only be marked as `requires_justification` if there's a scheduled change in storage.";
+
 			// If we are unable to enact an authority set it means our storage entry for scheduled
 			// changes is missing. Best to crash since this is likely a bug.
-			let _ = self.storage.enact_authority_set().expect(
-				"Headers must only be marked as `requires_justification` if there's a scheduled change in storage.",
-			);
-
-			// Clear the storage entry since we got a justification
-			self.storage.update_unfinalized_header(None);
+			let _ = self
+				.storage
+				.enact_authority_set(header.signal_hash.expect(SIGNAL_HASH_PROOF))
+				.expect(ENACT_SET_PROOF);
 		}
 
 		for header in finalized_headers.iter_mut() {
 			header.is_finalized = true;
 			header.requires_justification = false;
+			header.signal_hash = None;
 			self.storage.write_header(header);
 		}
 
@@ -329,17 +354,19 @@ mod tests {
 	use crate::justification::tests::*;
 	use crate::mock::helpers::*;
 	use crate::mock::*;
-	use crate::{BestFinalized, ImportedHeaders, PalletStorage};
+	use crate::{BestFinalized, BestHeight, HeaderId, ImportedHeaders, PalletStorage};
 	use codec::Encode;
 	use frame_support::{assert_err, assert_ok};
 	use frame_support::{StorageMap, StorageValue};
 	use sp_finality_grandpa::{AuthorityId, SetId};
+	use sp_runtime::{Digest, DigestItem};
 
 	fn unfinalized_header(num: u64) -> ImportedHeader<TestHeader> {
 		ImportedHeader {
 			header: test_header(num),
 			requires_justification: false,
 			is_finalized: false,
+			signal_hash: None,
 		}
 	}
 
@@ -354,6 +381,7 @@ mod tests {
 	}
 
 	// Useful for quickly writing a chain of headers to storage
+	// Input is expected in the form: vec![(num, requires_justification, is_finalized)]
 	fn write_headers<S: BridgeStorage<Header = TestHeader>>(
 		storage: &mut S,
 		headers: Vec<(u64, bool, bool)>,
@@ -363,6 +391,7 @@ mod tests {
 			header: test_header(0),
 			requires_justification: false,
 			is_finalized: true,
+			signal_hash: None,
 		};
 
 		<BestFinalized<TestRuntime>>::put(genesis.hash());
@@ -374,6 +403,7 @@ mod tests {
 				header: test_header(num),
 				requires_justification,
 				is_finalized,
+				signal_hash: None,
 			};
 
 			storage.write_header(&header);
@@ -383,6 +413,16 @@ mod tests {
 		imported_headers
 	}
 
+	// Given a block number will generate a chain of headers which don't require justification and
+	// are not considered to be finalized.
+	fn write_default_headers<S: BridgeStorage<Header = TestHeader>>(
+		storage: &mut S,
+		headers: Vec<u64>,
+	) -> Vec<ImportedHeader<TestHeader>> {
+		let headers = headers.iter().map(|num| (*num, false, false)).collect();
+		write_headers(storage, headers)
+	}
+
 	#[test]
 	fn fails_to_import_old_header() {
 		run_test(|| {
@@ -424,6 +464,7 @@ mod tests {
 				header: header.clone(),
 				requires_justification: false,
 				is_finalized: false,
+				signal_hash: None,
 			};
 			<ImportedHeaders<TestRuntime>>::insert(header.hash(), &imported_header);
 
@@ -444,6 +485,7 @@ mod tests {
 				header: parent,
 				requires_justification: false,
 				is_finalized: true,
+				signal_hash: None,
 			};
 			<ImportedHeaders<TestRuntime>>::insert(parent_hash, &imported_header);
 
@@ -457,7 +499,89 @@ mod tests {
 				.header_by_hash(header.hash())
 				.expect("Should have been imported successfully");
 			assert_eq!(stored_header.is_finalized, false);
-			assert_eq!(stored_header, storage.best_header());
+			assert_eq!(stored_header.hash(), storage.best_headers()[0].hash);
+		})
+	}
+
+	#[test]
+	fn successfully_imports_two_different_headers_at_same_height() {
+		run_test(|| {
+			let mut storage = PalletStorage::<TestRuntime>::new();
+
+			// We want to write the genesis header to storage
+			let _ = write_headers(&mut storage, vec![]);
+
+			// Both of these headers have the genesis header as their parent
+			let header_on_fork1 = test_header(1);
+			let mut header_on_fork2 = test_header(1);
+
+			// We need to change _something_ to make it a different header
+			header_on_fork2.state_root = [1; 32].into();
+
+			let mut verifier = Verifier {
+				storage: storage.clone(),
+			};
+
+			// It should be fine to import both
+			assert_ok!(verifier.import_header(header_on_fork1.clone()));
+			assert_ok!(verifier.import_header(header_on_fork2.clone()));
+
+			// We should have two headers marked as being the best since they're
+			// both at the same height
+			let best_headers = storage.best_headers();
+			assert_eq!(best_headers.len(), 2);
+			assert_eq!(
+				best_headers[0],
+				HeaderId {
+					number: *header_on_fork1.number(),
+					hash: header_on_fork1.hash()
+				}
+			);
+			assert_eq!(
+				best_headers[1],
+				HeaderId {
+					number: *header_on_fork2.number(),
+					hash: header_on_fork2.hash()
+				}
+			);
+			assert_eq!(<BestHeight<TestRuntime>>::get(), 1);
+		})
+	}
+
+	#[test]
+	fn correctly_updates_the_best_header_given_a_better_header() {
+		run_test(|| {
+			let mut storage = PalletStorage::<TestRuntime>::new();
+
+			// Write two headers at the same height to storage.
+			let imported_headers = write_default_headers(&mut storage, vec![1, 1]);
+
+			// The headers we manually imported should have been marked as the best
+			// upon writing to storage. Let's confirm that.
+			assert_eq!(storage.best_headers().len(), 2);
+			assert_eq!(<BestHeight<TestRuntime>>::get(), 1);
+
+			// Now let's build something at a better height.
+			let mut better_header = test_header(2);
+			better_header.parent_hash = imported_headers[1].hash();
+
+			let mut verifier = Verifier {
+				storage: storage.clone(),
+			};
+			assert_ok!(verifier.import_header(better_header.clone()));
+
+			// Since `better_header` is the only one at height = 2 we should only have
+			// a single "best header" now.
+			let best_headers = storage.best_headers();
+			assert_eq!(best_headers.len(), 1);
+			assert_eq!(
+				best_headers[0],
+				HeaderId {
+					number: *better_header.number(),
+					hash: better_header.hash()
+				}
+			);
+			assert_eq!(<BestHeight<TestRuntime>>::get(), 2);
 		})
 	}
 
@@ -465,9 +589,7 @@ mod tests {
 	fn related_headers_are_ancestors() {
 		run_test(|| {
 			let mut storage = PalletStorage::<TestRuntime>::new();
-
-			let headers = vec![(1, false, false), (2, false, false), (3, false, false)];
-			let mut imported_headers = write_headers(&mut storage, headers);
+			let mut imported_headers = write_default_headers(&mut storage, vec![1, 2, 3]);
 
 			for header in imported_headers.iter() {
 				assert!(storage.header_exists(header.hash()));
@@ -487,8 +609,7 @@ mod tests {
 		run_test(|| {
 			let mut storage = PalletStorage::<TestRuntime>::new();
 
-			let headers = vec![(1, false, false), (2, false, false), (3, false, false)];
-			let mut imported_headers = write_headers(&mut storage, headers);
+			let mut imported_headers = write_default_headers(&mut storage, vec![1, 2, 3]);
 			for header in imported_headers.iter() {
 				assert!(storage.header_exists(header.hash()));
 			}
@@ -501,6 +622,7 @@ mod tests {
 				header: bad_ancestor,
 				requires_justification: false,
 				is_finalized: false,
+				signal_hash: None,
 			};
 
 			let child = imported_headers.pop().unwrap();
@@ -514,8 +636,7 @@ mod tests {
 		run_test(|| {
 			let mut storage = PalletStorage::<TestRuntime>::new();
 
-			let headers = vec![(1, false, false), (2, false, false), (3, false, false)];
-			let mut imported_headers = write_headers(&mut storage, headers);
+			let mut imported_headers = write_default_headers(&mut storage, vec![1, 2, 3]);
 			for header in imported_headers.iter() {
 				assert!(storage.header_exists(header.hash()));
 			}
@@ -526,6 +647,7 @@ mod tests {
 				header: new_ancestor,
 				requires_justification: false,
 				is_finalized: false,
+				signal_hash: None,
 			};
 
 			let child = imported_headers.pop().unwrap();
@@ -536,12 +658,9 @@ mod tests {
 
 	#[test]
 	fn doesnt_import_header_which_schedules_change_with_invalid_authority_set() {
-		use sp_runtime::{Digest, DigestItem};
-
 		run_test(|| {
 			let mut storage = PalletStorage::<TestRuntime>::new();
-			let headers = vec![(1, false, false)];
-			let _imported_headers = write_headers(&mut storage, headers);
+			let _imported_headers = write_default_headers(&mut storage, vec![1]);
 			let mut header = test_header(2);
 
 			// This is an *invalid* authority set because the combined weight of the
@@ -568,8 +687,7 @@ mod tests {
 	fn finalizes_header_which_doesnt_enact_or_schedule_a_new_authority_set() {
 		run_test(|| {
 			let mut storage = PalletStorage::<TestRuntime>::new();
-			let headers = vec![(1, false, false)];
-			let _imported_headers = write_headers(&mut storage, headers);
+			let _imported_headers = write_default_headers(&mut storage, vec![1]);
 
 			// Nothing special about this header, yet Grandpa may have created a justification
 			// for it since it does that periodically
@@ -598,8 +716,7 @@ mod tests {
 	fn correctly_verifies_and_finalizes_chain_of_headers() {
 		run_test(|| {
 			let mut storage = PalletStorage::<TestRuntime>::new();
-			let headers = vec![(1, false, false), (2, false, false)];
-			let imported_headers = write_headers(&mut storage, headers);
+			let imported_headers = write_default_headers(&mut storage, vec![1, 2]);
 			let header = test_header(3);
 
 			let set_id = 1;
@@ -635,8 +752,16 @@ mod tests {
 	fn updates_authority_set_upon_finalizing_header_which_enacts_change() {
 		run_test(|| {
 			let mut storage = PalletStorage::<TestRuntime>::new();
-			let headers = vec![(1, false, false)];
-			let _imported_headers = write_headers(&mut storage, headers);
+			let genesis_hash = write_headers(&mut storage, vec![])[0].hash();
+
+			// We want this header to indicate that there's an upcoming set change on this fork
+			let parent = ImportedHeader {
+				header: test_header(1),
+				requires_justification: false,
+				is_finalized: false,
+				signal_hash: Some(genesis_hash),
+			};
+			storage.write_header(&parent);
 
 			let set_id = 1;
 			let authorities = authority_list();
@@ -654,21 +779,22 @@ mod tests {
 			let height = *header.number();
 			let authorities = vec![alice()];
 			let change = schedule_next_change(authorities, set_id, height);
-			storage.schedule_next_set_change(change.clone());
+			storage.schedule_next_set_change(genesis_hash, change.clone());
 
 			let mut verifier = Verifier {
 				storage: storage.clone(),
 			};
 
 			assert_ok!(verifier.import_header(header.clone()));
-			assert_eq!(storage.unfinalized_header(), Some(header.hash()));
+			assert_eq!(storage.missing_justifications().len(), 1);
+			assert_eq!(storage.missing_justifications()[0].hash, header.hash());
 
 			assert_ok!(verifier.import_finality_proof(header.hash(), justification.into()));
 			assert_eq!(storage.best_finalized_header().header, header);
 
 			// Make sure that we have updated the set now that we've finalized our header
 			assert_eq!(storage.current_authority_set(), change.authority_set);
-			assert_eq!(storage.unfinalized_header(), None);
+			assert!(storage.missing_justifications().is_empty());
 		})
 	}
 
@@ -682,6 +808,7 @@ mod tests {
 				header: genesis,
 				requires_justification: false,
 				is_finalized: true,
+				signal_hash: None,
 			};
 
 			// Make sure that genesis is the best finalized header
@@ -699,82 +826,4 @@ mod tests {
 			);
 		});
 	}
-
-	// We're supposed to enact a set change at header N. This means that when we import it we must
-	// remember that it requires a justification. We can continue importing headers past N but must
-	// not finalize any childen. At a later point in time we should be able to import the
-	// justification for N.
-	//
-	// Since N enacts a new authority set, when we finalize it we should see this change reflected
-	// correctly.
-	//
-	// [G] <- [N-1] <- [N] <- [N+1] <- [N+2]
-	//                  |                |- Import justification for N here
-	//                  |- Enacts change, needs justification
-	#[test]
-	fn allows_importing_justification_at_block_past_scheduled_change() {
-		run_test(|| {
-			let mut storage = PalletStorage::<TestRuntime>::new();
-			let headers = vec![(1, false, false)];
-			let imported_headers = write_headers(&mut storage, headers);
-
-			// Set up our initial authority set
-			let set_id = 1;
-			let authorities = authority_list();
-			let initial_authority_set = AuthoritySet::new(authorities.clone(), set_id);
-			storage.update_current_authority_set(initial_authority_set);
-
-			// This is header N
-			let header = test_header(2);
-
-			// Since we want to finalize N we need a justification for it
-			let grandpa_round = 1;
-			let justification = make_justification_for_header(&header, grandpa_round, set_id, &authorities).encode();
-
-			// Schedule a change at height N
-			let set_id = 2;
-			let height = *header.number();
-			let authorities = vec![alice()];
-			let change = schedule_next_change(authorities, set_id, height);
-			storage.schedule_next_set_change(change.clone());
-
-			// Import header N
-			let mut verifier = Verifier {
-				storage: storage.clone(),
-			};
-			assert!(verifier.import_header(header.clone()).is_ok());
-
-			// Header N should be marked as needing a justification
-			assert_eq!(
-				storage.header_by_hash(header.hash()).unwrap().requires_justification,
-				true
-			);
-
-			// Now we want to import some headers which are past N
-			let child = test_header(*header.number() + 1);
-			assert!(verifier.import_header(child.clone()).is_ok());
-
-			let grandchild = test_header(*child.number() + 1);
-			assert!(verifier.import_header(grandchild).is_ok());
-
-			// Even though we're a few headers ahead we should still be able to import
-			// a justification for header N
-			assert!(verifier
-				.import_finality_proof(header.hash(), justification.into())
-				.is_ok());
-
-			// Some checks to make sure that our header has been correctly finalized
-			let finalized_header = storage.header_by_hash(header.hash()).unwrap();
-			assert!(finalized_header.is_finalized);
-			assert_eq!(finalized_header.requires_justification, false);
-			assert_eq!(storage.best_finalized_header().header, header);
-
-			// Make sure we marked the parent of the header at N as finalized
-			assert!(storage.header_by_hash(imported_headers[1].hash()).unwrap().is_finalized);
-
-			// Since our header was supposed to enact a new authority set change when it got
-			// finalized let's make sure that the authority set actually changed
-			assert_eq!(storage.current_authority_set(), change.authority_set);
-		})
-	}
 }

bridges/primitives/millau/src/lib.rs

@@ -61,7 +61,7 @@ impl Chain for Millau {
 }
 
 /// Name of the `MillauHeaderApi::best_block` runtime method.
-pub const BEST_MILLAU_BLOCK_METHOD: &str = "MillauHeaderApi_best_block";
+pub const BEST_MILLAU_BLOCKS_METHOD: &str = "MillauHeaderApi_best_blocks";
 /// Name of the `MillauHeaderApi::finalized_block` runtime method.
 pub const FINALIZED_MILLAU_BLOCK_METHOD: &str = "MillauHeaderApi_finalized_block";
 /// Name of the `MillauHeaderApi::is_known_block` runtime method.
@@ -88,11 +88,13 @@ sp_api::decl_runtime_apis! {
 	/// This API is implemented by runtimes that are bridging with Millau chain, not the
 	/// Millau runtime itself.
 	pub trait MillauHeaderApi {
-		/// Returns number and hash of the best block known to the bridge module.
+		/// Returns number and hash of the best blocks known to the bridge module.
+		///
+		/// Will return multiple headers if there are many headers at the same "best" height.
 		///
 		/// The caller should only submit an `import_header` transaction that makes
 		/// (or leads to making) other header the best one.
-		fn best_block() -> (BlockNumber, Hash);
+		fn best_blocks() -> Vec<(BlockNumber, Hash)>;
 		/// Returns number and hash of the best finalized block known to the bridge module.
 		fn finalized_block() -> (BlockNumber, Hash);
 		/// Returns numbers and hashes of headers that require finality proofs.

bridges/primitives/rialto/src/lib.rs

@@ -79,11 +79,13 @@ sp_api::decl_runtime_apis! {
 	/// This API is implemented by runtimes that are bridging with Rialto chain, not the
 	/// Rialto runtime itself.
 	pub trait RialtoHeaderApi {
-		/// Returns number and hash of the best block known to the bridge module.
+		/// Returns number and hash of the best blocks known to the bridge module.
+		///
+		/// Will return multiple headers if there are many headers at the same "best" height.
 		///
 		/// The caller should only submit an `import_header` transaction that makes
 		/// (or leads to making) other header the best one.
-		fn best_block() -> (BlockNumber, Hash);
+		fn best_blocks() -> Vec<(BlockNumber, Hash)>;
 		/// Returns number and hash of the best finalized block known to the bridge module.
 		fn finalized_block() -> (BlockNumber, Hash);
 		/// Returns numbers and hashes of headers that require finality proofs.

bridges/primitives/runtime/src/chain.rs

@@ -40,7 +40,8 @@ pub trait Chain {
 		+ AtLeast32BitUnsigned
 		+ FromStr
 		+ MaybeMallocSizeOf
-		+ AsPrimitive<usize>;
+		+ AsPrimitive<usize>
+		+ Default;
 
 	/// A type that fulfills the abstract idea of what a Substrate hash is.
 	// Constraits come from the associated Hash type of `sp_runtime::traits::Header`

bridges/relays/substrate/src/headers_target.rs

@@ -90,10 +90,15 @@ where
 		let data = Bytes(Vec::new());
 
 		let encoded_response = self.client.state_call(call, data, None).await?;
-		let decoded_response: (P::Number, P::Hash) =
+		let decoded_response: Vec<(P::Number, P::Hash)> =
 			Decode::decode(&mut &encoded_response.0[..]).map_err(SubstrateError::ResponseParseFailed)?;
 
-		let best_header_id = HeaderId(decoded_response.0, decoded_response.1);
+		let best_header = decoded_response.last().ok_or_else(|| {
+			SubstrateError::ResponseParseFailed(
+				"Parsed an empty list of headers, we should always have at least one.".into(),
+			)
+		})?;
+		let best_header_id = HeaderId(best_header.0, best_header.1);
 		Ok(best_header_id)
 	}
 

bridges/relays/substrate/src/millau_headers_to_rialto.rs

@@ -24,7 +24,7 @@ use crate::{
 
 use async_trait::async_trait;
 use bp_millau::{
-	BEST_MILLAU_BLOCK_METHOD, FINALIZED_MILLAU_BLOCK_METHOD, INCOMPLETE_MILLAU_HEADERS_METHOD,
+	BEST_MILLAU_BLOCKS_METHOD, FINALIZED_MILLAU_BLOCK_METHOD, INCOMPLETE_MILLAU_HEADERS_METHOD,
 	IS_KNOWN_MILLAU_BLOCK_METHOD,
 };
 use codec::Encode;
@@ -65,7 +65,7 @@ impl HeadersSyncPipeline for MillauHeadersToRialto {
 
 #[async_trait]
 impl SubstrateHeadersSyncPipeline for MillauHeadersToRialto {
-	const BEST_BLOCK_METHOD: &'static str = BEST_MILLAU_BLOCK_METHOD;
+	const BEST_BLOCK_METHOD: &'static str = BEST_MILLAU_BLOCKS_METHOD;
 	const FINALIZED_BLOCK_METHOD: &'static str = FINALIZED_MILLAU_BLOCK_METHOD;
 	const IS_KNOWN_BLOCK_METHOD: &'static str = IS_KNOWN_MILLAU_BLOCK_METHOD;
 	const INCOMPLETE_HEADERS_METHOD: &'static str = INCOMPLETE_MILLAU_HEADERS_METHOD;