core: governance interface should return correct DKG master public keys (#85)

2 files changed, 21 insertions, 24 deletions

diff --git a/core/governance.go b/core/governance.go
index 4b625c0fa..ea0036a0e 100644
--- a/core/governance.go
+++ b/core/governance.go
@@ -118,15 +118,7 @@ func (g *Governance) DKGComplaints(round uint64) []*dkgTypes.Complaint {
 
 func (g *Governance) DKGMasterPublicKeys(round uint64) []*dkgTypes.MasterPublicKey {
 	headHelper := g.GetHeadHelper()
-	var dkgMasterPKs []*dkgTypes.MasterPublicKey
-	for _, pk := range headHelper.DKGMasterPublicKeys(big.NewInt(int64(round))) {
-		x := new(dkgTypes.MasterPublicKey)
-		if err := rlp.DecodeBytes(pk, x); err != nil {
-			panic(err)
-		}
-		dkgMasterPKs = append(dkgMasterPKs, x)
-	}
-	return dkgMasterPKs
+	return headHelper.UniqueDKGMasterPublicKeys(big.NewInt(int64(round)))
 }
 
 func (g *Governance) IsDKGFinal(round uint64) bool {
diff --git a/core/vm/governance.go b/core/vm/governance.go
index f4c2f4d93..8c6f70a04 100644
--- a/core/vm/governance.go
+++ b/core/vm/governance.go
@@ -1608,6 +1608,25 @@ func (s *GovernanceStateHelper) DKGMasterPublicKeys(round *big.Int) [][]byte {
 func (s *GovernanceStateHelper) PushDKGMasterPublicKey(round *big.Int, mpk []byte) {
 	s.appendTo2DByteArray(big.NewInt(dkgMasterPublicKeysLoc), round, mpk)
 }
+func (s *GovernanceStateHelper) UniqueDKGMasterPublicKeys(round *big.Int) []*dkgTypes.MasterPublicKey {
+	// Prepare DKGMasterPublicKeys.
+	var dkgMasterPKs []*dkgTypes.MasterPublicKey
+	existence := make(map[coreTypes.NodeID]struct{})
+	for _, mpk := range s.DKGMasterPublicKeys(round) {
+		x := new(dkgTypes.MasterPublicKey)
+		if err := rlp.DecodeBytes(mpk, x); err != nil {
+			panic(err)
+		}
+
+		// Only the first DKG MPK submission is valid.
+		if _, exists := existence[x.ProposerID]; exists {
+			continue
+		}
+		existence[x.ProposerID] = struct{}{}
+		dkgMasterPKs = append(dkgMasterPKs, x)
+	}
+	return dkgMasterPKs
+}
 
 // bytes[][] public dkgComplaints;
 func (s *GovernanceStateHelper) DKGComplaints(round *big.Int) [][]byte {
@@ -2214,21 +2233,7 @@ func (g *GovernanceContract) proposeCRS(nextRound *big.Int, signedCRS []byte) ([
 	prevCRS := g.state.CRS(round)
 
 	// Prepare DKGMasterPublicKeys.
-	var dkgMasterPKs []*dkgTypes.MasterPublicKey
-	existence := make(map[coreTypes.NodeID]struct{})
-	for _, mpk := range g.state.DKGMasterPublicKeys(round) {
-		x := new(dkgTypes.MasterPublicKey)
-		if err := rlp.DecodeBytes(mpk, x); err != nil {
-			panic(err)
-		}
-
-		// Only the first DKG MPK submission is valid.
-		if _, exists := existence[x.ProposerID]; exists {
-			continue
-		}
-		existence[x.ProposerID] = struct{}{}
-		dkgMasterPKs = append(dkgMasterPKs, x)
-	}
+	dkgMasterPKs := g.state.UniqueDKGMasterPublicKeys(round)
 
 	// Prepare DKGComplaints.
 	var dkgComplaints []*dkgTypes.Complaint