core: recover DKG master private shares (#487)

10 files changed, 255 insertions, 16 deletions

diff --git a/core/configuration-chain.go b/core/configuration-chain.go
index faea49f..9214cd9 100644
--- a/core/configuration-chain.go
+++ b/core/configuration-chain.go
@@ -106,6 +106,12 @@ func (cc *configurationChain) registerDKG(round uint64, threshold int) {
 		cc.recv,
 		round,
 		threshold)
+	err = cc.db.PutOrUpdateDKGMasterPrivateShares(round, *cc.dkg.prvShares)
+	if err != nil {
+		cc.logger.Error("Error put or update DKG master private shares", "error",
+			err)
+		return
+	}
 	go func() {
 		ticker := newTicker(cc.gov, round, TickerDKG)
 		defer ticker.Stop()
diff --git a/core/consensus.go b/core/consensus.go
index e3a7b1a..4201cbc 100644
--- a/core/consensus.go
+++ b/core/consensus.go
@@ -538,6 +538,11 @@ func newConsensusForRound(
 		logger:       logger,
 	}
 	cfgModule := newConfigurationChain(ID, recv, gov, nodeSetCache, db, logger)
+	dkg, err := recoverDKGProtocol(ID, recv, initRound, utils.GetDKGThreshold(initConfig), db)
+	if err != nil {
+		panic(err)
+	}
+	cfgModule.dkg = dkg
 	recv.cfgModule = cfgModule
 	appModule := app
 	if usingNonBlocking {
@@ -574,7 +579,6 @@ func newConsensusForRound(
 	baConfig := agreementMgrConfig{}
 	baConfig.from(initRound, initConfig, initCRS)
 	baConfig.SetRoundBeginHeight(initRoundBeginHeight)
-	var err error
 	con.baMgr, err = newAgreementMgr(con, initRound, baConfig)
 	if err != nil {
 		panic(err)
diff --git a/core/consensus_test.go b/core/consensus_test.go
index b758c77..d721bd8 100644
--- a/core/consensus_test.go
+++ b/core/consensus_test.go
@@ -203,6 +203,46 @@ func (s *ConsensusTestSuite) prepareConsensus(
 	return app, con
 }
 
+func (s *ConsensusTestSuite) prepareConsensusWithDB(
+	dMoment time.Time,
+	gov *test.Governance,
+	prvKey crypto.PrivateKey,
+	conn *networkConnection,
+	dbInst db.Database) (
+	*test.App, *Consensus) {
+
+	app := test.NewApp(0, nil)
+	nID := types.NewNodeID(prvKey.PublicKey())
+	network := conn.newNetwork(nID)
+	con := NewConsensus(
+		dMoment, app, gov, dbInst, network, prvKey, &common.NullLogger{})
+	conn.setCon(nID, con)
+	return app, con
+}
+
+func (s *ConsensusTestSuite) TestRegisteredDKGRecover() {
+	conn := s.newNetworkConnection()
+	prvKeys, pubKeys, err := test.NewKeys(1)
+	s.Require().NoError(err)
+	gov, err := test.NewGovernance(test.NewState(DKGDelayRound,
+		pubKeys, time.Second, &common.NullLogger{}, true), ConfigRoundShift)
+	s.Require().NoError(err)
+	gov.State().RequestChange(test.StateChangeRoundLength, uint64(200))
+	dMoment := time.Now().UTC()
+	dbInst, err := db.NewMemBackedDB()
+	s.Require().NoError(err)
+	_, con := s.prepareConsensusWithDB(dMoment, gov, prvKeys[0], conn, dbInst)
+
+	s.Require().Nil(con.cfgModule.dkg)
+
+	con.cfgModule.registerDKG(0, 10)
+
+	_, newCon := s.prepareConsensusWithDB(dMoment, gov, prvKeys[0], conn, dbInst)
+
+	s.Require().NotNil(newCon.cfgModule.dkg)
+	s.Require().True(newCon.cfgModule.dkg.prvShares.Equal(con.cfgModule.dkg.prvShares))
+}
+
 func (s *ConsensusTestSuite) TestDKGCRS() {
 	n := 21
 	lambda := 200 * time.Millisecond
diff --git a/core/crypto/dkg/dkg.go b/core/crypto/dkg/dkg.go
index f6b3e0e..425d96b 100644
--- a/core/crypto/dkg/dkg.go
+++ b/core/crypto/dkg/dkg.go
@@ -131,6 +131,33 @@ func (prvs *PrivateKeyShares) Equal(other *PrivateKeyShares) bool {
 	return true
 }
 
+// EncodeRLP implements rlp.Encoder
+func (prvs *PrivateKeyShares) EncodeRLP(w io.Writer) error {
+	mpks := make([][]byte, len(prvs.masterPrivateKey))
+	for i, m := range prvs.masterPrivateKey {
+		mpks[i] = m.GetLittleEndian()
+	}
+	return rlp.Encode(w, mpks)
+}
+
+// DecodeRLP implements rlp.Decoder
+func (prvs *PrivateKeyShares) DecodeRLP(s *rlp.Stream) error {
+	var dec [][]byte
+	if err := s.Decode(&dec); err != nil {
+		return err
+	}
+
+	for _, k := range dec {
+		var key bls.SecretKey
+		if err := key.SetLittleEndian(k); err != nil {
+			return err
+		}
+		prvs.masterPrivateKey = append(prvs.masterPrivateKey, key)
+	}
+
+	return nil
+}
+
 // PublicKeyShares represents a public key shares for DKG protocol.
 type PublicKeyShares struct {
 	shareCaches     []PublicKey
diff --git a/core/crypto/dkg/dkg_test.go b/core/crypto/dkg/dkg_test.go
index e097bdb..20dd908 100644
--- a/core/crypto/dkg/dkg_test.go
+++ b/core/crypto/dkg/dkg_test.go
@@ -328,6 +328,22 @@ func (s *DKGTestSuite) TestPublicKeySharesRLPEncodeDecode() {
 	s.Require().True(reflect.DeepEqual(b, bb))
 }
 
+func (s *DKGTestSuite) TestPrivateKeySharesRLPEncodeDecode() {
+	privShares, _ := NewPrivateKeyShares(10)
+
+	b, err := rlp.EncodeToBytes(privShares)
+	s.Require().NoError(err)
+
+	var newPrivShares PrivateKeyShares
+	err = rlp.DecodeBytes(b, &newPrivShares)
+	s.Require().NoError(err)
+
+	bb, err := rlp.EncodeToBytes(&newPrivShares)
+	s.Require().NoError(err)
+
+	s.Require().True(reflect.DeepEqual(b, bb))
+}
+
 func (s *DKGTestSuite) TestPublicKeySharesEquality() {
 	var req = s.Require()
 	IDs := s.genID(2)
diff --git a/core/db/interfaces.go b/core/db/interfaces.go
index ebbbbd4..e958611 100644
--- a/core/db/interfaces.go
+++ b/core/db/interfaces.go
@@ -50,6 +50,12 @@ var (
 	// ErrDKGPrivateKeyDoesNotExist raised when the DKG private key of the
 	// requested round does not exists.
 	ErrDKGPrivateKeyDoesNotExist = errors.New("dkg private key does not exists")
+	// ErrDKGMasterPrivateSharesExists raised when attempting to save DKG master private shares
+	// that already saved.
+	ErrDKGMasterPrivateSharesExists = errors.New("dkg master private shares exists")
+	// ErrDKGMasterPrivateSharesDoesNotExist raised when the DKG master private shares of the
+	// requested round does not exists.
+	ErrDKGMasterPrivateSharesDoesNotExist = errors.New("dkg master private shares does not exists")
 )
 
 // Database is the interface for a Database.
@@ -76,6 +82,7 @@ type Reader interface {
 	// DKG Private Key related methods.
 	HasDKGPrivateKey(round uint64) (bool, error)
 	GetDKGPrivateKey(round uint64) (dkg.PrivateKey, error)
+	GetDKGMasterPrivateShares(round uint64) (shares dkg.PrivateKeyShares, err error)
 }
 
 // Writer defines the interface for writing blocks into DB.
@@ -84,6 +91,7 @@ type Writer interface {
 	PutBlock(block types.Block) error
 	PutCompactionChainTipInfo(common.Hash, uint64) error
 	PutDKGPrivateKey(uint64, dkg.PrivateKey) error
+	PutOrUpdateDKGMasterPrivateShares(round uint64, shares dkg.PrivateKeyShares) error
 }
 
 // BlockIterator defines an iterator on blocks hold
diff --git a/core/db/level-db.go b/core/db/level-db.go
index 1fe29fa..efa1fec 100644
--- a/core/db/level-db.go
+++ b/core/db/level-db.go
@@ -29,9 +29,10 @@ import (
 )
 
 var (
-	blockKeyPrefix            = []byte("b-")
-	compactionChainTipInfoKey = []byte("cc-tip")
-	dkgPrivateKeyKeyPrefix    = []byte("dkg-prvs")
+	blockKeyPrefix               = []byte("b-")
+	compactionChainTipInfoKey    = []byte("cc-tip")
+	dkgPrivateKeyKeyPrefix       = []byte("dkg-prvs")
+	dkgMasterPrivateSharesPrefix = []byte("dkg-master-private-shares")
 )
 
 type compactionChainTipInfo struct {
@@ -188,6 +189,11 @@ func (lvl *LevelDBBackedDB) HasDKGPrivateKey(round uint64) (bool, error) {
 	return lvl.db.Has(lvl.getDKGPrivateKeyKey(round), nil)
 }
 
+// HasDKGMasterPrivateSharesKey check existence of DKG master private shares of one round.
+func (lvl *LevelDBBackedDB) HasDKGMasterPrivateSharesKey(round uint64) (bool, error) {
+	return lvl.db.Has(lvl.getDKGMasterPrivateSharesKey(round), nil)
+}
+
 // GetDKGPrivateKey get DKG private key of one round.
 func (lvl *LevelDBBackedDB) GetDKGPrivateKey(round uint64) (
 	prv dkg.PrivateKey, err error) {
@@ -221,6 +227,32 @@ func (lvl *LevelDBBackedDB) PutDKGPrivateKey(
 		lvl.getDKGPrivateKeyKey(round), marshaled, nil)
 }
 
+// GetDKGMasterPrivateShares get DKG master private shares of one round.
+func (lvl *LevelDBBackedDB) GetDKGMasterPrivateShares(round uint64) (
+	shares dkg.PrivateKeyShares, err error) {
+	queried, err := lvl.db.Get(lvl.getDKGMasterPrivateSharesKey(round), nil)
+	if err != nil {
+		if err == leveldb.ErrNotFound {
+			err = ErrDKGMasterPrivateSharesDoesNotExist
+		}
+		return
+	}
+
+	err = rlp.DecodeBytes(queried, &shares)
+	return
+}
+
+// PutOrUpdateDKGMasterPrivateShares save DKG master private shares of one round.
+func (lvl *LevelDBBackedDB) PutOrUpdateDKGMasterPrivateShares(
+	round uint64, shares dkg.PrivateKeyShares) error {
+	marshaled, err := rlp.EncodeToBytes(&shares)
+	if err != nil {
+		return err
+	}
+	return lvl.db.Put(
+		lvl.getDKGMasterPrivateSharesKey(round), marshaled, nil)
+}
+
 func (lvl *LevelDBBackedDB) getBlockKey(hash common.Hash) (ret []byte) {
 	ret = make([]byte, len(blockKeyPrefix)+len(hash[:]))
 	copy(ret, blockKeyPrefix)
@@ -236,3 +268,10 @@ func (lvl *LevelDBBackedDB) getDKGPrivateKeyKey(
 		ret[len(dkgPrivateKeyKeyPrefix):], round)
 	return
 }
+
+func (lvl *LevelDBBackedDB) getDKGMasterPrivateSharesKey(round uint64) (ret []byte) {
+	ret = make([]byte, len(dkgMasterPrivateSharesPrefix)+8)
+	copy(ret, dkgMasterPrivateSharesPrefix)
+	binary.LittleEndian.PutUint64(ret[len(dkgMasterPrivateSharesPrefix):], round)
+	return
+}
diff --git a/core/db/level-db_test.go b/core/db/level-db_test.go
index 40f9a60..4187d34 100644
--- a/core/db/level-db_test.go
+++ b/core/db/level-db_test.go
@@ -185,6 +185,44 @@ func (s *LevelDBTestSuite) TestDKGPrivateKey() {
 	s.Require().Equal(bytes.Compare(p.Bytes(), tmpPrv.Bytes()), 0)
 }
 
+func (s *LevelDBTestSuite) TestDKGMasterPrivateShares() {
+	dbName := fmt.Sprintf("test-db-%v-dkg-master-prv-shares.db", time.Now().UTC())
+	dbInst, err := NewLevelDBBackedDB(dbName)
+	s.Require().NoError(err)
+	defer func(dbName string) {
+		err = dbInst.Close()
+		s.NoError(err)
+		err = os.RemoveAll(dbName)
+		s.NoError(err)
+	}(dbName)
+
+	exists, err := dbInst.HasDKGMasterPrivateSharesKey(1)
+	s.Require().NoError(err)
+	s.Require().False(exists)
+
+	_, err = dbInst.GetDKGMasterPrivateShares(1)
+	s.Require().Equal(err.Error(), ErrDKGMasterPrivateSharesDoesNotExist.Error())
+
+	privShares, _ := dkg.NewPrivateKeyShares(10)
+
+	s.Require().NoError(dbInst.PutOrUpdateDKGMasterPrivateShares(1, *privShares))
+
+	tmpShares, err := dbInst.GetDKGMasterPrivateShares(1)
+	s.Require().NoError(err)
+	s.Require().True(tmpShares.Equal(privShares))
+
+	newPrivShares, _ := dkg.NewPrivateKeyShares(10)
+
+	// This privShare will override the old noe.
+	s.Require().NoError(dbInst.PutOrUpdateDKGMasterPrivateShares(1, *newPrivShares))
+
+	newTmpShares, err := dbInst.GetDKGMasterPrivateShares(1)
+	s.Require().NoError(err)
+	s.Require().True(newTmpShares.Equal(newPrivShares))
+
+	s.Require().False(newTmpShares.Equal(&tmpShares))
+}
+
 func TestLevelDB(t *testing.T) {
 	suite.Run(t, new(LevelDBTestSuite))
 }
diff --git a/core/db/memory.go b/core/db/memory.go
index 4bc08e7..548e41e 100644
--- a/core/db/memory.go
+++ b/core/db/memory.go
@@ -42,24 +42,27 @@ func (seq *blockSeqIterator) NextBlock() (types.Block, error) {
 
 // MemBackedDB is a memory backed DB implementation.
 type MemBackedDB struct {
-	blocksLock               sync.RWMutex
-	blockHashSequence        common.Hashes
-	blocksByHash             map[common.Hash]*types.Block
-	compactionChainTipLock   sync.RWMutex
-	compactionChainTipHash   common.Hash
-	compactionChainTipHeight uint64
-	dkgPrivateKeysLock       sync.RWMutex
-	dkgPrivateKeys           map[uint64]*dkg.PrivateKey
-	persistantFilePath       string
+	blocksLock                 sync.RWMutex
+	blockHashSequence          common.Hashes
+	blocksByHash               map[common.Hash]*types.Block
+	compactionChainTipLock     sync.RWMutex
+	compactionChainTipHash     common.Hash
+	compactionChainTipHeight   uint64
+	dkgPrivateKeysLock         sync.RWMutex
+	dkgPrivateKeys             map[uint64]*dkg.PrivateKey
+	dkgMasterPrivateSharesLock sync.RWMutex
+	dkgMasterPrivateShares     map[uint64]*dkg.PrivateKeyShares
+	persistantFilePath         string
 }
 
 // NewMemBackedDB initialize a memory-backed database.
 func NewMemBackedDB(persistantFilePath ...string) (
 	dbInst *MemBackedDB, err error) {
 	dbInst = &MemBackedDB{
-		blockHashSequence: common.Hashes{},
-		blocksByHash:      make(map[common.Hash]*types.Block),
-		dkgPrivateKeys:    make(map[uint64]*dkg.PrivateKey),
+		blockHashSequence:      common.Hashes{},
+		blocksByHash:           make(map[common.Hash]*types.Block),
+		dkgPrivateKeys:         make(map[uint64]*dkg.PrivateKey),
+		dkgMasterPrivateShares: make(map[uint64]*dkg.PrivateKeyShares),
 	}
 	if len(persistantFilePath) == 0 || len(persistantFilePath[0]) == 0 {
 		return
@@ -197,6 +200,34 @@ func (m *MemBackedDB) PutDKGPrivateKey(
 	return nil
 }
 
+// HasDKGMasterPrivateShares check existence of DKG master private shares of one round.
+func (m *MemBackedDB) HasDKGMasterPrivateShares(round uint64) (bool, error) {
+	m.dkgMasterPrivateSharesLock.RLock()
+	defer m.dkgMasterPrivateSharesLock.RUnlock()
+	_, exists := m.dkgMasterPrivateShares[round]
+	return exists, nil
+}
+
+// GetDKGMasterPrivateShares get DKG master private shares of one round.
+func (m *MemBackedDB) GetDKGMasterPrivateShares(round uint64) (
+	dkg.PrivateKeyShares, error) {
+	m.dkgMasterPrivateSharesLock.RLock()
+	defer m.dkgMasterPrivateSharesLock.RUnlock()
+	if shares, exists := m.dkgMasterPrivateShares[round]; exists {
+		return *shares, nil
+	}
+	return dkg.PrivateKeyShares{}, ErrDKGMasterPrivateSharesDoesNotExist
+}
+
+// PutOrUpdateDKGMasterPrivateShares save DKG master private shares of one round.
+func (m *MemBackedDB) PutOrUpdateDKGMasterPrivateShares(
+	round uint64, shares dkg.PrivateKeyShares) error {
+	m.dkgMasterPrivateSharesLock.Lock()
+	defer m.dkgMasterPrivateSharesLock.Unlock()
+	m.dkgMasterPrivateShares[round] = &shares
+	return nil
+}
+
 // Close implement Closer interface, which would release allocated resource.
 func (m *MemBackedDB) Close() (err error) {
 	// Save internal state to a pretty-print json file. It's a temporary way
diff --git a/core/dkg-tsig-protocol.go b/core/dkg-tsig-protocol.go
index 6c812db..259d07a 100644
--- a/core/dkg-tsig-protocol.go
+++ b/core/dkg-tsig-protocol.go
@@ -24,6 +24,7 @@ import (
 	"github.com/dexon-foundation/dexon-consensus/common"
 	"github.com/dexon-foundation/dexon-consensus/core/crypto"
 	"github.com/dexon-foundation/dexon-consensus/core/crypto/dkg"
+	"github.com/dexon-foundation/dexon-consensus/core/db"
 	"github.com/dexon-foundation/dexon-consensus/core/types"
 	typesDKG "github.com/dexon-foundation/dexon-consensus/core/types/dkg"
 	"github.com/dexon-foundation/dexon-consensus/core/utils"
@@ -164,6 +165,35 @@ func newDKGProtocol(
 	}
 }
 
+func recoverDKGProtocol(
+	ID types.NodeID,
+	recv dkgReceiver,
+	round uint64,
+	threshold int,
+	coreDB db.Database) (*dkgProtocol, error) {
+	shares, err := coreDB.GetDKGMasterPrivateShares(round)
+	if err != nil {
+		if err == db.ErrDKGMasterPrivateSharesDoesNotExist {
+			return nil, nil
+		}
+
+		return nil, err
+	}
+	return &dkgProtocol{
+		ID:                    ID,
+		recv:                  recv,
+		round:                 round,
+		threshold:             threshold,
+		idMap:                 make(map[types.NodeID]dkg.ID),
+		mpkMap:                make(map[types.NodeID]*dkg.PublicKeyShares),
+		masterPrivateShare:    &shares,
+		prvShares:             dkg.NewEmptyPrivateKeyShares(),
+		prvSharesReceived:     make(map[types.NodeID]struct{}),
+		nodeComplained:        make(map[types.NodeID]struct{}),
+		antiComplaintReceived: make(map[types.NodeID]map[types.NodeID]struct{}),
+	}, nil
+}
+
 func (d *dkgProtocol) processMasterPublicKeys(
 	mpks []*typesDKG.MasterPublicKey) (err error) {
 	d.idMap = make(map[types.NodeID]dkg.ID, len(mpks))