From 6cd4f5edf1bee4e46e460db6897526e7797a7f53 Mon Sep 17 00:00:00 2001
From: Hernando Castano <HCastano@users.noreply.github.com>
Date: Tue, 19 Jan 2021 11:03:26 -0500
Subject: [PATCH] Disallow Duplicate Best Headers (#653)

* Add test proving bug

* Add checks for duplicate headers

* Fix Clippy error
---
 bridges/modules/substrate/src/lib.rs      |  7 ++-
 bridges/modules/substrate/src/verifier.rs | 56 +++++++++++++++++++++--
 2 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/bridges/modules/substrate/src/lib.rs b/bridges/modules/substrate/src/lib.rs
index 2129f00049..e1c797c2ae 100644
--- a/bridges/modules/substrate/src/lib.rs
+++ b/bridges/modules/substrate/src/lib.rs
@@ -521,7 +521,12 @@ impl<T: Config> BridgeStorage for PalletStorage<T> {
 
 		match current_height.cmp(&best_height) {
 			Ordering::Equal => {
-				<BestHeaders<T>>::append(hash);
+				// Want to avoid duplicates in the case where we're writing a finalized header to
+				// storage which also happens to be at the best height the best height
+				let not_duplicate = !<ImportedHeaders<T>>::contains_key(hash);
+				if not_duplicate {
+					<BestHeaders<T>>::append(hash);
+				}
 			}
 			Ordering::Greater => {
 				<BestHeaders<T>>::kill();
diff --git a/bridges/modules/substrate/src/verifier.rs b/bridges/modules/substrate/src/verifier.rs
index edbc77d54a..5cb47cb656 100644
--- a/bridges/modules/substrate/src/verifier.rs
+++ b/bridges/modules/substrate/src/verifier.rs
@@ -24,6 +24,7 @@
 
 use crate::storage::{AuthoritySet, ImportedHeader, ScheduledChange};
 use crate::BridgeStorage;
+
 use bp_header_chain::justification::verify_justification;
 use finality_grandpa::voter_set::VoterSet;
 use sp_finality_grandpa::{ConsensusLog, GRANDPA_ENGINE_ID};
@@ -545,8 +546,23 @@ mod tests {
 		run_test(|| {
 			let mut storage = PalletStorage::<TestRuntime>::new();
 
+			// We want to write the genesis header to storage
+			let _ = write_headers(&mut storage, vec![]);
+
 			// Write two headers at the same height to storage.
-			let imported_headers = write_default_headers(&mut storage, vec![1, 1]);
+			let best_header = test_header(1);
+			let mut also_best_header = test_header(1);
+
+			// We need to change _something_ to make it a different header
+			also_best_header.state_root = [1; 32].into();
+
+			let mut verifier = Verifier {
+				storage: storage.clone(),
+			};
+
+			// It should be fine to import both
+			assert_ok!(verifier.import_header(best_header.hash(), best_header.clone()));
+			assert_ok!(verifier.import_header(also_best_header.hash(), also_best_header));
 
 			// The headers we manually imported should have been marked as the best
 			// upon writing to storage. Let's confirm that.
@@ -555,11 +571,8 @@ mod tests {
 
 			// Now let's build something at a better height.
 			let mut better_header = test_header(2);
-			better_header.parent_hash = imported_headers[1].hash();
+			better_header.parent_hash = best_header.hash();
 
-			let mut verifier = Verifier {
-				storage: storage.clone(),
-			};
 			assert_ok!(verifier.import_header(better_header.hash(), better_header.clone()));
 
 			// Since `better_header` is the only one at height = 2 we should only have
@@ -577,6 +590,39 @@ mod tests {
 		})
 	}
 
+	#[test]
+	fn doesnt_write_best_header_twice_upon_finalization() {
+		run_test(|| {
+			let mut storage = PalletStorage::<TestRuntime>::new();
+			let _imported_headers = write_default_headers(&mut storage, vec![1]);
+
+			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);
+
+			// Let's import our header
+			let header = test_header(2);
+			let mut verifier = Verifier {
+				storage: storage.clone(),
+			};
+			assert_ok!(verifier.import_header(header.hash(), header.clone()));
+
+			// Our header should be the only best header we have
+			assert_eq!(storage.best_headers()[0].hash, header.hash());
+			assert_eq!(storage.best_headers().len(), 1);
+
+			// Now lets finalize our best header
+			let grandpa_round = 1;
+			let justification = make_justification_for_header(&header, grandpa_round, set_id, &authorities).encode();
+			assert_ok!(verifier.import_finality_proof(header.hash(), justification.into()));
+
+			// Our best header should only appear once in the list of best headers
+			assert_eq!(storage.best_headers()[0].hash, header.hash());
+			assert_eq!(storage.best_headers().len(), 1);
+		})
+	}
+
 	#[test]
 	fn related_headers_are_ancestors() {
 		run_test(|| {