core: fix DKGProtocolInfo RLP decode/encode bug (#505)

2 files changed, 133 insertions, 8 deletions

diff --git a/core/db/level-db.go b/core/db/level-db.go
index c92a623..88f5801 100644
--- a/core/db/level-db.go
+++ b/core/db/level-db.go
@@ -59,6 +59,91 @@ type DKGProtocolInfo struct {
 	Reset                     uint64
 }
 
+// Equal compare with target DKGProtocolInfo.
+func (info *DKGProtocolInfo) Equal(target *DKGProtocolInfo) bool {
+	if !info.ID.Equal(target.ID) ||
+		info.Round != target.Round ||
+		info.Threshold != target.Threshold ||
+		info.IsMasterPrivateShareEmpty != target.IsMasterPrivateShareEmpty ||
+		info.IsPrvSharesEmpty != target.IsPrvSharesEmpty ||
+		info.Step != target.Step ||
+		info.Reset != target.Reset ||
+		!info.MasterPrivateShare.Equal(&target.MasterPrivateShare) ||
+		!info.PrvShares.Equal(&target.PrvShares) {
+		return false
+	}
+
+	if len(info.IDMap) != len(target.IDMap) {
+		return false
+	}
+	for k, v := range info.IDMap {
+		tV, exist := target.IDMap[k]
+		if !exist {
+			return false
+		}
+
+		if !v.IsEqual(&tV) {
+			return false
+		}
+	}
+
+	if len(info.MpkMap) != len(target.MpkMap) {
+		return false
+	}
+	for k, v := range info.MpkMap {
+		tV, exist := target.MpkMap[k]
+		if !exist {
+			return false
+		}
+
+		if !v.Equal(tV) {
+			return false
+		}
+	}
+
+	if len(info.PrvSharesReceived) != len(target.PrvSharesReceived) {
+		return false
+	}
+	for k := range info.PrvSharesReceived {
+		_, exist := target.PrvSharesReceived[k]
+		if !exist {
+			return false
+		}
+	}
+
+	if len(info.NodeComplained) != len(target.NodeComplained) {
+		return false
+	}
+	for k := range info.NodeComplained {
+		_, exist := target.NodeComplained[k]
+		if !exist {
+			return false
+		}
+	}
+
+	if len(info.AntiComplaintReceived) != len(target.AntiComplaintReceived) {
+		return false
+	}
+	for k, v := range info.AntiComplaintReceived {
+		tV, exist := target.AntiComplaintReceived[k]
+		if !exist {
+			return false
+		}
+
+		if len(v) != len(tV) {
+			return false
+		}
+		for kk := range v {
+			_, exist := tV[kk]
+			if !exist {
+				return false
+			}
+		}
+	}
+
+	return true
+}
+
 // NodeIDToNodeIDs the map with NodeID to NodeIDs.
 type NodeIDToNodeIDs map[types.NodeID]map[types.NodeID]struct{}
 
@@ -87,7 +172,8 @@ func (m NodeIDToNodeIDs) EncodeRLP(w io.Writer) error {
 }
 
 // DecodeRLP implements rlp.Encoder
-func (m NodeIDToNodeIDs) DecodeRLP(s *rlp.Stream) error {
+func (m *NodeIDToNodeIDs) DecodeRLP(s *rlp.Stream) error {
+	*m = make(NodeIDToNodeIDs)
 	var dec [][][]byte
 	if err := s.Decode(&dec); err != nil {
 		return err
@@ -111,7 +197,7 @@ func (m NodeIDToNodeIDs) DecodeRLP(s *rlp.Stream) error {
 			valueMap[value] = struct{}{}
 		}
 
-		m[key] = valueMap
+		(*m)[key] = valueMap
 	}
 
 	return nil
@@ -135,7 +221,8 @@ func (m NodeID) EncodeRLP(w io.Writer) error {
 }
 
 // DecodeRLP implements rlp.Encoder
-func (m NodeID) DecodeRLP(s *rlp.Stream) error {
+func (m *NodeID) DecodeRLP(s *rlp.Stream) error {
+	*m = make(NodeID)
 	var dec [][]byte
 	if err := s.Decode(&dec); err != nil {
 		return err
@@ -148,7 +235,7 @@ func (m NodeID) DecodeRLP(s *rlp.Stream) error {
 			return err
 		}
 
-		m[key] = struct{}{}
+		(*m)[key] = struct{}{}
 	}
 
 	return nil
@@ -178,7 +265,8 @@ func (m NodeIDToPubShares) EncodeRLP(w io.Writer) error {
 }
 
 // DecodeRLP implements rlp.Encoder
-func (m NodeIDToPubShares) DecodeRLP(s *rlp.Stream) error {
+func (m *NodeIDToPubShares) DecodeRLP(s *rlp.Stream) error {
+	*m = make(NodeIDToPubShares)
 	var dec [][]byte
 	if err := s.Decode(&dec); err != nil {
 		return err
@@ -197,7 +285,7 @@ func (m NodeIDToPubShares) DecodeRLP(s *rlp.Stream) error {
 			return err
 		}
 
-		m[key] = &value
+		(*m)[key] = &value
 	}
 
 	return nil
@@ -222,7 +310,8 @@ func (m NodeIDToDKGID) EncodeRLP(w io.Writer) error {
 }
 
 // DecodeRLP implements rlp.Encoder
-func (m NodeIDToDKGID) DecodeRLP(s *rlp.Stream) error {
+func (m *NodeIDToDKGID) DecodeRLP(s *rlp.Stream) error {
+	*m = make(NodeIDToDKGID)
 	var dec [][]byte
 	if err := s.Decode(&dec); err != nil {
 		return err
@@ -241,7 +330,7 @@ func (m NodeIDToDKGID) DecodeRLP(s *rlp.Stream) error {
 			return err
 		}
 
-		m[key] = value
+		(*m)[key] = value
 	}
 
 	return nil
diff --git a/core/db/level-db_test.go b/core/db/level-db_test.go
index c4f6b8a..d7f70bc 100644
--- a/core/db/level-db_test.go
+++ b/core/db/level-db_test.go
@@ -204,6 +204,42 @@ func (s *LevelDBTestSuite) TestDKGProtocol() {
 	s.Require().NoError(dbInst.PutOrUpdateDKGProtocol(DKGProtocolInfo{}))
 }
 
+func (s *LevelDBTestSuite) TestDKGProtocolInfoRLPEncodeDecode() {
+	protocol := DKGProtocolInfo{
+		ID:        types.NodeID{Hash: common.Hash{0x11}},
+		Round:     5,
+		Threshold: 10,
+		IDMap: NodeIDToDKGID{
+			types.NodeID{Hash: common.Hash{0x01}}: dkg.ID{},
+			types.NodeID{Hash: common.Hash{0x02}}: dkg.ID{},
+		},
+		MpkMap: NodeIDToPubShares{
+			types.NodeID{Hash: common.Hash{0x01}}: &dkg.PublicKeyShares{},
+			types.NodeID{Hash: common.Hash{0x02}}: &dkg.PublicKeyShares{},
+		},
+		AntiComplaintReceived: NodeIDToNodeIDs{
+			types.NodeID{Hash: common.Hash{0x01}}: map[types.NodeID]struct{}{
+				types.NodeID{Hash: common.Hash{0x02}}: {},
+			},
+			types.NodeID{Hash: common.Hash{0x03}}: map[types.NodeID]struct{}{
+				types.NodeID{Hash: common.Hash{0x04}}: {},
+			},
+		},
+		PrvSharesReceived: NodeID{
+			types.NodeID{Hash: common.Hash{0x01}}: struct{}{},
+		},
+	}
+
+	b, err := rlp.EncodeToBytes(&protocol)
+	s.Require().NoError(err)
+
+	newProtocol := DKGProtocolInfo{}
+	err = rlp.DecodeBytes(b, &newProtocol)
+	s.Require().NoError(err)
+
+	s.Require().True(protocol.Equal(&newProtocol))
+}
+
 func (s *LevelDBTestSuite) TestNodeIDToNodeIDsRLPEncodeDecode() {
 	m := NodeIDToNodeIDs{
 		types.NodeID{Hash: common.Hash{0x01}}: map[types.NodeID]struct{}{