core: add p2p whitelist for consortium

3 files changed, 117 insertions, 4 deletions

diff --git a/dex/handler.go b/dex/handler.go
index ca26bc6e8..4c1ed9a36 100644
--- a/dex/handler.go
+++ b/dex/handler.go
@@ -41,6 +41,7 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"math/big"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -55,6 +56,7 @@ import (
 	"github.com/tangerine-network/go-tangerine/consensus"
 	"github.com/tangerine-network/go-tangerine/core"
 	"github.com/tangerine-network/go-tangerine/core/types"
+	"github.com/tangerine-network/go-tangerine/core/vm"
 	"github.com/tangerine-network/go-tangerine/crypto"
 	dexDB "github.com/tangerine-network/go-tangerine/dex/db"
 	"github.com/tangerine-network/go-tangerine/dex/downloader"
@@ -89,6 +91,7 @@ const (
 
 	maxAgreementResultBroadcast = 3
 	maxFinalizedBlockBroadcast  = 3
+	checkPeerDuration           = 10 * time.Minute
 )
 
 // errIncompatibleConfig is returned if the requested protocols and configs are
@@ -349,6 +352,7 @@ func (pm *ProtocolManager) ReportBadPeerChan() chan<- interface{} {
 }
 
 func (pm *ProtocolManager) badPeerWatchLoop() {
+	go pm.checkPeerInWhitelist(pm.reportBadPeerChan)
 	for {
 		select {
 		case id := <-pm.reportBadPeerChan:
@@ -360,13 +364,42 @@ func (pm *ProtocolManager) badPeerWatchLoop() {
 	}
 }
 
+func (pm *ProtocolManager) checkPeerInWhitelist(reportBadPeerChan chan<- interface{}) {
+	for {
+		for id, p := range pm.peers.peers {
+			if !pm.inWhitelist(p) {
+				reportBadPeerChan <- id
+			}
+		}
+		time.Sleep(checkPeerDuration)
+	}
+}
+
 func (pm *ProtocolManager) newPeer(pv int, p *p2p.Peer, rw p2p.MsgReadWriter) *peer {
 	return newPeer(pv, p, newMeteredMsgWriter(rw))
 }
 
+func (pm *ProtocolManager) inWhitelist(p *peer) bool {
+	state, err := pm.blockchain.State()
+	if err != nil {
+		p.Log().Debug("get state fail in checking whitelist", "err", err)
+		return false
+	}
+	govState := vm.GovernanceState{StateDB: state}
+	if !govState.IsConsortium() {
+		return true
+	}
+	address := crypto.PubkeyToAddress(*p.Node().Pubkey())
+	return govState.WhitelistOffsetByAddress(address).Cmp(big.NewInt(0)) >= 0
+}
+
 // handle is the callback invoked to manage the life cycle of an eth peer. When
 // this function terminates, the peer is disconnected.
 func (pm *ProtocolManager) handle(p *peer) error {
+	if !pm.inWhitelist(p) {
+		p.Log().Debug("Peer disconnect: permission denied", "name", p.Name())
+		return p2p.DiscPermissionDenied
+	}
 	// Ignore maxPeers if this is a trusted peer
 	if pm.peers.Len() >= pm.maxPeers && !p.Peer.Info().Network.Trusted {
 		return p2p.DiscTooManyPeers
diff --git a/dex/handler_test.go b/dex/handler_test.go
index 80f36ae9e..f276e1e4a 100644
--- a/dex/handler_test.go
+++ b/dex/handler_test.go
@@ -17,9 +17,12 @@
 package dex
 
 import (
+	"fmt"
+	"io/ioutil"
 	"math"
 	"math/big"
 	"math/rand"
+	"net"
 	"testing"
 
 	"github.com/tangerine-network/go-tangerine/common"
@@ -30,6 +33,7 @@ import (
 	"github.com/tangerine-network/go-tangerine/dex/downloader"
 	"github.com/tangerine-network/go-tangerine/ethdb"
 	"github.com/tangerine-network/go-tangerine/p2p"
+	"github.com/tangerine-network/go-tangerine/p2p/enode"
 	"github.com/tangerine-network/go-tangerine/params"
 )
 
@@ -52,7 +56,7 @@ func TestProtocolCompatibility(t *testing.T) {
 	for i, tt := range tests {
 		ProtocolVersions = []uint{tt.version}
 
-		pm, _, err := newTestProtocolManager(tt.mode, 0, nil, nil)
+		pm, _, err := newTestProtocolManager(tt.mode, 0, nil, nil, params.TestChainConfig)
 		if pm != nil {
 			defer pm.Stop()
 		}
@@ -444,3 +448,74 @@ func testGetReceipt(t *testing.T, protocol int) {
 		t.Errorf("receipts mismatch: %v", err)
 	}
 }
+
+func TestHandlerWithConsoritum62(t *testing.T) { testHandlerWithConsoritum(t, 62) }
+func TestHandlerWithConsoritum63(t *testing.T) { testHandlerWithConsoritum(t, 63) }
+
+func testHandlerWithConsoritum(t *testing.T, version int) {
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		t.Errorf("gen key fail %v", err)
+	}
+
+	// set config and create protocol manager
+	config := params.NewTestChainConig()
+	config.ChainID = big.NewInt(int64(DefaultConfig.NetworkId))
+	config.Dexcon = params.NewTestDexonConfig()
+	config.Dexcon.IsConsortium = true
+	pm, _, err := newTestProtocolManager(downloader.FullSync, 0, nil, nil, config)
+	defer pm.Stop()
+	if err != nil {
+		t.Fatalf("Failed to create protocol manager: %v", err)
+	}
+
+	// create peer
+	pipenet1, pipenet2 := p2p.MsgPipe()
+	defer func() {
+		pipenet1.Close()
+		pipenet2.Close()
+	}()
+	node := enode.NewV4(&key.PublicKey, net.IP{}, 0, 0)
+	peer := pm.newPeer(version, p2p.NewPeerWithEnode(node, "handlerTest", nil), pipenet1)
+
+	// try to call handle, and should get permission denied error
+	if err := pm.handle(peer); err != p2p.DiscPermissionDenied {
+		t.Errorf("Expect get DiscPermissionDenied, but get %v", err)
+	}
+
+	// add address to whitelist, and create new pm
+	address := crypto.PubkeyToAddress(key.PublicKey)
+	config.Dexcon.AddressWhitelist = []common.Address{address}
+	pm2, _, err := newTestProtocolManager(downloader.FullSync, 0, nil, nil, config)
+	defer pm2.Stop()
+
+	if err != nil {
+		t.Fatalf("Failed to create protocol manager: %v", err)
+	}
+	handleErr := make(chan error)
+	go func() {
+		handleErr <- pm2.handle(peer)
+	}()
+	// do the handshake
+	msg, err := pipenet2.ReadMsg()
+	ioutil.ReadAll(msg.Payload)
+	p2p.Send(pipenet2, 0,
+		statusData{
+			ProtocolVersion: uint32(version),
+			NetworkId:       config.ChainID.Uint64(),
+			Number:          0,
+			CurrentBlock:    common.Hash{},
+			GenesisBlock:    pm2.blockchain.Genesis().Hash(),
+		},
+	)
+	// send status code to terminate the handleMsg loop
+	p2p.Send(pipenet2, StatusMsg, struct{}{})
+	err = <-handleErr
+	expectError := fmt.Errorf("%v - %v",
+		errorToString[ErrExtraStatusMsg],
+		"uncontrolled status message",
+	)
+	if err.Error() != expectError.Error() {
+		t.Errorf("err not match, expect: %s, but got: %s", expectError, err)
+	}
+}
diff --git a/dex/helper_test.go b/dex/helper_test.go
index 6e2cd3d77..5e171fef0 100644
--- a/dex/helper_test.go
+++ b/dex/helper_test.go
@@ -110,13 +110,18 @@ func (a *testApp) SubscribeNewFinalizedBlockEvent(
 // newTestProtocolManager creates a new protocol manager for testing purposes,
 // with the given number of blocks already known, and potential notification
 // channels for different events.
-func newTestProtocolManager(mode downloader.SyncMode, blocks int, generator func(int, *core.BlockGen), newtx chan<- []*types.Transaction) (*ProtocolManager, *ethdb.MemDatabase, error) {
+func newTestProtocolManager(mode downloader.SyncMode,
+	blocks int,
+	generator func(int, *core.BlockGen),
+	newtx chan<- []*types.Transaction,
+	chainConfig *params.ChainConfig,
+) (*ProtocolManager, *ethdb.MemDatabase, error) {
 	var (
 		evmux  = new(event.TypeMux)
 		engine = ethash.NewFaker()
 		db     = ethdb.NewMemDatabase()
 		gspec  = &core.Genesis{
-			Config: params.TestChainConfig,
+			Config: chainConfig,
 			Alloc:  core.GenesisAlloc{testBank: {Balance: big.NewInt(1000000), Staked: big.NewInt(0)}},
 		}
 		genesis       = gspec.MustCommit(db)
@@ -150,7 +155,7 @@ func newTestProtocolManager(mode downloader.SyncMode, blocks int, generator func
 // channels for different events. In case of an error, the constructor force-
 // fails the test.
 func newTestProtocolManagerMust(t *testing.T, mode downloader.SyncMode, blocks int, generator func(int, *core.BlockGen), newtx chan<- []*types.Transaction) (*ProtocolManager, *ethdb.MemDatabase) {
-	pm, db, err := newTestProtocolManager(mode, blocks, generator, newtx)
+	pm, db, err := newTestProtocolManager(mode, blocks, generator, newtx, params.TestChainConfig)
 	if err != nil {
 		t.Fatalf("Failed to create protocol manager: %v", err)
 	}