Remove future block heights from store when backtracking

Fixes #499

chain/src/chain.rs

@@ -81,11 +81,11 @@ impl Chain {
 		let head = match chain_store.head() {
 			Ok(tip) => tip,
 			Err(NotFoundErr) => {
+				let tip = Tip::new(genesis.hash());
 				chain_store.save_block(&genesis)?;
-				chain_store.setup_height(&genesis.header)?;
+				chain_store.setup_height(&genesis.header, tip)?;
 
 				// saving a new tip based on genesis
-				let tip = Tip::new(genesis.hash());
 				chain_store.save_head(&tip)?;
 				info!(
 					LOGGER,

chain/src/pipe.rs

@@ -349,7 +349,7 @@ fn update_head(b: &Block, ctx: &mut BlockContext) -> Result<Option<Tip>, Error>
 	if tip.total_difficulty > ctx.head.total_difficulty {
 		// update the block height index
 		ctx.store
-			.setup_height(&b.header)
+			.setup_height(&b.header, &ctx.head)
 			.map_err(|e| Error::StoreErr(e, "pipe setup height".to_owned()))?;
 
 		// in sync mode, only update the "body chain", otherwise update both the

chain/src/store.rs

@@ -135,18 +135,18 @@ impl ChainStore for ChainKVStore {
 	}
 
 	// lookup the block header hash by output commitment
- // lookup the block header based on this hash
- // to check the chain is correct compare this block header to
- // the block header currently indexed at the relevant block height (tbd if
- // actually necessary)
- //
- // NOTE: This index is not exhaustive.
- // This node may not have seen this full block, so may not have populated the
- // index.
- // Block headers older than some threshold (2 months?) will not necessarily be
- // included
- // in this index.
- //
+	// lookup the block header based on this hash
+	// to check the chain is correct compare this block header to
+	// the block header currently indexed at the relevant block height (tbd if
+	// actually necessary)
+	//
+	// NOTE: This index is not exhaustive.
+	// This node may not have seen this full block, so may not have populated the
+	// index.
+	// Block headers older than some threshold (2 months?) will not necessarily be
+	// included
+	// in this index.
+	//
 	fn get_block_header_by_output_commit(&self, commit: &Commitment) -> Result<BlockHeader, Error> {
 		let block_hash = self.db.get_ser(&to_key(
 			HEADER_BY_OUTPUT_PREFIX,
@@ -183,6 +183,10 @@ impl ChainStore for ChainKVStore {
 		option_to_not_found(self.db.get_ser(&u64_to_key(HEADER_HEIGHT_PREFIX, height)))
 	}
 
+	fn delete_header_by_height(&self, height: u64) -> Result<(), Error> {
+		self.db.delete(&u64_to_key(HEADER_HEIGHT_PREFIX, height))
+	}
+
 	fn get_output_by_commit(&self, commit: &Commitment) -> Result<Output, Error> {
 		option_to_not_found(
 			self.db
@@ -219,28 +223,36 @@ impl ChainStore for ChainKVStore {
 	}
 
 	/// Maintain consistency of the "header_by_height" index by traversing back
-	/// through the
-	/// current chain and updating "header_by_height" until we reach a
-	/// block_header
-	/// that is consistent with its height (everything prior to this will be
-	/// consistent).
-	/// We need to handle the case where we have no index entry for a given height to
-	/// account for the case where we just switched to a new fork and the height jumped
-	/// beyond current chain height.
-	fn setup_height(&self, header: &BlockHeader) -> Result<(), Error> {
+	/// through the current chain and updating "header_by_height" until we reach
+	/// a block_header that is consistent with its height (everything prior to
+	/// this will be consistent).
+	/// We need to handle the case where we have no index entry for a given
+	/// height to account for the case where we just switched to a new fork and
+	/// the height jumped beyond current chain height.
+	fn setup_height(&self, header: &BlockHeader, old_tip: &Tip)
+		-> Result<(), Error> {
+
+		// remove headers ahead if we backtracked
+		for n in header.height..old_tip.height {
+			self.delete_header_by_height(n)?;
+		}
+
 		self.db
 			.put_ser(&u64_to_key(HEADER_HEIGHT_PREFIX, header.height), header)?;
 
-		if header.height == 0 {
-			return Ok(());
-		}
+		if header.height > 0 {
+			let mut prev_header = self.get_block_header(&header.previous)?;
+			while prev_header.height > 0 {
+				if let Ok(_) = self.is_on_current_chain(&prev_header) {
+					break;
+				}
+				self.db
+					.put_ser(&u64_to_key(HEADER_HEIGHT_PREFIX, prev_header.height), &prev_header)?;
 
-		let prev_header = self.get_block_header(&header.previous)?;
-		if let Ok(_) = self.is_on_current_chain(&prev_header) {
-			Ok(())
-		} else {
-			self.setup_height(&prev_header)
+				prev_header = self.get_block_header(&prev_header.previous)?;
+			}
 		}
+		Ok(())
 	}
 }
 

chain/src/types.rs

@@ -226,6 +226,9 @@ pub trait ChainStore: Send + Sync {
 	/// Gets the block header at the provided height
 	fn get_header_by_height(&self, height: u64) -> Result<BlockHeader, store::Error>;
 
+	/// Delete the block header at the height
+	fn delete_header_by_height(&self, height: u64) -> Result<(), store::Error>;
+
 	/// Is the block header on the current chain?
 	/// Use the header_by_height index to verify the block header is where we think it is.
 	fn is_on_current_chain(&self, header: &BlockHeader) -> Result<(), store::Error>;
@@ -258,7 +261,7 @@ pub trait ChainStore: Send + Sync {
 	/// Saves the provided block header at the corresponding height. Also check
 	/// the consistency of the height chain in store by assuring previous
 	/// headers are also at their respective heights.
-	fn setup_height(&self, bh: &BlockHeader) -> Result<(), store::Error>;
+	fn setup_height(&self, bh: &BlockHeader, old_tip: &Tip) -> Result<(), store::Error>;
 }
 
 /// Bridge between the chain pipeline and the rest of the system. Handles

chain/tests/store_indices.rs

@@ -21,7 +21,7 @@ extern crate rand;
 
 use std::fs;
 
-use chain::ChainStore;
+use chain::{ChainStore, Tip};
 use core::core::hash::Hashed;
 use core::core::Block;
 use keychain::Keychain;
@@ -45,14 +45,14 @@ fn test_various_store_indices() {
 	global::set_mining_mode(ChainTypes::AutomatedTesting);
 	let genesis = pow::mine_genesis_block(None).unwrap();
 	chain_store.save_block(&genesis).unwrap();
-	chain_store.setup_height(&genesis.header).unwrap();
+	chain_store.setup_height(&genesis.header, &Tip::new(genesis.hash())).unwrap();
 
 	let block = Block::new(&genesis.header, vec![], &keychain, &key_id).unwrap();
 	let commit = block.outputs[0].commitment();
 	let block_hash = block.hash();
 
 	chain_store.save_block(&block).unwrap();
-	chain_store.setup_height(&block.header).unwrap();
+	chain_store.setup_height(&block.header, &Tip::from_block(&block.header)).unwrap();
 
 	let block_header = chain_store.get_block_header(&block_hash).unwrap();
 	assert_eq!(block_header.hash(), block_hash);