core, core: crypto: dkg: fix concurrent access bug (#551)

* core: crypto: dkg: fix concurrent access bug

* core: fix concurrent bug

3 files changed, 36 insertions, 10 deletions

diff --git a/core/consensus.go b/core/consensus.go
index e657c64..adaf51d 100644
--- a/core/consensus.go
+++ b/core/consensus.go
@@ -57,6 +57,8 @@ var (
 		"cannot verify block randomness")
 )
 
+type selfAgreementResult types.AgreementResult
+
 // consensusBAReceiver implements agreementReceiver.
 type consensusBAReceiver struct {
 	consensus         *Consensus
@@ -286,7 +288,8 @@ func (recv *consensusBAReceiver) ConfirmBlock(
 				IsEmptyBlock: isEmptyBlockConfirmed,
 				Randomness:   block.Randomness,
 			}
-			recv.consensus.baMgr.touchAgreementResult(result)
+			// touchAgreementResult does not support concurrent access.
+			recv.consensus.msgChan <- (*selfAgreementResult)(result)
 			recv.consensus.logger.Debug("Broadcast AgreementResult",
 				"result", result)
 			recv.consensus.network.BroadcastAgreementResult(result)
@@ -1217,6 +1220,8 @@ MessageLoop:
 			return
 		}
 		switch val := msg.(type) {
+		case *selfAgreementResult:
+			con.baMgr.touchAgreementResult((*types.AgreementResult)(val))
 		case *types.Block:
 			if ch, exist := func() (chan<- *types.Block, bool) {
 				con.lock.RLock()
diff --git a/core/crypto/dkg/dkg.go b/core/crypto/dkg/dkg.go
index b89ad10..ab43f51 100644
--- a/core/crypto/dkg/dkg.go
+++ b/core/crypto/dkg/dkg.go
@@ -450,11 +450,10 @@ func (prvs *PrivateKeyShares) Share(ID ID) (*PrivateKey, bool) {
 
 // NewEmptyPublicKeyShares creates an empty public key shares.
 func NewEmptyPublicKeyShares() *PublicKeyShares {
-	cache := &publicKeySharesCache{
-		index: make(map[ID]int),
-	}
 	pubShares := &PublicKeyShares{}
-	pubShares.cache.Store(cache)
+	pubShares.cache.Store(&publicKeySharesCache{
+		index: make(map[ID]int),
+	})
 	return pubShares
 }
 
@@ -481,9 +480,9 @@ func (pubs *PublicKeyShares) Share(ID ID) (*PublicKey, error) {
 }
 
 // AddShare adds a share.
-func (pubs *PublicKeyShares) AddShare(ID ID, share *PublicKey) error {
+func (pubs *PublicKeyShares) AddShare(shareID ID, share *PublicKey) error {
 	cache := pubs.cache.Load().(*publicKeySharesCache)
-	if idx, exist := cache.index[ID]; exist {
+	if idx, exist := cache.index[shareID]; exist {
 		if !share.publicKey.IsEqual(&cache.share[idx].publicKey) {
 			return ErrDuplicatedShare
 		}
@@ -492,9 +491,17 @@ func (pubs *PublicKeyShares) AddShare(ID ID, share *PublicKey) error {
 	pubs.lock.Lock()
 	defer pubs.lock.Unlock()
 	cache = pubs.cache.Load().(*publicKeySharesCache)
-	cache.index[ID] = len(cache.share)
-	cache.share = append(cache.share, *share)
-	pubs.cache.Store(cache)
+	newCache := &publicKeySharesCache{
+		index: make(map[ID]int, len(cache.index)+1),
+		share: make([]PublicKey, len(cache.share), len(cache.share)+1),
+	}
+	for k, v := range cache.index {
+		newCache.index[k] = v
+	}
+	copy(newCache.share, cache.share)
+	newCache.index[shareID] = len(newCache.share)
+	newCache.share = append(newCache.share, *share)
+	pubs.cache.Store(newCache)
 	return nil
 }
 
diff --git a/core/crypto/dkg/dkg_test.go b/core/crypto/dkg/dkg_test.go
index a679cf4..84dc2d4 100644
--- a/core/crypto/dkg/dkg_test.go
+++ b/core/crypto/dkg/dkg_test.go
@@ -394,6 +394,20 @@ func (s *DKGTestSuite) TestPublicKeySharesMove() {
 	req.True(pubShares2.Equal(pubShares3))
 }
 
+func (s *DKGTestSuite) TestPublicKeySharesConcurrent() {
+	t := 5
+	n := 10
+	IDs := make(IDs, n)
+	for i := range IDs {
+		id := common.NewRandomHash()
+		IDs[i] = NewID(id[:])
+	}
+	_, pubShare := NewPrivateKeyShares(t)
+	for _, id := range IDs {
+		go pubShare.Share(id)
+	}
+}
+
 func (s *DKGTestSuite) TestPrivateKeySharesEquality() {
 	var req = s.Require()
 	IDs := s.genID(2)