refactor: calculate balance directly from blockchain struct and verif…

…y tx with merkle proof

blockchain/blockchain.go

@@ -1,6 +1,7 @@
 package blockchain
 
 import (
+	"encoding/hex"
 	"encoding/json"
 	"log"
 	"strings"
@@ -20,6 +21,48 @@ func (bc *Blockchain) Airdrop(address string) {
 	bc.AddTransactionToTransactionPool(txn)
 }
 
+func (bc *Blockchain) CalculateBalance(address string) int {
+	log.Println("Fetch wallet transactions")
+	walletTx := bc.WalletIndex.GetWalletTransactions(address)
+	log.Println("teeest")
+	bal := 0
+
+	bb, _ := json.Marshal(walletTx)
+	log.Println(string(bb))
+
+	for _, txn := range walletTx {
+
+		block := bc.Blocks[txn.BlockIndex]
+		log.Println(block.ToJson())
+
+		proof, err := GenerateMerkleProof(block.Transactions, txn.TransactionIndex)
+		log.Printf("%x", proof)
+
+		if err != nil {
+			log.Panicf(err.Error())
+		}
+
+		merkleRootBytes, err := hex.DecodeString(block.MerkleRoot)
+		if err != nil {
+			log.Panic(err.Error())
+		}
+
+		isValidTx := VerifyTransaction(merkleRootBytes, txn.Transaction, proof)
+
+		if !isValidTx {
+			log.Panic("Invalid tx found while computing balance")
+		}
+
+		if txn.Transaction.From == address && txn.Transaction.Status == constants.STATUS_SUCCESS {
+			bal -= int(txn.Transaction.Value)
+		} else if txn.Transaction.To == address && txn.Transaction.Status == constants.STATUS_SUCCESS {
+			bal += int(txn.Transaction.Value)
+		}
+	}
+
+	return bal
+}
+
 func NewBlockchain(genesisBlock *Block) *Blockchain {
 	state, err := ReadFromDb()
 	bc := &Blockchain{}
@@ -74,8 +117,6 @@ func (bc *Blockchain) AddBlock(b *Block) {
 
 	for index, txn := range b.Transactions {
 		m[txn.TransactioHash] = true
-		balance := bc.WalletIndex.CalculateBalance(txn.From)
-		log.Printf("\n\nsender balance -> %d \n\n", balance)
 		bc.WalletIndex.AddTransaction(txn.From, nextBlockNumber, index, txn)
 		bc.WalletIndex.AddTransaction(txn.To, nextBlockNumber, index, txn)
 		bc.TransactionIndex.AddTransaction(txn.TransactioHash, index, nextBlockNumber)
@@ -105,8 +146,8 @@ func (bc *Blockchain) LastBlock() *Block {
 func (bc *Blockchain) CopyTransactionPool() []*Transaction {
 	t := make([]*Transaction, 0)
 	for _, txn := range bc.TransactionPool {
-		senderBalance := bc.WalletIndex.CalculateBalance(txn.From)
 		if txn.From != constants.BLOCKCHAIN_AIRDROP_ADDRESS {
+			senderBalance := bc.CalculateBalance(txn.From)
 			if senderBalance < int(txn.Value) {
 				txn.Status = constants.STATUS_FAILED
 			} else {

blockchain/merkle.go

@@ -72,6 +72,11 @@ func GenerateMerkleProof(transactions []*Transaction, index int) ([][]byte, erro
 	leaves := CreateLeafNodes(transactions)
 	proof := make([][]byte, 0)
 
+	if len(leaves) == 1 {
+		// Special case: only one leaf node
+		return proof, nil
+	}
+
 	for len(leaves) > 1 {
 		var newLevel []*Node
 		for i := 0; i < len(leaves); i += 2 {

blockchain/walletIndex.go

@@ -2,8 +2,6 @@ package blockchain
 
 import (
 	"log"
-
-	"github.com/stinkymonkeyph/gopher-blocks/constants"
 )
 
 type WalletTransactionIndex struct {
@@ -18,7 +16,7 @@ type WalletIndex struct {
 
 func NewWalletIndex() *WalletIndex {
 	return &WalletIndex{
-		Transactions: make(map[string][]*WalletTransactionIndex),
+		Transactions: make(map[string][]*WalletTransactionIndex, 0),
 	}
 }
 
@@ -38,28 +36,10 @@ func (w *WalletIndex) AddTransaction(address string, blockIndex int, txIndex int
 	log.Printf("Transactions for %s: %+v", address, w.Transactions[address])
 }
 
-func (w *WalletIndex) CalculateBalance(address string) int {
-	bal := 0
-
-	for _, txn := range w.Transactions[address] {
-		if txn.Transaction.From == address && txn.Transaction.Status == constants.STATUS_SUCCESS {
-			bal -= int(txn.Transaction.Value)
-		} else if txn.Transaction.To == address && txn.Transaction.Status == constants.STATUS_SUCCESS {
-			bal += int(txn.Transaction.Value)
-		}
-	}
-
-	return bal
-}
-
-func (w *WalletIndex) GetWalletTransactions(address string) []*Transaction {
-	t := make([]*Transaction, 0)
-
-	if w.Transactions[address] != nil {
-		for _, txi := range w.Transactions[address] {
-			t = append(t, txi.Transaction)
-		}
+func (w *WalletIndex) GetWalletTransactions(address string) []*WalletTransactionIndex {
+	if w.Transactions == nil {
+		w.Transactions = make(map[string][]*WalletTransactionIndex, 0)
 	}
 
-	return t
+	return w.Transactions[address]
 }

main.go

@@ -23,17 +23,10 @@ func main() {
 	bc.AddTransactionToTransactionPool(transaction1)
 	bc.Mining()
 	log.Print(bc.ToJSON())
-	senderBalance := bc.WalletIndex.CalculateBalance("0x1")
+	senderBalance := bc.CalculateBalance("0x1")
 	log.Printf("\n\n\nSender Balance: %d \n", senderBalance)
 
 	walletTxns := bc.WalletIndex.GetWalletTransactions("0x1")
 	wbytes, _ := json.Marshal(walletTxns)
 	log.Println(string(wbytes))
-
-	txnMetadata := bc.TransactionIndex.GetTransactionMetadata(walletTxns[0].TransactioHash)
-	txbytes, _ := json.Marshal(txnMetadata)
-	log.Println(string(txbytes))
-
-	retrievedBlock := bc.Blocks[txnMetadata.BlockHeight]
-	log.Println(retrievedBlock.ToJson())
 }