dex: rebuild connection if dkg reset, core/vm: fix gov (#308)

* dex: rebuild connection if dkg reset
* core: vm: fix governance contract

3 files changed, 21 insertions, 2 deletions

diff --git a/dex/handler.go b/dex/handler.go
index 8971ad500..6d0746d2f 100644
--- a/dex/handler.go
+++ b/dex/handler.go
@@ -1320,6 +1320,7 @@ func (pm *ProtocolManager) peerSetLoop() {
 		pm.peers.BuildConnection(CRSRound)
 		round = CRSRound
 	}
+	resetCount := uint64(0)
 
 	for {
 		select {
@@ -1334,9 +1335,12 @@ func (pm *ProtocolManager) peerSetLoop() {
 			if newRound == 0 {
 				break
 			}
+			reset := pm.gov.DKGResetCount(round)
 
-			log.Debug("ProtocolManager: new round", "round", newRound)
-			if newRound <= round {
+			log.Debug("ProtocolManager: new round",
+				"round", newRound,
+				"reset", reset)
+			if newRound <= round && resetCount == reset {
 				break
 			}
 
@@ -1345,6 +1349,10 @@ func (pm *ProtocolManager) peerSetLoop() {
 				if round >= 1 {
 					pm.peers.ForgetConnection(round - 1)
 				}
+			} else if newRound == round && resetCount+1 == reset {
+				pm.peers.ForgetConnection(newRound)
+				pm.gov.PurgeNotarySet(newRound)
+				pm.peers.BuildConnection(newRound)
 			} else {
 				// just forget all network connection and rebuild.
 				pm.peers.ForgetConnection(round)
@@ -1355,6 +1363,7 @@ func (pm *ProtocolManager) peerSetLoop() {
 				pm.peers.BuildConnection(newRound)
 			}
 			round = newRound
+			resetCount = reset
 		case <-pm.chainHeadSub.Err():
 			return
 		}
diff --git a/dex/helper_test.go b/dex/helper_test.go
index 1d903a907..1d267d9ae 100644
--- a/dex/helper_test.go
+++ b/dex/helper_test.go
@@ -221,6 +221,12 @@ func (g *testGovernance) CRSRound() uint64 {
 	return g.lenCRSFunc()
 }
 
+func (g *testGovernance) DKGResetCount(uint64) uint64 {
+	return 0
+}
+
+func (g *testGovernance) PurgeNotarySet(uint64) {}
+
 func (g *testGovernance) NotarySet(
 	round uint64) (map[string]struct{}, error) {
 	return g.notarySetFunc(round)
diff --git a/dex/protocol.go b/dex/protocol.go
index 2bcb57506..4da64b604 100644
--- a/dex/protocol.go
+++ b/dex/protocol.go
@@ -156,6 +156,10 @@ type governance interface {
 	CRSRound() uint64
 
 	NotarySet(uint64) (map[string]struct{}, error)
+
+	PurgeNotarySet(uint64)
+
+	DKGResetCount(uint64) uint64
 }
 
 type dexconApp interface {