make sure header PMMR is init correctly based on header_head from db (#3362)

* make sure header PMMR is init correctly based on header_head from db

* fix to ensure header PMMR is consistent with header_head in db on chain init

* change order of operations - validate header earlier
and avoid applying header to header PMMR before validation
as this potentially leaves the PMMR in a tricky state to rewind from

chain/src/chain.rs

@@ -1449,9 +1449,15 @@ fn setup_head(
 		}
 	}
 
-	// Setup our header_head if we do not already have one.
-	// Migrating back to header_head in db and some nodes may note have one.
-	if batch.header_head().is_err() {
+	// Make sure our header PMMR is consistent with header_head from db if it exists.
+	// If header_head is missing in db then use head of header PMMR.
+	if let Ok(head) = batch.header_head() {
+		header_pmmr.init_head(&head)?;
+		txhashset::header_extending(header_pmmr, &mut batch, |ext, batch| {
+			let header = batch.get_block_header(&head.hash())?;
+			ext.rewind(&header)
+		})?;
+	} else {
 		let hash = header_pmmr.head_hash()?;
 		let header = batch.get_block_header(&hash)?;
 		batch.save_header_head(&Tip::from_header(&header))?;

chain/src/pipe.rs

@@ -250,6 +250,11 @@ pub fn process_block_header(header: &BlockHeader, ctx: &mut BlockContext<'_>) ->
 		}
 	}
 
+	// We want to validate this individual header before applying it to our header PMMR.
+	validate_header(header, ctx)?;
+
+	// Apply the header to the header PMMR, making sure we put the extension in the correct state
+	// based on previous header first.
 	txhashset::header_extending(&mut ctx.header_pmmr, &mut ctx.batch, |ext, batch| {
 		rewind_and_apply_header_fork(&prev_header, ext, batch)?;
 		ext.validate_root(header)?;
@@ -260,7 +265,7 @@ pub fn process_block_header(header: &BlockHeader, ctx: &mut BlockContext<'_>) ->
 		Ok(())
 	})?;
 
-	validate_header(header, ctx)?;
+	// Add this new block header to the db.
 	add_block_header(header, &ctx.batch)?;
 
 	if has_more_work(header, &header_head) {

chain/src/txhashset/txhashset.rs

@@ -74,6 +74,41 @@ impl<T: PMMRable> PMMRHandle<T> {
 }
 
 impl PMMRHandle<BlockHeader> {
+	/// Used during chain init to ensure the header PMMR is consistent with header_head in the db.
+	pub fn init_head(&mut self, head: &Tip) -> Result<(), Error> {
+		let head_hash = self.head_hash()?;
+		let expected_hash = self.get_header_hash_by_height(head.height)?;
+		if head.hash() != expected_hash {
+			error!(
+				"header PMMR inconsistent: {} vs {} at {}",
+				expected_hash,
+				head.hash(),
+				head.height
+			);
+			return Err(ErrorKind::Other("header PMMR inconsistent".to_string()).into());
+		}
+
+		// 1-indexed pos and we want to account for subsequent parent hash pos.
+		// so use next header pos to find our last_pos.
+		let next_height = head.height + 1;
+		let next_pos = pmmr::insertion_to_pmmr_index(next_height + 1);
+		let pos = next_pos.saturating_sub(1);
+
+		debug!(
+			"init_head: header PMMR: current head {} at pos {}",
+			head_hash, self.last_pos
+		);
+		debug!(
+			"init_head: header PMMR: resetting to {} at pos {} (height {})",
+			head.hash(),
+			pos,
+			head.height
+		);
+
+		self.last_pos = pos;
+		Ok(())
+	}
+
 	/// Get the header hash at the specified height based on the current header MMR state.
 	pub fn get_header_hash_by_height(&self, height: u64) -> Result<Hash, Error> {
 		let pos = pmmr::insertion_to_pmmr_index(height + 1);