core, eth: roll back uncertain headers in failed fast syncs

2 files changed, 49 insertions, 3 deletions

diff --git a/core/blockchain.go b/core/blockchain.go
index 3e7dfa9ee..490552ea0 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -245,7 +245,21 @@ func (bc *BlockChain) SetHead(head uint64) {
 	if bc.currentBlock == nil {
 		bc.currentBlock = bc.genesisBlock
 	}
-	bc.insert(bc.currentBlock)
+	if bc.currentHeader == nil {
+		bc.currentHeader = bc.genesisBlock.Header()
+	}
+	if bc.currentFastBlock == nil {
+		bc.currentFastBlock = bc.genesisBlock
+	}
+	if err := WriteHeadBlockHash(bc.chainDb, bc.currentBlock.Hash()); err != nil {
+		glog.Fatalf("failed to reset head block hash: %v", err)
+	}
+	if err := WriteHeadHeaderHash(bc.chainDb, bc.currentHeader.Hash()); err != nil {
+		glog.Fatalf("failed to reset head header hash: %v", err)
+	}
+	if err := WriteHeadFastBlockHash(bc.chainDb, bc.currentFastBlock.Hash()); err != nil {
+		glog.Fatalf("failed to reset head fast block hash: %v", err)
+	}
 	bc.loadLastState()
 }
 
@@ -790,6 +804,27 @@ func (self *BlockChain) InsertHeaderChain(chain []*types.Header, checkFreq int)
 	return 0, nil
 }
 
+// Rollback is designed to remove a chain of links from the database that aren't
+// certain enough to be valid.
+func (self *BlockChain) Rollback(chain []common.Hash) {
+	for i := len(chain) - 1; i >= 0; i-- {
+		hash := chain[i]
+
+		if self.currentHeader.Hash() == hash {
+			self.currentHeader = self.GetHeader(self.currentHeader.ParentHash)
+			WriteHeadHeaderHash(self.chainDb, self.currentHeader.Hash())
+		}
+		if self.currentFastBlock.Hash() == hash {
+			self.currentFastBlock = self.GetBlock(self.currentFastBlock.ParentHash())
+			WriteHeadFastBlockHash(self.chainDb, self.currentFastBlock.Hash())
+		}
+		if self.currentBlock.Hash() == hash {
+			self.currentBlock = self.GetBlock(self.currentBlock.ParentHash())
+			WriteHeadBlockHash(self.chainDb, self.currentBlock.Hash())
+		}
+	}
+}
+
 // InsertReceiptChain attempts to complete an already existing header chain with
 // transaction and receipt data.
 func (self *BlockChain) InsertReceiptChain(blockChain types.Blocks, receiptChain []types.Receipts) (int, error) {
diff --git a/core/blockchain_test.go b/core/blockchain_test.go
index a614aaa2f..01667c21e 100644
--- a/core/blockchain_test.go
+++ b/core/blockchain_test.go
@@ -599,8 +599,8 @@ func testReorgBadHashes(t *testing.T, full bool) {
 			t.Errorf("last  block gasLimit mismatch: have: %x, want %x", ncm.GasLimit(), blocks[2].Header().GasLimit)
 		}
 	} else {
-		if ncm.CurrentHeader().Hash() != genesis.Hash() {
-			t.Errorf("last header hash mismatch: have: %x, want %x", ncm.CurrentHeader().Hash(), genesis.Hash())
+		if ncm.CurrentHeader().Hash() != headers[2].Hash() {
+			t.Errorf("last header hash mismatch: have: %x, want %x", ncm.CurrentHeader().Hash(), headers[2].Hash())
 		}
 	}
 }
@@ -775,6 +775,11 @@ func TestLightVsFastVsFullChainHeads(t *testing.T) {
 	height := uint64(1024)
 	blocks, receipts := GenerateChain(genesis, gendb, int(height), nil)
 
+	// Configure a subchain to roll back
+	remove := []common.Hash{}
+	for _, block := range blocks[height/2:] {
+		remove = append(remove, block.Hash())
+	}
 	// Create a small assertion method to check the three heads
 	assert := func(t *testing.T, kind string, chain *BlockChain, header uint64, fast uint64, block uint64) {
 		if num := chain.CurrentBlock().NumberU64(); num != block {
@@ -798,6 +803,8 @@ func TestLightVsFastVsFullChainHeads(t *testing.T) {
 		t.Fatalf("failed to process block %d: %v", n, err)
 	}
 	assert(t, "archive", archive, height, height, height)
+	archive.Rollback(remove)
+	assert(t, "archive", archive, height/2, height/2, height/2)
 
 	// Import the chain as a non-archive node and ensure all pointers are updated
 	fastDb, _ := ethdb.NewMemDatabase()
@@ -816,6 +823,8 @@ func TestLightVsFastVsFullChainHeads(t *testing.T) {
 		t.Fatalf("failed to insert receipt %d: %v", n, err)
 	}
 	assert(t, "fast", fast, height, height, 0)
+	fast.Rollback(remove)
+	assert(t, "fast", fast, height/2, height/2, 0)
 
 	// Import the chain as a light node and ensure all pointers are updated
 	lightDb, _ := ethdb.NewMemDatabase()
@@ -827,6 +836,8 @@ func TestLightVsFastVsFullChainHeads(t *testing.T) {
 		t.Fatalf("failed to insert header %d: %v", n, err)
 	}
 	assert(t, "light", light, height, 0, 0)
+	light.Rollback(remove)
+	assert(t, "light", light, height/2, 0, 0)
 }
 
 // Tests that chain reorganizations handle transaction removals and reinsertions.