From a96a5bc1e63b6a7fe04066da2e1786a227eb019f Mon Sep 17 00:00:00 2001
From: Ivan Kuznetsov <me@jeiwan.ru>
Date: Fri, 6 Oct 2017 12:29:55 +0700
Subject: [PATCH] Fix the main scenario

---
 blockchain.go |  6 +-----
 server.go     | 19 +++++++++++--------
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/blockchain.go b/blockchain.go
index 255d9e2..8436035 100644
--- a/blockchain.go
+++ b/blockchain.go
@@ -97,11 +97,6 @@ func NewBlockchain(nodeID string) *Blockchain {
 	return &bc
 }
 
-// Reset removes all blockchain data
-func (bc *Blockchain) Reset() {
-
-}
-
 // AddBlock saves the block into the blockchain
 func (bc *Blockchain) AddBlock(block *Block) {
 	err := bc.db.Update(func(tx *bolt.Tx) error {
@@ -127,6 +122,7 @@ func (bc *Blockchain) AddBlock(block *Block) {
 			if err != nil {
 				log.Panic(err)
 			}
+			bc.tip = block.Hash
 		}
 
 		return nil
diff --git a/server.go b/server.go
index 9653af9..b13f520 100644
--- a/server.go
+++ b/server.go
@@ -202,18 +202,16 @@ func handleBlock(request []byte, bc *Blockchain) {
 	fmt.Println("Recevied a new block!")
 	bc.AddBlock(block)
 
-	// TODO: how to update UTXOSet when the order of new blocks is not correct?
-	// UTXOSet := UTXOSet{bc}
-	// UTXOSet.Update(block)
 	fmt.Printf("Added block %x\n", block.Hash)
-	fmt.Printf("Added block %d\n", block.Height)
 
-	fmt.Println(blocksInTransit)
 	if len(blocksInTransit) > 0 {
 		blockHash := blocksInTransit[0]
 		sendGetData(payload.AddrFrom, "block", blockHash)
 
 		blocksInTransit = blocksInTransit[1:]
+	} else {
+		UTXOSet := UTXOSet{bc}
+		UTXOSet.Reindex()
 	}
 }
 
@@ -321,10 +319,11 @@ func handleTx(request []byte, bc *Blockchain) {
 		}
 	} else {
 		if len(mempool) >= 2 && len(miningAddress) > 0 {
+		MineTransactions:
 			var txs []*Transaction
 
-			for _, tx := range mempool {
-				// TODO: remove this check after improving MineBlock
+			for id := range mempool {
+				tx := mempool[id]
 				if bc.VerifyTransaction(&tx) {
 					txs = append(txs, &tx)
 				}
@@ -340,7 +339,7 @@ func handleTx(request []byte, bc *Blockchain) {
 
 			newBlock := bc.MineBlock(txs)
 			UTXOSet := UTXOSet{bc}
-			UTXOSet.Update(newBlock)
+			UTXOSet.Reindex()
 
 			fmt.Println("New block is mined!")
 
@@ -354,6 +353,10 @@ func handleTx(request []byte, bc *Blockchain) {
 					sendInv(node, "block", [][]byte{newBlock.Hash})
 				}
 			}
+
+			if len(mempool) > 0 {
+				goto MineTransactions
+			}
 		}
 	}
 }