M core/blockchain.go → core/blockchain.go

@@ -3,10 +3,13 @@
 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 @@ err = bc.addressdb.Put(PackPubK(tx.From), Uint64ToBytes(balance))
 		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
 
@@ -126,6 +134,16 @@ if err != nil {
 		return nil, err
 	}
 	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) {

M node/node_test.go → node/node_test.go

@@ -2,6 +2,8 @@ package node
 
 import (
 	"crypto/ecdsa"
+	"encoding/hex"
+	"fmt"
 	"io/ioutil"
 	"testing"
 
@@ -122,6 +124,7 @@ node, err := NewNode(privK, bc, true)
 	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 @@ 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(100))
+
 	// add another tx sending coins to the pubK1
 	privK1, err := core.NewKey()
 	assert.Nil(t, err)
@@ -199,4 +208,10 @@ block, err = node.NewBlock(txs)
 	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))
 }