Merge pull request #74 from getamis/refactor-eth-handler

Refactor protocol manager

Author: tailingchen

consensus/consensus.go

@@ -107,14 +107,8 @@ type PoW interface {
 type Istanbul interface {
 	Engine
 
-	// Add a peer
-	AddPeer(peerID string, publicKey *ecdsa.PublicKey) error
-
-	// Remove a peer
-	RemovePeer(peerID string) error
-
 	// Handle a message from peer
-	HandleMsg(peerID string, data []byte) error
+	HandleMsg(pubKey *ecdsa.PublicKey, data []byte) error
 
 	// Receive new chain head block
 	NewChainHead(block *types.Block)

consensus/istanbul/backends/simple/backend.go

@@ -36,7 +36,6 @@ import (
 func New(config *istanbul.Config, eventMux *event.TypeMux, privateKey *ecdsa.PrivateKey, db ethdb.Database) consensus.Istanbul {
 	backend := &simpleBackend{
 		config:           config,
-		peerSet:          newPeerSet(),
 		eventMux:         eventMux,
 		istanbulEventMux: new(event.TypeMux),
 		privateKey:       privateKey,
@@ -52,7 +51,6 @@ func New(config *istanbul.Config, eventMux *event.TypeMux, privateKey *ecdsa.Pri
 
 type simpleBackend struct {
 	config           *istanbul.Config
-	peerSet          *peerSet
 	valSet           istanbul.ValidatorSet
 	eventMux         *event.TypeMux
 	istanbulEventMux *event.TypeMux
@@ -83,33 +81,27 @@ func (sb *simpleBackend) Validators() istanbul.ValidatorSet {
 }
 
 func (sb *simpleBackend) Send(payload []byte, target common.Address) error {
-	peer := sb.peerSet.GetByAddress(target)
-	if peer == nil {
-		return errInvalidPeer
-	}
-
 	go sb.eventMux.Post(istanbul.ConsensusDataEvent{
-		PeerID: peer.ID(),
+		Target: target,
 		Data:   payload,
 	})
 	return nil
 }
 
 // Broadcast implements istanbul.Backend.Send
 func (sb *simpleBackend) Broadcast(payload []byte) error {
-	istanbulMsg := istanbul.MessageEvent{
-		Payload: payload,
-	}
-
-	// send to self
-	go sb.istanbulEventMux.Post(istanbulMsg)
-
-	// send to other peers
-	for _, peer := range sb.peerSet.List() {
-		go sb.eventMux.Post(istanbul.ConsensusDataEvent{
-			PeerID: peer.ID(),
-			Data:   payload,
-		})
+	for _, val := range sb.valSet.List() {
+		if val.Address() == sb.Address() {
+			// send to self
+			pbftMsg := istanbul.MessageEvent{
+				Payload: payload,
+			}
+			go sb.istanbulEventMux.Post(pbftMsg)
+
+		} else {
+			// send to other peers
+			sb.Send(payload, val.Address())
+		}
 	}
 	return nil
 }

consensus/istanbul/backends/simple/engine.go

@@ -30,6 +30,7 @@ import (
 	"github.com/ethereum/go-ethereum/consensus/istanbul/validator"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/crypto/sha3"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/rpc"
@@ -359,34 +360,12 @@ func (sb *simpleBackend) APIs(chain consensus.ChainReader) []rpc.API {
 	}}
 }
 
-// AddPeer implements consensus.Istanbul.AddPeer
-func (sb *simpleBackend) AddPeer(peerID string, publicKey *ecdsa.PublicKey) error {
-	peer := newPeer(peerID, publicKey)
-	// check is validator
-	if _, val := sb.valSet.GetByAddress(peer.Address()); val != nil {
-		// add to peer set
-		sb.peerSet.Add(peer)
-	}
-	return nil
-}
-
-// RemovePeer implements consensus.Istanbul.RemovePeer
-func (sb *simpleBackend) RemovePeer(peerID string) error {
-	sb.peerSet.Remove(peerID)
-	return nil
-}
-
 // HandleMsg implements consensus.Istanbul.HandleMsg
-func (sb *simpleBackend) HandleMsg(peerID string, data []byte) error {
-	peer := sb.peerSet.Get(peerID)
-	if peer == nil {
-		sb.logger.Error("Not in peer set", "peerID", peerID)
-		return errInvalidPeer
-	}
-
-	if _, val := sb.valSet.GetByAddress(peer.Address()); val == nil {
-		sb.logger.Error("Not in validator set", "peerAddr", peer.Address())
-		return errInvalidPeer
+func (sb *simpleBackend) HandleMsg(pubKey *ecdsa.PublicKey, data []byte) error {
+	addr := crypto.PubkeyToAddress(*pubKey)
+	if _, val := sb.valSet.GetByAddress(addr); val == nil {
+		sb.logger.Error("Not in validator set", "addr", addr)
+		return istanbul.ErrUnauthorizedAddress
 	}
 
 	go sb.istanbulEventMux.Post(istanbul.MessageEvent{