From 901bde1fd443e1375587a5b092160edf6bfe1f2c Mon Sep 17 00:00:00 2001
From: Ori Newman <orinewman1@gmail.com>
Date: Thu, 30 May 2019 18:14:27 +0300
Subject: [PATCH] [NOD-202] undo createDAGState if blockdag new fails (#318)

* [NOD-204] Add UTXOCommitment to GetBlockTemplateResult

* [NOD-204] Add UTXOCommitment to GetBlockTemplateResult

* [NOD-206] Avoid leaking blocks from previous miner when switching miners

* [NOD-202] Undo createDAGState if blockdag.New fails

* [NOD-202] Fix gofmt errors
---
 blockdag/dag.go   | 19 +++++++++++++++----
 blockdag/dagio.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/blockdag/dag.go b/blockdag/dag.go
index 09e85d663..acefc29f6 100644
--- a/blockdag/dag.go
+++ b/blockdag/dag.go
@@ -1864,14 +1864,23 @@ func New(config *Config) (*BlockDAG, error) {
 	// Initialize the chain state from the passed database.  When the db
 	// does not yet contain any DAG state, both it and the DAG state
 	// will be initialized to contain only the genesis block.
-	if err := dag.initDAGState(); err != nil {
+	err := dag.initDAGState()
+	if err != nil {
 		return nil, err
 	}
+	defer func() {
+		if err != nil {
+			err := dag.removeDAGState()
+			if err != nil {
+				panic(fmt.Sprintf("Couldn't remove the DAG State: %s", err))
+			}
+		}
+	}()
 
 	// Initialize and catch up all of the currently active optional indexes
 	// as needed.
 	if config.IndexManager != nil {
-		err := config.IndexManager.Init(dag.db, &dag, config.Interrupt)
+		err = config.IndexManager.Init(dag.db, &dag, config.Interrupt)
 		if err != nil {
 			return nil, err
 		}
@@ -1881,7 +1890,8 @@ func New(config *Config) (*BlockDAG, error) {
 
 	if genesis == nil {
 		genesisBlock := util.NewBlock(dag.dagParams.GenesisBlock)
-		isOrphan, err := dag.ProcessBlock(genesisBlock, BFNone)
+		var isOrphan bool
+		isOrphan, err = dag.ProcessBlock(genesisBlock, BFNone)
 		if err != nil {
 			return nil, err
 		}
@@ -1895,7 +1905,8 @@ func New(config *Config) (*BlockDAG, error) {
 	dag.genesis = genesis
 
 	// Initialize rule change threshold state caches.
-	if err := dag.initThresholdCaches(); err != nil {
+	err = dag.initThresholdCaches()
+	if err != nil {
 		return nil, err
 	}
 
diff --git a/blockdag/dagio.go b/blockdag/dagio.go
index aaef9420c..71b0cd00c 100644
--- a/blockdag/dagio.go
+++ b/blockdag/dagio.go
@@ -344,6 +344,49 @@ func (dag *BlockDAG) createDAGState() error {
 	return nil
 }
 
+func (dag *BlockDAG) removeDAGState() error {
+	err := dag.db.Update(func(dbTx database.Tx) error {
+		meta := dbTx.Metadata()
+
+		err := meta.DeleteBucket(blockIndexBucketName)
+		if err != nil {
+			return err
+		}
+
+		err = meta.DeleteBucket(utxoSetBucketName)
+		if err != nil {
+			return err
+		}
+
+		err = meta.DeleteBucket(utxoDiffsBucketName)
+		if err != nil {
+			return err
+		}
+
+		err = dbTx.Metadata().Delete(utxoSetVersionKeyName)
+		if err != nil {
+			return err
+		}
+
+		err = meta.DeleteBucket(subnetworksBucketName)
+		if err != nil {
+			return err
+		}
+
+		err = dbTx.Metadata().Delete(localSubnetworkKeyName)
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func dbPutLocalSubnetworkID(dbTx database.Tx, subnetworkID *subnetworkid.SubnetworkID) error {
 	if subnetworkID == nil {
 		return dbTx.Metadata().Put(localSubnetworkKeyName, []byte{})