From 85d0bcc5a3357821a9b8976fd1041017b5e3555e Mon Sep 17 00:00:00 2001
From: arnaucube <arnaucube@gmail.com>
Date: Thu, 11 Apr 2019 14:05:09 +0200
Subject: [PATCH] GetBalance

---
 core/blockchain.go | 20 +++++++++++++++++++-
 node/node_test.go  | 15 +++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/core/blockchain.go b/core/blockchain.go
index 9a56a4a..95d8856 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -3,10 +3,13 @@ package core
 import (
 	"bytes"
 	"crypto/ecdsa"
+	"encoding/hex"
 	"errors"
 	"fmt"
+	"strconv"
 
 	"github.com/arnaucube/slowlorisdb/db"
+	lvldberrors "github.com/syndtr/goleveldb/leveldb/errors"
 )
 
 type PoA struct {
@@ -99,18 +102,23 @@ func (bc *Blockchain) UpdateWalletsWithNewTx(tx *Tx) error {
 		if err != nil {
 			return err
 		}
+		fmt.Println("sent-->: balance of " + hex.EncodeToString(PackPubK(tx.From)[:10]) + ": " + strconv.Itoa(int(balance)))
 	}
 	for _, out := range tx.Outputs {
 		balanceBytes, err := bc.addressdb.Get(PackPubK(tx.To))
-		if err != nil {
+		if err != nil && err != lvldberrors.ErrNotFound {
 			return err
 		}
+		if err == lvldberrors.ErrNotFound {
+			balanceBytes = Uint64ToBytes(uint64(0))
+		}
 		balance := Uint64FromBytes(balanceBytes)
 		balance = balance + out.Value
 		err = bc.addressdb.Put(PackPubK(tx.To), Uint64ToBytes(balance))
 		if err != nil {
 			return err
 		}
+		fmt.Println("--> received: balance of " + hex.EncodeToString(PackPubK(tx.To)[:10]) + ": " + strconv.Itoa(int(balance)))
 	}
 	return nil
 
@@ -128,6 +136,16 @@ func (bc *Blockchain) GetBlock(hash Hash) (*Block, error) {
 	return block, nil
 }
 
+func (bc *Blockchain) GetBalance(pubK *ecdsa.PublicKey) (uint64, error) {
+	balanceBytes, err := bc.addressdb.Get(PackPubK(pubK))
+	if err != nil {
+		return uint64(0), err
+	}
+	balance := Uint64FromBytes(balanceBytes)
+	return balance, nil
+
+}
+
 func (bc *Blockchain) GetPrevBlock(hash Hash) (*Block, error) {
 	currentBlock, err := bc.GetBlock(hash)
 	if err != nil {
diff --git a/node/node_test.go b/node/node_test.go
index 1af916b..d8bdc97 100644
--- a/node/node_test.go
+++ b/node/node_test.go
@@ -2,6 +2,8 @@ package node
 
 import (
 	"crypto/ecdsa"
+	"encoding/hex"
+	"fmt"
 	"io/ioutil"
 	"testing"
 
@@ -122,6 +124,7 @@ func TestFromGenesisToTenBlocks(t *testing.T) {
 	assert.Nil(t, err)
 
 	// create the genesis block
+	// genesisBlock sends 100 to pubK
 	genesisBlock, err := node.CreateGenesis(&privK.PublicKey, uint64(100))
 	assert.Nil(t, err)
 	assert.NotEqual(t, genesisBlock.Signature, core.Signature{})
@@ -169,6 +172,12 @@ func TestFromGenesisToTenBlocks(t *testing.T) {
 	err = node.Bc.AddBlock(block)
 	assert.Nil(t, err)
 
+	balance, err := node.Bc.GetBalance(&pubK0)
+	assert.Nil(t, err)
+	fmt.Println(hex.EncodeToString(core.PackPubK(&pubK0)[:10]))
+	fmt.Println("balance in pubK0", balance)
+	assert.Equal(t, balance, uint64(100))
+
 	// add another tx sending coins to the pubK1
 	privK1, err := core.NewKey()
 	assert.Nil(t, err)
@@ -199,4 +208,10 @@ func TestFromGenesisToTenBlocks(t *testing.T) {
 	assert.Nil(t, err)
 	err = node.Bc.AddBlock(block)
 	assert.Nil(t, err)
+
+	balance, err = node.Bc.GetBalance(&pubK0)
+	assert.Nil(t, err)
+	fmt.Println(hex.EncodeToString(core.PackPubK(&pubK0)[:10]))
+	fmt.Println("balance in pubK0", balance)
+	assert.Equal(t, balance, uint64(90))
 }