From 5c4340f47d7ee44a4ae85ed467e2d1554fcf5f2f Mon Sep 17 00:00:00 2001
From: Ivan Kuznetsov <me@jeiwan.ru>
Date: Tue, 3 Oct 2017 16:21:15 +0700
Subject: [PATCH] Add -miner flag to 'startnode' command

---
 blockchain.go    |  1 +
 cli.go           |  5 +++--
 cli_send.go      |  3 +--
 cli_startnode.go | 16 +++++++++++++---
 server.go        | 29 +++++++++++++++++++++--------
 5 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/blockchain.go b/blockchain.go
index b6efe0f..255d9e2 100644
--- a/blockchain.go
+++ b/blockchain.go
@@ -275,6 +275,7 @@ func (bc *Blockchain) MineBlock(transactions []*Transaction) *Block {
 	var lastHeight int
 
 	for _, tx := range transactions {
+		// TODO: ignore transaction if it's not valid
 		if bc.VerifyTransaction(tx) != true {
 			log.Panic("ERROR: Invalid transaction")
 		}
diff --git a/cli.go b/cli.go
index 82a1127..b597b4c 100644
--- a/cli.go
+++ b/cli.go
@@ -20,7 +20,7 @@ func (cli *CLI) printUsage() {
 	fmt.Println("  printchain - Print all the blocks of the blockchain")
 	fmt.Println("  reindexutxo - Rebuilds the UTXO set")
 	fmt.Println("  send -from FROM -to TO -amount AMOUNT -mine - Send AMOUNT of coins from FROM address to TO. Mine on the same node, when -mine is set.")
-	fmt.Println("  startnode - Start a node with ID specified in NODE_ID env. var")
+	fmt.Println("  startnode -miner ADDRESS - Start a node with ID specified in NODE_ID env. var. -miner enables mining")
 }
 
 func (cli *CLI) validateArgs() {
@@ -55,6 +55,7 @@ func (cli *CLI) Run() {
 	sendTo := sendCmd.String("to", "", "Destination wallet address")
 	sendAmount := sendCmd.Int("amount", 0, "Amount to send")
 	sendMine := sendCmd.Bool("mine", false, "Mine immediately on the same node")
+	startNodeMiner := startNodeCmd.String("miner", "", "Enable mining mode and send reward to ADDRESS")
 
 	switch os.Args[1] {
 	case "getbalance":
@@ -149,6 +150,6 @@ func (cli *CLI) Run() {
 			startNodeCmd.Usage()
 			os.Exit(1)
 		}
-		cli.startNode(nodeID)
+		cli.startNode(nodeID, *startNodeMiner)
 	}
 }
diff --git a/cli_send.go b/cli_send.go
index cb607a8..75a5930 100644
--- a/cli_send.go
+++ b/cli_send.go
@@ -24,16 +24,15 @@ func (cli *CLI) send(from, to string, amount int, nodeID string, mineNow bool) {
 	wallet := wallets.GetWallet(from)
 
 	tx := NewUTXOTransaction(&wallet, to, amount, &UTXOSet)
-	cbTx := NewCoinbaseTX(from, "")
 
 	if mineNow {
+		cbTx := NewCoinbaseTX(from, "")
 		txs := []*Transaction{cbTx, tx}
 
 		newBlock := bc.MineBlock(txs)
 		UTXOSet.Update(newBlock)
 	} else {
 		sendTx(knownNodes[0], tx)
-		sendTx(knownNodes[0], cbTx)
 	}
 
 	fmt.Println("Success!")
diff --git a/cli_startnode.go b/cli_startnode.go
index 289ccd4..d390512 100644
--- a/cli_startnode.go
+++ b/cli_startnode.go
@@ -1,8 +1,18 @@
 package main
 
-import "fmt"
+import (
+	"fmt"
+	"log"
+)
 
-func (cli *CLI) startNode(nodeID string) {
+func (cli *CLI) startNode(nodeID, minerAddress string) {
 	fmt.Printf("Starting node %s\n", nodeID)
-	StartServer(nodeID)
+	if len(minerAddress) > 0 {
+		if ValidateAddress(minerAddress) {
+			fmt.Println("Mining is on. Address to receive rewards: ", minerAddress)
+		} else {
+			log.Panic("Wrong miner address!")
+		}
+	}
+	StartServer(nodeID, minerAddress)
 }
diff --git a/server.go b/server.go
index 7457cb0..3ccb0f3 100644
--- a/server.go
+++ b/server.go
@@ -16,6 +16,7 @@ const nodeVersion = 1
 const commandLength = 12
 
 var nodeAddress string
+var miningAddress string
 var knownNodes = []string{"localhost:3000"}
 var blocksInTransit = [][]byte{}
 var mempool = make(map[string]Transaction)
@@ -189,6 +190,10 @@ 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)
 
@@ -199,9 +204,6 @@ func handleBlock(request []byte, bc *Blockchain) {
 
 		blocksInTransit = blocksInTransit[1:]
 	}
-
-	UTXOSet := UTXOSet{bc}
-	UTXOSet.Update(block)
 }
 
 func handleInv(request []byte, bc *Blockchain) {
@@ -302,22 +304,32 @@ func handleTx(request []byte, bc *Blockchain) {
 
 	if nodeAddress == knownNodes[0] {
 		for _, node := range knownNodes {
-			if node != nodeAddress {
+			if node != nodeAddress && node != payload.AddFrom {
 				sendInv(node, "tx", [][]byte{tx.ID})
 			}
 		}
 	} else {
-		if len(mempool) >= 2 {
+		if len(mempool) >= 2 && len(miningAddress) > 0 {
 			var txs []*Transaction
 
 			for _, tx := range mempool {
-				txs = append(txs, &tx)
+				// TODO: remove this check after improving MineBlock
+				if bc.VerifyTransaction(&tx) {
+					txs = append(txs, &tx)
+				}
 			}
+
+			cbTx := NewCoinbaseTX(miningAddress, "")
+			txs = append(txs, cbTx)
+
 			newBlock := bc.MineBlock(txs)
 			UTXOSet := UTXOSet{bc}
 			UTXOSet.Update(newBlock)
 
-			mempool = make(map[string]Transaction)
+			for _, tx := range txs {
+				txID := hex.EncodeToString(tx.ID)
+				delete(mempool, txID)
+			}
 		}
 	}
 }
@@ -379,8 +391,9 @@ func handleConnection(conn net.Conn, bc *Blockchain) {
 }
 
 // StartServer starts a node
-func StartServer(nodeID string) {
+func StartServer(nodeID, minerAddress string) {
 	nodeAddress = fmt.Sprintf("localhost:%s", nodeID)
+	miningAddress = minerAddress
 	ln, err := net.Listen(protocol, nodeAddress)
 	if err != nil {
 		log.Panic(err)