core: add reset to db.PutDKGPrivateKey (#560)

6 files changed, 78 insertions, 53 deletions

diff --git a/core/configuration-chain.go b/core/configuration-chain.go
index 76917a3..e9e04a2 100644
--- a/core/configuration-chain.go
+++ b/core/configuration-chain.go
@@ -396,7 +396,8 @@ func (cc *configurationChain) runDKGPhaseNine(round uint64, reset uint64) error
 		return err
 	}
 	// Save private shares to DB.
-	if err = cc.db.PutDKGPrivateKey(round, *signer.privateKey); err != nil {
+	if err =
+		cc.db.PutDKGPrivateKey(round, reset, *signer.privateKey); err != nil {
 		return err
 	}
 	cc.dkgResult.Lock()
@@ -615,8 +616,9 @@ func (cc *configurationChain) recoverDKGInfo(
 		}()
 	}
 	if !signerExists && !ignoreSigner {
+		reset := cc.gov.DKGResetCount(round)
 		// Check if we have private shares in DB.
-		prvKey, err := cc.db.GetDKGPrivateKey(round)
+		prvKey, err := cc.db.GetDKGPrivateKey(round, reset)
 		if err != nil {
 			cc.logger.Warn("Failed to create DKGPrivateKey",
 				"round", round, "error", err)
@@ -638,7 +640,8 @@ func (cc *configurationChain) recoverDKGInfo(
 					"round", round, "error", err)
 				return err
 			}
-			if err = cc.db.PutDKGPrivateKey(round, *prvKeyRecover); err != nil {
+			if err = cc.db.PutDKGPrivateKey(
+				round, reset, *prvKeyRecover); err != nil {
 				cc.logger.Warn("Failed to save DKGPrivateKey",
 					"round", round, "error", err)
 			}
diff --git a/core/db/interfaces.go b/core/db/interfaces.go
index 2c32ebb..a571a80 100644
--- a/core/db/interfaces.go
+++ b/core/db/interfaces.go
@@ -80,8 +80,7 @@ type Reader interface {
 	GetCompactionChainTipInfo() (common.Hash, uint64)
 
 	// DKG Private Key related methods.
-	HasDKGPrivateKey(round uint64) (bool, error)
-	GetDKGPrivateKey(round uint64) (dkg.PrivateKey, error)
+	GetDKGPrivateKey(round, reset uint64) (dkg.PrivateKey, error)
 	GetDKGProtocol() (dkgProtocol DKGProtocolInfo, err error)
 }
 
@@ -90,7 +89,7 @@ type Writer interface {
 	UpdateBlock(block types.Block) error
 	PutBlock(block types.Block) error
 	PutCompactionChainTipInfo(common.Hash, uint64) error
-	PutDKGPrivateKey(uint64, dkg.PrivateKey) error
+	PutDKGPrivateKey(round, reset uint64, pk dkg.PrivateKey) error
 	PutOrUpdateDKGProtocol(dkgProtocol DKGProtocolInfo) error
 }
 
diff --git a/core/db/level-db.go b/core/db/level-db.go
index 5296833..da8bc0b 100644
--- a/core/db/level-db.go
+++ b/core/db/level-db.go
@@ -59,6 +59,11 @@ type DKGProtocolInfo struct {
 	Reset                     uint64
 }
 
+type dkgPrivateKey struct {
+	PK    dkg.PrivateKey
+	Reset uint64
+}
+
 // Equal compare with target DKGProtocolInfo.
 func (info *DKGProtocolInfo) Equal(target *DKGProtocolInfo) bool {
 	if !info.ID.Equal(target.ID) ||
@@ -478,13 +483,8 @@ func (lvl *LevelDBBackedDB) GetCompactionChainTipInfo() (
 	return
 }
 
-// HasDKGPrivateKey check existence of DKG private key of one round.
-func (lvl *LevelDBBackedDB) HasDKGPrivateKey(round uint64) (bool, error) {
-	return lvl.db.Has(lvl.getDKGPrivateKeyKey(round), nil)
-}
-
 // GetDKGPrivateKey get DKG private key of one round.
-func (lvl *LevelDBBackedDB) GetDKGPrivateKey(round uint64) (
+func (lvl *LevelDBBackedDB) GetDKGPrivateKey(round, reset uint64) (
 	prv dkg.PrivateKey, err error) {
 	queried, err := lvl.db.Get(lvl.getDKGPrivateKeyKey(round), nil)
 	if err != nil {
@@ -493,22 +493,32 @@ func (lvl *LevelDBBackedDB) GetDKGPrivateKey(round uint64) (
 		}
 		return
 	}
-	err = rlp.DecodeBytes(queried, &prv)
+	pk := dkgPrivateKey{}
+	err = rlp.DecodeBytes(queried, &pk)
+	if pk.Reset != reset {
+		err = ErrDKGPrivateKeyDoesNotExist
+		return
+	}
+	prv = pk.PK
 	return
 }
 
 // PutDKGPrivateKey save DKG private key of one round.
 func (lvl *LevelDBBackedDB) PutDKGPrivateKey(
-	round uint64, prv dkg.PrivateKey) error {
+	round, reset uint64, prv dkg.PrivateKey) error {
 	// Check existence.
-	exists, err := lvl.HasDKGPrivateKey(round)
-	if err != nil {
+	_, err := lvl.GetDKGPrivateKey(round, reset)
+	if err == nil {
+		return ErrDKGPrivateKeyExists
+	}
+	if err != ErrDKGPrivateKeyDoesNotExist {
 		return err
 	}
-	if exists {
-		return ErrDKGPrivateKeyExists
+	pk := &dkgPrivateKey{
+		PK:    prv,
+		Reset: reset,
 	}
-	marshaled, err := rlp.EncodeToBytes(&prv)
+	marshaled, err := rlp.EncodeToBytes(&pk)
 	if err != nil {
 		return err
 	}
diff --git a/core/db/level-db_test.go b/core/db/level-db_test.go
index 78e3b77..8b3c448 100644
--- a/core/db/level-db_test.go
+++ b/core/db/level-db_test.go
@@ -169,22 +169,31 @@ func (s *LevelDBTestSuite) TestDKGPrivateKey() {
 		s.NoError(err)
 	}(dbName)
 	p := dkg.NewPrivateKey()
-	// Check existence.
-	exists, err := dbInst.HasDKGPrivateKey(1)
-	s.Require().NoError(err)
-	s.Require().False(exists)
-	// We should be unable to get it, too.
-	_, err = dbInst.GetDKGPrivateKey(1)
+	// We should be unable to get it.
+	_, err = dbInst.GetDKGPrivateKey(1, 0)
 	s.Require().Equal(err.Error(), ErrDKGPrivateKeyDoesNotExist.Error())
 	// Put it.
-	s.Require().NoError(dbInst.PutDKGPrivateKey(1, *p))
+	s.Require().NoError(dbInst.PutDKGPrivateKey(1, 0, *p))
+	// We should be unable to get it because reset is different.
+	_, err = dbInst.GetDKGPrivateKey(1, 1)
+	s.Require().Equal(err.Error(), ErrDKGPrivateKeyDoesNotExist.Error())
 	// Put it again, should not success.
-	err = dbInst.PutDKGPrivateKey(1, *p)
+	err = dbInst.PutDKGPrivateKey(1, 0, *p)
 	s.Require().Equal(err.Error(), ErrDKGPrivateKeyExists.Error())
 	// Get it back.
-	tmpPrv, err := dbInst.GetDKGPrivateKey(1)
+	tmpPrv, err := dbInst.GetDKGPrivateKey(1, 0)
 	s.Require().NoError(err)
 	s.Require().Equal(bytes.Compare(p.Bytes(), tmpPrv.Bytes()), 0)
+	// Put it at different reset.
+	p2 := dkg.NewPrivateKey()
+	s.Require().NoError(dbInst.PutDKGPrivateKey(1, 1, *p2))
+	// We should be unable to get it because reset is different.
+	_, err = dbInst.GetDKGPrivateKey(1, 0)
+	// Get it back.
+	tmpPrv, err = dbInst.GetDKGPrivateKey(1, 1)
+	s.Require().NoError(err)
+	s.Require().Equal(bytes.Compare(p2.Bytes(), tmpPrv.Bytes()), 0)
+	s.Require().NotEqual(bytes.Compare(p2.Bytes(), p.Bytes()), 0)
 }
 
 func (s *LevelDBTestSuite) TestDKGProtocol() {
diff --git a/core/db/memory.go b/core/db/memory.go
index 971f758..6555de8 100644
--- a/core/db/memory.go
+++ b/core/db/memory.go
@@ -49,7 +49,7 @@ type MemBackedDB struct {
 	compactionChainTipHash   common.Hash
 	compactionChainTipHeight uint64
 	dkgPrivateKeysLock       sync.RWMutex
-	dkgPrivateKeys           map[uint64]*dkg.PrivateKey
+	dkgPrivateKeys           map[uint64]*dkgPrivateKey
 	dkgProtocolLock          sync.RWMutex
 	dkgProtocolInfo          *DKGProtocolInfo
 	persistantFilePath       string
@@ -61,7 +61,7 @@ func NewMemBackedDB(persistantFilePath ...string) (
 	dbInst = &MemBackedDB{
 		blockHashSequence: common.Hashes{},
 		blocksByHash:      make(map[common.Hash]*types.Block),
-		dkgPrivateKeys:    make(map[uint64]*dkg.PrivateKey),
+		dkgPrivateKeys:    make(map[uint64]*dkgPrivateKey),
 	}
 	if len(persistantFilePath) == 0 || len(persistantFilePath[0]) == 0 {
 		return
@@ -168,34 +168,29 @@ func (m *MemBackedDB) GetCompactionChainTipInfo() (
 	return m.compactionChainTipHash, m.compactionChainTipHeight
 }
 
-// HasDKGPrivateKey check existence of DKG private key of one round.
-func (m *MemBackedDB) HasDKGPrivateKey(round uint64) (bool, error) {
-	m.dkgPrivateKeysLock.RLock()
-	defer m.dkgPrivateKeysLock.RUnlock()
-	_, exists := m.dkgPrivateKeys[round]
-	return exists, nil
-}
-
 // GetDKGPrivateKey get DKG private key of one round.
-func (m *MemBackedDB) GetDKGPrivateKey(round uint64) (
+func (m *MemBackedDB) GetDKGPrivateKey(round, reset uint64) (
 	dkg.PrivateKey, error) {
 	m.dkgPrivateKeysLock.RLock()
 	defer m.dkgPrivateKeysLock.RUnlock()
-	if prv, exists := m.dkgPrivateKeys[round]; exists {
-		return *prv, nil
+	if prv, exists := m.dkgPrivateKeys[round]; exists && prv.Reset == reset {
+		return prv.PK, nil
 	}
 	return dkg.PrivateKey{}, ErrDKGPrivateKeyDoesNotExist
 }
 
 // PutDKGPrivateKey save DKG private key of one round.
 func (m *MemBackedDB) PutDKGPrivateKey(
-	round uint64, prv dkg.PrivateKey) error {
+	round, reset uint64, prv dkg.PrivateKey) error {
 	m.dkgPrivateKeysLock.Lock()
 	defer m.dkgPrivateKeysLock.Unlock()
-	if _, exists := m.dkgPrivateKeys[round]; exists {
+	if prv, exists := m.dkgPrivateKeys[round]; exists && prv.Reset == reset {
 		return ErrDKGPrivateKeyExists
 	}
-	m.dkgPrivateKeys[round] = &prv
+	m.dkgPrivateKeys[round] = &dkgPrivateKey{
+		PK:    prv,
+		Reset: reset,
+	}
 	return nil
 }
 
diff --git a/core/db/memory_test.go b/core/db/memory_test.go
index b0bfb16..33e0111 100644
--- a/core/db/memory_test.go
+++ b/core/db/memory_test.go
@@ -154,22 +154,31 @@ func (s *MemBackedDBTestSuite) TestDKGPrivateKey() {
 	s.Require().NoError(err)
 	s.Require().NotNil(dbInst)
 	p := dkg.NewPrivateKey()
-	// Check existence.
-	exists, err := dbInst.HasDKGPrivateKey(1)
-	s.Require().NoError(err)
-	s.Require().False(exists)
-	// We should be unable to get it, too.
-	_, err = dbInst.GetDKGPrivateKey(1)
+	// We should be unable to get it.
+	_, err = dbInst.GetDKGPrivateKey(1, 0)
 	s.Require().Equal(err.Error(), ErrDKGPrivateKeyDoesNotExist.Error())
 	// Put it.
-	s.Require().NoError(dbInst.PutDKGPrivateKey(1, *p))
+	s.Require().NoError(dbInst.PutDKGPrivateKey(1, 0, *p))
+	// We should be unable to get it because reset is different.
+	_, err = dbInst.GetDKGPrivateKey(1, 1)
+	s.Require().Equal(err.Error(), ErrDKGPrivateKeyDoesNotExist.Error())
 	// Put it again, should not success.
-	err = dbInst.PutDKGPrivateKey(1, *p)
+	err = dbInst.PutDKGPrivateKey(1, 0, *p)
 	s.Require().Equal(err.Error(), ErrDKGPrivateKeyExists.Error())
 	// Get it back.
-	tmpPrv, err := dbInst.GetDKGPrivateKey(1)
+	tmpPrv, err := dbInst.GetDKGPrivateKey(1, 0)
 	s.Require().NoError(err)
 	s.Require().Equal(bytes.Compare(p.Bytes(), tmpPrv.Bytes()), 0)
+	// Put it at different reset.
+	p2 := dkg.NewPrivateKey()
+	s.Require().NoError(dbInst.PutDKGPrivateKey(1, 1, *p2))
+	// We should be unable to get it because reset is different.
+	_, err = dbInst.GetDKGPrivateKey(1, 0)
+	// Get it back.
+	tmpPrv, err = dbInst.GetDKGPrivateKey(1, 1)
+	s.Require().NoError(err)
+	s.Require().Equal(bytes.Compare(p2.Bytes(), tmpPrv.Bytes()), 0)
+	s.Require().NotEqual(bytes.Compare(p2.Bytes(), p.Bytes()), 0)
 }
 
 func TestMemBackedDB(t *testing.T) {