core: fix PrivateShares RLP decode/encode bug (#510)

2 files changed, 62 insertions, 5 deletions

diff --git a/core/crypto/dkg/dkg.go b/core/crypto/dkg/dkg.go
index 425d96b..796609d 100644
--- a/core/crypto/dkg/dkg.go
+++ b/core/crypto/dkg/dkg.go
@@ -133,27 +133,79 @@ func (prvs *PrivateKeyShares) Equal(other *PrivateKeyShares) bool {
 
 // EncodeRLP implements rlp.Encoder
 func (prvs *PrivateKeyShares) EncodeRLP(w io.Writer) error {
+	data := make([][][]byte, 3)
+	shares := make([][]byte, len(prvs.shares))
+	for i, s := range prvs.shares {
+		shares[i] = s.Bytes()
+	}
+	data[0] = shares
+
+	shareIndex := make([][]byte, 0)
+	for k, v := range prvs.shareIndex {
+		shareIndex = append(shareIndex, k.GetLittleEndian())
+
+		vBytes, err := rlp.EncodeToBytes(uint64(v))
+		if err != nil {
+			return err
+		}
+		shareIndex = append(shareIndex, vBytes)
+	}
+	data[1] = shareIndex
+
 	mpks := make([][]byte, len(prvs.masterPrivateKey))
 	for i, m := range prvs.masterPrivateKey {
 		mpks[i] = m.GetLittleEndian()
 	}
-	return rlp.Encode(w, mpks)
+	data[2] = mpks
+	return rlp.Encode(w, data)
 }
 
 // DecodeRLP implements rlp.Decoder
 func (prvs *PrivateKeyShares) DecodeRLP(s *rlp.Stream) error {
-	var dec [][]byte
+	*prvs = PrivateKeyShares{}
+	var dec [][][]byte
 	if err := s.Decode(&dec); err != nil {
 		return err
 	}
 
-	for _, k := range dec {
+	var shares []PrivateKey
+	for _, bs := range dec[0] {
+		var key PrivateKey
+		err := key.SetBytes(bs)
+		if err != nil {
+			return err
+		}
+		shares = append(shares, key)
+	}
+	(*prvs).shares = shares
+
+	sharesIndex := map[ID]int{}
+	for i := 0; i < len(dec[1]); i += 2 {
+		var key ID
+		err := key.SetLittleEndian(dec[1][i])
+		if err != nil {
+			return err
+		}
+
+		var value uint64
+		err = rlp.DecodeBytes(dec[1][i+1], &value)
+		if err != nil {
+			return err
+		}
+
+		sharesIndex[key] = int(value)
+	}
+	(*prvs).shareIndex = sharesIndex
+
+	var mpks []bls.SecretKey
+	for _, bs := range dec[2] {
 		var key bls.SecretKey
-		if err := key.SetLittleEndian(k); err != nil {
+		if err := key.SetLittleEndian(bs); err != nil {
 			return err
 		}
-		prvs.masterPrivateKey = append(prvs.masterPrivateKey, key)
+		mpks = append(mpks, key)
 	}
+	(*prvs).masterPrivateKey = mpks
 
 	return nil
 }
diff --git a/core/crypto/dkg/dkg_test.go b/core/crypto/dkg/dkg_test.go
index 20dd908..f278932 100644
--- a/core/crypto/dkg/dkg_test.go
+++ b/core/crypto/dkg/dkg_test.go
@@ -330,6 +330,10 @@ func (s *DKGTestSuite) TestPublicKeySharesRLPEncodeDecode() {
 
 func (s *DKGTestSuite) TestPrivateKeySharesRLPEncodeDecode() {
 	privShares, _ := NewPrivateKeyShares(10)
+	privShares.shares = append(privShares.shares, PrivateKey{})
+	privShares.shareIndex = map[ID]int{
+		ID{}: 0,
+	}
 
 	b, err := rlp.EncodeToBytes(privShares)
 	s.Require().NoError(err)
@@ -342,6 +346,7 @@ func (s *DKGTestSuite) TestPrivateKeySharesRLPEncodeDecode() {
 	s.Require().NoError(err)
 
 	s.Require().True(reflect.DeepEqual(b, bb))
+	s.Require().True(privShares.Equal(&newPrivShares))
 }
 
 func (s *DKGTestSuite) TestPublicKeySharesEquality() {