core: Fix rlp encode/decode for DKGComplaint. (#441)

* types: Add RLP Encode/Decode to DKGComplaint

* Add test

* fix state test

3 files changed, 144 insertions, 1 deletions

diff --git a/core/test/state_test.go b/core/test/state_test.go
index ad3d1c4..79a7a85 100644
--- a/core/test/state_test.go
+++ b/core/test/state_test.go
@@ -66,7 +66,10 @@ func (s *StateTestSuite) newDKGComplaint(round uint64) *typesDKG.Complaint {
 			PrivateShare: *dkg.NewPrivateKey(),
 		},
 	}
-	s.Require().NoError(utils.NewSigner(prvKey).SignDKGComplaint(comp))
+	signer := utils.NewSigner(prvKey)
+	s.Require().NoError(signer.SignDKGComplaint(comp))
+	s.Require().NoError(signer.SignDKGPrivateShare(&comp.PrivateShare))
+	s.Require().False(comp.IsNack())
 	return comp
 }
 
diff --git a/core/types/dkg/dkg.go b/core/types/dkg/dkg.go
index efe8467..db58168 100644
--- a/core/types/dkg/dkg.go
+++ b/core/types/dkg/dkg.go
@@ -163,6 +163,64 @@ func (c *Complaint) Equal(other *Complaint) bool {
 		bytes.Compare(c.Signature.Signature, other.Signature.Signature) == 0
 }
 
+type rlpComplaint struct {
+	ProposerID   types.NodeID
+	Round        uint64
+	IsNack       bool
+	PrivateShare []byte
+	Signature    crypto.Signature
+}
+
+// EncodeRLP implements rlp.Encoder
+func (c *Complaint) EncodeRLP(w io.Writer) error {
+	if c.IsNack() {
+		return rlp.Encode(w, rlpComplaint{
+			ProposerID:   c.ProposerID,
+			Round:        c.Round,
+			IsNack:       true,
+			PrivateShare: c.PrivateShare.ProposerID.Hash[:],
+			Signature:    c.Signature,
+		})
+	}
+	prvShare, err := rlp.EncodeToBytes(&c.PrivateShare)
+	if err != nil {
+		return err
+	}
+	return rlp.Encode(w, rlpComplaint{
+		ProposerID:   c.ProposerID,
+		Round:        c.Round,
+		IsNack:       false,
+		PrivateShare: prvShare,
+		Signature:    c.Signature,
+	})
+}
+
+// DecodeRLP implements rlp.Decoder
+func (c *Complaint) DecodeRLP(s *rlp.Stream) error {
+	var dec rlpComplaint
+	if err := s.Decode(&dec); err != nil {
+		return err
+	}
+
+	var prvShare PrivateShare
+	if dec.IsNack {
+		copy(prvShare.ProposerID.Hash[:], dec.PrivateShare)
+		prvShare.Round = dec.Round
+	} else {
+		if err := rlp.DecodeBytes(dec.PrivateShare, &prvShare); err != nil {
+			return err
+		}
+	}
+
+	*c = Complaint{
+		ProposerID:   dec.ProposerID,
+		Round:        dec.Round,
+		PrivateShare: prvShare,
+		Signature:    dec.Signature,
+	}
+	return nil
+}
+
 // PartialSignature describe a partial signature in DKG protocol.
 type PartialSignature struct {
 	ProposerID       types.NodeID               `json:"proposer_id"`
diff --git a/core/types/dkg/dkg_test.go b/core/types/dkg/dkg_test.go
index 4eb17ac..7fe98cd 100644
--- a/core/types/dkg/dkg_test.go
+++ b/core/types/dkg/dkg_test.go
@@ -81,6 +81,88 @@ func (s *DKGTestSuite) TestRLPEncodeDecode() {
 	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())
+
+	// Test DKGPrivateShare.
+	p := PrivateShare{
+		ProposerID:   types.NodeID{Hash: common.Hash{1, 3, 5}},
+		Round:        10,
+		PrivateShare: *cryptoDKG.NewPrivateKey(),
+		Signature: crypto.Signature{
+			Type:      "123",
+			Signature: []byte{2, 4, 6},
+		},
+	}
+
+	b, err = rlp.EncodeToBytes(&p)
+	s.Require().NoError(err)
+
+	var pp PrivateShare
+	err = rlp.DecodeBytes(b, &pp)
+	s.Require().NoError(err)
+
+	bb, err = rlp.EncodeToBytes(&pp)
+	s.Require().NoError(err)
+	s.Require().True(reflect.DeepEqual(b, bb))
+	s.Require().True(p.ProposerID.Equal(pp.ProposerID))
+	s.Require().True(p.Round == pp.Round)
+	s.Require().True(reflect.DeepEqual(p.PrivateShare, pp.PrivateShare))
+	s.Require().True(reflect.DeepEqual(p.Signature, pp.Signature))
+
+	// Test DKG Nack Complaint.
+	c := Complaint{
+		ProposerID: d.ProposerID,
+		Round:      10,
+		PrivateShare: PrivateShare{
+			ProposerID: p.ProposerID,
+			Round:      10,
+		},
+		Signature: crypto.Signature{
+			Type:      "123",
+			Signature: []byte{3, 3, 3},
+		},
+	}
+	s.Require().True(c.IsNack())
+
+	b, err = rlp.EncodeToBytes(&c)
+	s.Require().NoError(err)
+
+	var cc Complaint
+	err = rlp.DecodeBytes(b, &cc)
+	s.Require().NoError(err)
+
+	bb, err = rlp.EncodeToBytes(&cc)
+	s.Require().NoError(err)
+	s.Require().True(reflect.DeepEqual(c, cc))
+	s.Require().True(c.ProposerID.Equal(cc.ProposerID))
+	s.Require().True(c.Round == cc.Round)
+	s.Require().True(reflect.DeepEqual(c.PrivateShare, cc.PrivateShare))
+	s.Require().True(reflect.DeepEqual(c.Signature, cc.Signature))
+
+	// Test DKG Complaint.
+	c = Complaint{
+		ProposerID:   d.ProposerID,
+		Round:        10,
+		PrivateShare: p,
+		Signature: crypto.Signature{
+			Type:      "123",
+			Signature: []byte{3, 3, 3},
+		},
+	}
+	s.Require().False(c.IsNack())
+
+	b, err = rlp.EncodeToBytes(&c)
+	s.Require().NoError(err)
+
+	err = rlp.DecodeBytes(b, &cc)
+	s.Require().NoError(err)
+
+	bb, err = rlp.EncodeToBytes(&cc)
+	s.Require().NoError(err)
+	s.Require().True(reflect.DeepEqual(c, cc))
+	s.Require().True(c.ProposerID.Equal(cc.ProposerID))
+	s.Require().True(c.Round == cc.Round)
+	s.Require().True(reflect.DeepEqual(c.PrivateShare, cc.PrivateShare))
+	s.Require().True(reflect.DeepEqual(c.Signature, cc.Signature))
 }
 
 func (s *DKGTestSuite) TestMasterPublicKeyEquality() {