core: add equality checker for dkg related structure (#237)

Besides adding equality, also renaming those fields.
- PublicKeyShares.shares -> shareCaches
- PublicKeyShares.shareIndex -> shareCacheIndex
- rlpPublicKeyShares.Shares -> ShareCaches
- rlpPublicKeyShares.ShareIndexK -> ShareCacheIndexK
- rlpPublicKeyShares.ShareIndexV -> ShareCahceIndexV

4 files changed, 379 insertions, 26 deletions

diff --git a/core/crypto/dkg/dkg.go b/core/crypto/dkg/dkg.go
index 066b070..184a94f 100644
--- a/core/crypto/dkg/dkg.go
+++ b/core/crypto/dkg/dkg.go
@@ -101,30 +101,84 @@ type PrivateKeyShares struct {
 	masterPrivateKey []bls.SecretKey
 }
 
+// Equal check equality between two PrivateKeyShares instances.
+func (prvs *PrivateKeyShares) Equal(other *PrivateKeyShares) bool {
+	// Check shares.
+	if len(prvs.shareIndex) != len(other.shareIndex) {
+		return false
+	}
+	for dID, idx := range prvs.shareIndex {
+		otherIdx, exists := other.shareIndex[dID]
+		if !exists {
+			return false
+		}
+		if !prvs.shares[idx].privateKey.IsEqual(
+			&other.shares[otherIdx].privateKey) {
+			return false
+		}
+	}
+	// Check master private keys.
+	if len(prvs.masterPrivateKey) != len(other.masterPrivateKey) {
+		return false
+	}
+	for idx, m := range prvs.masterPrivateKey {
+		if m.GetHexString() != other.masterPrivateKey[idx].GetHexString() {
+			return false
+		}
+	}
+	return true
+}
+
 // PublicKeyShares represents a public key shares for DKG protocol.
 type PublicKeyShares struct {
-	shares          []PublicKey
-	shareIndex      map[ID]int
+	shareCaches     []PublicKey
+	shareCacheIndex map[ID]int
 	masterPublicKey []bls.PublicKey
 }
 
 type rlpPublicKeyShares struct {
-	Shares          [][]byte
-	ShareIndexK     [][]byte
-	ShareIndexV     []uint32
-	MasterPublicKey [][]byte
+	ShareCaches      [][]byte
+	ShareCacheIndexK [][]byte
+	ShareCacheIndexV []uint32
+	MasterPublicKey  [][]byte
+}
+
+// Equal checks equality of two PublicKeyShares instance.
+func (pubs *PublicKeyShares) Equal(other *PublicKeyShares) bool {
+	// Check shares.
+	for dID, idx := range pubs.shareCacheIndex {
+		otherIdx, exists := other.shareCacheIndex[dID]
+		if !exists {
+			continue
+		}
+		if !pubs.shareCaches[idx].publicKey.IsEqual(
+			&other.shareCaches[otherIdx].publicKey) {
+			return false
+		}
+	}
+	// Check master public keys.
+	if len(pubs.masterPublicKey) != len(other.masterPublicKey) {
+		return false
+	}
+	for idx, m := range pubs.masterPublicKey {
+		if m.GetHexString() != other.masterPublicKey[idx].GetHexString() {
+			return false
+		}
+	}
+	return true
 }
 
 // EncodeRLP implements rlp.Encoder
 func (pubs *PublicKeyShares) EncodeRLP(w io.Writer) error {
 	var rps rlpPublicKeyShares
-	for _, share := range pubs.shares {
-		rps.Shares = append(rps.Shares, share.Serialize())
+	for _, share := range pubs.shareCaches {
+		rps.ShareCaches = append(rps.ShareCaches, share.Serialize())
 	}
 
-	for id, v := range pubs.shareIndex {
-		rps.ShareIndexK = append(rps.ShareIndexK, id.GetLittleEndian())
-		rps.ShareIndexV = append(rps.ShareIndexV, uint32(v))
+	for id, v := range pubs.shareCacheIndex {
+		rps.ShareCacheIndexK = append(
+			rps.ShareCacheIndexK, id.GetLittleEndian())
+		rps.ShareCacheIndexV = append(rps.ShareCacheIndexV, uint32(v))
 	}
 
 	for _, m := range pubs.masterPublicKey {
@@ -141,25 +195,25 @@ func (pubs *PublicKeyShares) DecodeRLP(s *rlp.Stream) error {
 		return err
 	}
 
-	if len(dec.ShareIndexK) != len(dec.ShareIndexV) {
+	if len(dec.ShareCacheIndexK) != len(dec.ShareCacheIndexV) {
 		return fmt.Errorf("invalid shareIndex")
 	}
 
 	ps := NewEmptyPublicKeyShares()
-	for _, share := range dec.Shares {
+	for _, share := range dec.ShareCaches {
 		var publicKey PublicKey
 		if err := publicKey.Deserialize(share); err != nil {
 			return err
 		}
-		ps.shares = append(ps.shares, publicKey)
+		ps.shareCaches = append(ps.shareCaches, publicKey)
 	}
 
-	for i, k := range dec.ShareIndexK {
+	for i, k := range dec.ShareCacheIndexK {
 		id, err := BytesID(k)
 		if err != nil {
 			return err
 		}
-		ps.shareIndex[id] = int(dec.ShareIndexV[i])
+		ps.shareCacheIndex[id] = int(dec.ShareCacheIndexV[i])
 	}
 
 	for _, k := range dec.MasterPublicKey {
@@ -205,6 +259,20 @@ func (pubs *PublicKeyShares) UnmarshalJSON(data []byte) error {
 	return nil
 }
 
+// Clone clones every fields of PublicKeyShares. This method is mainly
+// for testing purpose thus would panic when error.
+func (pubs *PublicKeyShares) Clone() *PublicKeyShares {
+	b, err := rlp.EncodeToBytes(pubs)
+	if err != nil {
+		panic(err)
+	}
+	pubsCopy := NewEmptyPublicKeyShares()
+	if err := rlp.DecodeBytes(b, pubsCopy); err != nil {
+		panic(err)
+	}
+	return pubsCopy
+}
+
 // NewID creates a ew ID structure.
 func NewID(id []byte) ID {
 	var blsID bls.ID
@@ -240,7 +308,7 @@ func NewPrivateKeyShares(t int) (*PrivateKeyShares, *PublicKeyShares) {
 			masterPrivateKey: msk,
 			shareIndex:       make(map[ID]int),
 		}, &PublicKeyShares{
-			shareIndex:      make(map[ID]int),
+			shareCacheIndex: make(map[ID]int),
 			masterPublicKey: mpk,
 		}
 }
@@ -329,15 +397,15 @@ func (prvs *PrivateKeyShares) Share(ID ID) (*PrivateKey, bool) {
 // NewEmptyPublicKeyShares creates an empty public key shares.
 func NewEmptyPublicKeyShares() *PublicKeyShares {
 	return &PublicKeyShares{
-		shareIndex: make(map[ID]int),
+		shareCacheIndex: make(map[ID]int),
 	}
 }
 
 // Share returns the share for the ID.
 func (pubs *PublicKeyShares) Share(ID ID) (*PublicKey, error) {
-	idx, exist := pubs.shareIndex[ID]
+	idx, exist := pubs.shareCacheIndex[ID]
 	if exist {
-		return &pubs.shares[idx], nil
+		return &pubs.shareCaches[idx], nil
 	}
 	var pk PublicKey
 	if err := pk.publicKey.Set(pubs.masterPublicKey, &ID); err != nil {
@@ -349,14 +417,14 @@ func (pubs *PublicKeyShares) Share(ID ID) (*PublicKey, error) {
 
 // AddShare adds a share.
 func (pubs *PublicKeyShares) AddShare(ID ID, share *PublicKey) error {
-	if idx, exist := pubs.shareIndex[ID]; exist {
-		if !share.publicKey.IsEqual(&pubs.shares[idx].publicKey) {
+	if idx, exist := pubs.shareCacheIndex[ID]; exist {
+		if !share.publicKey.IsEqual(&pubs.shareCaches[idx].publicKey) {
 			return ErrDuplicatedShare
 		}
 		return nil
 	}
-	pubs.shareIndex[ID] = len(pubs.shares)
-	pubs.shares = append(pubs.shares, *share)
+	pubs.shareCacheIndex[ID] = len(pubs.shareCaches)
+	pubs.shareCaches = append(pubs.shareCaches, *share)
 	return nil
 }
 
diff --git a/core/crypto/dkg/dkg_test.go b/core/crypto/dkg/dkg_test.go
index 44e68d8..cb167e4 100644
--- a/core/crypto/dkg/dkg_test.go
+++ b/core/crypto/dkg/dkg_test.go
@@ -25,6 +25,7 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/Spiderpowa/bls/ffi/go/bls"
 	"github.com/dexon-foundation/dexon/rlp"
 	"github.com/stretchr/testify/suite"
 
@@ -309,8 +310,8 @@ func (s *DKGTestSuite) TestPublicKeySharesRLPEncodeDecode() {
 	for i, id := range s.genID(1) {
 		privkey := NewPrivateKey()
 		pubkey := privkey.PublicKey().(PublicKey)
-		p.shares = append(p.shares, pubkey)
-		p.shareIndex[id] = i
+		p.shareCaches = append(p.shareCaches, pubkey)
+		p.shareCacheIndex[id] = i
 		p.masterPublicKey = append(p.masterPublicKey, pubkey.publicKey)
 	}
 
@@ -327,6 +328,66 @@ func (s *DKGTestSuite) TestPublicKeySharesRLPEncodeDecode() {
 	s.Require().True(reflect.DeepEqual(b, bb))
 }
 
+func (s *DKGTestSuite) TestPublicKeySharesEquality() {
+	var req = s.Require()
+	IDs := s.genID(2)
+	_, pubShares1 := NewPrivateKeyShares(4)
+	// Make a copy from an empty share.
+	pubShares2 := pubShares1.Clone()
+	req.True(pubShares1.Equal(pubShares2))
+	// Add two shares.
+	prvKey1 := NewPrivateKey()
+	pubKey1 := prvKey1.PublicKey().(PublicKey)
+	req.NoError(pubShares1.AddShare(IDs[0], &pubKey1))
+	prvKey2 := NewPrivateKey()
+	pubKey2 := prvKey2.PublicKey().(PublicKey)
+	req.True(pubShares1.Equal(pubShares2))
+	// Clone the shares.
+	req.NoError(pubShares2.AddShare(IDs[0], &pubKey1))
+	req.NoError(pubShares2.AddShare(IDs[1], &pubKey2))
+	// They should be equal now.
+	req.True(pubShares1.Equal(pubShares2))
+	req.True(pubShares2.Equal(pubShares1))
+}
+
+func (s *DKGTestSuite) TestPrivateKeySharesEquality() {
+	var req = s.Require()
+	IDs := s.genID(2)
+	prvShares1, _ := NewPrivateKeyShares(4)
+	// Make a copy of empty share.
+	prvShares2 := NewEmptyPrivateKeyShares()
+	req.False(prvShares1.Equal(prvShares2))
+	// Clone the master private key.
+	for _, m := range prvShares1.masterPrivateKey {
+		var key bls.SecretKey
+		req.NoError(key.SetLittleEndian(m.GetLittleEndian()))
+		prvShares2.masterPrivateKey = append(prvShares2.masterPrivateKey, key)
+	}
+	// Add two shares.
+	prvKey1 := NewPrivateKey()
+	req.NoError(prvShares1.AddShare(IDs[0], prvKey1))
+	prvKey2 := NewPrivateKey()
+	req.NoError(prvShares1.AddShare(IDs[1], prvKey2))
+	// They are not equal now.
+	req.False(prvShares1.Equal(prvShares2))
+	// Clone the shares.
+	req.NoError(prvShares2.AddShare(IDs[0], prvKey1))
+	req.NoError(prvShares2.AddShare(IDs[1], prvKey2))
+	// They should be equal now.
+	req.True(prvShares1.Equal(prvShares2))
+	req.True(prvShares2.Equal(prvShares1))
+}
+
+func (s *DKGTestSuite) TestPublicKeySharesClone() {
+	_, pubShares1 := NewPrivateKeyShares(4)
+	IDs := s.genID(2)
+	prvKey1 := NewPrivateKey()
+	pubKey1 := prvKey1.PublicKey().(PublicKey)
+	s.Require().NoError(pubShares1.AddShare(IDs[0], &pubKey1))
+	pubShares2 := pubShares1.Clone()
+	s.Require().True(pubShares1.Equal(pubShares2))
+}
+
 func TestDKG(t *testing.T) {
 	suite.Run(t, new(DKGTestSuite))
 }
diff --git a/core/types/dkg.go b/core/types/dkg.go
index 3bdb414..edd420d 100644
--- a/core/types/dkg.go
+++ b/core/types/dkg.go
@@ -18,6 +18,7 @@
 package types
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -38,6 +39,17 @@ type DKGPrivateShare struct {
 	Signature    crypto.Signature `json:"signature"`
 }
 
+// Equal checks equality between two DKGPrivateShare instances.
+func (p *DKGPrivateShare) Equal(other *DKGPrivateShare) bool {
+	return p.ProposerID.Equal(other.ProposerID) &&
+		p.ReceiverID.Equal(other.ReceiverID) &&
+		p.Round == other.Round &&
+		p.Signature.Type == other.Signature.Type &&
+		bytes.Compare(p.Signature.Signature, other.Signature.Signature) == 0 &&
+		bytes.Compare(
+			p.PrivateShare.Bytes(), other.PrivateShare.Bytes()) == 0
+}
+
 // DKGMasterPublicKey decrtibe a master public key in DKG protocol.
 type DKGMasterPublicKey struct {
 	ProposerID      NodeID              `json:"proposer_id"`
@@ -53,6 +65,16 @@ func (d *DKGMasterPublicKey) String() string {
 		d.Round)
 }
 
+// Equal check equality of two DKG master public keys.
+func (d *DKGMasterPublicKey) Equal(other *DKGMasterPublicKey) bool {
+	return d.ProposerID.Equal(other.ProposerID) &&
+		d.Round == other.Round &&
+		d.DKGID.GetHexString() == other.DKGID.GetHexString() &&
+		d.PublicKeyShares.Equal(&other.PublicKeyShares) &&
+		d.Signature.Type == other.Signature.Type &&
+		bytes.Compare(d.Signature.Signature, other.Signature.Signature) == 0
+}
+
 type rlpDKGMasterPublicKey struct {
 	ProposerID      NodeID
 	Round           uint64
@@ -126,6 +148,15 @@ func (c *DKGComplaint) String() string {
 		c.ProposerID.String()[:6], c.Round, c.PrivateShare)
 }
 
+// Equal checks equality between two DKGComplaint instances.
+func (c *DKGComplaint) Equal(other *DKGComplaint) bool {
+	return c.ProposerID.Equal(other.ProposerID) &&
+		c.Round == other.Round &&
+		c.PrivateShare.Equal(&other.PrivateShare) &&
+		c.Signature.Type == other.Signature.Type &&
+		bytes.Compare(c.Signature.Signature, other.Signature.Signature) == 0
+}
+
 // DKGPartialSignature describe a partial signature in DKG protocol.
 type DKGPartialSignature struct {
 	ProposerID       NodeID               `json:"proposer_id"`
@@ -148,6 +179,14 @@ func (final *DKGFinalize) String() string {
 		final.Round)
 }
 
+// Equal check equality of two DKGFinalize instances.
+func (final *DKGFinalize) Equal(other *DKGFinalize) bool {
+	return final.ProposerID.Equal(other.ProposerID) &&
+		final.Round == other.Round &&
+		final.Signature.Type == other.Signature.Type &&
+		bytes.Compare(final.Signature.Signature, other.Signature.Signature) == 0
+}
+
 // IsNack returns true if it's a nack complaint in DKG protocol.
 func (c *DKGComplaint) IsNack() bool {
 	return len(c.PrivateShare.Signature.Signature) == 0
diff --git a/core/types/dkg_test.go b/core/types/dkg_test.go
index 010d460..0558101 100644
--- a/core/types/dkg_test.go
+++ b/core/types/dkg_test.go
@@ -18,6 +18,7 @@
 package types
 
 import (
+	"math/rand"
 	"reflect"
 	"testing"
 
@@ -25,6 +26,7 @@ import (
 
 	"github.com/dexon-foundation/dexon-consensus-core/common"
 	"github.com/dexon-foundation/dexon-consensus-core/core/crypto"
+	"github.com/dexon-foundation/dexon-consensus-core/core/crypto/dkg"
 	"github.com/dexon-foundation/dexon/rlp"
 )
 
@@ -32,10 +34,32 @@ type DKGTestSuite struct {
 	suite.Suite
 }
 
+func (s *DKGTestSuite) genRandomBytes() []byte {
+	randomness := make([]byte, 32)
+	_, err := rand.Read(randomness)
+	s.Require().NoError(err)
+	return randomness
+}
+
+func (s *DKGTestSuite) genID() dkg.ID {
+	dID, err := dkg.BytesID(s.genRandomBytes())
+	s.Require().NoError(err)
+	return dID
+}
+
+func (s *DKGTestSuite) clone(src, dst interface{}) {
+	b, err := rlp.EncodeToBytes(src)
+	s.Require().NoError(err)
+	s.Require().NoError(rlp.DecodeBytes(b, dst))
+}
+
 func (s *DKGTestSuite) TestRLPEncodeDecode() {
+	dID := s.genID()
+	// Prepare master public key for testing.
 	d := DKGMasterPublicKey{
 		ProposerID: NodeID{common.Hash{1, 2, 3}},
 		Round:      10,
+		DKGID:      dID,
 		Signature: crypto.Signature{
 			Type:      "123",
 			Signature: []byte{4, 5, 6},
@@ -55,6 +79,167 @@ func (s *DKGTestSuite) TestRLPEncodeDecode() {
 	s.Require().True(d.ProposerID.Equal(dd.ProposerID))
 	s.Require().True(d.Round == dd.Round)
 	s.Require().True(reflect.DeepEqual(d.Signature, dd.Signature))
+	s.Require().Equal(d.DKGID.GetHexString(), dd.DKGID.GetHexString())
+}
+
+func (s *DKGTestSuite) TestDKGMasterPublicKeyEquality() {
+	var req = s.Require()
+	// Prepare source master public key.
+	master1 := &DKGMasterPublicKey{
+		ProposerID: NodeID{Hash: common.NewRandomHash()},
+		Round:      1234,
+		DKGID:      s.genID(),
+		Signature: crypto.Signature{
+			Signature: s.genRandomBytes(),
+		},
+	}
+	prvKey := dkg.NewPrivateKey()
+	pubKey := prvKey.PublicKey().(dkg.PublicKey)
+	_, pubShares := dkg.NewPrivateKeyShares(2)
+	req.NoError(pubShares.AddShare(s.genID(), &pubKey))
+	master1.PublicKeyShares = *pubShares
+	// Prepare another master public key by copying every field.
+	master2 := &DKGMasterPublicKey{}
+	s.clone(master1, master2)
+	// They should be equal.
+	req.True(master1.Equal(master2))
+	// Change round.
+	master2.Round = 2345
+	req.False(master1.Equal(master2))
+	master2.Round = 1234
+	// Change proposerID.
+	master2.ProposerID = NodeID{Hash: common.NewRandomHash()}
+	req.False(master1.Equal(master2))
+	master2.ProposerID = master1.ProposerID
+	// Change DKGID.
+	master2.DKGID = dkg.NewID(s.genRandomBytes())
+	req.False(master1.Equal(master2))
+	master2.DKGID = master1.DKGID
+	// Change signature.
+	master2.Signature = crypto.Signature{
+		Signature: s.genRandomBytes(),
+	}
+	req.False(master1.Equal(master2))
+	master2.Signature = master1.Signature
+	// Change public key share.
+	master2.PublicKeyShares = *dkg.NewEmptyPublicKeyShares()
+	req.False(master1.Equal(master2))
+}
+
+func (s *DKGTestSuite) TestDKGPrivateShareEquality() {
+	var req = s.Require()
+	share1 := &DKGPrivateShare{
+		ProposerID:   NodeID{Hash: common.NewRandomHash()},
+		ReceiverID:   NodeID{Hash: common.NewRandomHash()},
+		Round:        1,
+		PrivateShare: *dkg.NewPrivateKey(),
+		Signature: crypto.Signature{
+			Signature: s.genRandomBytes(),
+		},
+	}
+	// Make a copy.
+	share2 := &DKGPrivateShare{}
+	s.clone(share1, share2)
+	req.True(share1.Equal(share2))
+	// Change proposer ID.
+	share2.ProposerID = NodeID{Hash: common.NewRandomHash()}
+	req.False(share1.Equal(share2))
+	share2.ProposerID = share1.ProposerID
+	// Change receiver ID.
+	share2.ReceiverID = NodeID{Hash: common.NewRandomHash()}
+	req.False(share1.Equal(share2))
+	share2.ReceiverID = share1.ReceiverID
+	// Change round.
+	share2.Round = share1.Round + 1
+	req.False(share1.Equal(share2))
+	share2.Round = share1.Round
+	// Change signature.
+	share2.Signature = crypto.Signature{
+		Signature: s.genRandomBytes(),
+	}
+	req.False(share1.Equal(share2))
+	share2.Signature = share1.Signature
+	// Change private share.
+	share2.PrivateShare = *dkg.NewPrivateKey()
+	req.False(share1.Equal(share2))
+	share2.PrivateShare = share1.PrivateShare
+	// They should be equal after chaning fields back.
+	req.True(share1.Equal(share2))
+}
+
+func (s *DKGTestSuite) TestDKGComplaintEquality() {
+	var req = s.Require()
+	comp1 := &DKGComplaint{
+		ProposerID: NodeID{Hash: common.NewRandomHash()},
+		Round:      1,
+		PrivateShare: DKGPrivateShare{
+			ProposerID:   NodeID{Hash: common.NewRandomHash()},
+			ReceiverID:   NodeID{Hash: common.NewRandomHash()},
+			Round:        2,
+			PrivateShare: *dkg.NewPrivateKey(),
+			Signature: crypto.Signature{
+				Signature: s.genRandomBytes(),
+			},
+		},
+		Signature: crypto.Signature{
+			Signature: s.genRandomBytes(),
+		},
+	}
+	// Make a copy.
+	comp2 := &DKGComplaint{}
+	s.clone(comp1, comp2)
+	req.True(comp1.Equal(comp2))
+	// Change proposer ID.
+	comp2.ProposerID = NodeID{Hash: common.NewRandomHash()}
+	req.False(comp1.Equal(comp2))
+	comp2.ProposerID = comp1.ProposerID
+	// Change round.
+	comp2.Round = comp1.Round + 1
+	req.False(comp1.Equal(comp2))
+	comp2.Round = comp1.Round
+	// Change signature.
+	comp2.Signature = crypto.Signature{
+		Signature: s.genRandomBytes(),
+	}
+	req.False(comp1.Equal(comp2))
+	comp2.Signature = comp1.Signature
+	// Change share's round.
+	comp2.PrivateShare.Round = comp1.PrivateShare.Round + 1
+	req.False(comp1.Equal(comp2))
+	comp2.PrivateShare.Round = comp1.PrivateShare.Round
+	// After changing every field back, should be equal.
+	req.True(comp1.Equal(comp2))
+}
+
+func (s *DKGTestSuite) TestDKGFinalizeEquality() {
+	var req = s.Require()
+	final1 := &DKGFinalize{
+		ProposerID: NodeID{Hash: common.NewRandomHash()},
+		Round:      1,
+		Signature: crypto.Signature{
+			Signature: s.genRandomBytes(),
+		},
+	}
+	// Make a copy
+	final2 := &DKGFinalize{}
+	s.clone(final1, final2)
+	req.True(final1.Equal(final2))
+	// Change proposer ID.
+	final2.ProposerID = NodeID{Hash: common.NewRandomHash()}
+	req.False(final1.Equal(final2))
+	final2.ProposerID = final1.ProposerID
+	// Change round.
+	final2.Round = final1.Round + 1
+	req.False(final1.Equal(final2))
+	final2.Round = final1.Round
+	// Change signature.
+	final2.Signature = crypto.Signature{
+		Signature: s.genRandomBytes(),
+	}
+	req.False(final1.Equal(final2))
+	final2.Signature = final1.Signature
+	// After changing every field back, they should be equal.
+	req.True(final1.Equal(final2))
 }
 
 func TestDKG(t *testing.T) {