dex: add method to get NumChains, NotarySet, DKGSet easily

6 files changed, 80 insertions, 43 deletions

diff --git a/dex/governance.go b/dex/governance.go
index 37985cec4..32c2d79af 100644
--- a/dex/governance.go
+++ b/dex/governance.go
@@ -12,12 +12,11 @@ import (
 	coreTypes "github.com/dexon-foundation/dexon-consensus-core/core/types"
 
 	"github.com/dexon-foundation/dexon/common"
-	"github.com/dexon-foundation/dexon/core"
 	"github.com/dexon-foundation/dexon/core/types"
 	"github.com/dexon-foundation/dexon/core/vm"
 	"github.com/dexon-foundation/dexon/crypto"
-	"github.com/dexon-foundation/dexon/event"
 	"github.com/dexon-foundation/dexon/log"
+	"github.com/dexon-foundation/dexon/p2p/discover"
 	"github.com/dexon-foundation/dexon/params"
 	"github.com/dexon-foundation/dexon/rlp"
 	"github.com/dexon-foundation/dexon/rpc"
@@ -280,19 +279,64 @@ func (d *DexconGovernance) IsDKGFinal(round uint64) bool {
 	return count >= threshold
 }
 
-// TODO(sonic): finish these
-func (d *DexconGovernance) GetChainNum(uint64) uint32 {
-	return 3
+func (d *DexconGovernance) GetNumChains(round uint64) uint32 {
+	return d.Configuration(round).NumChains
 }
 
-func (d *DexconGovernance) GetNotarySet(uint32, uint64) map[string]struct{} {
-	return nil
-}
+func (d *DexconGovernance) NotarySet(chainID uint32, round uint64) map[string]struct{} {
+	id2Key := map[coreTypes.NodeID]coreCrypto.PublicKey{}
+
+	nodeSet := coreTypes.NewNodeSet()
+	for _, key := range d.NodeSet(round) {
+		id := coreTypes.NewNodeID(key)
+		id2Key[id] = key
+		nodeSet.Add(id)
+	}
+
+	cfg := d.Configuration(round)
+	crs := d.CRS(round)
 
-func (d *DexconGovernance) GetDKGSet(uint64) map[string]struct{} {
-	return nil
+	notarySet := nodeSet.GetSubSet(
+		int(cfg.NotarySetSize), coreTypes.NewNotarySetTarget(crs, chainID))
+
+	r := map[string]struct{}{}
+	for id := range notarySet {
+		compressed := id2Key[id]
+		// 33 bytes pubkey to 65 bytes pubkey
+		key, err := crypto.DecompressPubkey(compressed.Bytes())
+		if err != nil {
+			continue
+		}
+		r[discover.PubkeyID(key).String()] = struct{}{}
+	}
+	return r
 }
 
-func (d *DexconGovernance) SubscribeNewCRSEvent(ch chan core.NewCRSEvent) event.Subscription {
-	return nil
+func (d *DexconGovernance) DKGSet(round uint64) map[string]struct{} {
+	id2Key := map[coreTypes.NodeID]coreCrypto.PublicKey{}
+
+	nodeSet := coreTypes.NewNodeSet()
+	for _, key := range d.NodeSet(round) {
+		id := coreTypes.NewNodeID(key)
+		id2Key[id] = key
+		nodeSet.Add(id)
+	}
+
+	cfg := d.Configuration(round)
+	crs := d.CRS(round)
+
+	dkgSet := nodeSet.GetSubSet(
+		int(cfg.DKGSetSize), coreTypes.NewDKGSetTarget(crs))
+
+	r := map[string]struct{}{}
+	for id := range dkgSet {
+		compressed := id2Key[id]
+		// 33 bytes pubkey to 65 bytes pubkey
+		key, err := crypto.DecompressPubkey(compressed.Bytes())
+		if err != nil {
+			continue
+		}
+		r[discover.PubkeyID(key).String()] = struct{}{}
+	}
+	return r
 }
diff --git a/dex/handler.go b/dex/handler.go
index e013b9722..5613efdd8 100644
--- a/dex/handler.go
+++ b/dex/handler.go
@@ -26,7 +26,7 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/dexon-foundation/dexon-consensus-core/core/crypto"
+	coreCrypto "github.com/dexon-foundation/dexon-consensus-core/core/crypto"
 	coreTypes "github.com/dexon-foundation/dexon-consensus-core/core/types"
 
 	"github.com/dexon-foundation/dexon/common"
@@ -240,7 +240,7 @@ func (pm *ProtocolManager) Start(srvr p2pServer, maxPeers int) {
 
 	// run the peer set loop
 	pm.crsCh = make(chan core.NewCRSEvent)
-	pm.crsSub = pm.gov.SubscribeNewCRSEvent(pm.crsCh)
+	// pm.crsSub = pm.gov.SubscribeNewCRSEvent(pm.crsCh)
 	go pm.peerSetLoop()
 
 	// start sync handlers
@@ -680,7 +680,6 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 	case msg.Code == LatticeBlockMsg:
 		var rb rlpLatticeBlock
 		if err := msg.Decode(&rb); err != nil {
-			fmt.Println("decode lattice block error", err)
 			return errResp(ErrDecode, "msg %v: %v", msg, err)
 		}
 		pm.receiveCh <- fromRLPLatticeBlock(&rb)
@@ -815,7 +814,7 @@ func (pm *ProtocolManager) BroadcastLatticeBlock(block *coreTypes.Block) {
 
 // TODO(sonic): try to reduce traffic
 func (pm *ProtocolManager) SendDKGPrivateShare(
-	pub crypto.PublicKey, privateShare *coreTypes.DKGPrivateShare) {
+	pub coreCrypto.PublicKey, privateShare *coreTypes.DKGPrivateShare) {
 	id := discover.MustBytesID(pub.Bytes()[1:])
 	if p := pm.peers.Peer(id.String()); p != nil {
 		p.AsyncSendDKGPrivateShare(privateShare)
diff --git a/dex/helper_test.go b/dex/helper_test.go
index fc8053774..21803ed88 100644
--- a/dex/helper_test.go
+++ b/dex/helper_test.go
@@ -182,25 +182,21 @@ func newTestTransaction(from *ecdsa.PrivateKey, nonce uint64, datasize int) *typ
 
 // testGovernance is a fake, helper governance for testing purposes
 type testGovernance struct {
-	getChainNumFunc  func(uint64) uint32
-	getNotarySetFunc func(uint32, uint64) map[string]struct{}
-	getDKGSetFunc    func(uint64) map[string]struct{}
+	numChainsFunc func(uint64) uint32
+	notarySetFunc func(uint32, uint64) map[string]struct{}
+	dkgSetFunc    func(uint64) map[string]struct{}
 }
 
-func (g *testGovernance) GetChainNum(round uint64) uint32 {
-	return g.getChainNumFunc(round)
+func (g *testGovernance) GetNumChains(round uint64) uint32 {
+	return g.numChainsFunc(round)
 }
 
-func (g *testGovernance) GetNotarySet(chainID uint32, round uint64) map[string]struct{} {
-	return g.getNotarySetFunc(chainID, round)
+func (g *testGovernance) NotarySet(chainID uint32, round uint64) map[string]struct{} {
+	return g.notarySetFunc(chainID, round)
 }
 
-func (g *testGovernance) GetDKGSet(round uint64) map[string]struct{} {
-	return g.getDKGSetFunc(round)
-}
-
-func (g *testGovernance) SubscribeNewCRSEvent(ch chan core.NewCRSEvent) event.Subscription {
-	return nil
+func (g *testGovernance) DKGSet(round uint64) map[string]struct{} {
+	return g.dkgSetFunc(round)
 }
 
 // testPeer is a simulated peer to allow testing direct network calls.
diff --git a/dex/peer.go b/dex/peer.go
index e7c4f5d53..44f61354c 100644
--- a/dex/peer.go
+++ b/dex/peer.go
@@ -785,8 +785,8 @@ func (ps *peerSet) BuildNotaryConn(round uint64) {
 	ps.notaryHistory[round] = struct{}{}
 
 	selfID := ps.srvr.Self().ID.String()
-	for chainID := uint32(0); chainID < ps.gov.GetChainNum(round); chainID++ {
-		s := ps.gov.GetNotarySet(chainID, round)
+	for chainID := uint32(0); chainID < ps.gov.GetNumChains(round); chainID++ {
+		s := ps.gov.NotarySet(chainID, round)
 
 		// not in notary set, add group
 		if _, ok := s[selfID]; !ok {
@@ -825,8 +825,8 @@ func (ps *peerSet) ForgetNotaryConn(round uint64) {
 
 func (ps *peerSet) forgetNotaryConn(round uint64) {
 	selfID := ps.srvr.Self().ID.String()
-	for chainID := uint32(0); chainID < ps.gov.GetChainNum(round); chainID++ {
-		s := ps.gov.GetNotarySet(chainID, round)
+	for chainID := uint32(0); chainID < ps.gov.GetNumChains(round); chainID++ {
+		s := ps.gov.NotarySet(chainID, round)
 		if _, ok := s[selfID]; !ok {
 			ps.srvr.RemoveGroup(notarySetName(chainID, round))
 			continue
@@ -852,7 +852,7 @@ func (ps *peerSet) BuildDKGConn(round uint64) {
 	ps.lock.Lock()
 	defer ps.lock.Unlock()
 	selfID := ps.srvr.Self().ID.String()
-	s := ps.gov.GetDKGSet(round)
+	s := ps.gov.DKGSet(round)
 	if _, ok := s[selfID]; !ok {
 		return
 	}
@@ -882,7 +882,7 @@ func (ps *peerSet) ForgetDKGConn(round uint64) {
 
 func (ps *peerSet) forgetDKGConn(round uint64) {
 	selfID := ps.srvr.Self().ID.String()
-	s := ps.gov.GetDKGSet(round)
+	s := ps.gov.DKGSet(round)
 	if _, ok := s[selfID]; !ok {
 		return
 	}
diff --git a/dex/peer_test.go b/dex/peer_test.go
index 6e539e078..d0262e201 100644
--- a/dex/peer_test.go
+++ b/dex/peer_test.go
@@ -16,7 +16,7 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 	table := newNodeTable()
 
 	gov := &testGovernance{
-		getChainNumFunc: func(uint64) uint32 {
+		numChainsFunc: func(uint64) uint32 {
 			return 3
 		},
 	}
@@ -37,7 +37,7 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 		[]enode.ID{nodeID(0), nodeID(2), nodeID(6)},
 	}
 
-	gov.getNotarySetFunc = func(cid uint32, round uint64) map[string]struct{} {
+	gov.notarySetFunc = func(cid uint32, round uint64) map[string]struct{} {
 		m := map[uint64][][]enode.ID{
 			10: round10,
 			11: round11,
@@ -317,7 +317,7 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 
 	gov := &testGovernance{}
 
-	gov.getDKGSetFunc = func(round uint64) map[string]struct{} {
+	gov.dkgSetFunc = func(round uint64) map[string]struct{} {
 		m := map[uint64][]enode.ID{
 			10: []enode.ID{nodeID(0), nodeID(1), nodeID(2)},
 			11: []enode.ID{nodeID(1), nodeID(2), nodeID(5)},
diff --git a/dex/protocol.go b/dex/protocol.go
index 94241104b..3507965ae 100644
--- a/dex/protocol.go
+++ b/dex/protocol.go
@@ -124,13 +124,11 @@ type txPool interface {
 }
 
 type governance interface {
-	GetChainNum(uint64) uint32
+	GetNumChains(uint64) uint32
 
-	GetNotarySet(uint32, uint64) map[string]struct{}
+	NotarySet(uint32, uint64) map[string]struct{}
 
-	GetDKGSet(uint64) map[string]struct{}
-
-	SubscribeNewCRSEvent(ch chan core.NewCRSEvent) event.Subscription
+	DKGSet(uint64) map[string]struct{}
 }
 
 type p2pServer interface {