core: Add DKG final message. (#181)

3 files changed, 35 insertions, 0 deletions

diff --git a/simulation/governance.go b/simulation/governance.go
index 23e1884..f22abce 100644
--- a/simulation/governance.go
+++ b/simulation/governance.go
@@ -42,6 +42,7 @@ type simGovernance struct {
 	tsig               map[uint64]crypto.Signature
 	dkgComplaint       map[uint64][]*types.DKGComplaint
 	dkgMasterPublicKey map[uint64][]*types.DKGMasterPublicKey
+	dkgFinal           map[uint64]map[types.NodeID]struct{}
 	lambdaBA           time.Duration
 	lambdaDKG          time.Duration
 	roundInterval      time.Duration
@@ -64,6 +65,7 @@ func newSimGovernance(
 		tsig:               make(map[uint64]crypto.Signature),
 		dkgComplaint:       make(map[uint64][]*types.DKGComplaint),
 		dkgMasterPublicKey: make(map[uint64][]*types.DKGMasterPublicKey),
+		dkgFinal:           make(map[uint64]map[types.NodeID]struct{}),
 		lambdaBA: time.Duration(consensusConfig.LambdaBA) *
 			time.Millisecond,
 		lambdaDKG: time.Duration(consensusConfig.LambdaDKG) *
@@ -139,6 +141,12 @@ func (g *simGovernance) addNode(pubKey crypto.PublicKey) {
 
 // AddDKGComplaint adds a DKGComplaint.
 func (g *simGovernance) AddDKGComplaint(complaint *types.DKGComplaint) {
+	if g.IsDKGFinal(complaint.Round) {
+		return
+	}
+	if _, exist := g.dkgFinal[complaint.Round][complaint.ProposerID]; exist {
+		return
+	}
 	// TODO(jimmy-dexon): check if the input is valid.
 	g.dkgComplaint[complaint.Round] = append(
 		g.dkgComplaint[complaint.Round], complaint)
@@ -176,3 +184,20 @@ func (g *simGovernance) DKGMasterPublicKeys(
 	}
 	return masterPublicKeys
 }
+
+// AddDKGFinalize adds a DKG finalize message.
+func (g *simGovernance) AddDKGFinalize(final *types.DKGFinalize) {
+	// TODO(jimmy-dexon): check if the input is valid.
+	if _, exist := g.dkgFinal[final.Round]; !exist {
+		g.dkgFinal[final.Round] = make(map[types.NodeID]struct{})
+	}
+	g.dkgFinal[final.Round][final.ProposerID] = struct{}{}
+	if final.ProposerID == g.id {
+		g.network.broadcast(final)
+	}
+}
+
+// IsDKGFinal checks if DKG is final.
+func (g *simGovernance) IsDKGFinal(round uint64) bool {
+	return len(g.dkgFinal[round]) > int(g.Configuration(round).DKGSetSize)/3*2
+}
diff --git a/simulation/marshaller.go b/simulation/marshaller.go
index d677c1f..bb60432 100644
--- a/simulation/marshaller.go
+++ b/simulation/marshaller.go
@@ -99,6 +99,12 @@ func (m *jsonMarshaller) Unmarshal(
 			break
 		}
 		msg = psig
+	case "dkg-finalize":
+		final := &types.DKGFinalize{}
+		if err = json.Unmarshal(payload, final); err != nil {
+			break
+		}
+		msg = final
 	default:
 		err = fmt.Errorf("unrecognized message type: %v", msgType)
 	}
@@ -135,6 +141,8 @@ func (m *jsonMarshaller) Marshal(msg interface{}) (
 		msgType = "dkg-complaint"
 	case *types.DKGPartialSignature:
 		msgType = "dkg-partial-signature"
+	case *types.DKGFinalize:
+		msgType = "dkg-finalize"
 	default:
 		err = fmt.Errorf("unknwon message type: %v", msg)
 	}
diff --git a/simulation/node.go b/simulation/node.go
index f2ddd5a..b90c1f8 100644
--- a/simulation/node.go
+++ b/simulation/node.go
@@ -114,6 +114,8 @@ MainLoop:
 			n.gov.AddDKGComplaint(val)
 		case *types.DKGMasterPublicKey:
 			n.gov.AddDKGMasterPublicKey(val)
+		case *types.DKGFinalize:
+			n.gov.AddDKGFinalize(val)
 		default:
 			panic(fmt.Errorf("unexpected message from server: %v", val))
 		}