From ccd972806908e194952e0b86d1df0c063d2faf2d Mon Sep 17 00:00:00 2001
From: Leo Chen <leo@harmony.one>
Date: Fri, 18 Jan 2019 07:18:36 +0000
Subject: [PATCH] support the peerstore in beacon chain node

Signed-off-by: Leo Chen <leo@harmony.one>
---
 internal/beaconchain/libs/beaconchain.go | 15 +++++++++---
 internal/newnode/newnode.go              | 29 ++++++++++++++++++++----
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/internal/beaconchain/libs/beaconchain.go b/internal/beaconchain/libs/beaconchain.go
index b74c29f87..40ce153d9 100644
--- a/internal/beaconchain/libs/beaconchain.go
+++ b/internal/beaconchain/libs/beaconchain.go
@@ -18,6 +18,7 @@ import (
 	"github.com/harmony-one/harmony/p2p"
 	"github.com/harmony-one/harmony/p2p/host"
 	"github.com/harmony-one/harmony/p2p/p2pimpl"
+	peer "github.com/libp2p/go-libp2p-peer"
 )
 
 //BCState keeps track of the state the beaconchain is in
@@ -49,6 +50,7 @@ type BeaconChain struct {
 	state          BCState
 	rpcServer      *beaconchain.Server
 	Peer           p2p.Peer
+	Self           p2p.Peer // self Peer
 }
 
 //SaveFile is to store the file in which beaconchain info will be stored.
@@ -95,7 +97,8 @@ func New(numShards int, ip, port string) *BeaconChain {
 	bc := BeaconChain{}
 	bc.log = log.New()
 	bc.PubKey = generateBCKey()
-	bc.host = p2pimpl.NewHost(p2p.Peer{IP: ip, Port: port})
+	bc.Self = p2p.Peer{IP: ip, Port: port}
+	bc.host, _ = p2pimpl.NewHost(&bc.Self)
 	bcinfo := &BCInfo{NumberOfShards: numShards, NumberOfNodesAdded: 0,
 		IP:             ip,
 		Port:           port,
@@ -114,9 +117,10 @@ func generateBCKey() kyber.Point {
 //AcceptNodeInfo deserializes node information received via beaconchain handler
 func (bc *BeaconChain) AcceptNodeInfo(b []byte) *node.Info {
 	Node := bcconn.DeserializeNodeInfo(b)
-	bc.log.Info("New Node Connection", "IP", Node.IP, "Port", Node.Port)
-	bc.Peer = p2p.Peer{IP: Node.IP, Port: Node.Port}
+	bc.log.Info("New Node Connection", "IP", Node.IP, "Port", Node.Port, "PeerID", Node.PeerID)
+	bc.Peer = p2p.Peer{IP: Node.IP, Port: Node.Port, PeerID: Node.PeerID}
 	bc.host.AddPeer(&bc.Peer)
+
 	bc.BCInfo.NumberOfNodesAdded = bc.BCInfo.NumberOfNodesAdded + 1
 	shardNum, isLeader := utils.AllocateShard(bc.BCInfo.NumberOfNodesAdded, bc.BCInfo.NumberOfShards)
 	if isLeader {
@@ -189,3 +193,8 @@ func BCItoBC(bci *BCInfo) *BeaconChain {
 func SetSaveFile(path string) {
 	SaveFile = path
 }
+
+//GetID return ID
+func (bc *BeaconChain) GetID() peer.ID {
+	return bc.host.GetID()
+}
diff --git a/internal/newnode/newnode.go b/internal/newnode/newnode.go
index 2cc6d67be..e7b553871 100644
--- a/internal/newnode/newnode.go
+++ b/internal/newnode/newnode.go
@@ -18,6 +18,8 @@ import (
 	"github.com/harmony-one/harmony/p2p"
 	"github.com/harmony-one/harmony/p2p/host"
 	"github.com/harmony-one/harmony/p2p/p2pimpl"
+
+	multiaddr "github.com/multiformats/go-multiaddr"
 )
 
 //NewNode is ther struct for a candidate node
@@ -40,12 +42,17 @@ type NewNode struct {
 func New(ip string, port string) *NewNode {
 	priKey, pubKey := utils.GenKey(ip, port)
 	var node NewNode
+	var err error
 	node.PubK = pubKey
 	node.priK = priKey
 	node.Self = p2p.Peer{IP: ip, Port: port, PubKey: pubKey, ValidatorID: -1}
 	node.log = log.New()
 	node.SetInfo = false
-	node.host = p2pimpl.NewHost(node.Self)
+	node.host, err = p2pimpl.NewHost(&node.Self)
+	if err != nil {
+		node.log.Error("failed to create new host", "msg", err)
+		return nil
+	}
 	node.Leaders = map[uint32]p2p.Peer{}
 	return &node
 }
@@ -73,7 +80,7 @@ func (node *NewNode) requestBeaconChain(BCPeer p2p.Peer) (err error) {
 	if err != nil {
 		node.log.Error("Could not Marshall public key into binary")
 	}
-	nodeInfo := &proto_node.Info{IP: node.Self.IP, Port: node.Self.Port, PubKey: pubk}
+	nodeInfo := &proto_node.Info{IP: node.Self.IP, Port: node.Self.Port, PubKey: pubk, PeerID: node.host.GetID()}
 	msg := bcconn.SerializeNodeInfo(nodeInfo)
 	msgToSend := proto_identity.ConstructIdentityMessage(proto_identity.Register, msg)
 	gotShardInfo := false
@@ -108,9 +115,18 @@ func (node *NewNode) processShardInfo(msgPayload []byte) bool {
 	leaders := leadersInfo.Leaders
 	shardNum, isLeader := utils.AllocateShard(leadersInfo.NumberOfNodesAdded, leadersInfo.NumberOfShards)
 	for n, v := range leaders {
-		leaderPeer := p2p.Peer{IP: v.IP, Port: v.Port}
+		leaderPeer := p2p.Peer{IP: v.IP, Port: v.Port, PeerID: v.PeerID}
+
+		addr := fmt.Sprintf("/ip4/%s/tcp/%s", leaderPeer.IP, leaderPeer.Port)
+		targetAddr, err := multiaddr.NewMultiaddr(addr)
+		if err != nil {
+			log.Error("processShardInfo NewMultiaddr error", "error", err)
+			return false
+		}
+		leaderPeer.Addrs = append(leaderPeer.Addrs, targetAddr)
+
 		leaderPeer.PubKey = crypto.Ed25519Curve.Point()
-		err := leaderPeer.PubKey.UnmarshalBinary(v.PubKey[:])
+		err = leaderPeer.PubKey.UnmarshalBinary(v.PubKey[:])
 		if err != nil {
 			node.log.Error("Could not unmarshall leaders public key from binary to kyber.point")
 		}
@@ -144,3 +160,8 @@ func (node *NewNode) GetClientPeer() *p2p.Peer {
 func (node *NewNode) GetSelfPeer() p2p.Peer {
 	return node.Self
 }
+
+// AddPeer add new peer for newnode
+func (node *NewNode) AddPeer(p *p2p.Peer) error {
+	return node.host.AddPeer(p)
+}