From 470adef2c34879ac7f7c6c16291f40ef345a7414 Mon Sep 17 00:00:00 2001
From: Ivan Kuznetsov <me@jeiwan.ru>
Date: Sun, 1 Oct 2017 11:50:27 +0700
Subject: [PATCH] Implement 'block' and 'getdata' commands

---
 blockchain.go | 29 +++++++++++++++++
 server.go     | 89 +++++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 112 insertions(+), 6 deletions(-)

diff --git a/blockchain.go b/blockchain.go
index 0f310af..328cf06 100644
--- a/blockchain.go
+++ b/blockchain.go
@@ -97,6 +97,11 @@ func NewBlockchain(nodeID string) *Blockchain {
 	return &bc
 }
 
+// Reset removes all blockchain data
+func (bc *Blockchain) Reset() {
+
+}
+
 // FindTransaction finds a transaction by its ID
 func (bc *Blockchain) FindTransaction(ID []byte) (Transaction, error) {
 	bci := bc.Iterator()
@@ -169,6 +174,30 @@ func (bc *Blockchain) Iterator() *BlockchainIterator {
 	return bci
 }
 
+// GetBlock finds a block by its hash and returns it
+func (bc *Blockchain) GetBlock(blockHash []byte) (Block, error) {
+	var block Block
+
+	err := bc.db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte(blocksBucket))
+
+		blockData := b.Get(blockHash)
+
+		if blockData == nil {
+			return errors.New("Block is not found.")
+		}
+
+		block = *DeserializeBlock(blockData)
+
+		return nil
+	})
+	if err != nil {
+		return block, err
+	}
+
+	return block, nil
+}
+
 // GetBlockHashes returns a list of hashes of all the blocks in the chain
 func (bc *Blockchain) GetBlockHashes() [][]byte {
 	var blocks [][]byte
diff --git a/server.go b/server.go
index 0cbf599..83b649c 100644
--- a/server.go
+++ b/server.go
@@ -22,13 +22,25 @@ type addr struct {
 	AddrList []string
 }
 
+type block struct {
+	AddrFrom string
+	Block    []byte
+}
+
 type getblocks struct {
 	AddrFrom string
 }
 
+type getdata struct {
+	AddrFrom string
+	Type     string
+	ID       []byte
+}
+
 type inv struct {
-	Type  string
-	Items [][]byte
+	AddrFrom string
+	Type     string
+	Items    [][]byte
 }
 
 type verack struct {
@@ -81,6 +93,14 @@ func sendAddr(address string) {
 	sendData(address, request)
 }
 
+func sendBlock(addr string, b *Block) {
+	data := block{nodeAddress, b.Serialize()}
+	payload := gobEncode(data)
+	request := append(commandToBytes("block"), payload...)
+
+	sendData(addr, request)
+}
+
 func sendData(addr string, data []byte) {
 	conn, err := net.Dial(protocol, addr)
 	if err != nil {
@@ -95,7 +115,7 @@ func sendData(addr string, data []byte) {
 }
 
 func sendInv(address, kind string, items [][]byte) {
-	inventory := inv{kind, items}
+	inventory := inv{nodeAddress, kind, items}
 	payload := gobEncode(inventory)
 	request := append(commandToBytes("inv"), payload...)
 
@@ -109,6 +129,13 @@ func sendGetBlocks(address string) {
 	sendData(address, request)
 }
 
+func sendGetData(address, kind string, id []byte) {
+	payload := gobEncode(getdata{nodeAddress, kind, id})
+	request := append(commandToBytes("getdata"), payload...)
+
+	sendData(address, request)
+}
+
 func sendVersion(addr string) {
 	payload := gobEncode(verzion{nodeVersion, nodeAddress})
 
@@ -141,7 +168,25 @@ func handleAddr(request []byte) {
 	requestBlocks()
 }
 
-func handleInv(request []byte) {
+func handleBlock(request []byte, bc *Blockchain) {
+	var buff bytes.Buffer
+	var payload block
+
+	buff.Write(request[commandLength:])
+	dec := gob.NewDecoder(&buff)
+	err := dec.Decode(&payload)
+	if err != nil {
+		log.Panic(err)
+	}
+
+	blockData := payload.Block
+	block := DeserializeBlock(blockData)
+
+	fmt.Println("Recevied a new block!")
+	fmt.Println(block)
+}
+
+func handleInv(request []byte, bc *Blockchain) {
 	var buff bytes.Buffer
 	var payload inv
 
@@ -152,7 +197,14 @@ func handleInv(request []byte) {
 		log.Panic(err)
 	}
 
-	fmt.Printf("Recevied %d %s\n", len(payload.Items), payload.Type)
+	fmt.Printf("Recevied inventory with %d %s\n", len(payload.Items), payload.Type)
+	blocks := bc.GetBlockHashes()
+
+	if len(blocks) < len(payload.Items) {
+		for _, blockHash := range invResponse.Items {
+			sendGetData(sourceNode, "block", blockHash)
+		}
+	}
 }
 
 func handleGetBlocks(request []byte, bc *Blockchain) {
@@ -170,6 +222,27 @@ func handleGetBlocks(request []byte, bc *Blockchain) {
 	sendInv(payload.AddrFrom, "blocks", blocks)
 }
 
+func handleGetData(request []byte, bc *Blockchain) {
+	var buff bytes.Buffer
+	var payload getdata
+
+	buff.Write(request[commandLength:])
+	dec := gob.NewDecoder(&buff)
+	err := dec.Decode(&payload)
+	if err != nil {
+		log.Panic(err)
+	}
+
+	if payload.Type == "block" {
+		block, err := bc.GetBlock([]byte(payload.ID))
+		if err != nil {
+			return
+		}
+
+		sendBlock(payload.AddrFrom, &block)
+	}
+}
+
 func handleVersion(request []byte) {
 	var buff bytes.Buffer
 	var payload verzion
@@ -197,10 +270,14 @@ func handleConnection(conn net.Conn, bc *Blockchain) {
 	switch command {
 	case "addr":
 		handleAddr(request)
+	case "block":
+		handleBlock(request, bc)
 	case "inv":
-		handleInv(request)
+		handleInv(request, bc)
 	case "getblocks":
 		handleGetBlocks(request, bc)
+	case "getdata":
+		handleGetData(request, bc)
 	case "version":
 		handleVersion(request)
 	case "verack":